Commit Graph

690 Commits

Author SHA1 Message Date
0330c1da7a Fix README: multi-tile was layout bug not JIT bug, add example10, update status 2026-05-22 22:57:53 +00:00
79e35f10ad Add diag test with 8-mode TMA indexing from commit 2711611 2026-05-22 22:40:09 +00:00
9bba97158b auto: pre-test commit 2026-05-22 22:38:07 +00:00
dbd77f2bc4 DOCUMENT: TMA 8-mode indexing — the bug that cost us a full day. README + inline comments. 2026-05-22 21:28:58 +00:00
3a4524c318 Fix identity diag: same 8D TMA indexing fix 2026-05-22 21:21:52 +00:00
2fd6f02e8e FIX: Use full 8D indexing for tBgK/tVgV — mode 4 is the GMEM tile dim 2026-05-22 21:21:23 +00:00
cec5d1c313 Diagnostic: check tBgK/tVgV layout strides for degenerate dims 2026-05-22 21:20:46 +00:00
127c1b7a46 Test identity diag multi-tile 2026-05-22 21:14:42 +00:00
8df3f53d01 Minimal reference FMHA test: n=256 only 2026-05-22 21:13:37 +00:00
a49a2934f2 Test reference FMHA with proper API 2026-05-22 21:12:49 +00:00
a92054dd0e Test: CUTLASS reference FMHA on B200 multi-tile 2026-05-22 21:11:58 +00:00
51173b4805 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
8d5733bb4c Test: use kvh.index (pipeline state) as TMA GMEM coordinate 2026-05-22 20:36:21 +00:00
f2b1cc39d3 SMEM counter: separate allocate_tensor instead of struct field 2026-05-22 20:35:42 +00:00
4955fb5204 Fix SMEM counter type: cutlass.Int32 for MemRange 2026-05-22 20:35:17 +00:00
5f6e33c054 SMEM-backed kv_coord counter — JIT can't constant-fold SMEM reads 2026-05-22 20:34:52 +00:00
18a5c60e26 DEBUG: hardcoded Int32(1) to test if TMA can read tile 1 2026-05-22 20:34:21 +00:00
b7a1deed52 DEBUG: use Int32(kt) directly to test if coordinate matters 2026-05-22 20:34:03 +00:00
c01291f16a Test: kv_coord = warp_idx() * 0 — force SSA from runtime value 2026-05-22 20:33:40 +00:00
da4eebc9a4 DEBUG: add cute.printf for kv_coord runtime value 2026-05-22 20:33:03 +00:00
b8d4ac0ecb Test: Python range() instead of cutlass.range() for TMA loop 2026-05-22 20:32:44 +00:00
ae6197fc78 Test example9: drop try_acquire/pk, single loop-carried kv_coord 2026-05-22 20:32:25 +00:00
20290d7439 REVERT to working example7 (n=128 cos 0.999998). Example8 TMA fix didn't work. 2026-05-22 20:28:15 +00:00
72ad0eca9e Update stage_c test to example8: SSA kv_coord + per-tile O rescale 2026-05-22 20:27:58 +00:00
b1ee693998 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
578c7e58be Revert TMA to kt pattern (n=128 works), multi-tile TMA is separate bug 2026-05-22 20:10:35 +00:00
6a82036ebd TMA: use self.n_kv_tiles + kv_coord pattern from working diag test 2026-05-22 20:08:07 +00:00
74380ffc21 REVERT to 0bdcdc0 — the version that passed n=128 cos 0.999998 2026-05-22 20:07:18 +00:00
a3f0e73c82 DEBUG: disable O rescale + normalize, test if corr setup alone causes regression 2026-05-22 19:57:53 +00:00
aa43a902e1 Shared corr tensors for O rescale + final normalize, fix softmax loop 2026-05-22 19:55:06 +00:00
bd52e88ce1 Add O rescale with correction_rescale pattern + fix TMA to working diag pattern 2026-05-22 19:51:53 +00:00
93acac7e43 O normalize: exact CUTLASS correction_rescale pattern with 2D reg tensor 2026-05-22 19:50:37 +00:00
70349de6dc O normalize: TMEM round-trip with paired Ld/St atoms + standard epilogue_tma_store 2026-05-22 19:49:51 +00:00
02543b3d8d Fix: use NamedBarrier instead of mbarrier_arrive/wait 2026-05-22 19:49:07 +00:00
5f98eaec3a Fix: barrier_wait → mbarrier_wait, barrier_arrive → mbarrier_arrive 2026-05-22 19:48:31 +00:00
afe5d1ae21 Fix epilogue: corr_tile_size=16, proper epi_subtile tuple, match CUTLASS reference 2026-05-22 19:47:57 +00:00
8b93774d70 Fix example7: K slice (None,None,0,0) and softmax scale_log2 double-bug 2026-05-22 19:45:59 +00:00
6827862b3d Fix: fence_async_shared -> fence_view_async_shared 2026-05-22 19:25:05 +00:00
7806ac812a Clean v2: real softmax P, no O TMEM modify, standard epilogue. Baseline for custom epilogue work. 2026-05-22 19:24:25 +00:00
1176d75236 O normalize using tmem_ptr base (same as epilogue) + CUTLASS sub-tile pattern 2026-05-22 19:18:16 +00:00
86a4719afe O normalize with full layout (no sub-tiling), Repetition(64) 2026-05-22 19:17:03 +00:00
1d2111f772 Disable ALL O copies to verify baseline 2026-05-22 19:14:33 +00:00
d9f7443b77 Disable O rescale too for NO-OP test 2026-05-22 19:13:43 +00:00
62529df638 DEBUG: O load+store NO-OP to verify TMEM copy correctness 2026-05-22 19:12:43 +00:00
59575202a4 Re-enable O rescale + normalize with corr_tile_size=32 2026-05-22 19:11:17 +00:00
0589a14790 Disable O rescale + normalize, verify softmax P only 2026-05-22 19:09:33 +00:00
8db432ebed 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
b0498dd7ac Full pipeline: O rescale + final normalize with CUTLASS sub-tile approach 2026-05-22 19:05:59 +00:00
668c7e1a4e Test softmax P vs unnormalized reference (no O normalize) 2026-05-22 19:05:27 +00:00
46020baff1 O normalize using CUTLASS reference sub-tile approach 2026-05-22 19:04:37 +00:00