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
3cb339129b
FMHA SM100: Fix Phase 1 — single-thread reference for correctness
...
Use thread 0 for all computation (slow but correct).
SMEM for Q and O sharing across threads.
Online softmax with O rescale — correct D1.5 approach.
D3 SWA mask implemented.
Target: cos ~0.999998 then parallelize.
2026-05-28 05:32:47 +00:00
77fa34a9a6
fix: update launch wrapper for fmha_decode_ref
2026-05-28 05:28:49 +00:00
00ac46c9d3
FMHA SM100: Phase 1 — reference scalar implementation
...
Simpler approach first: scalar Q@K^T, softmax, P@V in registers.
No TMEM/MMA yet — verify correctness first, then replace with tcgen05.
- 192-thread CTA, all threads cooperate on one (batch, head)
- Online softmax with O rescale (correct D1.5 approach)
- D3 SWA mask, D4 causal (TODO), D5c sink (TODO)
- KV loaded in blocks of 128 for SMEM efficiency
- Correctness target: cos ~0.999998 against PyTorch reference
2026-05-28 05:27:36 +00:00
6f7449ce71
FMHA SM100: Fix tcgen05.mma PTX syntax — correct register constraints
...
- tcgen05.mma.cta_group::1.kind::f16 [tmem_c], desc_a, desc_b, idescE_hi, scaleC, {mask0..3}, pred
- idescE is upper 32 bits of the E descriptor
- scaleC is a float (1.0 for accumulate)
- mask is 4 uint32 values (0xFFFFFFFF for no masking)
2026-05-28 05:25:59 +00:00
a11a245307
fix: use unsigned short for BF16 storage, inline PTX for conversions
2026-05-28 05:24:32 +00:00
373900fa08
FMHA SM100: Fix launch wrapper to match new kernel API
2026-05-28 05:20:31 +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
09dfd4a41f
fix: rename .cpp to .cu for CUDA compilation
2026-05-28 05:16:41 +00:00
48baea7728
FMHA SM100: Remove CUTLASS includes, write raw PTX inline asm
...
CUTLASS headers transitively include cuda_bf16.h which has a CUDA 13.2
in_place_from bug. Writing tcgen05 PTX directly via inline asm instead.
No dependencies on CUTLASS C++ — pure PTX + CUDA runtime.
2026-05-28 05:15:07 +00:00
88d5995ec9
fix: define bf16_t using __bf16 built-in, avoid cuda_bf16.h bug
2026-05-28 05:14:01 +00:00
6bd3356582
fix: include cuda_bf16.h unconditionally, add --expt-relaxed-constexpr
2026-05-28 05:13:01 +00:00
c1266b5275
fix: include cuda_bf16.h only in device code
2026-05-28 05:12:30 +00:00
a64e55665b
fix: avoid cuda_bf16.h, use inline PTX for BF16 conversion
2026-05-28 05:12:08 +00:00
1734d13f60
fix: restore cuda_bf16.h include
2026-05-28 05:11:39 +00:00
8783a25deb
fix: guard cuda_bf16.h with __CUDA_ARCH__
2026-05-28 05:11:11 +00:00
5e389b5ed9
fix: remove duplicate desc declaration
2026-05-28 05:10:43 +00:00
7ac2499266
fix: defer UMMA descriptor — use placeholder for now
2026-05-28 05:10:15 +00:00
db17d8db9a
fix: cvta.to.shared PTX for SMEM address
2026-05-28 05:09:50 +00:00
e12a81ae36
fix: include cstdint
2026-05-28 05:09:28 +00:00
0c73a024ba
fix: guard CUTLASS includes with __CUDA_ARCH__ for host compilation
2026-05-28 05:09:07 +00:00
41e59a2423
FMHA SM100: Add SMEM descriptor construction for tcgen05.mma
2026-05-28 05:08:25 +00:00
230c350c77
FMHA SM100: Raw CUDA C++ decode kernel — initial skeleton
...
6-warp specialization using CUTLASS C++ atoms directly:
- tcgen05.mma for QK (SMEM→SMEM→TMEM) and PV (TMEM→SMEM→TMEM)
- TMEM accumulator with one-way correction epilogue (TMEM→regs→SMEM→GMEM)
- In-kernel O rescale via registers (fixes D1.5 TMEM round-trip!)
- D3/D4/D5c masks, NVFP4 quantize helpers, FP8 E4M3 encode
- PyTorch binding with head_dim template dispatch
This bypasses all CuTeDSL limitations: float→int, TMEM round-trip,
multi-CTA, hd=512 MLIR compilation hang.
2026-05-28 05:04:44 +00:00
b9f15c250f
Stage E: head-packed MQA/GQA, batch dim, custom_op, integration API
...
- production.py: head-packed M dimension for MQA/GQA (q_per_kv*T rows
in single launch per KV group, eliminating redundant K/V TMA loads)
- production.py: batch dimension support (outer Python loop)
- production.py: warmup_attention_kernels() for pre-compilation
- production.py: dsv4_attention_per_head() for exact per-head sink bias
- __init__.py: sparse_fmha_with_swa, dense_fmha_with_swa, swa_only_fmha
integration functions bridging AttentionSubBlock → production FMHA
- custom_ops.py: dsv4::sparse_fmha_with_swa custom_op registration
- test_production.py: comprehensive tests (MHA/MQA/GQA, head-packed vs
per-head parity, multi-segment KV, SWA+causal+sink, batch, edge cases)
2026-05-27 15:15:03 +00:00
2412a5431b
MQA/GQA: batch Q heads into kernel batch dim, shared K/V per KV group
2026-05-27 08:31:23 +00:00
778d9d4f4f
Compile with row_sums tensor so kernel writes per-row row_sums
2026-05-27 07:10:00 +00:00
0736a04d9b
Fix KV merge: use NORMALIZED O (O_unnorm/row_sum) with LSE
2026-05-27 07:07:51 +00:00
06e7f7ab48
Debug: print LSE values for 2-segment merge
2026-05-27 07:04:39 +00:00
8f8d14c300
Match tensor slicing exactly to test_d1_kv_merge (2D slices, 3D unsqueeze)
2026-05-27 06:58:28 +00:00
6ee61717c0
Match tensor shapes from working test_d1_kv_merge
2026-05-27 06:56:04 +00:00
36a6f07a7e
Fix: unsqueeze k/v when dim==2
2026-05-27 06:52:43 +00:00
fc4172937c
Clean production wrapper: always normalize=False + KV merge
2026-05-27 06:51:14 +00:00
8f87109f86
Single-segment: use normalize=False + per-row normalization from row_sums
2026-05-27 06:48:56 +00:00
fe55bf23a0
Split single-segment (normalized) and multi-segment (KV merge) paths
2026-05-27 06:46:30 +00:00
b70ab2a6ee
Return o_accum directly (un-normalized merge result)
2026-05-27 06:42:58 +00:00
6111db571c
Match working test: don't pass row_sums to kernel
2026-05-27 06:41:44 +00:00
312ac52d15
Normalize O_accum by exp(lse) before returning
2026-05-27 06:39:36 +00:00
ddc701af9b
Use exact merge formula from working test_d1_kv_merge.py
2026-05-27 06:38:04 +00:00
8321ccf9c1
Fix production KV merge: use normalized O for log-sum-exp merge
2026-05-27 06:36:24 +00:00
98c93c1cd8
Stage E: production attention wrapper + Python KV merge, clean fmha_smem_acc
2026-05-27 06:34:10 +00:00
51e456df44
Slice MMA tile coords from tOgO for TMA copy
2026-05-27 05:39:42 +00:00
1caa737b09
Move sC_flat_staged creation before const_expr guard
2026-05-27 05:38:39 +00:00
3c9dbc0c5d
Staged sC_flat with (128, pv_n_tile//2, 2) to match TMA atom
2026-05-27 05:37:05 +00:00
de2028b106
Split sC_flat into staged layout to match TMA atom decomposition
2026-05-27 05:35:56 +00:00
a0e9f7534b
Use tCgC_epi (transformed) for GMEM side of TMA partition
2026-05-27 05:34:40 +00:00
b02e103ac0
Add c_simple GMEM tensor (non-dynamic) for SMEM accumulator TMA store
2026-05-27 05:33:30 +00:00
2438826eee
Use tma_partition with group_modes on both sC_flat and gO
2026-05-27 05:31:47 +00:00
603f52de78
Fix gO creation: use slice_(pv_mma_tiler) like fmha.py
2026-05-27 05:30:50 +00:00
b39d7f1a14
Try cute.copy(tma_c, sC_flat, gO) directly
2026-05-27 05:29:51 +00:00
