Commit Graph

652 Commits

Author SHA1 Message Date
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
01621e1520 Diag: try runtime Int32(0+0) for kv_coord with cutlass.range 2026-05-22 17:57:58 +00:00
beecc4df47 Diag: use Python range() unrolling like stage C test 2026-05-22 17:56:59 +00:00
200430bd3f Fix diagnostic test: same Int32(kt) + n_kv_tiles fixes 2026-05-22 17:56:15 +00:00
c23ebd5b57 Try cutlass.range with Int32(kt) — now n_kv_tiles is Python int 2026-05-22 17:51:25 +00:00
4a41df51c4 FIX: n_kv_tiles as Python int (s_k//128) for range() unrolling
cute.size() returns a CuTeDSL symbol, not a Python int.
range() on a symbol can't iterate — the loop never unrolls.
Now n_kv_tiles is computed in __init__ as s_k // 128 (Python int).
2026-05-22 17:50:07 +00:00
70409636f7 Option 2: Python range() with Int32(kt) for TMA GMEM coord
cutlass.range traces once - kv_coord/kt are trace-time values,
not runtime loop-carried state. Python range() fully unrolls at
trace time, emitting distinct Int32(k) constants per iteration.
Int32(1) hardcoded already proved TMA CAN load from tile 1.
2026-05-22 17:47:43 +00:00
b55a38c4c3 Add example5: use cutlass.range induction variable as TMA GMEM coord 2026-05-22 17:47:10 +00:00
aacad257ea README: add fire_b200_test docs, update multi-tile blocker with real findings 2026-05-22 17:41:23 +00:00
93c28b9c29 Clean up debug prints, set kv_coord as Int32(0)
Key findings to relay to CUTLASS LLM:
- kv_coord=Int32(1) hardcode CHANGES the output (TMA CAN load from different tiles)
- kv_coord=Int32(0) + kv_coord += 1 does NOT increment at runtime
  (all multi-tile outputs identical to kv_coord=0)
- kv_coord=0 (plain Python int) also doesn't work
- Pipeline handle .count doesn't work either
- The TMA GMEM tile coordinate must be dynamic at kernel runtime,
  but CuTeDSL appears to constant-fold or not propagate the increment
2026-05-22 17:39:27 +00:00
1bba851911 DEBUG: try plain Python int kv_coord (like CUTLASS ref) 2026-05-22 17:34:30 +00:00
15b2a28d29 DEBUG: hardcode kv_coord=1 to test if TMA uses it 2026-05-22 17:32:53 +00:00
ff9ef6dcde DEBUG: try K slice (None,0,None,0) keeping mode 2 free 2026-05-22 17:30:06 +00:00
cec6f59d66 DEBUG: print tBgK/tVgV shapes before/after slice 2026-05-22 17:28:45 +00:00
8740ab5b27 Stage C: manual kv_coord + correct K GMEM slice + O rescale fence
Key fixes:
1. GMEM tile coord: manual Int32 kv_coord (not kvh.count)
2. K GMEM slice: (None,None,0,0) keeps mode 1 free (GMEM iter)
3. V GMEM slice: (None,0,None,0) keeps mode 2 free (GMEM iter)
4. Add fence_view_async_tmem_load before O rescale for visibility
2026-05-22 17:26:56 +00:00
f39c3a8e38 Add example4: manual kv_coord Int32 for GMEM tile indexing
Pipeline handle .count is NOT a GMEM coordinate — it's opaque pipeline
state. CUTLASS FMHA reference uses a manual Int32 counter (kv_coord)
for GMEM and handle.index for SMEM. .count is never used as a
coordinate anywhere in the reference.
2026-05-22 17:24:38 +00:00
711480a4fb README: add test harness instructions 2026-05-22 17:09:53 +00:00
83f8a13add run_test.sh: SIGKILL all children of screen session on cleanup
Deadlocked GPU processes ignore SIGHUP from screen -X quit.
Now kills the entire process group with SIGKILL, plus a catch-all
pkill for any python test_ processes.
2026-05-22 17:08:12 +00:00