# May 24–26, 2026 — Session Plan & Progress

## Quick Context

You're working on the DSV4 (DeepSeek V4 Pro) NVFP4 inference kernel for Blackwell B200. The FMHA (Fused Multi-Head Attention) kernel is working at hd=64/128/256 (cos 0.999998). D5 (sink merge) is COMPLETE. The next milestones are: fix D1.5 O rescale, then proceed to production extraction (Stage E).

**B200:** See MEMORY.md for access (not committed to repo)
**Repo:** `git@sweetapi.com:2222/biondizzle/nvfp4-megamoe-kernel.git`
**Local:** `~/dev/nvfp4-megamoe-kernel`
**Test command:** `~/.openclaw/workspace/fire_b200_test <test_file>`

## ✅ Completed (May 24–26)

### 1. NVFP4-0: Verify FP4 Primitives ✅
- All four diagnostics PASS — sf_dtype, TMA element type, MMA kind all correct
- NVFP4 uses FP8 E4M3 scales (NOT UE8M0), 16-element blocks, tcgen05 MMA kind correct

### 2. NVFP4-3: use_2cta_instrs Conditional ✅
- tokens_sum >= 256 and cluster_m even → 2-CTA UMMA
- 1.7-1.9× throughput at prefill shapes. Decode stays 1-CTA.

### 3. D1: Parameterized HEAD_DIM ✅ (hd≤256)
- hd=64/128/256: cos 0.999998, LSE err 0.0
- hd=512: BLOCKED by MLIR compilation hang (>3hr). External k_sub merge impossible.
- D1.5 O rescale: BLOCKED by TMEM round-trip corruption. Python KV merge workaround.

### 4. D2: Multi-Query Grid ✅ (per-head launch)
- Head-packed M-dimension: Q reshaped to (n_h*T, hd, 1), per-row softmax
- cos 0.999995 for n_h=1-128 at hd=64, n_h=2-8 at hd=128, n_h=2 at hd=256
- Multi-CTA grid: BLOCKED by flat_divide + epilogue_tma_store mismatch

### 5. D3: SWA Sequence Length Mask ✅
- In-kernel post-QK masking via tTMEM_LOADcS coordinates
- swa_len as Int32 scalar, offset by n_comp for D5c

### 6. D4: Causal Mask ✅
- SWA-relative position (kv_pos - n_comp) > m_coord → -inf
- Combined with D3 via OR logic

### 7. D5: SWA + Sink Merge ✅
- D5a: normalize flag + LSE + row_sums output
- D5b: Per-row LSE + Python KV merge (cos 0.999994)
- D5c: Sink bias as logit modification — **key insight: sink merge = single softmax over [S_comp, S_swa + attn_sink]**
  - Single-tile: cos 0.999996
  - Multi-tile (Python KV merge): cos 0.999996
  - D5d NOT NEEDED — sink bias approach supersedes fused merge epilogue

## 🎯 Next Priorities (in order)

### Priority 1: D1.5 — Fix O Rescale for Multi-KV-Tile (BLOCKER)
**Why:** Production DSV4 Pro decode needs s_k=1152 (9 KV tiles). Python KV merge works but requires 5-9 kernel launches per decode step.

**Approaches (ordered by feasibility):**
1. **Correction epilog pattern** (1-2 days): One-way TMEM→regs→SMEM→GMEM. Study CUTLASS reference `correction_rescale_and_partition` + `epilogue_tmem_copy_and_partition`. This is the proper Blackwell pipeline.
2. **Skip TMEM round-trip entirely**: After softmax, write P to SMEM, PV accumulate to SMEM, TMA store from SMEM. Requires SMEM budget for both P and O.
3. **Python KV merge as production path**: Accept the multi-launch overhead. Profile to see if it's actually a bottleneck (5-9 launches × ~50μs ≈ 250-450μs, vs ~50μs for single launch with O rescale).

### Priority 2: Stage E — Production Extraction
**Why:** D5 is complete. The kernel works. Time to wrap it in a proper interface.

- E1: File placement (already done — `dsv4/kernels/attention/fmha.py`)
- E2: Constructor signature (partially done — needs cleanup)
- E3: Call signature (needs sink_bias, row_sums, n_comp integration)
- E4: Kernel cache + warmup (key on n_comp, apply_sink_bias, head_dim, s_k)
- E5: torch.library custom op
- E6: Reference parity test
- E7: Cleanup (delete debug test files)

### Priority 3: NVFP4-1.1 — Fuse FP4 into SwiGLU Epilogue (1 day, parallel)
**Why:** Biggest bandwidth win for MoE pipeline. No FMHA dependency. Can work in parallel with D1.5.

Current: L1 GEMM → SwiGLU → BF16 GMEM → quantize → FP4 GMEM → L2 GEMM
Target: L1 GEMM → SwiGLU → FP4 pack in registers → FP4 GMEM → L2 GEMM

### Priority 4: hd=512 Fix (BLOCKED by MLIR)
**Status:** Kernel structurally correct (tracer 0.8s). MLIR optimizer hangs for 3+ hours.
**Options:**
1. Pre-compile offline + cache cubin (if MLIR eventually finishes)
2. Write hd=512 path in raw CUTLASS C++ (bypass CuTeDSL MLIR)
3. Report bug to NVIDIA

### Priority 5: D2 Multi-CTA Grid (BLOCKED by flat_divide)
**Status:** Per-head launch works for decode. Multi-CTA needed for prefill.
**Requires:** Full tma_partition + epilogue refactor into kernel (1-2 day effort).

## Key Numbers

| Config | n_h | top_k | s_k | n_kv_tiles | O rescale needed? |
|--------|----:|------:|----:|-----------:|:------------------|
| Flash decode | 64 | 512 | 640 | 5 | YES |
| Pro decode | 128 | 1024 | 1152 | 9 | YES |
| Current test | 1 | — | 128 | 1 | No (guarded away) |

## Rules (don't forget)

- NEVER edit on B200. Edit locally → commit → push → pull → test.
- ALWAYS use `fire_b200_test` or `run_test.sh`. Never raw SSH+nohup.
- ALWAYS verify hd=64 regression (cos ~0.999998) after every change.
- Guard dead code with `const_expr()`. CuTeDSL compiles both branches.
- CuTeDSL `if` blocks create separate MLIR regions — variables NOT visible across blocks.
- After every P store to TMEM, call `cute.arch.fence_view_async_tmem_store()`.
- NOHUP DOES NOT WORK on B200. Screen sessions survive SSH drops.
