133 lines
6.5 KiB
Markdown
133 lines
6.5 KiB
Markdown
# DSV4 NVFP4 Workspace
|
||
|
||
## Status (May 21, 2026 — 09:18 UTC)
|
||
|
||
### Stage A ✅ COMPLETE
|
||
Bare Q@K^T via tcgen05.mma → TMEM → GMEM. Cosine 0.999999.
|
||
|
||
### Stage B 🔨 IN PROGRESS — TMEM Alias Bug 4
|
||
|
||
Two MMAs chained: Q@K^T (SMEM source) → identity softmax in TMEM → P@V (TMEM source).
|
||
|
||
**Pipeline deadlock: ✅ FIXED. Softmax packing: ✅ CONFIRMED CORRECT.**
|
||
|
||
---
|
||
|
||
## Bug 4 (ACTIVE): Non-(128,128) PV MMA — V/B Staging or Output C/D Failure
|
||
|
||
### Summary
|
||
|
||
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.
|
||
|
||
**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.
|
||
|
||
### What Works / What Doesn't
|
||
|
||
- ✅ 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)
|
||
|
||
### 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
|
||
|
||
```
|
||
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
|
||
```
|
||
|
||
---
|
||
|
||
## Environment
|
||
|
||
- **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`
|