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

Update README, STAGE_D, STAGE_D2 with D1 rescale findings and D2 status

Browse Source
This commit is contained in:
biondizzle
2026-05-25 01:18:48 +00:00
parent 6cc151097e
commit 32850f6974
3 changed files with 31 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
README.md
View File

@@ -138,19 +138,20 @@ Summary
---
## Status (May 24, 2026 — 21:30 UTC)
## Status (May 25, 2026 — 01:10 UTC)
| 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 | ✅ COMPLETE | Real online softmax. Kernel outputs un-norm O + LSE (no TMEM round-trip). Migrated to `dsv4/kernels/attention/fmha.py` as `FmhaKernel`. |
| D1 | 🟡 hd≤256 DONE | Parameterized HEAD_DIM. qk_mma_tiler fix (hd=64/128/256 cos 0.999998). hd=512 SMEM fits but MLIR compilation hangs (>3hr). External k_sub merge proven impossible. |
| D2 | 🟡 Per-head DONE | Multi-query grid. Per-head launch works (cos 0.999998, n_h=64 hd=64). Multi-CTA grid deferred (requires tma_partition refactor). |
| D1 | 🟡 hd≤256 DONE | Parameterized HEAD_DIM. qk_mma_tiler fix (hd=64/128/256 cos 0.999998). hd=512 SMEM fits but MLIR compilation hangs (>3hr). External k_sub merge proven impossible. O rescale TMEM round-trip BROKEN (Ld32x32bOp/St32x32bOp corrupt data). Python KV merge workaround works. |
| D2 | 🟡 Per-head DONE | Multi-query grid. Per-head launch works (cos 0.999998, n_h=1-64 hd=64, n_h=2-8 hd=128, n_h=2 hd=256). Multi-CTA grid blocked: `flat_divide` + `epilogue_tma_store` layout mismatch. Requires full tma_partition refactor into kernel. |
| D3 | TODO | SWA sequence length mask (swa_lens per batch) |
| D4 | TODO | Causal mask on SWA branch only |
| D5 | 🟢 D5a+D5b DONE | D5a: normalize flag + LSE output (err=0.0). D5b: Python SWA+sink merge (cos 0.961). D5c/D5d: fused kernel merge TODO. |
| E1-E7 | TODO | Production extraction (class, custom op, cache, cleanup) |
| NVFP4-3 | ✅ DONE | `use_2cta_instrs` conditional in gemm_runner.py. 1.7-1.9× throughput at prefill shapes. |
---
@@ -389,6 +390,10 @@ Col 128+: O (PV acc, 64 FP32, rescale via Ld32x32bOp Repetition(16))
14. **Guard dead code with `const_expr`.** CuTeDSL compiles BOTH branches of Python `if` statements. Use `const_expr(condition)` to eliminate dead code at compile time. Critical for: O rescale (only when n_kv_tiles>1), LSE (only when normalize=False), SMEM-P path (only when use_smem_p=True), k_sub path (only when n_k_sub_tiles>1).
15. **External k_sub merge is mathematically impossible.** k_sub segments are additive in LOGIT space (S = S_0 + S_1), not attention weight space. You cannot recover softmax(S_0+S_1)@V from softmax(S_0)@V and softmax(S_1)@V. The D5 merge formula works for different token sets (additive in weight space), NOT for partial dot products. In-kernel k_sub accumulation before softmax is the only correct approach.
16. **`pv_n_tile` reduction is the easiest SMEM knob.** At hd>256, reducing pv_n_tile from 256 to 128 shrinks sV and sC by 2× each. Cost: 4 PV GEMM passes instead of 2. But PV is typically not the bottleneck, and this is simpler than SMEM overlap or Q tiling.
17. **O rescale TMEM round-trip with Ld32x32bOp/St32x32bOp is BROKEN.** Even a NO-OP round-trip (load O, multiply by 1.0, store back) corrupts data (cos 0.804 at s_k=256). The hand-constructed atoms don't preserve the C-fragment layout during round-trips. CUTLASS `correction_rescale` uses the same pattern — unclear why theirs works. **Workaround:** Python KV merge with per-segment LSE (cos 0.999998 for s_k up to 1024).
18. **KV merge formula uses NORMALIZED outputs, not un-normalized.** The correct D5 merge for different token sets: `O = sum_i [exp(lse_i) * O_i_norm] / sum_i [exp(lse_i)]`. Using `O_i_unnorm` instead of `O_i_norm` gives cos ~0.91. The un-norm merge only works when both segments share the same `row_max` (global max), which isn't the case for separate KV segments.
19. **`flat_divide` + `epilogue_tma_store` layout mismatch.** When using `cute.flat_divide` to create per-CTA GMEM views with runtime block coordinates (for multi-CTA grid), the resulting tensor layout is incompatible with CUTLASS's `epilogue_tma_store` pipeline, which expects the layout from `local_tile`. The tma_partition and epilogue must be refactored together to support multi-CTA grids.
20. **`local_tile` does not support runtime coordinates.** `cute.local_tile(mQ, tiler, (runtime_val, None))` fails at trace time. Must use `cute.flat_divide(mQ, tiler)` instead, which creates a tiled view with all rest dimensions accessible via runtime indexing.
---

