Skip to content

MoE: add a16w mix fused 2-stage kernels (bf16 A × mxfp4/int4/bf16 W) #948

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
coderfeli wants to merge 63 commits into
base: main
Choose a base branch
from
Open

MoE: add a16w mix fused 2-stage kernels (bf16 A × mxfp4/int4/bf16 W) #948

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
63 commits
Select commit Hold shift + click to select a range
1fe0ce0
[a16w4] Add cvt_scalef32_pk_bf16_fp4 Python wrapper
Jul 30, 2026
ba88a95
[a16w4] M1: add a16w4 (bf16 A x mxfp4 W) gemm1 body + builder
Jul 30, 2026
7e8b59e
[a16w4] M2: add a16w4 (bf16 intermediate x mxfp4 W2) gemm2 body + bui…
Jul 30, 2026
b44abb2
[a16w4] M3: sorted-position gemm1 intermediate + host a16w4 arm + e2e
Jul 30, 2026
963ddfd
[a16w4][perf] gemm1: default TILE_N 256->64 (kill VGPR spill, 5.7x s1)
Jul 30, 2026
759fb2a
[a16w4][perf][g2] bijective XCD round-robin grid for gemm2 (HBM chann…
Jul 30, 2026
b2abd1c
[a16w4][perf][g2] enable A-LDS XOR bank-conflict swizzle (gmem-source…
Jul 30, 2026
323fc5b
[a16w4][perf] gemm1: add SiTUv2 activation behind act= (SiLU default)
Jul 30, 2026
c9b4857
[a16w4][perf][g2] token-conditional B cache modifier (cached weight a…
Jul 30, 2026
d3023d2
[a16w4][perf][g2] extend cached-B to large M (U-shape cache threshold)
Jul 30, 2026
3afb739
[a16w4][perf] gemm1: cached B weight load (cache_modifier 2->0)
Jul 30, 2026
9eb81ed
[a16w4][perf] gemm1 pipeline (drop full drain)
Jul 30, 2026
de7207b
[a16w4] Merge gemm2 perf track (XCD grid, LDS swizzle, cached-B) into…
Jul 30, 2026
f019ad1
[a16w4][perf] gemm1 wave-split N (kill all-wave redundancy)
Jul 31, 2026
03833cb
[a16w4] M4: wire a16w4 into test matrices + regression
Jul 31, 2026
9e9fb9c
[a16w4][perf] gemm1 default TILE_N 64->256 (fatter N tile)
Jul 31, 2026
ec5ce0e
[a16w4][perf] gemm1 phase-separate A-LDS read from next-tile A-DMA
Jul 31, 2026
ab03f95
[a16w4][doc] gemm1 re-diff at f019 + landed levers verdict
Jul 31, 2026
ebc63bf
[a16w4] Merge gemm1 config-match (TILE_N=256 + phase-sep A-LDS) into …
Jul 31, 2026
d941462
[a16w4] gemm1 TILE_N adaptive (256 if inter%256==0 else 128); drop sc…
Jul 31, 2026
a6c3970
[a16w4] host: aiter tile-config interface + CSV loader; gemm1 b_nt U-…
Jul 31, 2026
b2241c0
[a16w4] Merge host aiter tile-config interface + CSV loader + gemm1 b…
Jul 31, 2026
678ad90
[a16w4] Extract a16w4 into standalone kernels/moe/moe_2stage_a16wmix …
Jul 31, 2026
7cc8c89
[a16w4] a8w4-interleave combined dispatch + comparison harness
Jul 31, 2026
cb082cd
[a16wi4] Add int4-weight path to moe_2stage_a16wmix + route int4_bf16…
Jul 31, 2026
802e210
[a16w4] Merge a8w4-interleave combined dispatch + comparison harness
Jul 31, 2026
82a61ab
[test] Fix masked mxfp_moe/a16w4 correctness gates (assert + strict c…
Jul 31, 2026
6579eb4
[repro] Add reproducer + issue draft for pre-existing mxfp_moe a4w4/a…
Jul 31, 2026
9f05b5c
[a16w4][perf] k_wave (kw2/kw4) intra-block slice-K + LDS reduce
Jul 31, 2026
4e88a69
[a16wi4][perf] tile-config align to legacy int4 (tile_n=128)
Jul 31, 2026
f3c9a5a
[a16w4] Merge k_wave (kw2/kw4) param (interface parity; kw1 byte-iden…
Jul 31, 2026
c367701
[a16w4] Merge strict correctness gates + a4w4/a8w4 reproducer + issue…
Jul 31, 2026
4acba3b
[a16w16] Add w_dtype=bf16 (raw bf16 W) path to moe_2stage_a16wmix
Jul 31, 2026
f9530d1
[a16w4][perf] persistent CU-limited gemm2 grid (opt-in; fix E896 over…
Jul 31, 2026
c327d0a
[a16w4][perf] CSV-driven per-token config (tile_n=64+kw small-M)
Jul 31, 2026
8680655
[a16w4] Make CSV per-token config opt-in (default uses fast tile_n; n…
Jul 31, 2026
d111db6
[a16w16] Fix bf16-W path correctness at large E (expert stride/layout)
Jul 31, 2026
0b739aa
[a16w4] Merge persistent CU-limited gemm2 grid (opt-in, default-off b…
Jul 31, 2026
8d29248
[a16w4] Tighten moe_2stage_a16wmix: cut comments/dead logic (behavior…
Jul 31, 2026
c2b12c1
[a16w4] Modernize moe_2stage_a16wmix to fx.* surface (dialect)
Jul 31, 2026
c59ea6e
[a16w4] Modernize moe_2stage_a16wmix to fx.* surface (buffer_ops)
Jul 31, 2026
1a64a52
[a16w4] Fix gemm2 s2 regression (restore @flyc.jit on _atomic_bf16_ep…
Jul 31, 2026
72dfd46
[a16w4] CSV loader: accept int4/bf16 rows + map aiter split-K (_kb) t…
Jul 31, 2026
cfa5cfc
[a16w4] gemm1 mxfp4: default tile_n=128 to fix occupancy cliff at hig…
Jul 31, 2026
0450f82
[a16w4] gemm1: clamp A-load buffer resource to real activation size (…
Jul 31, 2026
451dc23
[a16w4] gemm1 mxfp4: high-token default TILE_K=128 + xcd_swizzle=1
Jul 31, 2026
8d9cffe
[a16w4] gemm1 mxfp4: very-low-M (tok<=2) default TILE_N=64 (under-fil…
Jul 31, 2026
e75f0f9
[MoE a16wi4] gemm1/gemm2: VALU-lean gfx950 int4->bf16 dequant
Jul 31, 2026
47a20f5
[MoE a16wi4] gemm1: block_m=64 at the tok2048-class L2-fill point
Jul 31, 2026
49e7666
[a16w4] Remove stray scratch bench committed with block_m helper
Jul 31, 2026
712c28f
Revert "[MoE a16wi4] gemm1/gemm2: VALU-lean gfx950 int4->bf16 dequant"
Jul 31, 2026
6603353
[MoE a16wi4] gemm1/gemm2: fix + re-enable VALU-lean int4->bf16 dequant
Jul 31, 2026
08825ab
[MoE a16wi4] gemm1: TILE_N=64 at BM=64 fill point (1->2 wave occupancy)
Jul 31, 2026
cbfe7c1
[a16w4] Remove scratch interleave-sweep driver from the kernel package
Aug 1, 2026
411f761
[a16w4] gemm1 tok<=2: 4-way slice-K (TILE_K=128 + k_wave=4)
Aug 1, 2026
9e19ebf
[a16w4] CI fix: gate int4_bf16 new-kernel routing to gfx950 (legacy o…
Aug 1, 2026
54368cb
[a16wmix] gfx942 (CDNA3) support: arch-gate K=16 MFMA + scalar int4 d…
Aug 1, 2026
68621dc
[a16wmix] Condense verbose comments in moe_2stage_a16wmix (comment-on…
Aug 1, 2026
51c143a
[a16w4] CI: fall back int4_bf16 to legacy on gfx942 (K=16 path aborts…
Aug 1, 2026
ff61d32
[a16wmix] gfx942: fix a16wi4 A-tile LDS staging (direct-to-LDS 128b u…
Aug 1, 2026
c762d5b
[a16w4] CI: xfail fp4/a8w4 in test_moe_stage2_standalone (pre-existin…
Aug 1, 2026
cd3f125
[a16wmix cleanup] Remove a4w4/a8w4 artifacts; consolidate common+host…
Aug 1, 2026
3d1e03d
[a16wmix cleanup] Modernize scale-load buffer_ops to the fx.* layout API
Aug 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Loading
Loading