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

Commit Graph

  • d36b727898 D1.1: Add SMEM-P path behind use_smem_p flag (stub: zero sP) biondizzle 2026-05-23 06:01:02 +00:00
  • bd0b56dddd D1.0: Replace HEAD_DIM=64 with self.head_dim constructor parameter biondizzle 2026-05-23 05:55:03 +00:00
  • 249a581d8a Add STAGE_D.md: step-by-step runbook and todo list for D1-D5 biondizzle 2026-05-23 05:52:03 +00:00
  • bfacfeca7b Rename FmhaV3StageC → FmhaKernel — no dev stage artifacts in production API biondizzle 2026-05-23 05:45:58 +00:00
  • 787a25516d Update README: reflect Stage C migration, built indexer/router/compressor, SMEM-P path, CuTeDSL scoping lesson biondizzle 2026-05-23 05:42:44 +00:00
  • b39301ebc6 Migrate Stage C kernel (proven cos 0.97) into module - exact copy, no modifications biondizzle 2026-05-23 05:36:22 +00:00
  • 6c9a9d72f1 Fix TMEM-P offset calc: match Stage C with p_cols_fp32 from pv_mma_tiler[2] biondizzle 2026-05-23 05:18:37 +00:00
  • 03748c4215 Add missing TMEM fence after P store in TMEM-P path biondizzle 2026-05-23 05:17:45 +00:00
  • 9c93d655de Fix p_cols_fp32: use pv_mma_tiler[2] (K-dim) not [1] (N-dim) biondizzle 2026-05-23 05:16:19 +00:00
  • 32ae44d97d Fix PV A-operand major mode: K for TMEM-P, a_major for SMEM-P biondizzle 2026-05-23 05:14:08 +00:00
  • 7df67d5237 Fix CuTeDSL scoping: hoist P store vars out of if block biondizzle 2026-05-23 05:12:30 +00:00
  • 2addbeed7d Fix O rescale: use Stage C proven correction_rescale pattern biondizzle 2026-05-23 05:10:46 +00:00
  • 86f3e9cf32 Fix tOrP0 indexing: 3-dim slice (None,None,kb) not 4-dim biondizzle 2026-05-23 05:09:19 +00:00
  • 75fec90eef Fix CuTeDSL scoping: unconditionally define tOrP0 and tCrP biondizzle 2026-05-23 05:08:10 +00:00
  • fd157cf38e Fix CuTeDSL variable scoping: define tOrP0 and tCrP in both branches biondizzle 2026-05-23 05:07:30 +00:00
  • 2cc9786491 Fix p_tmem_s: use ComposedLayout from make_smem_layout_a, pass as kernel arg biondizzle 2026-05-23 05:06:45 +00:00
  • d0869054d5 Consolidate FMHA stages A/B/C into unified kernel module with SMEM-P stub biondizzle 2026-05-23 05:04:43 +00:00
  • 5af491cd73 WIP: make_tiled_copy_C for P→SMEM biondizzle 2026-05-23 03:56:56 +00:00
  • cf080ccf00 fix: cpasync.CopyOp for reg→SMEM biondizzle 2026-05-23 03:54:49 +00:00
  • a18c639021 fix: CopyAtomUniversalOp biondizzle 2026-05-23 03:52:47 +00:00
  • f869f31625 WIP: tiled copy for P→SMEM (zero fill) biondizzle 2026-05-23 03:51:58 +00:00
  • 6ac61b69c8 fix: cute.copy(dst, src) order biondizzle 2026-05-23 03:51:00 +00:00
  • 99f13cf52e fix: BFloat16 not Float32 for bf16 reg biondizzle 2026-05-23 03:50:09 +00:00
  • 748873a58c WIP: P→SMEM write stub (zero fill, proper mapping TODO) biondizzle 2026-05-23 03:49:05 +00:00
  • 469665f69a fix: partition_A not partition_S biondizzle 2026-05-23 03:47:53 +00:00
  • 3ee330a84c fix: make_smem_layout_epi not make_epilogue_smem_layout biondizzle 2026-05-23 03:47:09 +00:00
  • 1d8b4aaab3 WIP: SMEM P path for PV (compiles but P write not implemented) biondizzle 2026-05-23 03:46:01 +00:00
  • 4c3239b9c0 debug: hd=64 with CUDA_LAUNCH_BLOCKING biondizzle 2026-05-23 03:42:53 +00:00
  • 2e2ac4eb1f D1: P store as BF16 using PV A-fragment layout (tOrP0) biondizzle 2026-05-23 03:42:07 +00:00
  • 902bda5c31 D1: align P store and PV A-fragment layouts via tP biondizzle 2026-05-23 03:40:10 +00:00
  • 95cf4159f2 D1: P store as BF16 using PV A-fragment layout biondizzle 2026-05-23 03:38:24 +00:00
  • e33f582436 D1: P store uses tOrP0.layout (PV A-fragment TMEM layout) biondizzle 2026-05-23 03:36:40 +00:00
  • 063783dd6b D1 fix: P store uses PV A-fragment layout (p_tmem_s.outer) biondizzle 2026-05-23 03:35:49 +00:00
  • 4df5dafcc9 D1: test raw unnormalized output via epilogue_tma_store biondizzle 2026-05-23 03:33:59 +00:00
  • 6dfc06cdef test: paired atoms epilog from old commit 6ee28d8 biondizzle 2026-05-23 03:32:53 +00:00
  • e1fc4cee60 D1: paired atoms epilogue (no TMEM round-trip) biondizzle 2026-05-23 03:29:51 +00:00
  • 32481f8a2b d1: sweep hd=64,128,256 biondizzle 2026-05-23 03:26:10 +00:00
  • 3ec7f36e62 fix: use mV.iterator biondizzle 2026-05-23 03:25:29 +00:00
  • 73f359b565 fix: use mQ not q for LayoutEnum biondizzle 2026-05-23 03:24:58 +00:00
  • 4203d6ad20 d1: add diagnostic script biondizzle 2026-05-23 03:24:16 +00:00
  • b249b8f135 D1: N-tile support for HEAD_DIM>256 biondizzle 2026-05-23 03:22:23 +00:00
  • f2dced88a3 d1: add hd=512 test biondizzle 2026-05-23 03:20:46 +00:00
  • 06c3df8aad d1: add quick regression test (hd=64 only) biondizzle 2026-05-23 03:20:12 +00:00
  • 0102294d29 D1: Parameterize HEAD_DIM in FmhaKernel (64→512) biondizzle 2026-05-23 03:19:52 +00:00
  • 6dd71aaf56 docs: revised Stage D/E plan — indexer removes paged TMA, one kernel for CSA/HCA/SWA, sink merge biondizzle 2026-05-23 03:10:41 +00:00
  • b1fe18acbb cleanup: remove archive/ (240 stale files), stale example9/10, fix test table, add Stage D plan biondizzle 2026-05-23 03:05:08 +00:00
  • 6eb9729b06 docs: update README with Stage C TMEM layout mismatch findings and status biondizzle 2026-05-23 03:01:04 +00:00
  • a7038de611 fix: revert to composition layout for hand-constructed atoms (matching CUTLASS) biondizzle 2026-05-23 02:54:54 +00:00
  • de7cc61480 fix: use logical_divide (not composition) for O rescale/normalize atoms to match get_tmem_load_op layout biondizzle 2026-05-23 02:53:59 +00:00
  • 71934f78b2 fix: add NO-OP TMEM round-trip to re-map O from MMA to epilog layout biondizzle 2026-05-23 02:50:53 +00:00
  • ac1feeb895 fix: use TMEM round-trip normalize + epilogue_tma_store (known ~3% error) biondizzle 2026-05-23 02:49:46 +00:00
  • 8ac05f2b26 fix: correct bSG_gC indexing (6 modes) biondizzle 2026-05-23 02:45:30 +00:00
  • c403b57617 diag: print bSG shapes for TMA store indexing biondizzle 2026-05-23 02:44:47 +00:00
  • 1d3c169805 fix: typo from_dlcap -> from_dlpack biondizzle 2026-05-23 02:44:00 +00:00
  • 4b9949092e fix: correction_epilog with paired atoms + pre-partitioned TMA store outside if block biondizzle 2026-05-23 02:41:07 +00:00
  • 24412df914 test: NO-OP round-trip + normalize at n=128 and n=256 biondizzle 2026-05-23 02:37:50 +00:00
  • 22f8870fbd fix: correction_epilog with paired atoms + pre-partitioned TMA store biondizzle 2026-05-23 02:34:33 +00:00
  • fb4888079a diag: NO-OP round-trip before normalize on 2D pattern biondizzle 2026-05-23 02:32:40 +00:00
  • 3416c46a05 fix: O rescale uses 2D register tensor pattern, remove fence_view_async_tmem_load biondizzle 2026-05-23 02:31:28 +00:00
  • b457d196af fix: use paired atoms for correction_epilog + cute.copy TMA store biondizzle 2026-05-23 02:26:57 +00:00
  • e3d6d6eebf diag: add CUDA_LAUNCH_BLOCKING for crash debug biondizzle 2026-05-23 02:25:46 +00:00
  • 1ee99cbef6 fix: inline epilogue_tma_store with inv_row_sum multiply using paired atoms biondizzle 2026-05-23 02:24:36 +00:00
  • 09568b46ba fix: use cute.copy instead of cpasync.copy for TMA store biondizzle 2026-05-23 02:23:16 +00:00
  • 23d77286ff fix: correction_epilog with get_tmem_load_op paired atoms + direct TMA store biondizzle 2026-05-23 02:19:41 +00:00
  • b10a8af198 diag: NO-OP TMEM round-trip test — load+store back unchanged biondizzle 2026-05-23 02:15:28 +00:00
  • 641034217d fix: inline epilogue with paired atoms + inv_row_sum normalize, no TMEM round-trip biondizzle 2026-05-23 02:13:52 +00:00
  • f9bb2c18e3 fix: all epilogue warps do TMA store, no dynamic if inside method biondizzle 2026-05-23 01:41:36 +00:00
  • 7123c214bf fix: correction_epilog with get_tmem_load_op paired atoms, no TMEM round-trip biondizzle 2026-05-23 01:40:13 +00:00
  • 9d29441690 fix: use attn_raw (not softmax'd) for unnorm computation biondizzle 2026-05-23 01:36:27 +00:00
  • 8f65bcd02f diag: skip kernel normalize, do Python-side normalize to isolate TMEM round-trip issue biondizzle 2026-05-23 01:35:18 +00:00
  • 9347119ade diag: print expected unnorm P@V for comparison with raw kernel output biondizzle 2026-05-23 01:28:32 +00:00
  • e5b4da7274 diag: skip final normalize, test raw PV output via epilogue_tma_store biondizzle 2026-05-23 01:27:03 +00:00
  • e7b2cb5a32 fix: O rescale uses 2D register tensor pattern (matching CUTLASS correction_rescale) biondizzle 2026-05-23 01:25:53 +00:00
  • c1f531c0e6 fix: pre-compute tmem_load_epi_atom in __call__, pass to kernel biondizzle 2026-05-23 01:24:33 +00:00
  • 9c9eb18c73 fix: index into TMA partitioned tensors for copy biondizzle 2026-05-23 01:23:04 +00:00
  • 6503dc8659 fix: use flat_divide+group_modes(0,2) for TMA store, matching CUTLASS biondizzle 2026-05-23 01:22:22 +00:00
  • 5630398498 fix: use gC not tCgC for TMA partition, group modes 0-3 biondizzle 2026-05-23 01:20:52 +00:00
  • 6b2dda02ec fix: use tma_partition for TMA store in correction_epilog biondizzle 2026-05-23 01:20:09 +00:00
  • e70230825a fix: replace TMEM round-trip normalize with CUTLASS correction_epilog pattern biondizzle 2026-05-23 01:18:56 +00:00
  • ed8855285c diag: inv_row_sum=1.0 to test raw PV, n=128 only biondizzle 2026-05-23 01:17:14 +00:00
  • dc8d95a5ce diag: test original code n=128+256 to confirm baseline biondizzle 2026-05-23 01:13:29 +00:00
  • 58c472ef73 diag: disable O rescale properly, test n=128+256 baseline biondizzle 2026-05-23 01:12:50 +00:00
  • 88b0b2dada diag: test n=128 and n=256 both with rescale disabled biondizzle 2026-05-23 01:12:00 +00:00
  • 4712a0f90d fix: indentation error in diag disable biondizzle 2026-05-23 01:11:25 +00:00
  • 26fb90bbba diag: disable O rescale to isolate the issue (n=256 only) biondizzle 2026-05-23 01:11:00 +00:00
  • d1514a4167 debug: add wide-search diagnostics for n=256 O rescale biondizzle 2026-05-23 01:02:33 +00:00
  • a1f08f9488 🚀 MULTI-TILE SOFTMAX + O RESCALE WORKING: n=128 cos 0.999998, n=256 cos 0.80 biondizzle 2026-05-23 00:35:42 +00:00
  • 62ef6d6ae9 Debug: add row_sum/inv_row_sum printf at final normalize biondizzle 2026-05-23 00:34:38 +00:00
  • 910183e7aa Fix ALL loops: use self.n_kv_tiles everywhere biondizzle 2026-05-23 00:33:38 +00:00
  • 5c12c72c05 Fix softmax loop: use self.n_kv_tiles not cute.size(gK, mode=[3]) biondizzle 2026-05-23 00:30:49 +00:00
  • 355d02f56c auto: pre-test commit biondizzle 2026-05-23 00:29:49 +00:00
  • aed0e50946 Add O rescale with pre-built paired atoms (corr_tile_size=16) biondizzle 2026-05-23 00:28:44 +00:00
  • e236e91def Reset to working_softmax_maybe.py + TMA fix only biondizzle 2026-05-23 00:27:41 +00:00
  • f4b7bed463 Fix final normalize: use working 2D register tensor pattern from working_softmax_maybe.py biondizzle 2026-05-23 00:25:16 +00:00
  • 764d3aaf2f Fix: add self.n_kv_tiles to __init__ biondizzle 2026-05-23 00:23:38 +00:00
  • ba31a1cedf Add per-tile O rescale (O *= acc_scale) to softmax loop biondizzle 2026-05-23 00:22:12 +00:00
  • 0ced79ab37 Clean up: archive diagnostics and superseded tests biondizzle 2026-05-23 00:17:07 +00:00
  • b61df94706 auto: pre-test commit biondizzle 2026-05-23 00:05:07 +00:00
  • eadd870d80 auto: pre-test commit biondizzle 2026-05-23 00:02:33 +00:00
  • 79faa14cef auto: pre-test commit biondizzle 2026-05-23 00:00:47 +00:00