nvfp4-megamoe-kernel/ROADMAP.md

# ROADMAP

Living document. Current state, active blockers, priority order, and what to build next. Architecture and lessons live in README.md — this file is for "what now."

**Last updated:** 2026-05-26 (revised after correcting D1.5 fix-path analysis)

---

## Current status

### Working

| Component | hd | n | cos | Status |
|---|---:|---:|---:|---|
| FMHA TMEM-P | 64 | 128 | 0.999998 | ✅ |
| FMHA TMEM-P / SMEM-P | 128 | 128 | 0.999997 | ✅ |
| FMHA TMEM-P | 256 | 128 | 0.999998 | ✅ |
| FMHA multi-tile (Python KV merge) | 64 | up to 1024 | 0.999998 | ✅ Workaround — see below |
| D3 SWA length mask (in-kernel) | 128 | 128 | 0.999996 | ✅ |
| D4 causal mask on SWA (in-kernel) | 128 | 128 | 0.999996 | ✅ |
| D5c sink merge single-tile | 64 | 128 | 0.999996 | ✅ |
| D5c sink merge multi-tile (Python KV merge) | 64 | 256 | 0.999996 | ✅ |
| Per-head multi-head launch | 64 | 128 | 0.999995 | ✅ n_h=1–128 |
| MoE fused SwiGLU (NVFP4) | — | — | matches ref | ✅ Clamping in kernel |
| Dense router (sqrt-softplus) | — | — | matches ref | ✅ |
| Hash router | — | — | matches ref | ✅ |
| `use_2cta_instrs` conditional | — | — | 1.7–1.9× speedup | ✅ M≥256 prefill |
| NVFP4 primitives | — | — | E4M3 SF / mxf4nvf4 / 16-elem | ✅ Verified |

### Known blockers

| Blocker | Impact | Status |
|---|---|---|
| **Per-kt O rescale in TMEM (D1.5)** | Multi-tile attention requires 5–9 kernel launches per decode step instead of 1 (Python KV merge) | Workaround is correct (cos 0.999998). Whether to fix depends on profiling — see Priority 1. |
| **TMEM final-normalize round-trip** | Cannot do in-kernel `O /= row_sum` cleanly | **Already worked around** — emit un-normalized O + LSE, external divide is exact. Not a blocker for shipping. |
| **`epilogue_tma_store` blocks D2 multi-CTA + NVFP4-1.2** | Per-head Python launch wastes 128 launches per Pro decode step; FMHA output forces BF16 GMEM round-trip before wo_a | Unblocked by Priority 2 (the one-way final-epilogue rewrite — which is the **only** part of the MoE pattern that legitimately ports to FMHA). |
| **hd=512 MLIR backend hang** | Cannot compile single-kernel hd=512 (>3hr optimizer time, structurally correct) | Decode works via head-packed M with hd≤256 chunks. Single-kernel hd=512 only needed for prefill efficiency. Low priority. |

---

## What the MoE epilogue pattern actually buys us (and what it doesn't)

This deserves stating explicitly because the previous version of this document had it wrong.

**The MoE pattern is one-way only.** `dsv4/kernels/gemm/fused_swiglu.py` uses `epilogue_tmem_copy_and_partition` + `epilogue_smem_copy_and_partition` to construct a TMEM → registers → SMEM → GMEM pipeline. There is **no** corresponding store-back-to-TMEM helper, and no inverse pairing of the t2r atom. The MoE epilogue runs *once*, *after* all MMA K-tiles are accumulated. It never needs to mutate the TMEM accumulator mid-loop.

**FMHA's per-kt O rescale is structurally different.** PV uses `tcgen05.mma` with `ACCUMULATE=True` across the kt loop. The accumulator must live in TMEM because that's where MMA reads it and writes it. When row_max changes between kt iterations, the running O accumulator has to be multiplied by `acc_scale` *in TMEM* before the next PV — load to registers, multiply, store back. The store-back is the part that's broken: `Ld32x32bOp` and `St32x32bOp` built as separate atoms have hardware column mappings that don't match, producing ~3% corruption even on NO-OP round-trip. No software layout transformation in CuTeDSL Python has, so far, made them pair correctly.

**Therefore:** porting the MoE pattern to FMHA **only fixes the one-way paths**:

1. The final epilogue (after the last kt, when O is being written to GMEM for good).
2. Any FP4 amax + pack fusion into that final epilogue.

It does **not** fix the per-kt rescale. That is a separate, harder problem with three possible paths laid out below.

