Commit Graph

5 Commits

Author SHA1 Message Date
2eb44a00bf fix(tmem): warp-collective TMEM ops + one-way correction epilogue
Key fixes for fmha_epilogue_sm100.cuh hang:
- tcgen05.ld/st are WARP-COLLECTIVE: ALL 32 lanes must execute
- Old code guarded TMEM ops with if(tid==0) = warp divergence = HANG
- tmem_dealloc now uses tmem_base (value from alloc), not SMEM pointer
- Compute attention in SMEM, then do one-way TMEM pipeline:
  SMEM → TMEM (warp-collective store) → regs (warp-collective load)
  → normalize in regs → BF16 cast → GMEM
- This proves the MoE-style one-way correction epilogue on FMHA

Also: enable TMEM kernel test + hd=128 in standalone test
2026-05-28 07:27:25 +00:00
bcc5d0b6cb FMHA SM100: Add TMEM+correction epilogue kernel (Priority 2)
New file: fmha_epilogue_sm100.cuh
- TMEM alloc/dealloc/load/store via tcgen05 PTX
- One-way correction epilogue: TMEM→regs→normalize→BF16→GMEM
- D1.5 fix: O rescale in REGISTERS (TMEM→regs→multiply→TMEM)
- Same pattern as MoE epilogue but with normalize instead of SwiGLU
- Unblocks D2 multi-CTA and NVFP4-1.2 (register slot for FP4 pack)

Test: hd=64 + hd=128, reference vs TMEM kernels
2026-05-28 06:27:56 +00:00
8eb735618f fix: use expf for softmax (not exp2f with scale) 2026-05-28 05:34:03 +00:00
2d4e2c57e0 auto: pre-test commit 2026-05-28 05:22:23 +00:00
a30ebfb197 FMHA SM100: Full kernel with TMET PTX, UMMA descriptors, softmax loop
- TMEM alloc/dealloc/load/store via inline PTX (tcgen05.*)
- UMMA SMEM descriptor construction (make_umma_desc)
- QK GEMM via tcgen05.mma.kind::f16 inline asm
- Online softmax with D3/D4/D5c masks
- O rescale in REGISTERS (D1.5 fix — no TMEM round-trip!)
- FP4 quantize helpers (hs2e2m1, fp8_e4m3_encode)
- Still needs: PV GEMM, proper P staging, TMEM O load/store
2026-05-28 05:19:34 +00:00