diff --git a/README.md b/README.md index 6485d2d1..bd78b435 100644 --- a/README.md +++ b/README.md @@ -1,68 +1,73 @@ # DSV4 Inference Kernel -## ⚠️⚠️⚠️ CRITICAL: TMA Partition Tensor Coordinate Space ⚠️⚠️⚠️ +## ⚠️⚠️⚠️ CRITICAL: TMA Partition Tensor Mode Ordering ⚠️⚠️⚠️ **THIS BUG COST US AN ENTIRE DAY. READ THIS. BURN IT INTO YOUR BRAIN.** -After `cpasync.tma_partition()`, the output GMEM tensor has **8 TMA coordinate dimensions**: +After `cpasync.tma_partition()`, the output GMEM tensor has **4 modes** (verified on B200): ``` -tBgK TMA coord space: (1, 1, 1, 1, n_kv_tiles, 1, 1, 1) - 0 1 2 3 4 5 6 7 +tBgK shape: (((64, 128), 1), ?, KV_tiles, ?) + mode 0 1 2 3 ``` -**Mode 4 is the GMEM tile dimension.** The dimension you index with `kt` to load different K/V tiles. - -The Python-visible shape only shows 4 modes, but the TMA coordinate space is 8-dimensional. You MUST apply an 8-None no-op pre-slice to open the full coordinate space before indexing. +**Mode 2 is the GMEM tile dimension.** The dimension you index with `kt` to load different K/V tiles. ### THE WRONG WAY (what we did — silently loads from tile 0 forever): ```python -# ❌❌❌ 4-MODE PRE-SLICE COLLAPSES THE GMEM TILE AXIS ❌❌❌ -# The (None, None, 0, 0) slice only addresses 4 of 8 TMA coord dims. -# Modes 4-7 get collapsed to coordinate 0. TMA ALWAYS reads tile 0. -tBgK = tBgK[(None, None, 0, 0)] # ← WRONG! +# ❌❌❌ (None,None,0,0) KEEPS MODES 0,1 FREE, SETS MODE 2 TO 0 ❌❌❌ +# Mode 2 (the KV tile dim) gets collapsed to coordinate 0. +# TMA ALWAYS reads from tile 0. +tBgK = tBgK[(None, None, 0, 0)] # ← WRONG! Mode 2 pinned to 0! -# The copy "works" but kv_coord indexes mode 1 (size 1), so -# every coordinate maps to the same TMA descriptor. -cute.copy(tma_k, tBgK[(None, kv_coord)], ...) # ← kv_coord is ignored! +# The copy "works" but kv_coord indexes mode 1 (inner GEMM K, not KV tiles). +cute.copy(tma_k, tBgK[(None, kv_coord)], ...) # ← kv_coord indexes wrong mode! ``` -### THE RIGHT WAY (what actually works — confirmed on B200 at n=256): +### THE RIGHT WAY (verified on B200 at n=128 and n=256): ```python -# ✅ 8-None no-op pre-slice opens the full TMA coordinate space -tBgK = tBgK[(None, None, None, None, None, None, None, None)] +# ✅ (None,0,None,0) keeps modes 0 and 2 free → 2D tensor +# Mode 2 (KV tiles) survives as the second mode. +tBgK = tBgK[(None, 0, None, 0)] -# ✅ Index mode 4 (the GMEM tile dim) in the copy call -cute.copy(tma_k, tBgK[None, None, None, None, kt, None, None, None], ...) -# ^^ MODE 4 — THE GMEM TILE DIM +# ✅ [None, kt] indexes the surviving mode 1 (originally mode 2 = KV tiles) +cute.copy(tma_k, tBgK[None, kt], ...) +# ^^ THIS IS THE KV TILE DIM +``` + +**Verified shapes on B200 (May 22, n=256, inside @cute.kernel):** +``` +Before slice: tBgK = (((64,128),1), Int32(?), Int32(?), Int32(?)) — 4 modes +After (None,0,None,0): tBgK = (((64,128),1), Int32(?)) — 2 modes ``` ### WHY THIS IS SO INSIDIOUS -1. **No error, no warning.** The slice `tBgK[(None,None,0,0)]` silently collapses modes 4-7. -2. **Single-tile (n=128) works perfectly.** With only 1 KV tile, mode 4 is size 1, so the bug is invisible. -3. **All multi-tile tests produce "reasonable" output.** The TMA loads from tile 0 every time, so you get a valid (but wrong) attention computation. Cosine similarity is 0.7-0.9, not NaN. +1. **No error, no warning.** The slice `tBgK[(None,None,0,0)]` silently sets mode 2 to 0. +2. **Single-tile (n=128) works perfectly.** With only 1 KV tile, mode 2 is size 1, so the bug is invisible. +3. **Multi-tile tests produce "reasonable" output.** The TMA loads from tile 0 every time, so you get a valid (but wrong) attention computation. Cosine similarity is 0.7-0.9, not NaN. 4. **The strides are all 0.** Printing `tBgK.layout.stride` shows all zeros for TMA tensors. You can't detect the bug from strides alone. 5. **`cute.printf` shows `kv_coord=0`.** We thought the JIT was constant-folding the variable. It wasn't — the variable was fine, but it was indexing the wrong mode. +6. **The 8-mode theory was wrong.** We assumed tma_partition produced 8 TMA coordinate dimensions. It produces 4. The 8-None no-op slice fails with "weakly congruent" at JIT compile. ### THE LESSON -**TMA tensors use a coordinate space, not a pointer space.** The TMA instruction at PTX level takes integer coordinates (`crd0, crd1, crd2, ...`), not pointers. CuTeDSL's `tma_partition` produces a tensor whose layout maps logical coordinates to TMA coordinate tuples. When you pre-slice with fewer dimensions than the TMA descriptor expects, the extra coordinate dimensions get collapsed to 0. +**PRINT THE SHAPES. ALWAYS.** Run `print(f"tBgK: shape={cute.shape(tBgK)}")` inside `@cute.kernel` at trace time. The shapes are your ground truth. Reasoning about mode counts without evidence is how we wasted a day. -**The 8-None no-op pre-slice is mandatory for multi-tile TMA.** Without it, the GMEM tile axis (mode 4) is invisible and unindexable. +**The correct pre-slice depends on which mode is the GMEM tile iteration axis.** For our `local_tile` + `partition_B` + `group_modes(0,3)` pattern, mode 2 is the KV tile axis. `(None,0,None,0)` keeps it free. `(None,None,0,0)` collapses it to 0. ```python -# After tma_partition, ALWAYS apply the 8-None no-op pre-slice: -tBgK = tBgK[(None, None, None, None, None, None, None, None)] -tVgV = tVgV[(None, None, None, None, None, None, None, None)] +# ALWAYS verify the shape at trace time: +print(f"tBgK shape: {cute.shape(tBgK)}") # 4 modes +print(f"tBgK after slice: {cute.shape(tBgK[(None,0,None,0)])}") # 2 modes -# Then index mode 4 in cute.copy: -cute.copy(tma_k, tBgK[None, None, None, None, kt, None, None, None], ...) +# Then index the 2D tensor: +cute.copy(tma_k, tBgK[None, kt], ...) ``` -**IF YOU SKIP THE 8-NONE PRE-SLICE, MULTI-TILE TMA WILL BE SILENTLY BROKEN.** +**IF YOU USE (None,None,0,0) INSTEAD OF (None,0,None,0), MULTI-TILE TMA WILL BE SILENTLY BROKEN.** --- @@ -140,7 +145,7 @@ Summary |-------|--------|-------------| | A | ✅ COMPLETE | Q@K^T via tcgen05.mma → TMEM → GMEM | | B | ✅ COMPLETE | QK → identity softmax → P@V pipeline (TMEM alias, KV-tile interleaving) | -| C | ⚠️ MULTI-TILE IN PROGRESS | Single-tile cos 0.999998. TMA fix: n=256 cos 0.9956. Need O rescale + pipeline cycling. | +| C | ⚠️ MULTI-TILE TMA FIXED | n=128 cos 0.999998 ✅. TMA fix: n=256 loads 2 tiles. Pipeline cycling needed for n≥384. O rescale needed. | | C' | 🔨 IN PROGRESS | Multi-tile TMA indexing fix + correction warps. See below. | | D | TODO | Full decode attention: paged KV cache, multi-query, causal mask | | E | TODO | Production kernel: extract into dsv4/kernels/attention/, PyTorch custom op, vLLM bridge | @@ -294,30 +299,39 @@ What it does: ### Multi-Tile TMA Fix (RESOLVED — was a LAYOUT bug, not a JIT bug) -After `cpasync.tma_partition()`, the output GMEM tensor has **8 TMA coordinate dimensions**: +After `cpasync.tma_partition()`, the output GMEM tensor has **4 modes**: `(((64,128),1), ?, KV_tiles, ?)`. -``` -tBgK TMA coord space: (1, 1, 1, 1, n_kv_tiles, 1, 1, 1) - 0 1 2 3 4 5 6 7 -``` +**Mode 2 is the GMEM tile dimension.** Our old pre-slice `tBgK[(None, None, 0, 0)]` kept modes 0,1 free and set mode 2 to 0, so TMA always read tile 0. The bug looked like "JIT constant-folding" but was purely a layout error. -**Mode 4 is the GMEM tile dimension.** Our old pre-slice `tBgK[(None, None, 0, 0)]` collapsed modes 4-7 to coordinate 0, so TMA always read tile 0. The bug looked like "JIT constant-folding" but was purely a layout error. - -**The fix:** 8-None no-op pre-slice + 8-mode indexing with `kt` at mode 4: +**The fix:** `(None,0,None,0)` keeps modes 0,2 free, then `[None, kt]` indexes KV tiles: ```python -tBgK = tBgK[(None, None, None, None, None, None, None, None)] -cute.copy(tma_k, tBgK[None, None, None, None, kt, None, None, None], ...) +tBgK = tBgK[(None, 0, None, 0)] +cute.copy(tma_k, tBgK[None, kt], ...) ``` -**Results after fix:** +**Results after TMA fix (verified on B200, May 22):** - n=128: cos 0.999998 ✅ -- n=256: cos 0.9956 ✅ (lower because no O rescale yet) +- n=256: cos 0.71 (TMA loads 2 tiles correctly, needs O rescale for 0.9999) +- n=512/1024: output identical to n=256 — pipeline not cycling past kv_stage=2 + +**Verified tensor shapes (diag prints inside @cute.kernel on B200, n=256):** +``` +Before (None,0,None,0) pre-slice: + tAgQ: (((64,128),1), Int32(?), Int32(?), Int32(?)) — 4 modes + tBgK: (((64,128),1), Int32(?), Int32(?), Int32(?)) — 4 modes + tVgV: (((64,128),1), 1, 1, 1) — 4 modes + +After (None,0,None,0) pre-slice: + tAgQ: (((64,128),1), Int32(?)) — 2 modes, mode 1 = KV tiles + tBgK: (((64,128),1), Int32(?)) — 2 modes, mode 1 = KV tiles + tVgV: (((64,128),1), 1) — 2 modes, mode 1 = 1 (static) +``` ### Remaining for Multi-Tile 1. O rescale between tiles: `O *= exp2(old_max - new_max)` — needed for n=256+ to hit 0.9999 -2. Pipeline state cycling for n≥384 (3+ tiles with 2 pipeline stages) +2. Pipeline state cycling for n≥384 (3+ tiles with 2 pipeline stages) — output identical for all n>256, meaning only 2 KV tiles are loaded 3. Correction warps for production (separate softmax/correction/epilogue) 4. 12-warp layout @@ -325,7 +339,7 @@ cute.copy(tma_k, tBgK[None, None, None, None, kt, None, None, None], ...) | File | Status | Notes | |------|--------|-------| -| `fmha_v3_stage_c_example10.py` | 🔨 CURRENT | 8-mode TMA, combined K+V pipeline, O rescale, final normalize | +| `fmha_v3_stage_c_example10.py` | 🔨 CURRENT | (None,0,None,0) TMA, combined K+V pipeline, O rescale, final normalize | | `test_fmha_v3_stage_c_full.py` | OK n=128 | Working real softmax + O normalization | | `fmha_v3_stage_c_example1.py` | BROKEN multi-tile | First fix attempt, TMA still loads tile 0 | | `fmha_v3_stage_c_example2.py` | DEADLOCK | Combined K+V barrier, compiles but deadlocks | @@ -347,7 +361,7 @@ Warps 0-3: Softmax, Warps 4-7: Correction, Warp 8: MMA, Warp 9: TMA, Warp 10: Ep 1. `vectorize=True` loops: ONLY load/store/print 2. `.reduce(cute.ReductionOp.MAX)`: reduces ENTIRE C-fragment to scalar — global max, not per-row 3. `cute.arch.fmax`: impure for vectorizer — use plain `range()` loop -4. `tBgK`/`tVgV` have 8 TMA coord dims after tma_partition — 8-None no-op pre-slice required, mode 4 is GMEM tile dim +4. `tBgK`/`tVgV` have 4 modes after tma_partition — (None,0,None,0) keeps mode 2 (KV tiles) free, [None, kt] indexes it 5. `tBgK[(None, 0, None, 0)]` hardcodes GMEM iteration to tile 0 6. `softmax_done_bar` NamedBarrier is reusable across tiles