---

## Priority 1: Profile production decode to determine if the rescale fix is needed

Before investing days in fixing the per-kt rescale, measure whether the 5–9 launch overhead from Python KV merge is actually a bottleneck.

At Pro decode (s_k=1152, n_kv_tiles=9), Python merge dispatches 9 kernels per step. Conservative launch overhead ~50 μs per kernel ≈ 450 μs/step in launch overhead alone. If a full decode step (all 61 layers, MoE, embedding, sampler) takes ~30 ms, that's ~1.5% of latency. If it takes ~10 ms, it's ~4.5%. Whether that's worth a 1–2 week refactor depends on the actual measurement.

**Action:**
- [ ] Profile Pro decode at s_k=1152 with current Python KV merge. Measure: total step latency, launch overhead from FMHA dispatches, FMHA compute time per launch.
- [ ] Measure CPU dispatch overhead on the host (Python loop + kernel cache lookup).
- [ ] Decision rule: if Python merge overhead is < 5% of total decode latency, defer Priority 8 indefinitely. Ship Python merge as production path.

**Done when:** there's a profiled number that justifies (or doesn't) the engineering investment in Priority 8.

---

## Priority 2: One-way final-epilogue rewrite

**What:** Replace the `utils.gemm.sm100.epilogue_tma_store(...)` call at the end of FMHA (`dsv4/kernels/attention/fmha.py` lines 565–577) with the MoE-style explicit pipeline:

```
transform_partitioned_tensor_layout → epilogue_tmem_copy_and_partition →
[register slot — optional normalize/cast/FP4 pack] →
epilogue_smem_copy_and_partition → flat_divide → cpasync.tma_partition → TMA store
```

This is **strictly one-way**. The kt-loop rescale code stays exactly as it is (using the broken hand-built atoms — see Priority 8 for the fix path).

**What this enables:**

1. **Optional in-kernel normalize.** Adds an `if const_expr(self.normalize):` block at the register slot to multiply by `inv_row_sum`. Currently external code does the divide on the un-normalized output. Tiny perf win, not the main reason to do this.
2. **Unblocks NVFP4-1.2** (Priority 6) — gives a register-level modification slot in the FMHA output path where FP4 amax + pack can live, eliminating the BF16 GMEM materialization between FMHA and wo_a.
3. **Likely unblocks D2 multi-CTA grid** (Priority 4) — the current `epilogue_tma_store` is what couldn't accept the `flat_divide`-based GMEM coordinate system. Switching to the explicit `cpasync.tma_partition(tma_c, ..., cute.flat_divide(tCgC_transformed, epi_tile))` path puts FMHA on the same TMA pattern MoE uses successfully, which should accept multi-CTA block coordinates.

**Caveats to verify on B200 before assuming this works:**

- Whether `tma_partition` survives inside the `if warp_idx < self.mma_warp_id` block. MoE calls it from inside its epilogue warp's `if`, but that has not been tested in FMHA's region tree. Previous attempts at the full pattern triggered 20+ minute MLIR compile times before reaching a verdict.
- Whether `transform_partitioned_tensor_layout` accepts FMHA's `tOtO0` (TMEM iterator with offset) directly, or whether it needs a fresh tensor built at `tmem_ptr + self.tmem_o0_offset` with `tCtO_fake.layout`.
- The `epilogue_tmem_copy_and_partition` helper signature on the current CuTeDSL version — print on B200 before coding.

**Failure mode to watch for:** if compile hangs as it did previously, this rewrite is genuinely blocked and the chain of follow-ons (Priorities 4 and 6) need alternative paths.

**Effort:** 1–2 days if the helpers cooperate. Multiple days if the MLIR hang reappears.

**Done when:**
- hd=64/128/256 regression cos ≥ 0.999998 holds with both `normalize=True` and `normalize=False` paths.
- LSE output still matches reference for `normalize=False` callers.
- Compile time is reasonable (< 5 min) — if not, document the hang and fall back.

---

## Priority 3: NVFP4-1.1 — Fuse FP4 quant into MoE SwiGLU epilogue

**Independent of FMHA. Can run in parallel with Priority 2.** Biggest bandwidth win in the codebase.

Current:
```
padded_x_fp4 → L1 GEMM → SwiGLU → BF16 GMEM
                                  ↓
                          quantize_activation_nvfp4 (separate kernel)
                                  ↓
                         padded_activated_fp4 → L2 GEMM
```