6
STAGE_D.md
View File

@@ -166,11 +166,11 @@ pv_n_tile shown in parens; hd>256 uses pv_n_tile=128 (4 PV GEMM passes) to fit S
**Priority:** MEDIUM. Not a correctness blocker (external normalization is exact). Would enable in-kernel normalization for D5c/d. Also blocks NVFP4-1.2 (inverse RoPE FP4 fuse).
**Issue 2: O rescale for kt>0 (multi-KV-tile).** When s_k > 128 (n_kv_tiles > 1), the O accumulator must be rescaled after each KV tile. The current rescale code uses hand-constructed TMEM atoms and is guarded away with `const_expr(n_kv_tiles > 1)` at n=128. **Untested and likely broken for s_k > 128.**
**Issue 2: O rescale for kt>0 (multi-KV-tile).** CONFIRMED BROKEN (May 24). Even a NO-OP round-trip (load O, multiply by 1.0, store back) produces cos 0.804 at s_k=256. The `Ld32x32bOp`/`St32x32bOp` atoms corrupt data regardless of the rescale factor. The same atoms in CUTLASS `correction_rescale` use the same pattern — unclear why theirs works with 12-warp layout but ours fails with 6-warp.
**Priority:** HIGH. DSV4 Pro uses top_k=1024 (s_k=1024, n_kv_tiles=8). D2 multi-head decode WILL exercise s_k>128. If O rescale is broken, all D2 correctness at production configs will fail.
**Workaround (VERIFIED):** Python KV merge with per-segment LSE. Run kernel with s_k=128 (1 KV tile, no rescale) per segment. Merge using: `O = sum_i [exp(lse_i) * O_i_norm] / sum_i [exp(lse_i)]`. Verified cos 0.999998 for s_k=256, 384, 512, 1024 at hd=64. **Caveat:** requires per-row LSE output (currently only row 0 is written; per-row verified correct with max err 0.000001 but CuTe tensor indexing needs work for full per-row output).
**Fix approach:** Replace hand-constructed TMEM round-trip with the correction_rescale pattern from CUTLASS (`correction_rescale_and_partition`). This is the same one-way TMEM→SMEM path as D1.5 Issue 1, but applied to the O accumulator rescale between KV tiles.
**Priority:** HIGH for production (DSV4 Pro needs s_k=1024). The Python merge works but adds kernel launch overhead (8 launches for s_k=1024). Fused in-kernel rescale requires fixing the TMEM round-trip or using a different accumulator strategy.
---

58
STAGE_D2.md
View File

