782de32644
FIX: tma_partition tensors have 4 modes, not 8. Mode 2 is GMEM tile dim.
...
The 8-mode indexing (tBgK[None,None,None,None,kt,None,None,None]) fails at
JIT compilation with 'coord and shape are weakly congruent' error. The actual
MLIR tensor shape is (((64,128),1),?,?,?) — 4 modes, not 8.
The working fix from commit 845ad98
2026-05-22 23:08:27 +00:00
265c4583b7
Fix test_fmha_v3_stage_c.py: 8-mode TMA indexing (mode 4 = GMEM tile dim)
2026-05-22 22:58:10 +00:00
25054b6b09
Fix README: multi-tile was layout bug not JIT bug, add example10, update status
2026-05-22 22:57:53 +00:00
a3873a4f7d
Add diag test with 8-mode TMA indexing from commit 2711611
2026-05-22 22:40:09 +00:00
5201fbfe20
auto: pre-test commit
2026-05-22 22:38:07 +00:00
0466eaa407
DOCUMENT: TMA 8-mode indexing — the bug that cost us a full day. README + inline comments.
2026-05-22 21:28:58 +00:00
8fd2ac48d8
Fix identity diag: same 8D TMA indexing fix
2026-05-22 21:21:52 +00:00
dfed61dbf3
FIX: Use full 8D indexing for tBgK/tVgV — mode 4 is the GMEM tile dim
2026-05-22 21:21:23 +00:00
a9c0114879
Diagnostic: check tBgK/tVgV layout strides for degenerate dims
2026-05-22 21:20:46 +00:00
da6df02aec
Test identity diag multi-tile
2026-05-22 21:14:42 +00:00
5834db1f75
Minimal reference FMHA test: n=256 only
2026-05-22 21:13:37 +00:00
8b0f6bc0cd
Test reference FMHA with proper API
2026-05-22 21:12:49 +00:00
f94fd52a08
Test: CUTLASS reference FMHA on B200 multi-tile
2026-05-22 21:11:58 +00:00
a1a1c34a4f
REVERT to working baseline (n=128 cos 0.999998). Multi-tile TMA is a CuTeDSL JIT limitation.
2026-05-22 20:37:21 +00:00
20d517a08b
Test: use kvh.index (pipeline state) as TMA GMEM coordinate
2026-05-22 20:36:21 +00:00
0822b3476c
SMEM counter: separate allocate_tensor instead of struct field
2026-05-22 20:35:42 +00:00
4ff379fc3d
Fix SMEM counter type: cutlass.Int32 for MemRange
2026-05-22 20:35:17 +00:00
8d6fe8075d
SMEM-backed kv_coord counter — JIT can't constant-fold SMEM reads
2026-05-22 20:34:52 +00:00
e90e25c6da
DEBUG: hardcoded Int32(1) to test if TMA can read tile 1
2026-05-22 20:34:21 +00:00
f522490d3a
DEBUG: use Int32(kt) directly to test if coordinate matters
2026-05-22 20:34:03 +00:00
8889fe5766
Test: kv_coord = warp_idx() * 0 — force SSA from runtime value
2026-05-22 20:33:40 +00:00
8a7446a2c1
DEBUG: add cute.printf for kv_coord runtime value
2026-05-22 20:33:03 +00:00
cef4bcf686
Test: Python range() instead of cutlass.range() for TMA loop
2026-05-22 20:32:44 +00:00
4a102a9953
Test example9: drop try_acquire/pk, single loop-carried kv_coord
2026-05-22 20:32:25 +00:00
2aa6e4d234
REVERT to working example7 (n=128 cos 0.999998). Example8 TMA fix didn't work.
2026-05-22 20:28:15 +00:00
fa1f5b8ef6
Update stage_c test to example8: SSA kv_coord + per-tile O rescale
2026-05-22 20:27:58 +00:00
594efdcbc0
Clean up tests: archive superseded files, keep only essential unit tests
...
Kept in tests/unit/:
- test_fmha_v3.py (stages A+B)
- test_fmha_v3_diag.py (identity softmax, n=128+256)
- test_fmha_v3_stage_c.py (real softmax, n=128 cos 0.999998)
- layertest.py + cudagraph_test.py (required for every change)
- infrastructure: cache, custom_op, cutedsl, router, fp4, fused, interleave
Archived: 19 superseded unit tests + 10 root-level scratch files
Root level: only fmha_v3_stage_c_example7.py remains (now in unit/)
2026-05-22 20:25:27 +00:00
52857aee16
Revert TMA to kt pattern (n=128 works), multi-tile TMA is separate bug
2026-05-22 20:10:35 +00:00
dc9a5bc499
TMA: use self.n_kv_tiles + kv_coord pattern from working diag test
2026-05-22 20:08:07 +00:00
e96eb789a0
REVERT to 0bdcdc0 — the version that passed n=128 cos 0.999998
2026-05-22 20:07:18 +00:00
ad6ea439a4
DEBUG: disable O rescale + normalize, test if corr setup alone causes regression
2026-05-22 19:57:53 +00:00
4fe88a6994
Shared corr tensors for O rescale + final normalize, fix softmax loop
2026-05-22 19:55:06 +00:00
7a2c62c7fc
Add O rescale with correction_rescale pattern + fix TMA to working diag pattern
2026-05-22 19:51:53 +00:00
f31f4210eb
O normalize: exact CUTLASS correction_rescale pattern with 2D reg tensor
2026-05-22 19:50:37 +00:00
ea3de238b3
O normalize: TMEM round-trip with paired Ld/St atoms + standard epilogue_tma_store
2026-05-22 19:49:51 +00:00
642f715091
Fix: use NamedBarrier instead of mbarrier_arrive/wait
2026-05-22 19:49:07 +00:00
b1bd8bc8ac
Fix: barrier_wait → mbarrier_wait, barrier_arrive → mbarrier_arrive
2026-05-22 19:48:31 +00:00
d3b49c9998
Fix epilogue: corr_tile_size=16, proper epi_subtile tuple, match CUTLASS reference
2026-05-22 19:47:57 +00:00
633a272988
Fix example7: K slice (None,None,0,0) and softmax scale_log2 double-bug
2026-05-22 19:45:59 +00:00
1212b663c3
Fix: fence_async_shared -> fence_view_async_shared
2026-05-22 19:25:05 +00:00
e229155d01
Clean v2: real softmax P, no O TMEM modify, standard epilogue. Baseline for custom epilogue work.
2026-05-22 19:24:25 +00:00
06941c317d
O normalize using tmem_ptr base (same as epilogue) + CUTLASS sub-tile pattern
2026-05-22 19:18:16 +00:00
877ebb7584
O normalize with full layout (no sub-tiling), Repetition(64)
2026-05-22 19:17:03 +00:00
e0a604c508
Disable ALL O copies to verify baseline
2026-05-22 19:14:33 +00:00
64d174fdf6
Disable O rescale too for NO-OP test
2026-05-22 19:13:43 +00:00
5a0b575bfe
DEBUG: O load+store NO-OP to verify TMEM copy correctness
2026-05-22 19:12:43 +00:00
c2fe1bffb5
Re-enable O rescale + normalize with corr_tile_size=32
2026-05-22 19:11:17 +00:00
95d3d1bf03
Disable O rescale + normalize, verify softmax P only
2026-05-22 19:09:33 +00:00
69ec2ce7e5
FIX: O sub-tile count should be HEAD_DIM/corr_tile_size, not 128/corr_tile_size
2026-05-22 19:06:52 +00:00
c4de6aa5e7
Full pipeline: O rescale + final normalize with CUTLASS sub-tile approach
2026-05-22 19:05:59 +00:00