Target:
```
padded_x_fp4 → L1 GEMM → SwiGLU → online amax → FP8 scale + FP4 pack → FP4 GMEM → L2 GEMM
```

The SwiGLU + clamp result already lives in registers at `tRS_rC.store(acc_vec_bf16)` (line 2207 of `fused_swiglu.py`). That's the slot for amax + FP4 pack.

**Per-microblock amax (16 contiguous elements):**
1. `shfl_xor` butterfly reduction across the 4 threads holding the 16 elements.
2. FP8 E4M3 scale = amax / 6 (FP4 e2m1 max).
3. Per-element FP4 pack: `sign_bit << 3 | (clamped_val / scale).to(uint3)`. Two elements → one byte.
4. 16 packed nibbles → 64-bit word → SMEM stage → TMA store.
5. FP8 scale → separate scale-factor SMEM stage → TMA store to the L2 SFA buffer.

**Done when:** `padded_activated_fp4` and `padded_activated_x_sf` scratch buffers go away, `quantize_activation_nvfp4` between L1 and L2 disappears, L1→L2 cosine matches reference.

---

## Priority 4: D2 multi-CTA grid

**Depends on Priority 2.** Currently per-head Python launch dispatches 128 kernels per Pro decode step. Multi-CTA grid collapses that to 1.

**Grid:** `(num_M_tiles, num_query_heads, batch)` — at decode T=1: `(1, 128, batch)`.

**Q tensor layout:** Option 1 — `(batch, n_h, T, head_dim)` with head as a TMA mode. Matches CUTLASS reference, allows per-head LSE output, generalizes to GQA later.

**MQA K/V sharing:** start with independent K/V loads per CTA (each loads its own copy). At decode hd=512, K/V per CTA is ~128 KB; 128 CTAs × 128 KB = 16 MB total, comfortably within HBM bandwidth. Cluster-wide sharing via `cluster_shape_mn=(1, num_query_heads, 1)` is a future optimization once profiling shows it matters.

**Done when:** `n_h=128, batch=4, T=1` at hd=512 produces correct output with single launch, per-head LSE writes to `mLSE[batch, head, m_row]` correctly.

---

## Priority 5: Stage E — Production extraction

D5 is complete. Wrap the kernel in a proper interface.

| Step | What | Status |
|---|---|---|
| E1 | File placement: `dsv4/kernels/attention/fmha.py` | ✅ Done |
| E2 | Constructor signature (`head_dim`, `num_query_heads`, `sliding_window`, `top_k`, sink/causal flags, dtypes) | ⚠️ Partial — needs cleanup |
| E3 | Call signature: `q`, `compressed_kv`, `swa_kv`, `swa_lens`, `sink_logits`, `request_ids`, `o`, `stream` | ⚠️ Needs sink_bias / row_sums integration |
| E4 | Kernel cache + warmup, keyed on `(head_dim, num_query_heads, top_k, n_comp, apply_sink_bias, is_causal, ...)` | TODO |
| E5 | `torch.library.custom_op("dsv4::sparse_fmha_with_swa", mutates_args=("o",))` | TODO |
| E6 | Reference parity test against FP32 oracle in `dsv4/reference/attention.py` | TODO |
| E7 | Cleanup: delete debug test files, keep only `tests/unit/test_fmha_kernel.py` | TODO |

Block table, paged KV, FP8 dequant, inv_scale — all handled upstream by the indexer + gather chain. FMHA sees a dense BF16 `[T, top_k, head_dim]` tile.

---

## Priority 6: NVFP4-1.2 — Fuse FP4 quant into FMHA output → wo_a path

**Depends on Priority 2** (uses the register slot in the new final epilogue).

Currently: FMHA emits BF16 → inverse RoPE → BF16 GMEM → wo_a quantizes to FP4.

Target: register slot in FMHA's new final epilogue does `O / row_sum` *and* inverse RoPE rotation *and* per-microblock amax *and* FP4 pack. wo_a reads FP4 directly with no GMEM materialization.

Same pattern as Priority 3, different home (FMHA final epilogue, not MoE epilogue).

---

## Priority 7: NVFP4-2 — FP4 KV pipeline depth in FMHA

**Depends on Priority 2 being solid at BF16 KV first.** FP4 KV shrinks tiles ~4×; same SMEM budget supports more pipeline stages.

| KV dtype | Tile size (hd=512) | Stages fitting 192 KB |
|---|---:|---:|
| BF16 | 128 KB | 2 |
| FP8 | 64 KB | 4 |
| FP4 | ~36 KB | 6 |

