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

---

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
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	Workaround	Fix path
D1.5 TMEM round-trip corruption	Hand-built atoms produce 3% error on NO-OP round-trip; blocks in-kernel multi-tile O rescale and in-kernel normalize	Emit un-normalized O + LSE; Python KV merge for s_k>128	Priority 1: correction-epilog rewrite (sketch below)
hd=512 MLIR backend hang	Cannot compile hd=512 kernel (>3hr optimizer time, structurally correct)	Run hd=512 via head-packed M with hd≤256 chunks; ship without hd=512 if needed for D2	Pre-compile cubin / raw CUTLASS C++ / report NVIDIA bug
D2 multi-CTA grid (flat_divide + epilogue_tma_store)	Per-head Python launch wastes 128 launches per decode step at Pro	Per-head launch (works, just slow)	Unblocked by correction-epilog rewrite (uses flat_divide + tma_partition like MoE does)

---

Priority 1: Correction epilog rewrite (unblocks D1.5 + a chain of follow-ons)

Why this first: Every downstream item needs the kernel to have a register-level slot in the epilogue for modification. The current epilogue_tma_store path with hand-built atoms doesn't have one. The correction-epilog pattern does.

The pattern is already in the codebase — dsv4/kernels/gemm/fused_swiglu.py uses it for the MoE SwiGLU epilogue (lines 2021, 2064–2229). Library helpers, paired atoms, one-way TMEM → registers → SMEM → GMEM. SwiGLU + clamping math sits between the t2r and r2s copies. That's the exact slot FMHA needs.

What changes:

Replace the FMHA epilogue (dsv4/kernels/attention/fmha.py lines 549–597 — the epilogue_tma_store call) with:

1. Setup (run once per kernel, outside the kt loop):

   • tCtO_transformed = utils.gemm.sm100.transform_partitioned_tensor_layout(tOtO0)
   • tCgC_transformed = utils.gemm.sm100.transform_partitioned_tensor_layout(tCgC)
   • tiled_copy_t2r, tTR_tO_base, tTR_rO = utils.gemm.sm100.epilogue_tmem_copy_and_partition(...)
   • tiled_copy_r2s, tRS_rC, tRS_sC = utils.gemm.sm100.epilogue_smem_copy_and_partition(...)
   • TMA partition for C via flat_divide(tCgC_transformed, epi_tile) + cpasync.tma_partition(...)

2. Final epilogue (replaces the round-trip normalize):

   • Subtile loop, in each subtile: cute.copy(tiled_copy_t2r, tTR_tO_mn, tTR_rO) → multiply by inv_row_sum in registers if normalize=True → cast to BF16 → cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[...]) → TMA SMEM → GMEM.

3. Per-kt O rescale (replaces the broken hand-built round-trip on lines 524–544):

   • Inside the kt loop, when kt > 0: same t2r → multiply by acc_scale in registers → store back via the paired atom (tiled_copy_t2r.retile_to_S() or equivalent — verify exact API on B200).

What unblocks:

• D1.5 issue 1 (round-trip corruption): gone.
• D1.5 issue 2 (per-kt rescale): gone.
• In-kernel multi-tile attention (single launch for s_k=1152, not 9).
• NVFP4-1.2 (fuse FP4 quant into FMHA output → wo_a path): the register slot is where amax + FP4 pack go.
• D2 multi-CTA grid: flat_divide + tma_partition path is the same one MoE uses successfully. The flat_divide vs local_tile mismatch resolves.

Caveats to print and verify on B200:

• Exact CUTLASS helper API for the "store back to TMEM" direction (retile_to_S form vs separate helper vs same-base-tensor pattern).
• Whether transform_partitioned_tensor_layout accepts tOtO0 (TMEM iterator with offset) or needs a fresh tensor built at tmem_ptr + self.tmem_o0_offset with tCtO_fake.layout.
• Whether tma_partition inside if warp_idx < self.mma_warp_id works in this kernel's region tree. The MoE kernel does it; if FMHA hits "weakly congruent," hoist the partition call above the warp branch.

Done when:

