tile-ai / tile-ai/tilelang

[BUG] Unexpected CTA-wide __syncthreads() inserted between split T.ws producer/consumer scopes

Open
#2,192 0 comments 0 reactions 1 assignee Claimed by @LeiWang1999 View on GitHub
bug
Dominant language
Python
Stars
7.4k
Forks
742
Avg merge
1d 2h
Merged PRs (30d)
108

Description

### Required prerequisites

- [x] I have read the documentation .
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) that this hasn't already been reported. (comment there if it has.)

### What version of TileLang are you using?

0.1.9

### System information

3.12.3 (main, Feb 4 2025, 14:48:35) [GCC 13.3.0] linux
0.1.9
2.8.0a0+5228986c39.nv25.05

PyTorch version: 2.8.0a0+5228986c39.nv25.05
CUDA used to build PyTorch: 12.9
OS: Ubuntu 24.04 x86_64
Python version: 3.12.3
Is CUDA available: True
GPU: NVIDIA H200
Nvidia driver version: 570.86.15
CUDA runtime version: 12.9.41
cuDNN version: 9.x

CPU: x86_64, Intel Xeon, 2 sockets, 192 logical CPUs

Relevant packages:
numpy==1.26.4
torch==2.8.0a0+5228986c39.nv25.5
torchvision==0.22.0a0
triton==3.1.0
pytorch-triton==3.3.0+
torch_tensorrt==2.8.0a0

### Problem description

`T.ws` producer/consumer scopes appear to be serialized by an unexpected CTA-wide `__syncthreads()`.

This can be reproduced directly with the existing TileLang example:

`python3 tilelang/examples/warp_specialize/example_warp_specialize_gemm_barrierpipe_stage2.py
`

The example contains a producer scope and a consumer scope inside the same loop:

```python
for ko in range(T.ceildiv(K, block_K)):
with T.ws(1):
T.mbarrier_wait_parity(...)
T.tma_copy(...)
T.tma_copy(...)
T.mbarrier_arrive(...)

with T.ws(0):
T.mbarrier_wait_parity(...)
T.gemm(...)
T.mbarrier_arrive(...)
```

However, the generated CUDA source contains CTA-wide synchronization between these scopes:

```cpp
for (int ko = 0; ko < ...; ++ko) {
__syncthreads();

if (128 <= ((int)threadIdx.x)) {
// producer: wait empty, issue TMA loads, arrive ready
}

__syncthreads();

if (((int)threadIdx.x) < 128) {
// consumer: wait ready, run WGMMA, arrive empty
}
}
```

This serializes the producer and consumer phases within each pipeline iteration. As a result, the generated code does not appear to implement an overlapped producer/consumer pipeline, even though the source uses separate `T.ws` scopes and explicit mbarrier synchronization.

### Reproducible example code

The Python snippets:

```python
python3 tilelang/examples/warp_specialize/example_warp_specialize_gemm_barrierpipe_stage2.py
```

### Traceback

```pytb

```

### Expected behavior

The producer and consumer warp-specialized roles should be able to run concurrently, synchronized only by the explicit mbarriers in the program.

In other words, the generated code should be closer to a persistent role split such as:

```cpp
if (producer_threads) {
for (...) {
// producer pipeline
}
} else {
for (...) {
// consumer pipeline
}
}
```

without CTA-wide `__syncthreads()` inserted between producer and consumer work in each iteration.

### Additional context

If the producer/consumer split is written manually using `threadIdx.x` / `T.get_thread_binding()` instead of separate `T.ws` scopes, the generated CUDA source does not contain the extra per-iteration `__syncthreads()`.

For example, writing the roles as a single thread-id branch:

```python
tx = T.get_thread_binding()

if tx >= 128:
for ko in range(T.ceildiv(K, block_K)):
T.mbarrier_wait_parity(...)
T.tma_copy(...)
T.tma_copy(...)
T.mbarrier_arrive(...)
else:
for ko in range(T.ceildiv(K, block_K)):
T.mbarrier_wait_parity(...)
T.gemm(...)
T.mbarrier_arrive(...)
```

generates a persistent producer/consumer role split without the extra CTA-wide synchronization between the roles.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.