At 1M-context decode where KV reads dominate, deeper pipelines hide more TMA latency.

**Implementation:**
- TMA loads FP4 NoPE dims (`e2m1_x2` packed) to SMEM slot 0.
- TMA loads BF16 RoPE dims to SMEM slot 1.
- TMA loads FP8 scale factors to SMEM slot 2.
- SMEM dequant FP4 → BF16 vectorized (`* FP8_scale`, 16-element microblocks).
- Concatenate `[NoPE, RoPE]` in SMEM.
- MMA reads contiguous BF16 from SMEM.

**Test:** FP4+BF16 split input → identical output to pure BF16 input. Dequant must be transparent.

---

## Priority 8 (conditional on P1 profile): Per-kt O rescale fix

**Only justified if Priority 1 profiling shows Python KV merge overhead > 5% of decode latency.** Otherwise defer indefinitely — the current correct workaround ships.

There are three possible paths. None is a small change.

### Path A: CUTLASS atom replication

The CUTLASS C++ Blackwell FMHA does a TMEM round-trip for O rescale using `SM100_TMEM_LOAD_32dp32b_4x_atom` and `SM100_TMEM_STORE_32dp32b_4x_atom`, which are paired by hardware design. So in principle TMEM round-trip is possible. The question is whether CuTeDSL Python exposes the specific atom variants and layout configuration CUTLASS uses, and whether they can be paired correctly through `make_tmem_copy`.

**Steps:**
- [ ] Read `/root/cutlass/.../blackwell/kernel/attention/fmha/fmha.py` (or equivalent C++ reference) and document the exact atom + repetition + tensor layout used for `correction_rescale`.
- [ ] Enumerate what CuTeDSL Python exposes: `dir(tcgen05.copy)`, available `LdNxNbOp` / `StNxNbOp` variants, what `Repetition(N)` controls.
- [ ] Identify the difference between current `Ld32x32bOp(Repetition(16))` + `St32x32bOp(Repetition(16))` and whatever CUTLASS uses.
- [ ] Build a minimal NO-OP round-trip test (load O, store back unchanged) with the candidate atom configuration. Verify cos = 1.0.
- [ ] If NO-OP passes, retest with `* acc_scale` modification.

**Risk:** CuTeDSL Python may not expose the necessary atom variants, or may not allow the layout configuration CUTLASS uses. In that case, escalate to Path B or C.

**Effort if it works:** 2–4 days investigation + 1–2 days porting.

### Path B: O accumulator in registers, manual PV

Restructure FMHA so the PV accumulator is register-resident, not TMEM-resident. Each kt: read V from SMEM, read P from TMEM/SMEM, compute one-shot PV (no accumulate) writing to a temporary TMEM region, then load to registers and add to register-resident running O (with acc_scale applied to the running O before the add).

**Implications:**
- Register pressure is severe at hd=512. 512 FP32 per row × 1 row per thread × 128 threads = 128 KB of registers just for O. Possible but tight.
- PV without TMEM accumulate is a non-standard MMA usage. May need to use a smaller PV tile and accumulate in registers across sub-tiles.
- Loses the natural MMA-accumulator pipeline overlap.

**Effort:** 1–2 weeks. High risk of regressing hd=64/128/256 paths during the refactor.

### Path C: O accumulator in SMEM

Variant of Path B with O in SMEM instead of registers. PV writes to a temporary TMEM region, gets loaded to registers, applies acc_scale * existing_SMEM_O + new_PV, stores back to SMEM. Final epilogue reads from SMEM.

**Implications:**
- SMEM budget tightens significantly (need O in SMEM = ~64 KB at hd=512).
- Adds SMEM read/write pressure on every kt.
- May require dropping kv_stage to 1 across the board.

**Effort:** ~1 week. Lower risk than Path B but bigger perf impact (more SMEM traffic).

### Recommended order if profiling demands a fix

1. Try Path A first — least invasive, may just work with the right atom config.
2. If Path A confirmed impossible in CuTeDSL Python, try Path C (SMEM-resident O) before Path B (register-resident O). SMEM gives more headroom at hd=512.
3. Path B only if both fail and the perf gap is truly critical.

---

## Priority 9: hd=512 single-kernel fix

**Currently blocked.** MLIR optimizer hangs > 3 hours on the hd=512 kernel. Tracer completes in 0.8s — kernel is structurally correct.

