biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Commit Graph

  • 8ccbdec1ed 🚀🚀🚀 TMA MULTI-TILE FIX VERIFIED ON B200 🚀🚀🚀 biondizzle 2026-05-22 23:51:29 +00:00
  • 0dc1fe5f71 FIX: (None,0,None,0) for ALL tma_partition outputs — verified shapes on B200 biondizzle 2026-05-22 23:35:55 +00:00
  • 0534f19aaf auto: pre-test commit biondizzle 2026-05-22 23:34:03 +00:00
  • ea7c414480 auto: pre-test commit biondizzle 2026-05-22 23:30:43 +00:00
  • 2236823293 auto: pre-test commit biondizzle 2026-05-22 23:29:15 +00:00
  • fa6388dbf6 auto: pre-test commit biondizzle 2026-05-22 23:28:16 +00:00
  • 9c5122f180 auto: pre-test commit biondizzle 2026-05-22 23:27:33 +00:00
  • 00e103ab91 FIX: (None,0,None,0) pre-slice keeps KV tile axis (mode 2) free biondizzle 2026-05-22 23:25:40 +00:00
  • 30eaba39aa FIX: 8-None no-op pre-slice opens full TMA coordinate space (8 dims) biondizzle 2026-05-22 23:18:40 +00:00
  • 9c5adcee46 FIX: tma_partition tensors have 4 modes, not 8. Mode 2 is GMEM tile dim. biondizzle 2026-05-22 23:08:27 +00:00
  • d7cdf63c58 Fix test_fmha_v3_stage_c.py: 8-mode TMA indexing (mode 4 = GMEM tile dim) biondizzle 2026-05-22 22:58:10 +00:00
  • 0330c1da7a Fix README: multi-tile was layout bug not JIT bug, add example10, update status biondizzle 2026-05-22 22:57:53 +00:00
  • 79e35f10ad Add diag test with 8-mode TMA indexing from commit 2711611 biondizzle 2026-05-22 22:40:09 +00:00
  • 9bba97158b auto: pre-test commit biondizzle 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. biondizzle 2026-05-22 21:28:58 +00:00
  • 3a4524c318 Fix identity diag: same 8D TMA indexing fix biondizzle 2026-05-22 21:21:52 +00:00
  • 2fd6f02e8e FIX: Use full 8D indexing for tBgK/tVgV — mode 4 is the GMEM tile dim biondizzle 2026-05-22 21:21:23 +00:00
  • cec5d1c313 Diagnostic: check tBgK/tVgV layout strides for degenerate dims biondizzle 2026-05-22 21:20:46 +00:00
  • 127c1b7a46 Test identity diag multi-tile biondizzle 2026-05-22 21:14:42 +00:00
  • 8df3f53d01 Minimal reference FMHA test: n=256 only biondizzle 2026-05-22 21:13:37 +00:00
  • a49a2934f2 Test reference FMHA with proper API biondizzle 2026-05-22 21:12:49 +00:00
  • a92054dd0e Test: CUTLASS reference FMHA on B200 multi-tile biondizzle 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. biondizzle 2026-05-22 20:37:21 +00:00
  • 8d5733bb4c Test: use kvh.index (pipeline state) as TMA GMEM coordinate biondizzle 2026-05-22 20:36:21 +00:00
  • f2b1cc39d3 SMEM counter: separate allocate_tensor instead of struct field biondizzle 2026-05-22 20:35:42 +00:00
  • 4955fb5204 Fix SMEM counter type: cutlass.Int32 for MemRange biondizzle 2026-05-22 20:35:17 +00:00
  • 5f6e33c054 SMEM-backed kv_coord counter — JIT can't constant-fold SMEM reads biondizzle 2026-05-22 20:34:52 +00:00
  • 18a5c60e26 DEBUG: hardcoded Int32(1) to test if TMA can read tile 1 biondizzle 2026-05-22 20:34:21 +00:00
  • b7a1deed52 DEBUG: use Int32(kt) directly to test if coordinate matters biondizzle 2026-05-22 20:34:03 +00:00
  • c01291f16a Test: kv_coord = warp_idx() * 0 — force SSA from runtime value biondizzle 2026-05-22 20:33:40 +00:00
  • da4eebc9a4 DEBUG: add cute.printf for kv_coord runtime value biondizzle 2026-05-22 20:33:03 +00:00
  • b8d4ac0ecb Test: Python range() instead of cutlass.range() for TMA loop biondizzle 2026-05-22 20:32:44 +00:00
  • ae6197fc78 Test example9: drop try_acquire/pk, single loop-carried kv_coord biondizzle 2026-05-22 20:32:25 +00:00
  • 20290d7439 REVERT to working example7 (n=128 cos 0.999998). Example8 TMA fix didn't work. biondizzle 2026-05-22 20:28:15 +00:00
  • 72ad0eca9e Update stage_c test to example8: SSA kv_coord + per-tile O rescale biondizzle 2026-05-22 20:27:58 +00:00
  • b1ee693998 Clean up tests: archive superseded files, keep only essential unit tests biondizzle 2026-05-22 20:25:27 +00:00
  • 578c7e58be Revert TMA to kt pattern (n=128 works), multi-tile TMA is separate bug biondizzle 2026-05-22 20:10:35 +00:00
  • 6a82036ebd TMA: use self.n_kv_tiles + kv_coord pattern from working diag test biondizzle 2026-05-22 20:08:07 +00:00
  • 74380ffc21 REVERT to 0bdcdc0 — the version that passed n=128 cos 0.999998 biondizzle 2026-05-22 20:07:18 +00:00
  • a3f0e73c82 DEBUG: disable O rescale + normalize, test if corr setup alone causes regression biondizzle 2026-05-22 19:57:53 +00:00
  • aa43a902e1 Shared corr tensors for O rescale + final normalize, fix softmax loop biondizzle 2026-05-22 19:55:06 +00:00
  • bd52e88ce1 Add O rescale with correction_rescale pattern + fix TMA to working diag pattern biondizzle 2026-05-22 19:51:53 +00:00
  • 93acac7e43 O normalize: exact CUTLASS correction_rescale pattern with 2D reg tensor biondizzle 2026-05-22 19:50:37 +00:00
  • 70349de6dc O normalize: TMEM round-trip with paired Ld/St atoms + standard epilogue_tma_store biondizzle 2026-05-22 19:49:51 +00:00
  • 02543b3d8d Fix: use NamedBarrier instead of mbarrier_arrive/wait biondizzle 2026-05-22 19:49:07 +00:00
  • 5f98eaec3a Fix: barrier_wait → mbarrier_wait, barrier_arrive → mbarrier_arrive biondizzle 2026-05-22 19:48:31 +00:00
  • afe5d1ae21 Fix epilogue: corr_tile_size=16, proper epi_subtile tuple, match CUTLASS reference biondizzle 2026-05-22 19:47:57 +00:00
  • 8b93774d70 Fix example7: K slice (None,None,0,0) and softmax scale_log2 double-bug biondizzle 2026-05-22 19:45:59 +00:00
  • 6827862b3d Fix: fence_async_shared -> fence_view_async_shared biondizzle 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. biondizzle 2026-05-22 19:24:25 +00:00
  • 1176d75236 O normalize using tmem_ptr base (same as epilogue) + CUTLASS sub-tile pattern biondizzle 2026-05-22 19:18:16 +00:00
  • 86a4719afe O normalize with full layout (no sub-tiling), Repetition(64) biondizzle 2026-05-22 19:17:03 +00:00
  • 1d2111f772 Disable ALL O copies to verify baseline biondizzle 2026-05-22 19:14:33 +00:00
  • d9f7443b77 Disable O rescale too for NO-OP test biondizzle 2026-05-22 19:13:43 +00:00
  • 62529df638 DEBUG: O load+store NO-OP to verify TMEM copy correctness biondizzle 2026-05-22 19:12:43 +00:00
  • 59575202a4 Re-enable O rescale + normalize with corr_tile_size=32 biondizzle 2026-05-22 19:11:17 +00:00
  • 0589a14790 Disable O rescale + normalize, verify softmax P only biondizzle 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 biondizzle 2026-05-22 19:06:52 +00:00
  • b0498dd7ac Full pipeline: O rescale + final normalize with CUTLASS sub-tile approach biondizzle 2026-05-22 19:05:59 +00:00
  • 668c7e1a4e Test softmax P vs unnormalized reference (no O normalize) biondizzle 2026-05-22 19:05:27 +00:00
  • 46020baff1 O normalize using CUTLASS reference sub-tile approach biondizzle 2026-05-22 19:04:37 +00:00
  • e56ffbbe32 Fix O normalize: use 2D register tensor indexing biondizzle 2026-05-22 18:59:15 +00:00
  • 43e9eb2f30 Add O normalization with sub-tile TMEM read-modify-write biondizzle 2026-05-22 18:58:35 +00:00
  • 603385a88c Simplify: softmax P only, no O rescale/normalize yet biondizzle 2026-05-22 18:57:00 +00:00
  • 46462b7c8b Real softmax test built on working identity diag biondizzle 2026-05-22 18:55:31 +00:00
  • 4fef047f5c CRITICAL FIX: remove extra scale_log2 in softmax (minus_row_max and acc_scale) biondizzle 2026-05-22 18:52:58 +00:00
  • fe58d07b8c FIX: K slice (None,None,0,0) like working diag biondizzle 2026-05-22 18:51:22 +00:00
  • 12defe923d Diag: identity softmax on example6 pipeline to isolate softmax bug biondizzle 2026-05-22 18:49:33 +00:00
  • d7e0aa6168 Quick test: working v3 with n=256 multi-tile biondizzle 2026-05-22 18:47:54 +00:00
  • 6a7b955c19 DEBUG: add version marker to confirm code changes are running biondizzle 2026-05-22 18:47:05 +00:00
  • fb1057eb5f CRITICAL FIX: TMA pre-slice (None,0,None,0) → (None,None,0,0) to keep GMEM tile dim free biondizzle 2026-05-22 18:45:55 +00:00
  • 2deb28827a Diag: TMA shapes with hardcoded major modes biondizzle 2026-05-22 18:43:55 +00:00
  • d2a3b83aa2 Diag: simplified TMA shape analysis biondizzle 2026-05-22 18:43:30 +00:00
  • 73109dde85 Diag: print TMA partition shapes for multi-tile debugging biondizzle 2026-05-22 18:43:05 +00:00
  • 375b24304b FIX: Use Python range() in TMA warp for concrete per-iteration GMEM coords biondizzle 2026-05-22 18:26:31 +00:00
  • 0dd6fefd66 FIX: Force SSA GMEM coord via n_kv_tiles - n_kv_tiles instead of cutlass.range kt biondizzle 2026-05-22 18:25:13 +00:00
  • 2a4731426e FIX: TMEM offset bug in O rescale/normalize — use tOtO0.iterator not tOtO.iterator biondizzle 2026-05-22 18:23:46 +00:00
  • 1e0805ad60 Diag: test n=384 (3 tiles) to find crash boundary biondizzle 2026-05-22 18:07:07 +00:00
  • 1aa4a91d01 Diag: test all sizes 128-1024 biondizzle 2026-05-22 18:06:28 +00:00
  • 8586603280 DEBUG: disable O rescale to isolate NaN cause biondizzle 2026-05-22 18:05:46 +00:00
  • 45a0fb9971 Add NaN/inf checking to stage C test biondizzle 2026-05-22 18:01:11 +00:00
  • 147d85a617 CRITICAL FIX: K GMEM slice (None,None,0,0) not (None,0,None,0) biondizzle 2026-05-22 17:59:57 +00:00
  • bc3e94ff45 Diag: try K slice (None,None,0,0) keeping mode 1 (CUTLASS ref style) biondizzle 2026-05-22 17:59:01 +00:00
  • 601e662dd4 Diag: try runtime Int32(0+0) for kv_coord with cutlass.range biondizzle 2026-05-22 17:57:58 +00:00
  • e5030cbea5 Diag: use Python range() unrolling like stage C test biondizzle 2026-05-22 17:56:59 +00:00
  • 5ee34d925b Fix diagnostic test: same Int32(kt) + n_kv_tiles fixes biondizzle 2026-05-22 17:56:15 +00:00
  • d2bbdd59f6 Try cutlass.range with Int32(kt) — now n_kv_tiles is Python int biondizzle 2026-05-22 17:51:25 +00:00
  • bf80fbee99 FIX: n_kv_tiles as Python int (s_k//128) for range() unrolling biondizzle 2026-05-22 17:50:07 +00:00
  • 0b3bc3a16d Option 2: Python range() with Int32(kt) for TMA GMEM coord biondizzle 2026-05-22 17:47:43 +00:00
  • 083e205aae Add example5: use cutlass.range induction variable as TMA GMEM coord biondizzle 2026-05-22 17:47:10 +00:00
  • 56769cdbf5 README: add fire_b200_test docs, update multi-tile blocker with real findings biondizzle 2026-05-22 17:41:23 +00:00
  • f80f8eb38f Clean up debug prints, set kv_coord as Int32(0) biondizzle 2026-05-22 17:39:27 +00:00
  • 36cf1a363b DEBUG: try plain Python int kv_coord (like CUTLASS ref) biondizzle 2026-05-22 17:34:30 +00:00
  • d95e2221c2 DEBUG: hardcode kv_coord=1 to test if TMA uses it biondizzle 2026-05-22 17:32:53 +00:00
  • 59746b46fc DEBUG: try K slice (None,0,None,0) keeping mode 2 free biondizzle 2026-05-22 17:30:06 +00:00
  • b6cefba31c DEBUG: print tBgK/tVgV shapes before/after slice biondizzle 2026-05-22 17:28:45 +00:00
  • ba2cefb668 Stage C: manual kv_coord + correct K GMEM slice + O rescale fence biondizzle 2026-05-22 17:26:56 +00:00
  • 5a6c4d2cd2 Add example4: manual kv_coord Int32 for GMEM tile indexing biondizzle 2026-05-22 17:24:38 +00:00
  • 81eee05018 README: add test harness instructions biondizzle 2026-05-22 17:09:53 +00:00
  • a950f978d3 run_test.sh: SIGKILL all children of screen session on cleanup biondizzle 2026-05-22 17:08:12 +00:00