• hd=64/128/256 regression cos ≥ 0.999998 holds with normalize=True and normalize=False.
• New multi-tile s_k=256 test with kt > 0 rescale gives cos ≥ 0.999998 (not the current 0.997 Python-merge workaround, the real in-kernel rescale).
• Existing Python KV merge tests continue to pass (test_d15_multi_kv.py).

---

Priority 2: Stage E — Production extraction

D5 is complete. The kernel works. Wrap it 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

Notably absent from the call signature: block_table, paged KV, inv_scale, FP8 dequant. All handled upstream by the indexer + gather kernel chain. FMHA sees a dense BF16 [T, top_k, head_dim] tile.

---

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

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

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 that hold 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.

Subtlety: NVFP4 microblock = 16 elements. Port the same 16-element logic from dsv4/ops/quantize.py. Don't accidentally use the 32-element MXFP4 block.

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 (no regression from BF16 intermediate).
• L2 GEMM reads FP4 scales produced by L1 epilogue.

---

Priority 4: D2 multi-CTA grid

Currently per-head Python launch (works, cos 0.999995, but 128 launches per decode step at Pro).

Multi-CTA grid is unblocked by Priority 1 — the flat_divide + tma_partition path becomes available once the epilogue uses the MoE pattern.

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

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

Q tensor layout: Option 1 — (batch, n_h, T, head_dim) with head as a TMA mode (matches CUTLASS reference and allows per-head LSE output). Picked over Option 2 (heads packed into M) because it generalizes better to GQA later.

Done when:

• n_h=128, batch=4, T=1 at hd=512 produces correct output with single launch.
• Per-head LSE writes to correct mLSE[batch, head, m_row] position.

---

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

Depends on Priority 1 (correction epilog gives the register slot).

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

Target: register slot in FMHA epilogue does the divide-by-row_sum and inverse RoPE rotation and per-microblock amax + FP4 pack. wo_a reads FP4 directly.

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

---

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

Depends on Priority 1 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 that fit (192 KB budget)
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 (packed e2m1_x2) 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 in vectorized form (* 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 7: hd=512 fix

Blocked. Per Priority 4, multi-CTA grid + head-packed M means decode at hd=512 can route through pv_n_tile=128 and n_k_sub_tiles=2, which compiles fine for hd=256. The hd=512 single-kernel compile is the missing piece for prefill efficiency, not correctness.

Options:

1. Pre-compile hd=512 cubin offline (accept 1–2 hour compile if MLIR ever finishes — uncertain).
2. Add no-softmax mode emitting raw S to GMEM, call twice for k_sub=0/1, accumulate in Python, softmax once. Two launches but no MLIR hang.
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.

Lower priority than the chain above — at decode T=1, n_h=128, hd=512 the head-packed approach already works without needing a single hd=512 kernel.

---

Priority 8: 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 the indexer 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 per paper (2× speedup on top-k selector, 99.7% recall).

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

---

Build order — recommended sequencing

Now ─┬─ Priority 1 (correction epilog rewrite)
     │     │
     │     └─→ unblocks D1.5, D2 multi-CTA, NVFP4-1.2
     │
     ├─ Priority 3 (NVFP4-1.1 fuse FP4 in SwiGLU) ← parallel, independent
     │
     ↓
     Verify hd=64/128/256 regressions hold
     │
     ↓
     Priority 2 (Stage E production extraction)
     │
     ↓
     Priority 4 (D2 multi-CTA grid)
     │
     ↓
     Priority 5 (NVFP4-1.2 fuse FP4 in FMHA output)
     │
     ↓
     Priority 6 (NVFP4-2 FP4 KV pipeline)
     │
     ↓
     Priority 7 (hd=512 fix — only if prefill efficiency demands it)
     │
     ↓
     Priority 8 (indexer FP4 tensor-core scoring) — Stage F

Priority 3 has no dependency on Priorities 1 or 2 and can run on a parallel branch.

---

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 for compressed KV; 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 (Priority 1) working in-kernel. Today's Python KV merge ships correct results at the cost of 5–9 launches per step.