biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

README: full roadmap — Stage C (real softmax), D (paged KV), E (production kernel)

Browse Source 
Document canonical test files, obsolete test sprawl, and the path from
test_fmha_v3.py → cutedsl/kernel/attention/fmha_kernel.py → vLLM integration.
Also: TMEM layout for Stage C, key lessons from A&B.
This commit is contained in:
biondizzle
2026-05-21 15:43:01 +00:00
parent ad24792fc7
commit 20564425ec
1 changed files with 168 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

228
README.md
View File

@@ -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.