Files

biondizzle 4bb0e063cc D1.5: Replace broken TMEM round-trip with correction epilogue (paired atoms)

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.

2026-05-26 19:11:19 +00:00

5.2 KiB

Raw Blame History

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):

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.
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.
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:

Pre-compile offline + cache cubin (if MLIR eventually finishes)
Write hd=512 path in raw CUTLASS C++ (bypass CuTeDSL MLIR)
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.

5.2 KiB Raw Blame History Unescape Escape