- **hd=512: MLIR compilation hangs.** SMEM budget fixed (192KB ✅), kernel structure correct (tracer 0.8s), but MLIR→PTX backend optimizer cannot process the IR in reasonable time (>3 hours). Both `range()` unrolled and `cutlass.range(unroll=1)` runtime loops trigger this. This is a CuTeDSL/MLIR toolchain limitation.
- **External k_sub merge doesn't work.** k_sub segments are additive in logit space (S = S_0 + S_1), not attention weight space. The D5 merge formula does not apply. In-kernel k_sub accumulation is the only correct approach.
- **O rescale (kt>0):** Uses hand-constructed TMEM atoms. May corrupt data for n>128 (multi-KV-tile). At n=128 (1 KV tile, kt=0), no rescale needed. Guarded with `const_expr(n_kv_tiles > 1)`.
This caused the QK GEMM to only compute 64 of 128 (or 256, 512) dimensions at hd>64. The QK dot products were half the correct length, producing wrong attention scores.
The MLIR→PTX backend optimizer has exponential-or-worse behavior for kernels with TMA pipeline acquire/release inside loops. Both unrolled (Python `range`) and runtime (`cutlass.range unroll=1`) loops trigger this. The Python tracer is fast (0.8s) because it just generates IR. The MLIR optimizer then chews on that IR for hours. **Workaround:** keep pipeline loops as simple as possible. Consider raw CUDA C++ for complex kernels.
### 2. External k_sub Merge is Mathematically Impossible
You CANNOT merge the outputs of two attention calls that compute softmax(Q_k0 @ K_k0^T)@V and softmax(Q_k1 @ K_k1^T)@V into softmax(Q @ K^T)@V. The k_sub segments are additive in LOGIT space (S = S_0 + S_1), but softmax is nonlinear. The D5 merge formula works because sparse and SWA attend over DIFFERENT token sets (additive in weight space). k_sub attends over the SAME tokens with PARTIAL dot products. These are fundamentally different operations. **The only correct approach is in-kernel accumulation (S_0 + S_1 before softmax).**
### 3. 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. The cost is 4 PV GEMM passes instead of 2. But PV is typically not the bottleneck. This is simpler than SMEM overlap (which requires CuTeDSL SmemAllocator changes) or Q tiling (which adds pipeline complexity).
### 4. Guard Dead Code with const_expr
CuTeDSL compiles BOTH branches of Python `if` statements, generating IR for code that will never execute at a given head_dim. Use `const_expr(condition)` to eliminate dead code at compile time. This is critical for:
- O rescale code (only needed when n_kv_tiles > 1)
- LSE computation (only needed when normalize=False)
- SMEM-P path (only needed when use_smem_p=True)
### 5. Don't Mix Python Loops and CuTeDSL Pipeline Operations
Python `for` loops unroll at trace time, creating N copies of the loop body in the IR. For pipeline acquire/release + TMA copy + GEMM, each copy is substantial. `cutlass.range(unroll=1)` creates a runtime loop with one copy of the body. **For pipeline operations, prefer `cutlass.range(unroll=1)` to reduce IR size**, even though the MLIR optimizer may still struggle with it.
### 6. The k_tile Parameter is the Key to hd=512
At hd=512, the kernel splits Q and K into sub-tiles of size `k_tile=256` along the head_dim. Each sub-tile is loaded via TMA, processed by MMA, and accumulated. `n_k_sub_tiles = head_dim // k_tile = 2`. The k_tile parameter controls the sub-tile size and the number of iterations. **k_tile must be ≤ 256** (MMA instruction K-dim limit) and must evenly divide head_dim.
**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).** 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.
**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).
**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.
**Problem:** hd=512 exceeds the MMA instruction's max K-dim (256). Must split Q and K into 2 sub-tiles along head_dim (k_tile=256, n_k_sub_tiles=2). The QK dot product is S = Q_k0 @ K_k0^T + Q_k1 @ K_k1^T (additive in logit space).
**SMEM budget: SOLVED.** pv_n_tile=128 for hd>256 reduces sV from 64KB→32KB, sC from 64KB→32KB. Total 192KB ✅.
**Compilation: BLOCKED.** The CuTeDSL MLIR→PTX backend optimizer cannot compile the hd=512 kernel in reasonable time. Both Python `range()` (unrolled IR) and `cutlass.range(unroll=1)` (runtime loop) produce IR that the optimizer chews on for 3+ hours without finishing. The Python tracer completes in 0.8s — the kernel is structurally correct. This is a toolchain limitation.
**External merge: IMPOSSIBLE.** The D5 online softmax merge formula assumes separate attention distributions over different token sets (additive in weight space). k_sub segments are additive in LOGIT space (S = S_0 + S_1), not weight space. You cannot recover softmax(S_0 + S_1)@V from softmax(S_0)@V and softmax(S_1)@V. In-kernel accumulation before softmax is the only correct approach.
1.**Pre-compile hd=512 kernel offline.** Accept 1-2 hour compilation during build. Cache the cubin. This works if the MLIR optimizer eventually finishes (it might just be slow, not stuck — but 3+ hours is excessive even for pre-compilation).
2.**Add no-softmax mode to the kernel.** Output raw S (QK scores) without softmax. Call twice for k_sub=0 and k_sub=1. Accumulate S_0+S_1 in Python. Apply softmax once. This requires modifying the softmax warp to optionally skip normalization and output S to GMEM instead of P to TMEM/SMEM.
3.**Write hd=512 kernel in CUTLASS C++.** Bypass CuTeDSL's MLIR backend entirely. Use raw CUTLASS C++ with tcgen05 MMA intrinsics. More work but compilation is fast (seconds).
4.**Report CuTeDSL MLIR optimizer bug.** The optimizer should handle this IR in reasonable time. File an issue with NVIDIA.
Three honest buckets. A fourth speculative bucket flagged at the end.
### NVFP4-0: Verify Right Blackwell FP4 Primitives ⚡ DO FIRST
**No correctness or quality risk. Pure correctness of implementation.** If these are wrong, we're running wrong MMA shapes silently.
#### NVFP4-0.1 — sf_dtype tracing
**What:** Trace the SF dtype through the full pipeline: `gemm_runner.py` → `dense.py` → `blockscaled_utils` → TMEM layout.
**The problem:** `dense.py` line 137 says NVF4 supports `Float8E8M0FNU/Float8E4M3FN` at sf_vec_size=16. But UE8M0 is the MXFP4/MXFP8 scale format. NVFP4 uses **FP8 E4M3**. The examples on lines 90/100 show `Float8E8M0FNU` at sf_vec_size=16 which is the MXFP4 path. **Need to verify the runner is passing E4M3, not E8M0.**
Action:
- [ ] Print `sf_dtype` in `gemm_runner.py` at construction: `print(f"sf_dtype={sf_dtype}, sf_vec_size={SF_VEC_SIZE}")`
- [ ] Print `self.sf_dtype` in `dense.py``BlockScaledGEMM.__init__`
- [ ] Print `self.sf_vec_size` in `dense.py`
- [ ] Trace through `blockscaled_utils.make_sm100_sf_layout` — does it produce E4M3 packing (4 FP8 E4M3 → 1 int32) or UE8M0 packing?
- [ ]**If wrong sf_dtype is found:** fix in `gemm_runner.py` SF_DTYPE constant, retest MoE cosine
#### NVFP4-0.2 — SF TMEM layout verification
**What:** NVFP4 expects scale factors in TMEM in a specific transposed-packed layout. UE4M3 for NVFP4 (4 packed FP8 E4M3 per int32 word). The comment in `dense.py` about "SM100 requires scaling factors in packed UE8M0 format" is for **MXFP8**, not NVFP4.
**What:** `float4_e2m1fn_x2` must survive all the way into TMA descriptor creation. Blackstone TMA supports `e2m1_x2` packed-FP4 element type directly. Loading as `uint8` works but loses tensor-core awareness.
Action:
- [ ] Trace `float4_e2m1fn_x2` through `quantize.py` → TMA atom creation in `fmha.py`
- [ ] Print the GMEM tensor dtype at FMHA kernel input
- [ ] Print the TMA atom dtype at construction
- [ ] Verify `cpasync.tma_partition` receives `float4_e2m1fn_x2` element type, not uint8
- [ ]**If uint8 fallback:** fix TMA atom creation in `fmha.py`
#### NVFP4-0.4 — MMA kind is mxf4nvf4
**What:** Blackwell has a single MMA kind for both MXFP4 and NVFP4. NVFP4 = scales are FP8 E4M3, 16-element block. MXFP4 = scales are UE8M0, 32-element block. The MMA kind is determined by scale-factor type at runtime. Need to confirm tcgen05 is inferring NVFP4.
Action:
- [ ] Print `tcgen05.mma.kind` at GEMM construction (if accessible)
- [ ] Print the MMA instruction shape `(M, N, K)` confirmed by JIT compile
- [ ] Verify it matches Blackwell MMA shape for NVFP4 (not MXFP4)
**Execution:** These are 5-minute print jobs. Do all 4 NVFP4-0 items before touching any code. If any of them reveals a wrong dtype, fix it FIRST before anything else. A wrong sf_dtype poisons every FP4 GEMM result.
---
### NVFP4-1: Eliminate BF16 Round-Trips After FP4 GEMMs 🔴 PURE-WIN, NO QUALITY RISK
**These are pure bandwidth/compute wins. The math doesn't change — we just avoid precision loss and kernel launch overhead.**
The amax reduction is in-registers: for an epi tile with 16 contiguous elements per thread, each tile produces one FP8 E4M3 scale and 64 bits of packed FP4 nibbles. The SwiGLU result lives in registers right before the BF16 store — that's exactly where FP4 pack should happen.
**What you save:**
- ~2× GMEM bandwidth between L1 and L2 (FP4 instead of BF16)
4. Write packed nibbles to GMEM as `float4_e2m1fn_x2`
5. Write FP8 scale to SF TMA buffer
**The amax subtlety:** For NVFP4 the microblock is 16 elements. Port the same 16-element logic from `quantize.py` into the epilogue. Do NOT use 32-element MXFP4 microblocks.
**Note:** NVFP4-1.2 and NVFP4-1.3 depend on D1.5 (correction epilogue fix) because those epilogues need the clean one-way TMEM path. NVFP4-1.1 (MoE SwiGLU) is independent.
---
### NVFP4-2: FP4 KV Pipeline Depth in FMHA 🔴 STAGE D, DEPENDS ON D1
**FP4 KV shrinks tiles 4×, same SMEM budget buys 3× more pipeline stages.**
- [ ]**Prerequisite:** D1 working at BF16 first. Cannot skip.
- [ ]**Test:** FP4+BF16 split input → identical output to pure BF16 input (dequant is transparent)
---
### NVFP4-3: use_2cta_instrs for Production MoE 🟢 30 MINUTES, PURE PERF
**This is the single biggest single-knob perf win for FP4 GEMMs on B200.**
**What:** `FusedSwiGLUScaledGroupedGemmKernel` supports 2-CTA UMMA but defaults to `False`. With 2-CTA, the B operand is TMA-multicast: each CTA reads half of B, peers cross the Infiniband link. Effective MMA tile M doubles (128→256, 256→512).
**Measured win:** 1.7–1.9× throughput over single-CTA at prefill/batch shapes.
**Decision tree:**
- M < 128 (decode single-token): 1-CTA is correct. 2-CTA wastes hardware.
- M ≥ 256 (prefill or batched decode): 2-CTA is free perf.
- [ ]**Test:** throughput comparison at M=256, 512, 1024
- [ ]**Scope:** MoE-side, `gemm_runner.py`. Does not affect FMHA.
---
### ⚠️ Speculative: Beyond V4 Paper Validation
The following are real potential wins but go beyond what the V4 paper explicitly validated for FP4. Listed for completeness, do NOT implement without explicit sign-off from Mike.
1.**Indexer FP4 tensor-core scoring (paper §5.2.1 "QK path in the indexer... cached, loaded, and multiplied entirely in FP4")**
- Paper says the indexer SHOULD do QK in FP4 with tensor cores
- Current: scalar FP32 dot products with no tensor cores
- Huge scope: 2-3 weeks minimum
- **Risk:** FP4 dot product precision for index selection needs recall validation
- **Verdict:** Track for Stage F. Do NVFP4-0.4 first.
2.**MXFP4 vs NVFP4 for indexer scoring** — not validated in the paper for indexer specifically. Evaluate after NVFP4-0.
3.**NVFP4 for full attention Q×K^T GEMM** — Already closed. NVFP4 Q×K^T is too lossy (cos 0.86 vs FP32). Attention stays FP16/FP32.
4.**Per-token FP8 activation scaling in FMHA** — Different precision model, not validated. Out of scope.
---
## NVFP4 Execution Order
| # | Task | Scope | Risk | Blocks | Est. |
|---|------|-------|------|--------|------|
| NVFP4-0.1 | sf_dtype tracing | Both | NONE — print only | D1 if wrong | 5 min |
| NVFP4-0.2 | SF TMEM layout | Both | NONE — print only | D1 if wrong | 5 min |
| NVFP4-0.3 | FP4 TMA element type | FMHA | NONE — print only | D1 if wrong | 5 min |
| NVFP4-0.4 | MMA kind verification | GEMM | NONE — print only | everything | 5 min |
| NVFP4-3 | use_2cta_instrs conditional | MoE | NONE — perf only | nothing | 30 min |
| NVFP4-1.1 | Fuse FP4 quant into SwiGLU epilogue | MoE | NONE | nothing | 1 day |
| NVFP4-1.2 | Fuse FP4 quant into invRoPE→wo_a | Attention | NONE | D1.5 | 1 day |
| NVFP4-1.3 | Fuse FP4 quant into mHC mixing | Attention | NONE | post-D5 | 2 days |
| NVFP4-2 | FP4 KV pipeline depth | FMHA | NONE — perf only | D1 | 1 day |
**NVFP4-0 results gate the critical path.** If NVFP4-0.1–0.4 find a wrong sf_dtype, the fix comes before D2. Everything else is either parallel or post-D1.
**NVFP4-3 (use_2cta_instrs) is the fastest win and has no dependencies.** Do it immediately after the NVFP4-0 prints.
**NVFP4-1.1 (fuse FP4 into SwiGLU) is the next-biggest win.** No FMHA dependency. Do it in parallel with D2.
**NVFP4-2 (FP4 KV) depends on D1 being solid.** Do after D2 or alongside hd=512 fix.
