ProducerConsumerWarpSpecialized: extend WS+TMA to sibling pipelined loops
- Dominant language
- Python
- Stars
- 7.4k
- Forks
- 742
- Avg merge
- 1d 50m
- Merged PRs (30d)
- 104
Description
## Summary
`tilelang.transform.ProducerConsumerWarpSpecialized` (defined in
`src/transform/producer_consumer_ws.cc`) only wraps the **first**
pipelined loop per block in a WS+TMA scope. Subsequent sibling
`T.Pipelined` loops in the same block fall back to per-thread
`cp_async_gs` copies even when their copies would otherwise be
TMA-eligible.
This is acknowledged in the source as a v1 limitation:
```cpp
* Limitations (v1):
* - Pure TMA pipelines only (no mixed TMA + cp.async)
* - No conditionally guarded loop bodies (phase counters)
* - Single pipelined loop per block
* - No pre-loop TMA prefetch / prologue optimizations
```
Filing this as a tracking issue + a record of what doesn't work, in
case it's useful for whoever picks this up.
## Motivating workload
5-pass sparsemax-attention forward kernel (AdaSplash2-style) on sm_90a:
each pass is its own `T.Pipelined` K-scan. Pass 1's K_shared is loaded
via TMA; passes 2-5 emit `cp_async_gs` per thread instead.
In the generated CUDA, passes 2-5 collectively account for ~3x as many
SMEM-bytes-loaded as pass 1, so leaving them on cp.async meaningfully
caps achievable HBM bandwidth for the kernel.
## What we tried (negative result)
A small C++ patch:
1. `FindPipelineLoop` skips subtrees inside
`AttrStmt(kWarpSpecializationScope)` (otherwise iteration re-finds the
already-wrapped first loop):
```cpp
if (auto *attr = stmt.as()) {
if (attr->attr_key == attr::kWarpSpecializationScope) {
return nullptr;
}
return FindPipelineLoop(attr->body);
}
```
2. Driver iterates `ProducerConsumerWSRewriter::Substitute` until
`TiledWSCandidate::Check` returns false.
Built TileLang v0.1.9 from source with the patch, dropped the rebuilt
`libtilelang.so` + `libtvm.so` into the venv, ran a minimal 2-sibling
GEMM kernel + the 5-pass sparsemax-attention kernel.
**Result**: identical codegen to unpatched. Both kernels still generate
exactly **one** `if (256 <= threadIdx.x)` WS partition. The micro-test
2-sibling GEMM produced `tma_load=3 cp_async_gs=6 ws_partitions=1`
(same as baseline).
Most likely the rewriter's stateful assumptions (`num_threads_`,
mbarrier allocation, `ws_transformed_`) cause iter-2's `Substitute` to
either bail silently or no-op. Fixing this properly would require:
- Per-loop disjoint mbarrier ranges (not one shared `barrier_buf` per
pass invocation),
- Idempotent `num_threads_` updates,
- Possibly a unified producer warp that serves all sibling loops
sequentially rather than one WS scope per loop.
That's a bigger refactor than I have bandwidth for right now, but
recording the dead end so the next attempt doesn't repeat it.
## Reproducer
Patch + build script + micro-test: experiment/tilelang-patched-ws
branch of the project that hit this:
- `patches/tilelang_iterative_ws.patch` – the .cc diff
- `scripts/build_patched_tilelang.py` – Modal-based source rebuild
- `scripts/micro_test_patched_ws.py` – 2-sibling-pipelined-GEMM smoke
test
## Environment
- TileLang: v0.1.9 (commit 441c3b06acb23b09d68639532d3a21f427370ced)
- Modal H100 (sm_90a), CUDA 12.8
- nvidia/cuda:12.8.0-devel-ubuntu22.04 + Python 3.12 + uv venv
Contributor guide
Assessment
This issue has not been assessed yet.