biondizzle
4bb0e063cc
Replace hand-constructed Ld32x32bOp/St32x32bOp TMEM round-trip with the proven correction epilogue pattern from fused_swiglu.py: 1. O rescale (kt>0): TMEM→REGS (paired load), multiply by acc_scale, REGS→TMEM (paired store via retile_to_S). No layout mismatch. 2. Final O output: One-way TMEM→REGS→SMEM→GMEM using epilogue_tmem_copy_and_partition + epilogue_smem_copy_and_partition + TMA partition. Register-level normalization (divide by row_sum) or raw BF16 cast for D5a path. This fixes both D1.5 issues: - Issue 1: TMEM round-trip corruption (hand-constructed atoms) - Issue 2: O rescale for multi-KV-tile (kt>0) Supports normalize=True (in-kernel) and normalize=False (D5a external). Uses epilog_sync_bar + c_pipe for SMEM→GMEM, replacing epilogue_tma_store.
103 lines
5.2 KiB
Markdown
103 lines
5.2 KiB
Markdown
# 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.
|