diff --git a/README.md b/README.md index 257311ba..b4932212 100644 --- a/README.md +++ b/README.md @@ -1,59 +1,178 @@ # DSV4 NVFP4 Kernel -## Status (May 21, 2026 — 15:35 UTC) +## Status (May 21, 2026 — 15:40 UTC) -### Stage A ✅ COMPLETE -Bare Q@K^T via tcgen05.mma → TMEM → GMEM. Cosine 0.999999. - -### Stage B ✅ COMPLETE — QK → Softmax → PV pipeline working for (128,64) PV -Cosine 0.999999 with identity softmax and random V. - -### Stage C 🔨 NEXT -Real softmax (exp, row max, row sum, rescale). Multi-tile with proper accumulation. +| Stage | Status | Description | +|-------|--------|-------------| +| A | ✅ COMPLETE | Q@K^T via tcgen05.mma → TMEM → GMEM | +| B | ✅ COMPLETE | QK → identity softmax → P@V pipeline (TMEM alias, KV-tile interleaving) | +| C | 🔨 NEXT | Real softmax: row max, exp, rescale, row sum | +| D | TODO | Full decode attention: paged KV cache, multi-query, causal mask | +| E | TODO | Production kernel: integrate into `cutedsl/`, PyTorch custom op, vLLM bridge | --- -## Stage B — Bug 4b Root Cause & Fix +## Where Things Live -### The Bug: TMEM P/O Region Overlap +### Canonical Test Files (current working code) -**Symptom:** (128,64) PV produces NaN or zeros. (128,128) PV works fine. +| File | Stage | What It Tests | +|------|-------|---------------| +| `tests/test_fmha_v3.py` | A+B | Full QK→softmax→PV with KV-tile interleaving. **This is the canonical Stage B kernel.** | +| `tests/test_pv64_with_softmax.py` | B | Single AB pipeline variant. Simpler, also passing. | +| `tests/test_128_128_vdiag.py` | A+B | (128,128) PV baseline. Square case. | +| `tests/test_qkonly.py` | A | QK with split Q/KV pipelines, no softmax, no PV. | +| `tests/test_qk_softmax.py` | A+partial B | QK + softmax writes P to TMEM. No PV. | -**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. +### Reference Implementations (non-CuTe) -For (128,128) PV, `find_tmem_tensor_col_offset(tOtO)` returns 128, so O starts after P — no overlap. It worked by accident. +| File | Description | +|------|-------------| +| `cutedsl/blackwell_attention.py` | Pure PyTorch reference: RoPE, KV cache, causal attention, SWA | +| `cutedsl/csa_attention.py` | CSA/HCA sparse attention reference | +| `cutedsl/native_swa_decode.py` | SWA decode reference | -### The Fix +### CuTeDSL Kernel Modules (production code) -Place O after both S and P regions: +| File | Description | +|------|-------------| +| `cutedsl/kernel/moe/` | NVFP4 MoE GEMM kernels (fused SwiGLU, grouped MM) | +| `cutedsl/kernel/blockscaled_gemm/` | Block-scaled GEMM | +| `cutedsl/nvfp4_linear.py` | NVFP4 linear layer wrapper | +| `cutedsl/runner.py` | MoE runner | +| `cutedsl/moe_pipeline.py` | MoE pipeline orchestration | -```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 +### Obsolete Tests (do not use, can delete) + +~100+ test files from debugging stages A and B. Key patterns: +- `test_stage_b_v1.py` through `test_stage_b_v30.py` — incremental Bug 4b debugging +- `test_128_16_*.py` — early (128,16) PV attempts with wrong head_dim +- `test_tmem_*.py`, `test_bf16_*.py` — standalone TMEM/copy debugging +- `test_pv64_no_softmax.py`, `test_pv64_fmha_v.py`, `test_pv64_kmajor_v.py` — Bug 4b root cause isolation + +These can be deleted once the canonical tests are stable and the kernel is extracted. + +--- + +## Stage C: Real Softmax + +### What We Have Now + +Identity softmax: load S FP32 from TMEM, convert to BF16, store P back to TMEM. This proves the TMEM pipeline works but isn't a real softmax. + +### What We Need + +FMHA-style online softmax per KV-tile: + +``` +For each KV tile: + 1. QK → S (FP32 in TMEM) + 2. Load S row-max for this tile: tile_max[j] = max(S[j,:]) + 3. Compute new row max: new_max[j] = max(old_max[j], tile_max[j]) + 4. Rescale O: O[j,:] *= exp(old_max[j] - new_max[j]) + 5. Compute P: P[j,i] = exp(S[j,i] - new_max[j]) + 6. Store P to TMEM (BF16, same C-fragment composition store) + 7. Update row sum: row_sum[j] = row_sum[j] * exp(old_max[j] - new_max[j]) + sum(P[j,:]) + 8. PV: O[j,:] += P[j,:] @ V[i,:] +After all tiles: + 9. O[j,:] /= row_sum[j] (final normalization) ``` -### Secondary Fix: FMHA-Style V Reconstruction +### Key Challenges -V from DLPack has logical shape (n, hd) but PV B-operand expects (hd, n). Reconstruct inside CuTe: +1. **Row max across tiles** — Must track per-row maximum across KV-tiles and rescale O when a new max is found. This is the core of online softmax. +2. **Row sum accumulation** — Must accumulate exp(sum) across tiles with proper rescaling. +3. **FP32 precision** — Row max, rescale, and row sum must stay in FP32 for numerical stability. Only P (the exp values) get cast to BF16 for TMEM store. +4. **O rescale in TMEM** — When a new row max is found, the existing O in TMEM must be multiplied by `exp(old_max - new_max)`. This requires loading O, rescaling, and storing back. Same TMEM load/store machinery as softmax P. +5. **Final normalization** — After all KV-tiles, divide O by row_sum. Can be done as part of the epilogue. -```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 +### Expected Structure + +The softmax epilogue warps will expand significantly: +- Currently: load S → convert BF16 → store P (identity softmax) +- After Stage C: load S → compute tile_max → compare with old_max → rescale O → compute exp → store P → update row_sum + +The MMA loop remains the same (QK → softmax → PV per tile). The softmax just does more work between QK completion and PV start. + +--- + +## Stage D: Full Decode Attention + +### What We Have After Stage C + +A working QK → real softmax → PV kernel for a single query sequence against a contiguous KV block. Fixed dimensions (128×128 QK, 128×64 PV), single CTA. + +### What We Need + +1. **Paged KV cache** — KV comes from a paged cache (fp8 or bf8 with per-token inverse scale), not a contiguous tensor. TMA loads must follow page tables. +2. **Multi-query** — Multiple query sequences in flight, each with different KV lengths. Requires grid dimensions > 1, possibly persistent kernel. +3. **Causal masking** — QK must mask future positions. For decode (1 query vs N KVs), this is trivial (no mask needed). For prefill, need a causal mask. +4. **Variable sequence length** — Each CTA handles a different number of KV tiles. The loop bound `n_kv_tiles` becomes dynamic. +5. **Multiple head dimensions** — HEAD_DIM=16, 64, 128 all need to work. Currently only HEAD_DIM=64 is tested. +6. **CSA/HCA sparse attention** — For compress_ratio > 1, KV is read from compressed cache instead of full KV cache. Different TMEM layouts, different attention patterns. + +### Key Question: Do We Need Stage D As A Separate Stage? + +Stage D is really about *scaling* the Stage C kernel, not adding fundamentally new compute. The core pipeline (QK → softmax → PV) doesn't change. What changes is: +- Where the data comes from (paged cache vs contiguous tensor) +- How many CTAs run (grid size) +- Whether we need causal masking + +This could be folded into the production kernel directly rather than being a separate test stage. + +--- + +## Stage E: Production Kernel + +### Goal + +Replace `cutedsl/blackwell_attention.py` (pure PyTorch) with a CuTeDSL kernel that runs on the Blackwell tensor cores. + +### Steps + +1. **Extract kernel from `test_fmha_v3.py`** → `cutedsl/kernel/attention/fmha_kernel.py` + - Class `FmhaKernel` with `@cute.jit __call__` + - Clean parameter interface: Q, K, V, O tensors + config + - No hardcoded dimensions — all derived from MMA shapes + +2. **Add real softmax (Stage C)** to the extracted kernel + +3. **Add paged KV cache support (Stage D)** + - Page table TMA or gather-style loads + - Per-sequence KV length tracking + +4. **Wrap as PyTorch custom op** → `cutedsl/custom_ops.py` + - `blackwell_fmha_forward(q, k, v, ...) -> o` + - Autograd support (or torch.compile integration) + - torch.library custom op registration + +5. **Integrate with vLLM** → `vllm/attention/ops/blackwell_fmha.py` + - Replace the broken FlashMLA Blackwell path + - Hook into vLLM's paged attention interface + - Support both prefill and decode modes + +6. **Benchmark and tune** + - Profile against PyTorch SDPA baseline + - Tune tile sizes, pipeline stages, SMEM usage + - Verify numerical accuracy vs reference across head dims and sequence lengths + +### File Structure (target) + +``` +cutedsl/ + kernel/ + attention/ + __init__.py + fmha_kernel.py ← extracted, clean CuTeDSL kernel + fmha_softmax.py ← real softmax (Stage C) + fmha_epilogue.py ← row sum normalization, O output + blackwell_attention.py ← PyTorch reference (keep for testing) + custom_ops.py ← PyTorch custom op wrappers ``` --- -## TMEM Layout +## TMEM Layout (Current) ``` Col: 0 32 64 96 128 192 256 @@ -62,38 +181,27 @@ Col: 0 32 64 96 128 192 256 | 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. +For Stage C, we'll need additional TMEM regions: +- `row_max` — per-row FP32 max (128 rows × 1 col = 128 FP32 values, can use 4 TMEM columns) +- `row_sum` — per-row FP32 sum (128 rows × 1 col, 4 TMEM columns) +- `old_max` — per-row FP32 previous max (4 TMEM columns) + +These are tiny (4-8 TMEM columns each). They can go in the gap at columns 96-128 or after O. --- -## Softmax P Store (FMHA Pattern) +## Key Lessons From Stages A & B -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. +1. **NEVER use `find_tmem_tensor_col_offset()` as a TMEM placement decision.** It returns footprint size, not a safe column offset. The P/O overlap bug cost the entire Bug 4b debugging session. -```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) +2. **FMHA never trusts DLPack tensor layouts.** Reconstruct V as (hd, s_k) MN-major inside CuTe. The DLPack shape (n, hd) has wrong logical modes for PV B-operand. -# 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) -``` +3. **TMEM allocation must be power of 2.** `TmemAllocator.allocate()` asserts this. -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) -``` +4. **P/A alias works.** QK C-fragment composition store + PV A-fragment read alias the same physical TMEM columns. Proven for (128,64) and (128,128). ---- +5. **Square hides bugs.** (128,128) PV worked for every wrong approach because both dims are 128. Always test non-square cases. -## Test Files +6. **`St32x32bOp` MUST use Float32, NOT BFloat16.** BFloat16 causes illegal memory access. -- **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. +7. **First PV ACCUMULATE=False.** Otherwise adds uninitialized TMEM to output.