3c0451a3e5
REVERT to working example7 (n=128 cos 0.999998). Example8 TMA fix didn't work.
2026-05-22 20:28:15 +00:00
880bd9ef81
Update stage_c test to example8: SSA kv_coord + per-tile O rescale
2026-05-22 20:27:58 +00:00
c395b279d2
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
a54a241052
Revert TMA to kt pattern (n=128 works), multi-tile TMA is separate bug
2026-05-22 20:10:35 +00:00
956ca1ecfd
TMA: use self.n_kv_tiles + kv_coord pattern from working diag test
2026-05-22 20:08:07 +00:00
242ffebcd9
REVERT to 0bdcdc0 — the version that passed n=128 cos 0.999998
2026-05-22 20:07:18 +00:00
3814838107
DEBUG: disable O rescale + normalize, test if corr setup alone causes regression
2026-05-22 19:57:53 +00:00
2c51f82382
Shared corr tensors for O rescale + final normalize, fix softmax loop
2026-05-22 19:55:06 +00:00
f165257c50
Add O rescale with correction_rescale pattern + fix TMA to working diag pattern
2026-05-22 19:51:53 +00:00
0bdcdc0efd
O normalize: exact CUTLASS correction_rescale pattern with 2D reg tensor
2026-05-22 19:50:37 +00:00
009cf9f80d
O normalize: TMEM round-trip with paired Ld/St atoms + standard epilogue_tma_store
2026-05-22 19:49:51 +00:00
f821dd00fe
Fix: use NamedBarrier instead of mbarrier_arrive/wait
2026-05-22 19:49:07 +00:00
3bd406e925
Fix: barrier_wait → mbarrier_wait, barrier_arrive → mbarrier_arrive
2026-05-22 19:48:31 +00:00
544f0ca52b
Fix epilogue: corr_tile_size=16, proper epi_subtile tuple, match CUTLASS reference
2026-05-22 19:47:57 +00:00
d526aa04fb
Fix example7: K slice (None,None,0,0) and softmax scale_log2 double-bug
2026-05-22 19:45:59 +00:00
558aac0581
Fix: fence_async_shared -> fence_view_async_shared
2026-05-22 19:25:05 +00:00
d62e6fc9ca
Clean v2: real softmax P, no O TMEM modify, standard epilogue. Baseline for custom epilogue work.
2026-05-22 19:24:25 +00:00
183292d919
O normalize using tmem_ptr base (same as epilogue) + CUTLASS sub-tile pattern
2026-05-22 19:18:16 +00:00
365e8f53af
O normalize with full layout (no sub-tiling), Repetition(64)
2026-05-22 19:17:03 +00:00
7bc94e610a
Disable ALL O copies to verify baseline
2026-05-22 19:14:33 +00:00
590d6e9fba
Disable O rescale too for NO-OP test
2026-05-22 19:13:43 +00:00
51cec1405d
DEBUG: O load+store NO-OP to verify TMEM copy correctness
2026-05-22 19:12:43 +00:00
e3e2668192
Re-enable O rescale + normalize with corr_tile_size=32
2026-05-22 19:11:17 +00:00
d77c965646
Disable O rescale + normalize, verify softmax P only
2026-05-22 19:09:33 +00:00
dcc64dd14d
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
48b24ba005
Full pipeline: O rescale + final normalize with CUTLASS sub-tile approach
2026-05-22 19:05:59 +00:00
a85894df89
Test softmax P vs unnormalized reference (no O normalize)
2026-05-22 19:05:27 +00:00
c0b39fc2bf
O normalize using CUTLASS reference sub-tile approach
2026-05-22 19:04:37 +00:00
3dbda0eebb
Fix O normalize: use 2D register tensor indexing
2026-05-22 18:59:15 +00:00
6b61d5274c
Add O normalization with sub-tile TMEM read-modify-write
2026-05-22 18:58:35 +00:00
b936c6220d
Simplify: softmax P only, no O rescale/normalize yet
2026-05-22 18:57:00 +00:00
e2fad84205
Real softmax test built on working identity diag
2026-05-22 18:55:31 +00:00
d3b662d3a8
CRITICAL FIX: remove extra scale_log2 in softmax (minus_row_max and acc_scale)
2026-05-22 18:52:58 +00:00
32869c7378
FIX: K slice (None,None,0,0) like working diag
2026-05-22 18:51:22 +00:00
4b1fc7ee1f
Diag: identity softmax on example6 pipeline to isolate softmax bug
2026-05-22 18:49:33 +00:00
912f92c6b5
Quick test: working v3 with n=256 multi-tile
2026-05-22 18:47:54 +00:00
5b6392beaa
DEBUG: add version marker to confirm code changes are running
2026-05-22 18:47:05 +00:00
c7d55a5f49
CRITICAL FIX: TMA pre-slice (None,0,None,0) → (None,None,0,0) to keep GMEM tile dim free
2026-05-22 18:45:55 +00:00
f734610268
Diag: TMA shapes with hardcoded major modes
2026-05-22 18:43:55 +00:00
18a589347c
Diag: simplified TMA shape analysis
2026-05-22 18:43:30 +00:00
7ad4ddb6ba
Diag: print TMA partition shapes for multi-tile debugging
2026-05-22 18:43:05 +00:00
67c5a0928d
FIX: Use Python range() in TMA warp for concrete per-iteration GMEM coords
2026-05-22 18:26:31 +00:00
54de81985f
FIX: Force SSA GMEM coord via n_kv_tiles - n_kv_tiles instead of cutlass.range kt
2026-05-22 18:25:13 +00:00
f3503fc1ee
FIX: TMEM offset bug in O rescale/normalize — use tOtO0.iterator not tOtO.iterator
2026-05-22 18:23:46 +00:00
0b7ae7c969
Diag: test n=384 (3 tiles) to find crash boundary
2026-05-22 18:07:07 +00:00
640ec3e96e
Diag: test all sizes 128-1024
2026-05-22 18:06:28 +00:00
02d993ecac
DEBUG: disable O rescale to isolate NaN cause
2026-05-22 18:05:46 +00:00
1c3970fe58
Add NaN/inf checking to stage C test
2026-05-22 18:01:11 +00:00
d7a0fc2bc2
CRITICAL FIX: K GMEM slice (None,None,0,0) not (None,0,None,0)
...
K from QK MMA B-partition has GMEM iter at mode 1, NOT mode 2.
(None,0,None,0) hardcodes mode 1 to 0 → TMA always loads tile 0.
(None,None,0,0) keeps mode 1 free → correct multi-tile loading.
Proof: diag n=256 went from cos 0.711 → 0.999999 with this one change.
2026-05-22 17:59:57 +00:00
b6a2904e93
Diag: try K slice (None,None,0,0) keeping mode 1 (CUTLASS ref style)
2026-05-22 17:59:01 +00:00
