biondizzle
1e6adf5e01
- fmha_multihead_launch.cu: PyTorch launch wrapper for fmha_6warp_multihead_kernel
(c10::BFloat16 boundary, uint16_t bf16_t inside kernel, zero-cost casts)
- fmha_multihead_op.py: torch.utils.cpp_extension JIT loader + custom_op registration
(dsv4::fmha_multihead_decode for torch.compile)
- production.py: fast path dispatch for T=1, n_segments==1, hd in {64,128,256}
Falls through to CuTeDSL slow path for multi-segment/prefill
- test_p3_fast_decode.py: integration test (MHA/MQA/GQA, cosine >= 0.999998)
Architecture:
Grid: dim3(1, n_h, batch_size) — one CTA per (head, batch)
MQA: k_head_stride=0 so all Q heads share same K/V
Single kernel launch, zero cudaDeviceSynchronize on hot path
Normalized output for single-segment decode
5.6 KiB
OK. Here's the full status.
STATUS UPDATE — DSV4 NVFP4 Kernel (as of 2026-05-28)
What's DONE and PROVEN (cos 0.999+)
| Component | Status | Notes |
|---|---|---|
| FMHA core (hd=64/128/256) | ✅ | TMEM-P (hd≤64), SMEM-P (hd>64). cos 0.999997-0.999998 |
| D1: multi-KV-tile via Python KV merge | ✅ | Correct, cos 0.999998. 5–9 kernel launches per decode step |
| D2: per-head multi-head launch | ✅ | n_h=1–128, head-packed MQA/GQA |
| D3: SWA length mask (in-kernel) | ✅ | cos 0.999996 |
| D4: causal mask on SWA (in-kernel) | ✅ | cos 0.999996, combined with D3 |
| D5c: sink merge | ✅ | Single softmax over [S_comp, S_swa + attn_sink]. Single-tile and multi-tile both pass |
| D5b: per-row LSE output | ✅ | All 128 rows write correctly now |
| Stage E: production wrapper | ✅ | production.py — kernel cache, batch dim, MQA/GQA head-packed, torch.library.custom_op |
| MoE fused SwiGLU (NVFP4) | ✅ | Clamping in kernel, 7-warp specialization |
| Router (hash + dense) | ✅ | Both modes |
| use_2cta_instrs conditional | ✅ | 1.7–1.9× prefill speedup |
| NVFP4 primitives | ✅ | sf_dtype, TMA, MMA kind all verified |
| GPU-only NVFP4 quantize | ✅ | Byte-exact match with Python |
| KV cache infra | ✅ | allocator, paged_cache, state_cache, flush, schema, handle, manager |
| Compressor (CSA/HCA) | ✅ | flush_write kernels, FP8/FP4 quantize |
| Indexer | ✅ | gather_kv, score_topk (FP32 scalar). Compiles and runs on B200 |
| Model assembly | ✅ | config, layer_schedule, all 43 Flash + 61 Pro layers construct and validate |
What's BLOCKED / UNDONE
| Item | Status | What's needed |
|---|---|---|
| D1.5: in-kernel O rescale (TMEM round-trip) | ❌ FUNDAMENTALLY BROKEN | TMEM load→store atoms have mismatched column mappings. NO-OP round-trip corrupts data. Closed issue — do not re-attempt. Production path is Python KV merge. |
| Priority 1: Profile production decode | ❌ Not done | Need to measure if 5–9 launch Python KV merge overhead actually matters. Gates Priority 8. |
| Priority 2: One-way final-epilogue rewrite | ❌ Not done | Replace epilogue_tma_store with MoE-style TMEM→REGS→SMEM→GMEM. Unlocks P4 (multi-CTA) and P6 (FP4 fuse). Attempted in fmha_smem_acc.py — many commits, unclear if it's working. |
| Priority 3: NVFP4-1.1 FP4 quant in MoE epilogue | ❌ Not done | Fuse amax+FP4 pack into L1→L2 path. Independent of FMHA. |
| Priority 4: D2 multi-CTA grid | ❌ Blocked on P2 | epilogue_tma_store can't accept flat_divide coordinates. |
| Priority 5: Stage E cleanup | ⚠️ Partial | production.py exists. E1✅ E2⚠️ E3⚠️ E4-E7 TODO. Many debug test files still in tests/unit/ (66 test files, lots of diagnostic/debug artifacts). |
| Priority 6: NVFP4-1.2 FP4 in FMHA output | ❌ Blocked on P2 | Needs register slot in new final epilogue. |
| Priority 7: NVFP4-2 FP4 KV pipeline | ❌ Blocked on P2 | FP4 KV dequant in SMEM for deeper pipeline stages. |
| Priority 8: Per-kt rescale fix | ❌ Conditional on P1 | Three paths (A/B/C) if profiling shows Python merge overhead >5%. |
| Priority 9: hd=512 single-kernel | ❌ MLIR hang | CuTeDSL MLIR optimizer can't handle it. Decode works via head-packed M + hd≤256 chunks. |
| Priority 10: Indexer FP4 tensor-core scoring | ❌ Stage F | Scalar FP32 scoring today. Needs FP4 MMA + warp-level top-k. 2–3 weeks. |
MAY_24_2026_PLAN_NEW.md (TMEM Round-Trip Investigation Plan)
This is a 5-phase investigation plan for Priority 8, Path A (CUTLASS atom replication). Status: not executed. The plan was written but MEMORY.md already contains the conclusion: TMEM round-trip is FUNDAMENTALLY BROKEN and will NEVER work. The plan's Phase 1 (read CUTLASS reference) and Phase 2 (NO-OP round-trip tests) were effectively done through the many test_d15_* and test_tmem_roundtrip_minimal.py test files. All variants failed. The plan is historical — its conclusion (escalate to Path C / Python KV merge) is what's already in production.
ROADMAP.md
The roadmap is current and accurate as of 2026-05-26. The priority ordering and dependency chain are correct. No priorities have been completed since the last update — the latest work was Stage E production extraction (commit b9f15c2).
Attention Kernel Folder
fmha.py(592 lines) — main production kernelfmha_smem_acc.py(592 lines) — SMEM accumulator variant, attempted but unclear status (likely broken/in-progress for P2)production.py(427 lines) — working production wrapper with Python KV mergefmha_backup_*.py— historical backups- 66 test files in
tests/unit/— heavy diagnostic residue from the D1.5 investigation
Summary: Where We're At
The kernel works. FMHA at hd=64/128/256 with all masks (SWA, causal, sink), multi-head, multi-KV-tile via Python merge — all proven on B200. The MoE, router, indexer, compressor, KV cache — all working.
The next move is Priority 1 (profile production decode) to determine if the Python KV merge overhead matters. If it's <5% of decode latency, you ship as-is and the D1.5 TMEM round-trip problem becomes irrelevant.
If profiling says fix it, the path is Priority 2 first (one-way final epilogue rewrite using MoE pattern), which unblocks the multi-CTA and FP4 fusion chain. The per-kt rescale (P8) is a separate, harder problem with three possible approaches.
The MAY_24_2026_PLAN_NEW.md is dead — its investigation was run, and the answer was "TMEM round-trip doesn't work." Should be archived or deleted.