Update both READMEs: Stage B complete, document TMEM overlap root cause
- Workspace README: full rewrite with Stage B ✅, Bug 4b root cause (P/O overlap),
FMHA V reconstruction, TMEM layout diagram, softmax store pattern, updated footguns
- Kernel README: focused on the bug, fix, and current test status
- Key lesson documented: NEVER use find_tmem_tensor_col_offset() as O placement
This commit is contained in:
185
README.md
185
README.md
@@ -1,132 +1,99 @@
|
||||
# DSV4 NVFP4 Workspace
|
||||
# DSV4 NVFP4 Kernel
|
||||
|
||||
## Status (May 21, 2026 — 09:18 UTC)
|
||||
## Status (May 21, 2026 — 15:35 UTC)
|
||||
|
||||
### Stage A ✅ COMPLETE
|
||||
Bare Q@K^T via tcgen05.mma → TMEM → GMEM. Cosine 0.999999.
|
||||
|
||||
### Stage B 🔨 IN PROGRESS — TMEM Alias Bug 4
|
||||
### Stage B ✅ COMPLETE — QK → Softmax → PV pipeline working for (128,64) PV
|
||||
Cosine 0.999999 with identity softmax and random V.
|
||||
|
||||
Two MMAs chained: Q@K^T (SMEM source) → identity softmax in TMEM → P@V (TMEM source).
|
||||
|
||||
**Pipeline deadlock: ✅ FIXED. Softmax packing: ✅ CONFIRMED CORRECT.**
|
||||
### Stage C 🔨 NEXT
|
||||
Real softmax (exp, row max, row sum, rescale). Multi-tile with proper accumulation.
|
||||
|
||||
---
|
||||
|
||||
## Bug 4 (ACTIVE): Non-(128,128) PV MMA — V/B Staging or Output C/D Failure
|
||||
## Stage B — Bug 4b Root Cause & Fix
|
||||
|
||||
### Summary
|
||||
### The Bug: TMEM P/O Region Overlap
|
||||
|
||||
The softmax writes P to TMEM using the **QK C-fragment layout**. The PV MMA reads P from TMEM using the **PV A-fragment layout**. For (128,128) PV these layouts agree. For (128,16) PV they disagree — the PV A-fragment reads from different TMEM columns than where the softmax wrote, producing zero output.
|
||||
**Symptom:** (128,64) PV produces NaN or zeros. (128,128) PV works fine.
|
||||
|
||||
**FMHA uses (128,16) PV with the same construction and works.** The root difference is not yet identified despite exhaustive comparison. FMHA references: `p_tmem_layout_staged = make_smem_layout_a(pv_mma, pv_mma_tiler, q_dtype, 1)` — same call we make.
|
||||
**Root cause:** PV output accumulator O was placed at `find_tmem_tensor_col_offset(tOtO)`, which returns 64 for (128,64) PV. P occupies TMEM columns [32, 96). O at [64, 128) overlaps P at [64, 96). While PV MMA reads P (A-operand), it simultaneously writes O (D-operand) to overlapping TMEM columns. The A-operand gets corrupted mid-computation.
|
||||
|
||||
### What Works / What Doesn't
|
||||
For (128,128) PV, `find_tmem_tensor_col_offset(tOtO)` returns 128, so O starts after P — no overlap. It worked by accident.
|
||||
|
||||
- ✅ PV (128,128) output, V=I or random → cosine 1.0 / 0.999999
|
||||
- ✅ PV (128,128) with zero-padded V (head_dim=16) → cosine 1.0 **WORKAROUND**
|
||||
- ✅ PV (128,64), all-ones V → cosine 0.999999 (uniform hides bug)
|
||||
- ✅ PV (128,64), single-element V → cosine 1.0 (sparse hides bug)
|
||||
- ❌ PV (128,64), truncated identity V → cosine 0.02
|
||||
- ❌ PV (128,16), V=I(128,128) → cosine 0.0 (all zeros)
|
||||
- ❌ PV (128,16) with P at S offset (no softmax) → NaN (FP32→BF16 reinterpret)
|
||||
### The Fix
|
||||
|
||||
### Root Cause (CONFIRMED May 21 09:50 UTC)
|
||||
|
||||
**Bug is NOT the TMEM alias.** The PV A-fragment layout is identical for all PV sizes (confirmed by C++ source and diagnostics): all PV sizes produce tOrP2_s = (2048, 1, 8), size=16384.
|
||||
|
||||
**The real bug: V SMEM only holds 1 K-tile (2048 BF16), but the PV MMA iterates 8 K-phases.** For non-(128,128) V, most K-phases read wrong V data.
|
||||
|
||||
- (128,128) PV + V=I works by coincidence (V=I makes the projection self-consistent)
|
||||
- (128,32) PV + V=(32,128) fails because V SMEM only has V[0:16,:], K-phases 1-7 read wrong data
|
||||
- Zero-padded V works because V=(128,128) covers all 8 K-phases; rows beyond head_dim are zero
|
||||
|
||||
**How FMHA avoids this:** FMHA interleaves QK and PV per KV-tile. Each tile loads 16 K-rows of V, and PV processes only that tile. This ensures V SMEM always has the correct data.
|
||||
|
||||
**Workaround:** Zero-pad V to 128 K-rows (2-4x compute waste, but correct). Proper fix: FMHA-style KV-tile interleaving.
|
||||
|
||||
### Current Workaround
|
||||
|
||||
Use **(128,128) PV with zero-padded V**. This wastes compute (8× for head_dim=16, 2× for head_dim=64) but produces correct results (cosine 1.0). For the production kernel, we'll use this initially and optimize to (128,16) PV once the TMEM alias is resolved.
|
||||
|
||||
### Required Fixes (Not Yet Applied)
|
||||
|
||||
1. **Primary**: Softmax must write P using the PV A-fragment TMEM layout, not the QK C-fragment layout. Requires constructing a `make_tmem_copy` with `tP` (PV layout) as the destination, and rearranging register data from QK partition to PV partition.
|
||||
|
||||
2. **Secondary**: `epi_tile` must use PV's cta tile, and `self.cta_tile_shape_mnk` must be swapped before `epilogue_tma_store`. FMHA sets `self.epi_tile = self.pv_mma_tiler[:2]` directly.
|
||||
|
||||
3. **Alternative (for later)**: Investigate using `composition()` to create a hybrid layout that both the QK softmax write and PV A-fragment read can agree on.
|
||||
|
||||
---
|
||||
|
||||
## Bugs 1–3: ✅ FIXED
|
||||
|
||||
### Bug 1: V B-Operand Must Be MN-Major
|
||||
|
||||
FMHA requires V to be **MN-major** for the PV MMA B-operand. V must be shaped (head_dim, seq) = (64, 128) with strides (1, 64) via `as_strided`.
|
||||
|
||||
### Bug 2: C-Fragment Composition Store for P — CONFIRMED CORRECT
|
||||
|
||||
FP32→BF16 packing via C-fragment composition store works. ⛔ `St32x32bOp` MUST use Float32, NOT BFloat16.
|
||||
|
||||
### Bug 3: First PV Must Use ACCUMULATE=False
|
||||
|
||||
If ACCUMULATE=True on the first PV, `O = P@V + old_O` adds uninitialized TMEM. FMHA: `pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, kphase_idx != 0)`.
|
||||
|
||||
---
|
||||
|
||||
## Pipeline Deadlock — ✅ FIXED (May 21)
|
||||
|
||||
Three root causes found and fixed:
|
||||
1. `PipelineUmmaAsync` for mma_si must NOT pass `cta_layout_vmnk`
|
||||
2. TMA warp must NOT call `tmem.wait_for_alloc()`
|
||||
3. `pipeline.PipelineTmaStore` (not `TmaStorePipeline`)
|
||||
|
||||
---
|
||||
|
||||
## ⛔ FOOTGUNS — CUTLASS CuTeDSL Landmines
|
||||
|
||||
1. **St32x32bOp with BFloat16 → ILLEGAL MEMORY ACCESS** — Must use Float32 + `cute.recast_ptr`
|
||||
2. **V major ≠ K major** — V must be MN-major, use `as_strided`
|
||||
3. **C-fragment → A-fragment TMEM alias only works when N_MMA matches** — (128,128) works, (128,64) breaks
|
||||
4. **PipelineUmmaAsync consumer = thread count, NOT warp count** — `32 * len(warp_ids)`
|
||||
5. **mma_si pipeline must NOT pass cta_layout_vmnk**
|
||||
6. **TMA warp excluded from tmem barrier**
|
||||
7. **First PV ACCUMULATE=False**
|
||||
8. **TMEM offset: FP32 ptr + 32 = BF16 ptr + 64** (width scaling)
|
||||
9. **epi_tile must use PV cta_tile, not QK**
|
||||
10. **CuTe nested layout modes flatten sequentially** — `((128,16),1,(4,2)):((65536,1),0,(16,64))` is sequential
|
||||
|
||||
---
|
||||
|
||||
## Architecture: Per-Tile Flow
|
||||
Place O after both S and P regions:
|
||||
|
||||
```python
|
||||
p_cols_fp32 = pv_mma_tiler[2] * q_dtype.width // qk_acc_dtype.width # 128*16/32 = 64
|
||||
p_end = tmem_p0_offset + p_cols_fp32 # 32 + 64 = 96
|
||||
s_cols = qk_mma_tiler[1] # 128
|
||||
o_after = max(s_cols, p_end) # 128
|
||||
tmem_o0_offset = ((o_after + 31) // 32) * 32 # align to 32 = 128
|
||||
```
|
||||
For each KV tile:
|
||||
1. Load warp writes sKV[stage] (paged FP8 gather via indexed cp.async)
|
||||
2. MMA warp issues MMA1: sQ @ sKV[stage]^T → tmem_scores (accumulate=False)
|
||||
Signals scores_full_mbar (via PipelineUmmaAsync commit)
|
||||
3. Epilogue warps wait on mma_si consumer (scores ready), then:
|
||||
a. tcgen05.ld scores from TMEM → register fragments
|
||||
b. Compute tile_max, new_max, rescale = exp(old_max - new_max)
|
||||
c. Apply rescale to tmem_output IN PLACE (tmem_output *= rescale)
|
||||
d. tcgen05.st exp(scores - new_max) back to TMEM → P operand
|
||||
e. Release mma_si (softmax_done — MMA warp can re-acquire and issue PV MMA)
|
||||
4. MMA warp waits on mma_si acquire (softmax done), MMA2: P @ sV → tmem_output (accumulate=True)
|
||||
5. Stage released, load warp can refill it
|
||||
|
||||
After all tiles: epilogue warps tcgen05.ld tmem_output, divide by row_sum, cast to BF16, store to GMEM
|
||||
### Secondary Fix: FMHA-Style V Reconstruction
|
||||
|
||||
V from DLPack has logical shape (n, hd) but PV B-operand expects (hd, n). Reconstruct inside CuTe:
|
||||
|
||||
```python
|
||||
v_fmha = cute.make_tensor(
|
||||
v.iterator,
|
||||
cute.make_layout(
|
||||
(HEAD_DIM, s_k, 1),
|
||||
stride=(1, HEAD_DIM, HEAD_DIM * s_k),
|
||||
),
|
||||
)
|
||||
v_major = LayoutEnum.from_tensor(v_fmha).mma_major_mode() # MN
|
||||
# Use v_fmha in make_tiled_tma_atom_B, NOT the DLPack v
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Environment
|
||||
## TMEM Layout
|
||||
|
||||
- **Server**: root@45.76.247.107 (B200, 180 GiB HBM3e per GPU)
|
||||
- **venv**: `source /root/dsv4-nvfp4-workspace/venv/bin/activate`
|
||||
- **PYTHONPATH**: `/root/dsv4-nvfp4-workspace/kernel`
|
||||
- **Model**: `/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4`
|
||||
- **vLLM repo**: `/root/dsv4-nvfp4-workspace/vllm` (modified for Blackwell)
|
||||
- **Pseudocode**: `/root/fragile-kernel-example/README.md`
|
||||
- **fmha.py reference**: `/root/cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py`
|
||||
- **fmha_bwd.py reference**: `/root/cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha_bwd.py`
|
||||
```
|
||||
Col: 0 32 64 96 128 192 256
|
||||
|---- S ----|---- P ----| |---- O ----|
|
||||
| QK acc | Softmax P | (gap) | PV acc |
|
||||
| 128 FP32 | 64 FP32 | 32 col | 64 FP32 |
|
||||
```
|
||||
|
||||
P aliases part of S (softmax overwrites S columns 32-95 with P). O must not overlap P or S.
|
||||
|
||||
---
|
||||
|
||||
## Softmax P Store (FMHA Pattern)
|
||||
|
||||
Store uses QK C-fragment composition. Read uses PV A-fragment. These are two separate aliases of the same physical TMEM — the P/A alias works (proven by no-softmax test) because both layouts depend on M=128 and K, not on PV output N.
|
||||
|
||||
```python
|
||||
# Store (softmax writes P)
|
||||
tStP = cute.make_tensor(tStS.iterator + tmem_p0_offset,
|
||||
cute.composition(tStS.layout, cute.make_layout((128, p_cols_fp32))))
|
||||
tiled_tmem_store = tcgen05.make_tmem_copy(store_atom, tStP)
|
||||
|
||||
# Read (PV MMA reads P)
|
||||
tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer)
|
||||
tOrP = pv_thr.make_fragment_A(tP)[None,None,None,0]
|
||||
tOrP0 = cute.make_tensor(tOrP.iterator + width_scale * tmem_p0_offset, tOrP.layout)
|
||||
```
|
||||
|
||||
Register bridge (FP32 backing + BF16 view):
|
||||
```python
|
||||
rP_words = cute.make_rmem_tensor(tScP.shape, qk_acc_dtype)
|
||||
rP_bf16 = cute.make_tensor(recast_ptr(rP_words.iterator, dtype=q_dtype), tTMEM_LOADrS.layout)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Test Files
|
||||
|
||||
- **tests/test_fmha_v3.py** — Full pipeline with KV-tile interleaving. PASS.
|
||||
- **tests/test_pv64_with_softmax.py** — Single AB pipeline. PASS.
|
||||
- **tests/test_128_128_vdiag.py** — (128,128) PV baseline. PASS.
|
||||
- **tests/test_qkonly.py** — QK only. PASS.
|
||||
- **tests/test_qk_softmax.py** — QK + softmax (no PV). PASS.
|
||||
|
||||
Reference in New Issue
Block a user