Decode works via head-packed M with `pv_n_tile=128` and `n_k_sub_tiles=2`, so this is only a prefill efficiency issue. Lower priority than the chain above.

**Options:**
1. Pre-compile cubin offline (accept 1–2 hour compile, cache result).
2. Add no-softmax mode emitting raw S to GMEM; call twice for k_sub=0/1, accumulate in Python, softmax once externally.
3. Write hd=512 path in raw CUTLASS C++. Bypasses CuTeDSL MLIR entirely. Most realistic if NVIDIA can't fix the optimizer.
4. Report CuTeDSL MLIR optimizer bug to NVIDIA.

---

## Priority 10: Indexer FP4 tensor-core scoring (Stage F)

Paper §5.2.1: *"the QK path in the indexer of CSA, where QK activations are cached, loaded, and multiplied entirely in FP4."*

Current indexer (`dsv4/kernels/cuda/indexer_score_topk.cu`): scalar FP32 dot products, no tensor cores, spinlock-protected shared-memory heap. Single largest perf gap in the codebase. At 1M-context decode it scores ~250K compressed entries per query token — the spinlock heap will not scale to top_k=1024.

**Target:** port DeepGEMM `fp8_paged_mqa_logits` to FP4 inputs with `tcgen05.mma.kind=mxf4nvf4`. Plus per-warp partial top-k merged with a final reduction tree (or radix-select). Plus FP32→BF16 score quantization (paper claims 2× speedup on top-k selector at 99.7% recall).

**Scope:** 2–3 weeks. Stage F. Do not start until the FP4 epilogue patterns from Priorities 3 and 6 are established — they inform the indexer's FP4 load + score paths.

---

## Build order — recommended sequencing

```
Priority 1 (PROFILE) ──► gates Priority 8
                        │
Priority 2 (one-way    ─┼─► unblocks Priority 4 (multi-CTA)
final epilogue)         │   unblocks Priority 6 (FP4 fuse in FMHA)
                        │
Priority 3 (NVFP4-1.1) ─┴── parallel, independent
                        
[verify hd regressions]
        │
        ▼
Priority 4 (D2 multi-CTA grid)
        │
        ▼
Priority 5 (Stage E production extraction)
        │
        ▼
Priority 6 (NVFP4-1.2 FP4 fuse in FMHA output)
        │
        ▼
Priority 7 (NVFP4-2 FP4 KV pipeline)
        │
        ▼
Priority 8 (per-kt rescale fix — ONLY if P1 says it matters)
        │
        ▼
Priority 9 (hd=512 fix — only if prefill efficiency demands)
        │
        ▼
Priority 10 (indexer FP4 tensor-core scoring) — Stage F
```

**Key change from the previous version:** Priority 1 is now a profiling task, not an engineering task. Priority 2 is scoped honestly — it's the one-way path only, and it does *not* fix the per-kt rescale. The per-kt rescale is Priority 8, conditional, and three paths deep because there is no easy fix.

---

## Speculative — beyond what the V4 paper validated

Listed for completeness. **Do not implement without explicit sign-off.**

1. **NVFP4 compressed KV NoPE dims.** Paper validated FP8; FP4 would halve cache again. Risk: compounds quantization noise on already-lossy compressed KV.
2. **MXFP4 vs NVFP4 for indexer scoring.** Not validated for indexer specifically.
3. **NVFP4 for full attention Q×K^T GEMM.** Closed. Cos 0.86 vs FP32 in earlier tests. Attention stays BF16/FP32.
4. **Per-token FP8 activation scaling in FMHA.** Not validated. Out of scope.
5. **2:4 structured sparsity on FP4 expert weights.** V4 not trained with structured sparsity. Off the table for the released checkpoint.
6. **NVFP4 LM head + MTP head.** Big VRAM win (~1.4 GB saved on Pro). Modest quality risk on rare-token logits. Test against held-out eval before shipping.

---

## Key numbers to remember

| Config | n_h | top_k | s_k decode | n_kv_tiles | Multi-tile? |
|---|---:|---:|---:|---:|:---|
| Flash decode | 64 | 512 | 640 | 5 | YES |
| Pro decode | 128 | 1024 | 1152 | 9 | YES |
| Current single-tile test | 1 | — | 128 | 1 | NO |

Production decode needs the multi-tile path. Today's Python KV merge ships correct results at the cost of 5–9 launches per step. Whether that cost matters is what Priority 1 measures.