@@ -192,51 +192,33 @@ O has shape `(batch, n_h, T, head_dim)`. Each CTA writes its head's output. The
## To-Do List
- [ ] **D2.1:** Define multi-head tensor shapes in `FmhaKernel.__init__`
- Add `num_query_heads` constructor parameter
- Add `batch_size` parameter (or infer from tensor shapes at launch time)
- Keep `num_kv_heads=1` hardcoded for now (MQA)
### ✅ Completed (Per-Head Launch Approach)
- [ ] **D2.2:** Create test file `tests/unit/test_d2_multihead.py`
- Test with n_h=2, batch=1, T=1, hd=64 (minimal multi-head)
- Test with n_h=8, batch=2, T=1, hd=64 (small batch)
- Reference: PyTorch SDPA with MQA (all heads share K/V)
- Start with single-head regression (n_h=1, batch=1) to verify nothing breaks
- [x] **D2.2:** Create test file `tests/unit/test_d2_perhead.py`
- Per-head launch: kernel with grid=(1,1,1), Python iterates over heads/batches
- Verified: n_h=1,2,8,16,64 at hd=64; n_h=2,8 at hd=128; n_h=2 at hd=256
- All cos ≥0.999998
- [ ] **D2.3:** Modify `__call__` to construct multi-head TMA descriptors
- Q tensor: `(batch, n_h, T, head_dim)` — TMA over (batch, n_h, T, head_dim)
- K/V tensors: `(batch, 1, s_k, head_dim)` — TMA over (batch, s_k, head_dim) (squeeze kv_heads)
- O tensor: `(batch, n_h, T, head_dim)` — TMA over (batch, n_h, T, head_dim)
- LSE tensor: `(batch, n_h, T)` — per-head LSE
- [x] **D2.7 + D2.8:** Multi-head correctness across configs
- n_h=64, batch=1, hd=64 → Flash decode config — PASS
- hd=128, n_h=8 — PASS
- hd=256, n_h=2 — PASS
- [ ] **D2.4:** Compute grid shape and pass to launch
- `grid = (ceil_div(T, 128), num_query_heads, batch)`
- The `cta_tiler` remains `(128, 128, head_dim)` for QK, `(128, pv_n_tile, head_dim)` for PV
### 🟡 Blocked (Multi-CTA Grid)
- [ ] **D2.5:** Add block coordinate mapping inside `_kernel`
- `bidx, bidy, bidz = cute.arch.block_idx()`
- Map to `m_tile_idx, head_idx, batch_idx`
- Use head_idx to index the Q TMA tensor's head mode
- Use batch_idx for batch mode
- [ ] **D2.1:** Add `num_query_heads` and `batch_size` to `FmhaKernel.__init__`
- Simple to add, but the grid change is blocked (see below)
- [ ] **D2.6:** Adjust TMA loads for per-head Q and shared K/V
- Q load: index by `(batch_idx, head_idx, m_tile_idx, ...)`
- K/V load: index by `(batch_idx, ...)` (no head index)
- O store: index by `(batch_idx, head_idx, m_tile_idx, ...)`
- [ ] **D2.7:** Test multi-head correctness
- n_h=2, batch=1, T=1, hd=64 → cos ≥0.999 per head
- n_h=8, batch=2, T=1, hd=64 → cos ≥0.999 per head
- n_h=64, batch=1, T=1, hd=64 → Flash decode config
- [ ] **D2.8:** Test at hd=128, hd=256
- Same tests as D2.7 but with larger head dims
- Verify SMEM budget still fits (no per-head SMEM change — each CTA has its own SMEM)
- [ ] **D2.3–D2.6:** Multi-CTA grid with runtime block coordinates
- **BLOCKED:** `cute.local_tile` does not support runtime coordinates. Must use `cute.flat_divide` instead.
- **BLOCKED:** `flat_divide` + `epilogue_tma_store` layout mismatch. The epilogue pipeline expects `tCgC` from `local_tile`, but `flat_divide` produces a different coordinate system.
- **Requires:** Full refactor of `tma_partition` + `epilogue_tma_store` to work with `flat_divide`-based GMEM views. This means moving ALL GMEM tensor partitioning into the kernel (like CUTLASS reference does).
- **CUTLASS reference approach:** Uses `flat_divide` + `tma_partition` inside the TMA warp block, and a custom epilogue that handles the flat_divide coordinate system. Estimated 1-2 day effort.
- [ ] **D2.9:** LSE output for multi-head
- LSE shape: `(batch, n_h, T)` or `(batch, n_h, ceil_div(T, 128) * 128)` padded
- Each CTA writes its head's LSE to the correct position
- Needed for D5 merge
- Per-row LSE verified correct (max err 0.000001) but CuTe tensor indexing needs work
- Currently only row 0 is written (sfw_idx==0 guard)
- Full per-row output needed for D5 KV merge
---