MoE: add a16w mix fused 2-stage kernels (bf16 A × mxfp4/int4/bf16 W) - #948
Open
coderfeli wants to merge 59 commits into
Open
MoE: add a16w mix fused 2-stage kernels (bf16 A × mxfp4/int4/bf16 W)#948coderfeli wants to merge 59 commits into
coderfeli wants to merge 59 commits into
Conversation
Mirrors cvt_scalef32_pk_f32_fp4; the mxfp4->bf16 upconvert primitive for the a16w4 (bf16 x mxfp4) MoE path (e8m0 block scale folded into the cvt). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
New dedicated a16w4 stage1 path in mxfp_moe/gemm1.py (separate from the a4w4/a8w4 scaled-f8f6f4 body; that code is byte-identical, diff is purely additive). bf16 activations (no A-scale), in-kernel mxfp4->bf16 W upconvert via cvt_scalef32_pk_bf16_fp4 (e8m0 scale folded in), non-scaled MFMA(16,16,32,BFloat16) K=32 via fx.gemm, SiLU epilogue -> bf16 intermediate [tokens*topk, inter_dim]. A/B/scale addressing ported from aiter compile_mixed_moe_gemm1_a16w4 (load_b_raw_mxfp4_dwordx4 layout_b, make_preshuffle_scale_layout, CK A sub-lane read). Un-pipelined (perf is M5). A-LDS bank-conflict XOR swizzle disabled for M1 (linear layout) pending an M5 re-derivation; both the DMA write and LDS read go through one helper so they stay consistent. Isolated stage1 vs torch_moe_gemm1 (bf16 A, scale=None; mxfp4 W): cos=0.999999 on tokens=128/K=2048/INTER=512/NE=8/topk=2/BM=32 and tokens=64/INTER=256/BM=16 (cold). a4w4/a8w4 variants still green (AC-2). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…lder New dedicated a16w4 stage2 path in mxfp_moe/gemm2.py (separate from the a4w4/a8w4 scaled-f8f6f4 body; that code is byte-identical, diff purely additive). Reuses the M1-proven a16w4 primitives (shared helpers imported from gemm1: _a16w4_swizzle_xor16, _e8m0_byte_to_f32): bf16 A intermediate -> LDS (indexed by sorted position, like the a4w4 path), mxfp4 W2 dwordx4 load + cvt_scalef32_pk_bf16_fp4 upconvert (N=model_dim, K=inter_dim), non-scaled MFMA(16,16,32,BFloat16) K=32 via fx.gemm. Output is bf16 [tokens, model_dim] via the existing atomic (routing-weighted) scatter epilogue -- reused unchanged; accm[mi][ni] laid out to match its (mchunk, wave-J) indexing (wave owns TILE_N/4 = 64 cols, num_acc_n=4). Linear A-LDS (XOR swizzle deferred to M5, as in M1). Isolated stage2 vs torch_moe_gemm2 (bf16 A intermediate, scale=None; mxfp4 W2, doweight_stage2=True): cos=0.999997 on tokens=128/model_dim=2048/ INTER=512/NE=8/topk=2/BM=32 and tokens=64/model_dim=1024/INTER=256/BM=16 (cold + warm). Mini e2e chain (kernel gemm1 -> host re-sort -> kernel gemm2) vs full torch e2e: cos=0.999997. a4w4/a8w4 variants still green (AC-2). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1. gemm1 a16w4 epilogue now stores the bf16 intermediate by SORTED POSITION (row = bx_m + row_in_tile), shape [sorted_size, inter_dim], matching the a4w4 fused inter_sorted_* contract so gemm2_a16w4 consumes it drop-in (no host gather). Padding rows masked (token >= tokens). out_rsrc sized from cumsum0 (padded sorted rows). Localized to _gemm1_body_a16w4 + its call site; a4w4/a8w4 gemm1 byte-identical (diff touches only a16w4 lines). 2. host.py: new flydsl_a16w4_gemm1 / flydsl_a16w4_gemm2 arms (+ cached compile helpers) mirroring the a4w4 arms: bf16 A (no A-scale), bf16 [sorted_size, inter] intermediate (no intermediate scale), reuse sorting/cumsum/m_indices, grid, and standard shuffle_weight + e8m0_shuffle W layout. a4w4/a8w4 host arms unchanged (only the import line expanded). Full e2e cold + warm vs torch_moe_gemm1+torch_moe_gemm2 (bf16 A scale=None, mxfp4 W1/W2, doweight_stage2): small (128, md2048, i512, E8, k2, BM32): cos 0.999997 kimi512(128, md3584, i512, E896,k16, BM32): cos 0.999986 gemm1-alone (sorted-position) still cos 0.999999. test_mxfp_moe_variants (fp4+fp8) green (AC-2). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The a16w4 gemm1 default TILE_N=256 gave num_acc_n=16 -> 2*(2*16) f32[4] accumulators = 256 VGPR for accs alone, hitting the 256-VGPR cap with a 2764B scratch spill and 1 wave/SIMD occupancy. rocprofv3 kernel-trace: gemm1_a16w4 tok128 3922us, VGPR=256, Scratch=2764. TILE_N=64 (matching aiter's tuned t32x64x256 gemm1 tile for Kimi-K3 3584x512) drops accumulators 4x: VGPR 256->128, Scratch 2764->0 (spill gone), occupancy doubles. Measured s1 (median-of-3, GPU5 gfx950, warm, Kimi-K3 3584x512 E896 k16): tok1 408->47, tok16 1106->210, tok128 3955->680, tok1024 4449->836, tok8192 20512->4051 us. Wins every token. a16w4 e2e cos unchanged (kimi512 0.999986 cold); a4w4/a8w4 byte-identical off-path (test_mxfp_moe_variants 6/6 green); no spill (Scratch=0). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…el balance) a16w4 gemm2 grid was plain m-major linear (no XCD remap), unlike the a4w4 gemm2 path which already round-robins the launch index across the 8 XCDs. flyprof on tok16 (bundle runs/g2_tok16, mapped 99.9%) shows the kernel is HBM-bandwidth-bound (L2 hit 14.9%, achieved 2705/8000 GB/s, MFMA only 4.6% of instrs, total_stall 82.5%) with the m-major grid clustering consecutive tiles onto the same XCD/HBM channels. Port the same bijective 8-XCD round-robin (+ optional M-group swizzle) into compile_gemm2_a16w4_port, gated by xcd_swizzle (default 1). a16w4-only; the a4w4/a8w4 gemm2 code is byte-identical (36 ins, 1 del, all inside the a16w4 builder). host.py needs no change (the default flows through). Perf (s2 median-of-5, iters=40, GPU6, Kimi-K3 3584x512 E896 k16): tok1 7.59->7.29 (-4%), tok16 49.70->44.78 (-10%), tok128 160.61->140.80 (-12%), tok1024 200.92->182.83 (-9%), tok4096 555.44->503.26 (-9%). Now WIN tok1/1024/ 4096 vs aiter, tok128 1.07x. AC-P1 cos 0.99921 held (== base). AC-P2 test_mxfp_moe_variants 6/6 green. AC-P3 VGPR 141->144, spill=0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… variant) The a16w4 gemm2 A-LDS was LINEAR: M1 deferred the XOR swizzle because flipping enable=True on the old scheme (swizzled LDS *dest* of the buffer_load_lds DMA) NaN'd -- the DMA hardware does not honor an arbitrary swizzled per-lane LDS dest. Root-caused and fixed by matching kernels/gemm/mxfp4_preshuffle.py: keep the LDS dest LINEAR and swizzle the GMEM *source* column instead, so linear LDS slot [row][col] holds A[row][swz(row,col)]; the LDS read applies the SAME swizzle to its offset to recover the right logical K. Both sides go through _a16w4_swizzle_xor16(..., enable=True). a16w4-gemm2-only (gemm1 and the shared helper default stay enable=False -> byte-identical off-path). flyprof tok16 (bundle runs/g2_tok16_swz vs runs/g2_tok16): SQ_LDS_BANK_CONFLICT 36.77M -> 7.35M (-80%); L2 hit 14.9% -> 32.8%; lds stall 19.45% -> 3.6%; lgkmcnt 12.1% -> 6.3%; total_stall 82.5% -> 76.2%. Perf (s2 median, iters=40, GPU6, Kimi-K3 3584x512 E896 k16), R1(XCD) -> R2(+swz): tok16 44.78->38.66 (-14%, gap 1.31x->1.13x), tok128 140.80->125.59 (-11%, now 0.95x WIN), tok1024 182.83->168.47 (-8% WIN), tok4096 503.26->472.97 (-6% WIN), tok1 ~7.3->7.7 (noise, WIN). Now BEAT aiter on 4 of 5 tokens. AC-P1 cold cos 0.99960/0.99959 (== base, NaN gate cleared). AC-P2 test_mxfp_moe_variants 6/6 green. AC-P3 VGPR 144->156, spill=0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds aiter's SiTUv2 activation as an opt-in epilogue for the a16w4 gemm1 so the perf-parity run against aiter (which uses SiTUv2 for Kimi-K3) is apples-to-apples. Default stays "silu" so the a16w4 correctness tests (torch ref = SiLU) and all existing callers are byte-identical. situ(g) = beta * tanh(g/beta) * sigmoid(g) (beta=1) situ_up(u) = linear_beta * tanh(u/linear_beta) (linear_beta=1) out = situ(clamp(g, max=+7)) * situ_up(clamp(u, [-7,+7])) New _situ_mul_batch / _tanh_f32 / _sigmoid_f32 in mxfp4_gemm_common.py (tanh via the exp2/rcp expansion, sign-restored, matching aiter mixed_moe tanh_elem). act= threads through compile_gemm1_a16w4_port (baked into the kernel name for cache keying) and flydsl_a16w4_gemm1. Verified: SiTUv2 stage1 vs a torch SiTUv2 reference cos=0.999999 (max_abs 0.0020, small shape, cold). SiLU default e2e cos 0.999986 (kimi512) unchanged. test_mxfp_moe_variants 6/6 green. Both act paths VGPR=128, Scratch=0 (no spill). a4w4/a8w4 byte-identical off-path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…t small M) gemm2's B (mxfp4 weight) load used cache_modifier=2 (SLC / non-temporal stream) for all tokens. At small M the whole expert's B is reused across few M-blocks, so CACHED loads (cache_modifier=0, nt off) win on L2 reuse -- matching aiter's bnt0-at-small-M vs bnt2 gemm2 selection. Thread b_cache_mod through _gemm2_body_a16w4 / compile_gemm2_a16w4_port (default 2 = prior behavior, name-tagged only when != 2 -> byte-identical off-path). Host flydsl_a16w4_gemm2 picks 0 (cached) for M_logical<=16, else 2 (nt). a16w4-only; a4w4/a8w4 gemm2 bodies + host arms untouched. Same-session median-of-5 A/B (3584x512, GPU6), nt(bcm2) vs cached(bcm0): tok8 23.5->22.0 (-6%), tok16 43.0->38.6 (-7..-10%, cached selected); crossover at tok32 (nt +8%), so threshold=16. tok128/1024 keep nt (unchanged from R2). tok1 is a wash (both ~7, already beats aiter). Timing is ground truth; aggregate flyprof L2 counters at tok16 are capture-noisy. AC-P1 cold cos 0.99960/0.99959 (tok16 cached + tok128 nt paths). AC-P2 test_mxfp_moe_variants 6/6 green. AC-P3 VGPR 156, spill=0 (aux bit only). gemm2 now beats aiter on 4/5 tokens; tok16 ~1.08x (down from 1.45x at base). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
flyprof tok16384 (bundle runs/g2_tok16384): the a16w4 gemm2 is NOT HBM-bound at large M (L2 hit 65.1%, HBM only 603/8000 GB/s) -- it is atomic-epilogue-bound (vmcnt bubble 45%, top hotspots the shared _atomic_bf16_epilog global_atomic_pk_ add_bf16 + stids/sweights loads) plus K-loop B-load vmcnt. The atomic epilog is shared with a4w4/a8w4 and already matches aiter, so it stays untouched. The one gated lever: B cache modifier is a U-shape in M. Cached loads (cache_modifier=0) win at BOTH ends -- small M (B reused across few M-blocks) and large M (>=2048; L2 hit ~65%, so nt streaming would bypass a reusable cache) -- while nt (cache_modifier=2) wins the middle band (32..1024). A/B median-of-5: tok2048 -5%, tok4096 -5.4%, tok8192 -2.7%, tok16384 -1.9%; tok1024 keeps nt. Host threshold M<=16 -> (M<=16 or M>=2048) cached, else nt. Host-only change; gemm2.py kernel unchanged (bcm0 already supported, spill=0 verified in R3). a16w4 gemm2 host arm only -- a4w4/a8w4 unaffected. Final large-M s2 (median iters=40, GPU6, 3584x512): tok8192 800->774 (0.965x WIN vs aiter 802.1), tok16384 1535->1508 (1.061x, down from 1.08x, vs aiter 1421.0). AC-P1 cold cos 0.99959 (tok4096 cached + tok128 nt). AC-P2 test_mxfp_moe_variants 6/6 green. AC-P3 no new kernel variant (bcm0 spill=0). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gemm1's B (mxfp4 W1 gate+up) load used cache_modifier=2 (nt / stream) for all tokens. A/B (median-of-3, GPU5, Kimi-K3 3584x512, TILE_N=64) shows the CACHED load (cache_modifier=0, nt off) wins across the ENTIRE token range -- unlike gemm2 (where nt won large M), gemm1's B has enough cross-M-block L2 reuse (E896, 2*inter cols) that streaming pessimizes it everywhere: s1 nt(2) -> cached(0): tok1 46.7->42.0 (-10%), tok2 48.3->44.1, tok4 57.6->54.0, tok16 219.8->214.9 (-2%), tok128 681.8->653.4 (-4%), tok1024 815.8->773.8 (-5%), tok4096 2269->2076 (-9%), tok8192 4032->3695, tok16384 7648->6939 (-9%). flyprof (bubbles, tok128) confirms the mechanism: vmem_load 23.4%->17.4% (less global-load work from L2 reuse), occupancy mean_waves 2.04->2.19. (The aggregate PMC L2-hit counter is capture-noisy per gemm2's ledger note; median-of-3 timing + the vmem_load bubble are the reliable evidence.) Thread b_cache_mod through _gemm1_body_a16w4 / compile_gemm1_a16w4_port (name-tagged only when != 2 -> off-path byte-identical) and flydsl_a16w4_gemm1 (default now 0 = cached, the measured optimum). a16w4-only; a4w4/a8w4 bodies + host arms untouched. AC-P1 cold cos 0.999986 (kimi512) / 0.999997 (small) unchanged. AC-P2 test_mxfp_moe_variants 6/6 green. AC-P3 VGPR=128, Scratch=0 (no spill). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
ISA diff (isa-diff-gemm1.md) proved our un-pipelined hot loop ran load-all -> s_waitcnt vmcnt(0) expcnt(0) lgkmcnt(0) -> s_barrier -> compute-all, 14x (full VMEM drain + 2 barriers/tile, max 3 loads in flight) vs aiter's software pipeline (partial lgkmcnt waits, 1 barrier, ~9 loads in flight). The earlier R2/R3 prefetch attempts FAILED because they kept the full vmcnt(0) drain -- THE key miss the ISA diff identified. Fix in _gemm1_body_a16w4: - A-LDS double-buffered (A_LDS_STAGES=2): tile K+1's DMA writes the pong slot while tile K reads ping. - Prologue loads tile-0 A(->LDS)+B+B-scale; each iter issues tile K+1's A-DMA + B/B-scale AFTER the wait/barrier but BEFORE tile K's MFMA, so they stay in flight across the MFMA cluster. - Replace the full s_waitcnt(0) drain with a PARTIAL s_waitcnt(lgkmcnt=0) (the buffer_load..lds A-DMA completes on lgkmcnt; B vmem is NOT drained) + exactly ONE barrier per tile. lgkmcnt(0) targets tile K's A-DMA exactly because K+1's DMA is issued only after the barrier. s1 median-of-3 (GPU5, iters=50, TILE_N=64, 3584x512), base 3afb739 -> pipe: tok16 209.1->197.8 (-5.4%), tok128 656.8->605.5 (-7.8%), tok1024 791.6->766.5 (-3.2%), tok4096 2073.5->2022.5 (-2.5%), tok16384 6856->6685 (-2.5%). Wins every token. AC-P1 cold cos 0.999986 (kimi512) / 0.999997 (small) unchanged. AC-P2 test_mxfp_moe_variants 6/6 green (a4w4/a8w4 byte-identical off-path). AC-P3 VGPR=128, AGPR=0, Scratch=0 (no spill); LDS 16KB->32KB (2 slots). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… main Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
_gemm1_body_a16w4 never used `wave`: all 4 waves of the 256-thread workgroup redundantly computed the FULL TILE_N (num_acc_n=TILE_N/16) for both gate and up, loaded the same B, upconverted the same mxfp4, ran the same MFMA, and wrote the same results to the same global rows (idempotent, so correctness held but ~4x wasted work). This is the 128-MFMA/tile vs aiter's 32-MFMA/tile gap. Fix mirrors _gemm2_body_a16w4's proven wave-split: _n_per_wave = TILE_N/4, num_acc_n = _n_per_wave/16, and n_tile_base = wave*_n_per_wave shifts each wave's gate/up N addressing (col_g / scale ng) so each wave owns a distinct TILE_N/4 column slice. The epilogue already writes col_g_list[ni], now per-wave-distinct -> full non-overlapping coverage. Verified (TILE_N=64, K=3584, cold cache): MFMA 128->32/tile (448 total = 14*32, == aiter), VGPR 256->94 (no spill, AGPR/scratch 0). s1 median-of-3 ~2x faster (tok16 201.8->119.8us, 128 606.2->345.2, 1024 772.1->410.6, 4096 2024->1013, 16384 6735->3341). Cold e2e cos kimi512 0.999986 / small 0.999997 (== fp4 bar). a4w4/a8w4 untouched (variants test green). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Wires the a16w4 (bf16 A x mxfp4 W1/W2) fused MoE into the pytest matrices:
- Export flydsl_a16w4_gemm1/2 from kernels.moe.mxfp_moe (__init__).
- _run_a16w4_moe_e2e: dedicated e2e harness -- A stays raw bf16 (no quant,
x_scale=None), W1/W2 mxfp4 (standard shuffle + e8m0), bf16 [sorted_size,inter]
intermediate (no ascaleout), routes flydsl_a16w4_gemm1/2 (atomic gemm2), ref =
torch_moe_gemm1(bf16 A, scale=None) -> ref1.to(bf16) -> torch_moe_gemm2 (no fp4
re-quant). _run_mxfp_moe_e2e gains an a_dtype=="a16" branch that delegates to it.
- test_moe_gemm_2stage: add "a16w4" to the in_dtype parametrize + CLI --in_dtype
choices + "all" expansion; encode a16w4 constraints (gemm1 inter%64==0 &
2*inter%256==0; gemm2 model_dim%256==0 & inter%256==0; atomic-only -> skip reduce
& tile_m must be in {16,32,64}); route a16w4 -> a_dtype="a16".
- test_mxfp_moe_variants: add "a16w4" a_dtype; only bm32_atomic applies (bf16 A has
no inline_quant/interleave gemm1 variant -> other two skip).
- New test_a16w4_moe_e2e: small (128x1024x256 E8 k2) + Kimi-K3 (3584x512 E896 k16,
large_shape) shapes.
Cold (FLYDSL_RUNTIME_ENABLE_CACHE=0, gfx950): a16w4 e2e small + kimi512 + variants
bm32_atomic + matrix FP4-M all PASS at the fp4 verify bar. a4w4/a8w4 not regressed
(fp4/a8w4/fp8/bf16 rows 23 passed / 0 failed; variants 6/6). Style gate clean.
Shape constraints encoded: a16w4 needs tile_k>=256, model_dim>512, inter%256==0,
model_dim%256==0, tile_m in {16,32,64}, atomic epilog only (no reduce).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Re-diff at HEAD (isa-diff-gemm1-f019.md) confirmed the wave-split (f019ad1) already brought MFMA to 32/tile and VGPR to 94 (below aiter's 134), and measured occupancy at ~3.8-4.0 waves/CU -- so occupancy was NOT the residual 1.5x. The dominant remaining lever is N-tile geometry: aiter's tuned CSV (kimik3 abf16_wfp4) uses t32x128x256 for the mid/large tokens while gemm1 was pinned at tile_n=64. Swept tile_n 64/128/256 (median-of-3, dev7, cold): tile_n=256 wins or ties at every token except a ~3% wash at tok128 (within gfx950 clock noise), and beats aiter at tok16. With the wave-split N-split, tile_n=256 -> num_acc_n=4, VGPR 256 (no spill, AGPR/scratch 0), LDS 32 KB unchanged (A-LDS only). s1 median-of-3 (us), before(tn64)/after(tn256) vs aiter s1: tok tn64 tn256 aiter new gap 16 117.1 89.1 92.6 0.96x (BEAT) 128 345.7 289.3 228.6 1.27x 1024 406.8 341.0 276.6 1.23x 4096 1005.4 770.9 - - 16384 3334.5 1953.5 1313.3 1.49x (was 2.53x) tok16384 gap 2.53x -> 1.49x; mid-band ~1.5x -> ~1.25x. Cold e2e cos kimi512 0.999986 / small 0.999997 (== fp4 bar). variants test green. a4w4/a8w4 host arms + builders untouched (byte-identical off-path). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Re-diff (isa-diff-gemm1-f019.md) found ours interleaved the next tile's buffer_load..lds (A DMA, a VMEM op) among the current tile's ds_read cluster, forcing per-read s_waitcnt vmcnt(0) drains (2/iter avg, 3 on the pong iter) that also flushed the in-flight B mxfp4-weight loads. Aiter instead reads the whole resident A-LDS buffer FIRST, THEN issues the next A-DMA, keeping one vmcnt(0) only at the buffer-swap and reaching vmcnt(8) in-flight depth. Mirror it: hoist all m_repeat*k_unroll ds_read_b128 into preload_a() before the kt+1 A-DMA + B loads; compute_tile now consumes preloaded fragments. ISA at tn256: vmcnt(0) 24->23, vmcnt(N>0) 44->65 (deeper overlap), VGPR 256 (no spill, AGPR/scratch 0), MFMA 1792 unchanged. Marginal at tn256 (the fat N tile already amortizes the drains) but a consistent small win at large M with no downside. s1 median (dev7, cold, same-session head-to-head), Lever1-only/+phase-sep: tok1024 326.4 -> 324.7 (-0.5%) tok4096 772.1 -> 774.8 (+0.3%, wash) tok16384 1957.3 -> 1917.5 (-2.0%, seen twice) Cold e2e cos kimi512 0.999986 / small 0.999997 (== fp4 bar). variants green. a4w4/a8w4 use _gemm1_body (untouched); this only edits _gemm1_body_a16w4. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…main Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ratch ISA doc Keeps the tn256 win on inter%256==0 shapes (kimi512, small) while preserving support for inter%256!=0 shapes (e.g. 3584x384) instead of a hard TILE_N=256. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…shape s1 win Refactor flydsl_a16w4_gemm1/gemm2 to aiter's tile-config param interface (tile_m/tile_n/tile_k, waves_per_eu, k_batch, b_nt, xcd_swizzle, gate_mode, act), removing the home-grown adaptive TILE_N. tile_m->BM, tile_n->TILE_N, tile_k->TILE_K, b_nt->W-load cache modifier. Kernel builders: add xcd_swizzle + waves_per_eu to compile_gemm1_a16w4_port (xcd = bijective XCD/HBM-channel grid remap, same _xcd pattern as a4w4 gemm1 / a16w4 gemm2; waves_per_eu = rocdl.waves_per_eu min-occupancy value_attr) and waves_per_eu to compile_gemm2_a16w4_port (xcd already present). a4w4/a8w4 builders + host untouched (byte-identical off-path). CSV loader: pick_a16w4_config()/_load_a16w4_csv()/_decode_a16w4_kname() parse aiter's kimik3_fp4_tuned_fmoe.csv abf16_wfp4 rows -> per-token tile-config (tile_m/n/k, waves_per_eu, xcd, b_nt, k_wave) as a candidate-geometry source. Perf (the one portable lever): gemm1 b_nt default is now the measured per-M W-load U-shape (nt streaming b_nt=2 for 16<=tok<=1024, cached b_nt=0 at the ends), same mechanism as gemm2. Measured s1 (median-of-3, 3584x512 E896 k16, gfx950): tok16 90.3->82.3 (-8.8%), tok128 287.0->254.7 (-11.2%), tok1024 338.0->313.7 (-7.2%); tok4096/16384 unchanged (cached). cos held (kimi512 0.999986, small 0.999997), no VGPR/SGPR/scratch spill. test_a16w4_moe_e2e updated to the new interface (waves_per_eu left None: our fixed 4-wave tile_n=256 body is 256-VGPR-bound, so aiter's tuned waves_per_eu 3/4 spills and regresses ~4-5x here -- that knob is for aiter's tile_n=64 body). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…_nt win Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…(unmix from mxfp_moe) Move the a16w4 (bf16 A x mxfp4 W) fused 2-stage MoE kernel out of kernels/moe/mxfp_moe/ into a new self-contained package kernels/moe/moe_2stage_a16wmix/ so it is no longer mixed with the a4w4/a8w4 (mxfp4-activation) code. New package (gemm1/gemm2/host/common/__init__): moves _gemm1_body_a16w4, compile_gemm1_a16w4_port, gemm1_a16w4_grid, _gemm2_body_a16w4, compile_gemm2_a16w4_port, gemm2_a16w4_grid, the a16w4 host glue + CSV loader, and the a16w4-only helpers (_a16w4_swizzle_xor16, _e8m0_byte_to_f32). common.py is a curated copy of the subset of shared helpers a16w4 uses (raw/udiv/umod, buffer_rsrc, global_i32 helpers, gep/ lds ptr helpers, silu/situ activation, kmchunks/lds_acc), plus a private copy of _atomic_bf16_epilog, so the package has no import back into mxfp_moe. mxfp_moe: pure deletion of the a16w4 additions -- gemm1.py/gemm2.py have zero added lines (a4w4/a8w4 execution paths byte-identical); removed the now-unused _situ_mul_batch/_tanh_f32/_sigmoid_f32 from mxfp4_gemm_common and trimmed the a16w4-only imports/exports. tests: import flydsl_a16w4_gemm1/2 from the new package; _run_a16w4_moe_e2e behavior unchanged. cvt_scalef32_pk_bf16_fp4 stays in expr/rocdl (shared primitive). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add compare_interleave_sweep.py: a combined-dispatch comparison harness mirroring aiter's gfx950 per-1x32 (mxfp4-weight) INTERLEAVE/Swiglu production path. Routes token<256 -> our a16w4 (bf16 A x mxfp4 W, moe_2stage_a16wmix) and token>=256 -> our a8w4 (fp8 A x mxfp4 W, mxfp_moe interleave). Both regimes consume the identical shuffle_weight + e8m0_shuffle W layout (interleave is purely an internal kernel N-column addressing choice, so no boundary re-shuffle). Reports per-token s1/s2/total (median-of-3) and, alongside the loose shipped e2e gate, a TRUSTWORTHY cos vs a bf16-dense MoE reference (aiter's yardstick). Additive only; no kernel changes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… to it Extend the standalone a16wmix MoE package with a second weight format, a16wi4 (bf16 A x signed-int4 W, groupwise scale, group_size=32), alongside the existing a16w4 (mxfp4). int4 W reuses the a16w4 mxfp4 body verbatim: packed 2 nibbles/byte in the SAME preshuffle byte layout, dequant via v_cvt_off_f32_i4 (folding the x16 correction into the per-group scale) -> the identical non-scaled MFMA(16,16,32,bf16). Adds a w_dtype="mxfp4"|"int4" param to compile_gemm1/2_a16w4_port, the bodies, and the host wrappers; w_dtype="mxfp4" is byte-identical to the original a16w4 (kernel names and const_expr-gated paths unchanged). Groupwise bf16 scale uses (E, N, G//2, 2) packed layout via a16wi4_scale_to_kernel_layout; the K->group map (adj_ku) matches the a16w4 e8m0 addressing and is verified against the weight preshuffle. Route the legacy moe_gemm_2stage int4_bf16 test entry (run_moe_stage1/2) to the new path: stage1 always, stage2 for bf16 output (the a16w4 atomic epilog is bf16-only; f16-out stage2 keeps the legacy builder). Per-row (group_size=-1) is expressed as all-equal groups. Correctness: cold cos ~0.99999 vs a torch int4-groupwise reference across shapes; the legacy int4_bf16 groupwise + 2stage atomic/reduce tests pass via the new path. No VGPR/SGPR spill (gemm1/gemm2 spill_count=0). a16w4 mxfp4 e2e/variants and a4w4/a8w4 regression green. Style gate clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…os); xfail broken a4w4/a8w4 The two fused-MoE e2e gates (_run_a16w4_moe_e2e, _run_mxfp_moe_e2e) called verify_output() with the return value discarded (never asserted) AND with logits_diff_threshold=1, so they validated nothing. Even the logits threshold was dead: verify_output early-returns True when <5% of elements exceed the allclose tol, so the loose rtol/atol=0.5 masked the check regardless. Make the gates honest: assert verify_output(..., rtol=2e-3, atol=2e-3, logits_diff_threshold=2e-3). Independent strict-cosine verification (real cosine of e2e output vs the torch reference, cold) confirms: a16w4 e2e cos ~0.9999 (control, faithful) -> strict-asserted, passes a4w4 e2e cos ~0.12 (broken) -> xfail a8w4 e2e cos ~0.07 (broken) -> xfail Stage split localizes it: a4w4 stage1 cos = 1.0 but gemm2-isolated cos = 0.12 (shared down-proj broken); a8w4 stage1 cos = 0.16 (fp8-gemm1 A-path broken too). The broken kernels are also memory-unsafe: launching them and unwinding an xfail under pytest teardown cascades an illegal-address crash into other tests, so the a4w4/a8w4 xfails use run=False (variants) / pytest.xfail() before launch (2stage) to document the expected failure without executing the kernel. a16w4/a16wi4 stay strict-asserted-and-passing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…8w4 breakage The fused MX-FP4 MoE kernels in kernels/moe/mxfp_moe/ are numerically broken: a4w4 e2e cos ~0.12, a8w4 ~0.07, vs the a16w4 control (separate kernels) at ~0.9999. The shared flydsl_mxfp4_gemm2 down-proj is broken and, for a8w4, the fp8 gemm1 A-path additionally. The kernels are also memory-unsafe (illegal address under teardown). This is documentation/tracking only -- no kernel fix. - tools/repro/repro_mxfp_moe_a4w4_a8w4_broken.py: deterministic, cold-cache, self-validating (a16w4 control) strict-cosine reproducer with PASS/FAIL. - docs/issues/mxfp_moe_a4w4_a8w4_broken.md: ready-to-file ROCm/FlyDSL issue draft (gh CLI absent) with evidence table, localization, and repro command. - tests/kernels/test_moe_gemm.py: point the a4w4/a8w4 xfail reasons at the issue doc + repro, with a TODO(issue #NNN) placeholder. xfail(run=False) retained for CI stability. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add aiter-style intra-block slice-K (k_wave) to the a16w4/a16wi4 gemm1:
repartition the 4 waves into (4/k_wave) N-waves x k_wave K-waves, each K-wave
computes a K-slice (klen=K/k_wave) of a wider N-slice, and the k_wave partial
f32 accumulators are reduced in LDS (gate/up in separate rounds to bound peak
scratch) before the SiLU epilogue. Only the primary k-group (wave_k_id==0)
writes the intermediate; peers hold the identical reduced sum.
Everything is const_expr(k_wave>1)-gated: the k_wave=1 path is byte-identical
to before (verified: identical final ISA, 260 VGPR / 0 spill). Threaded through
compile_gemm1_a16w4_port and the host flydsl_a16w4_gemm1 (the CSV loader already
decodes _kw{N}); k_wave in {1,2,4} with 4%k_wave==0 and K%(k_wave*TILE_K)==0.
Verdict: slice-K does NOT help this kernel's fixed-256-thread / TILE_N=256
geometry. Because the workgroup stays at 4 waves, k_wave>1 widens each wave's
N-tile by k_wave x (num_acc_n 4 -> 8 -> 16), which blows the register file:
kw2 = 512 VGPR / 282 spill (~2x slower), kw4 = 512 VGPR / 969 spill (~3x
slower) vs kw1 = 260 VGPR / 0 spill. aiter benefits because its kw configs use
tile_n=64 and add real waves (num_waves_total up to 8 = 512 threads); a fixed
256-thread block cannot add parallelism, so slice-K only redistributes the same
MFMA work worse plus LDS-reduce overhead. Real cosine ~0.999999 for kw1/kw2/kw4
(kw4 needs tile_k=128 on Kimi K=3584: klen=896 not divisible by 4*256). The
param is landed for aiter-interface parity, not as a small-M win.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The a16wi4 (bf16 A x signed-int4 W groupwise) path that now serves int4_bf16
inherited the a16w4-mxfp4 fat-wave tile_n=256 geometry, which is a consistent
loss for the int4 (bf16-A) kernel: measured median-of-3 on gfx950, tile_n=128
beats 256 on BOTH stages across every int4 MoE shape/token tested (model_dim/
inter in {256,512,1024,2048}, tokens 16..1024) -- e.g. gemm1 inter=1024 tok=256
15.8 vs 30.7 us; gemm2 similar. int4 is bandwidth/grid-fill bound, so the wide
256-col tile just halves the N-tiles per expert (fewer waves to hide latency)
with no MFMA-density payoff, unlike the fp4 path where 256 is the tuned tile.
- host: add _default_tile_n(N, w_dtype); int4 -> 128 (fallback when N%128!=0),
mxfp4 unchanged (256 when N%256==0). gemm1/gemm2 tile_n=None now delegate to
it (gemm2 previously had no None handling).
- gemm2 _atomic_bf16_epilog: derive num_acc_n / _n_per_wave / s-count from BN so
the tile_n=128 geometry lowers (was hardcoded to the tile_n=256 4-way split).
Byte-identical for mxfp4 tile_n=256.
- test: int4_bf16 new-path stages pass tile_n=None (host default = 128) instead
of hardcoding 256; add inert, documented FLYDSL_INT4_FORCE_LEGACY /
FLYDSL_A16WI4_TILE_N perf-measurement escape hatches.
Correctness (cold, FLYDSL_RUNTIME_ENABLE_CACHE=0): int4_bf16 new-path + legacy-
routed (reduce/f16-out) tests pass vs torch int4-groupwise reference. a16w4-mxfp4
non-regressed (test_a16w4_moe_e2e + FP4-M cosine pass). Residual vs legacy int4
at the caller's BM=32 is a stage2 tile_m effect (the a16wi4 kernel prefers BM=16,
which reaches parity/win) controlled by the caller, not the host tile-selection.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…tical) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… draft Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add a16w16 (bf16 A x raw bf16 W, no scale, no upconvert) as a third weight format alongside a16w4 (mxfp4) and a16wi4 (int4). B is loaded straight into the bf16 MFMA(16,16,32,bf16) fragment: each dwordx4 (8 bf16) IS one K32 ku fragment. No dequant, no scale buffer. The bf16 W preshuffle (shuffle_weight layout=(16,16), elem_bytes=2) is K-reindexed so the (klane,ku)->K map matches the fp4 path. mxfp4/int4 paths stay byte-identical via const_expr gates. Cold real cos (bf16-dense MoE ref) ~= 1.0 (0.999997, unquantized). a16w4/ a16wi4 non-regression green. a16w16 is 4x the W bytes of a16w4 (bf16 2B vs fp4 0.5B) so ~4.8x slower on the small BW-bound shape (reported, expected). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…-launch investigation) Add an opt-in persistent, CU-limited launch grid to the shared a16w4/a16wi4 gemm2 (down-proj), mirroring aiter's persist path: cap the launch to min(total_work, NUM_CU=256) CTAs and have each CTA round-robin over the REAL work-tiles [0, bound) (bound from the cumsum, so padded empty m-blocks are never executed), with the existing XCD swizzle applied per visited tile so each real tile is computed exactly once. Threaded through gemm2_a16w4_grid, compile_gemm2_a16w4_port, and flydsl_a16w4_gemm2. Default is OFF and byte-identical to the prior grid (persist tag not appended, no loop emitted). Correctness verified bit-exact (maxdiff=0.0) between persist-on and persist-off across tok 1/16/128/1024 at E896, and the strict-cos e2e (test_a16w4_moe_e2e, incl. kimi 3584x512 E896 k16) + a16wi4 int4_bf16 tests pass. Finding: the persistent grid does NOT close the E896 perf gap and is a small (~1-7%) regression, so it is left off by default. Measured on gfx950: launching only the real tiles (zero empty CTAs) is the same latency as the padded launch (477us both), i.e. empty-CTA launch overhead is ~0 and is not the bottleneck. The E896 cost is real-tile compute -- the sort emits ~1 padded m-block per distinct expert (234 blocks for 256 active slots at tok16), each computing a full BM x TILE_N x K tile for ~1 token. aiter's fast a16w4 gemm2 uses a tile_k=128 REDUCE epilogue (not atomic); the persistent grid cannot address that. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Wire aiter's tuned kimik3_fp4_tuned_fmoe.csv per-token tile config into the
a16w4 (bf16 A x mxfp4 W) dispatch. resolve_a16w4_gemm{1,2}_config parse
(tile_m/n/k, waves_per_eu, xcd_swizzle, b_nt, k_wave) verbatim from the CSV row
for (model_dim, inter_dim, E, topk, tokens) and pass them to the builders.
Default-enabled (use_csv_config) for mxfp4; falls back to the adaptive default
when no CSV row matches (non-kimi shapes, missing CSV). a16wi4/a16w16/a4w4/a8w4
untouched (guarded on w_dtype=="mxfp4").
At small M the CSV yields tile_n=64 + k_wave in {2,4}, which -- unlike a wide
tile_n=256+kw (spills to 512 VGPR) -- keeps per-wave num_acc_n = tile_n*kw/64 at
or below the no-spill default footprint (<=4) while adding intra-block K-slice
parallelism. Measured (gfx950, 3584x512, median-of-3): s1 speedup at tok
1/2/4/8 = 1.86/1.46/1.21/1.12x; all configs 0 spill (tn64 kw1=122, kw2=175,
kw4=236 VGPR vs default tn256=260). e2e real cos 0.99999 (>= fp4 bar).
kw4 divisibility: the builder requires K % (k_wave*tile_k)==0; aiter names kw4
rows tile_k=256 but relies on K-padding we lack (K=3584: 3584 % 1024 != 0).
_kw_tile_k_for keeps the CSV k_wave and picks the largest tile_k satisfying the
constraint (kw4 -> tile_k=128), else drops to k_wave=1. gemm2 CSV rows violating
D_INTER%tile_k or model_dim%tile_n (e.g. inter=384, tile_k=256) return None ->
adaptive default.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…o mid-M regress) aiter's CSV tile_n=64 configs are tuned for aiter's thin-wave body and regress our fat-wave gemm1 at mid tokens (16-128). Default to our tuned tile_n; CSV params remain available (use_csv_config=True / aiter head-to-head) and win at small M (tok1-8). No custom tile tables added; CSV is parsed verbatim when used. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The a16w16 (bf16 A x raw bf16 W) path produced garbage at large expert count (kimi512 E896: e2e logits_diff 0.35, stage1 cos 0.66) while small-E passed and mxfp4/int4 passed at E896. Root cause: the bf16 W buffer offset overflowed 32 bits at large E. The bf16 W [E, N_OUT, K] is 2 B/elem (4x the fp4 0.5 B/elem), and the kernel folded the expert offset into the load index (n_blk = (e*N_OUT + col)//16) against a whole-tensor buffer resource. buffer_load computes its byte offset as offset*elem_bytes in i32, so at E896 the byte offset reached ~6.6 GB (gemm1, K=3584) / ~3.3 GB (gemm2), overflowing i32; the whole-tensor num_records (6.6 GB) also exceeded the 32-bit buffer num_records field (clamped to 4 GB, making high experts read OOB->0). fp4/int4 escape both (~4x smaller bytes). Fix (bf16 path only, const_expr-gated): fold the per-expert base into the i64 buffer resource address (base = arg_bq + e*N_OUT*K*2) and index W WITHIN the expert (n_blk carries no expert_off). Per-expert byte offset stays < 8 MB and per-expert num_records < 8 MB, both well under 2^31. mxfp4/int4 keep the whole-tensor resource + expert_off-in-index path byte-identical. a16w16 kimi512 e2e cos 0.657 -> 0.99999, stage1 cos 0.658 -> 1.0; small-E still passes. a16w4/a16wi4 non-regression green (byte-identical, const_expr gates). No spill. Gate: existing strict kimi512-a16w16 e2e assert (E896). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…uilding block) Kept as an opt-in (persist=True) building block; default-off is byte-identical. Does not close the E896 small-M gap (that's atomic-vs-reduce+tile_k, not grid), but a correct/tested option for grid-capped launches. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> # Conflicts: # kernels/moe/moe_2stage_a16wmix/host.py
…-preserving) Reduce the incremental-build/ledger-style comment bloat accumulated over many tuning rounds to the density of the neighboring mxfp_moe kernels, while keeping the load-bearing ABI/layout notes (fp4/int4->bf16 fragment mapping, the bf16 expert-offset-in-i64 overflow fix, the LDS buffer_load_lds swizzle hazard) and all intentional opt-in features (k_wave, persist, CSV config, mxfp4/int4/bf16). Strictly comment/docstring-only: the executable AST is byte-identical modulo docstrings across all five files. Verified behavior-preserving on gfx950: a16w4/a16w16/a16wi4 tests unchanged (7 passed / 2 skipped / 6 xfailed), real cos ~0.999995 for a16w4 and a16w16 at the kimi512 shape. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace the clean, behavior-preserving raw-arith call-sites with the current fx.* operator/method surface, matching the mxfp_moe/attention idiom: - arith.cmpi(eq)+arith.select -> (a == b).select(t, f) (B-scale lo/hi, x4) - arith.cmpi(ult) -> (token < i32_ntok); arith.andi -> & (epilogue store mask) - arith.TruncFOp(T.bf16, y) -> y.to(fx.BFloat16) (epilogue + int4 upconvert) Kept as raw (layout-API ceiling / reference-aligned, noted in code): the preshuffled-W buffer_ops.buffer_load (crd2idx layout) and per-lane scalar B-scale loads, the masked buffer_store output, and the llvm.AtomicRMWOp bf16 scatter epilogue -- all match mxfp_moe's justified raw machinery. XCD-swizzle arith.minsi kept (no fx signed-int-min operator). The A-DMA already uses make_buffer_tensor+fx.copy and the MMA already uses make_mma_atom+fx.gemm. Verified COLD (FLYDSL_RUNTIME_ENABLE_CACHE=0, gfx950): full gate 7 passed / 2 skipped / 6 xfailed (unchanged); real cos a16w4/a16w16 = 0.999995 (unchanged); s1/s2 median-of-3 within noise (tok128 s1 332/s2 1662us, tok1024 s1 343/s2 965us vs baseline 333/1665, 342/964); no new spill (gemm1 260 VGPR/0 spill, gemm2 156 VGPR/0 spill, 0 scratch mem-ops); style clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Convert the hot-loop W dwordx4 weight load from raw buffer_ops.buffer_load over a
create_buffer_resource_from_addr V# to the current BufferCopy128b copy-atom idiom
used by the mxfp_moe reference:
- add _global_i32_buffer_tiles helper (curated from mxfp4_gemm_common) to common
- W load: fx.copy(BufferCopy128b(b_cache_mod) atom, slice(w_tiles, tile_idx), r)
replacing buffer_load(w_rsrc, ..., cache_modifier=b_cache_mod); the runtime cache
modifier is carried in the atom aux field. Both load_b_raw (mxfp4/int4) and
load_b_raw_bf16 (a16w16) in gemm1 and gemm2.
Kept raw (layout-API ceiling, noted in code): the per-lane scalar (vec_width=1)
e8m0 B-scale buffer_load (no layout form, dict-cached), and the masked predicated
bf16 output buffer_store + create_buffer_resource_from_addr (no copy-atom form for
a masked scatter-store). The llvm.AtomicRMWOp bf16 scatter epilogue and XCD-swizzle
arith.minsi also stay (match mxfp_moe's justified raw machinery).
Verified COLD (FLYDSL_RUNTIME_ENABLE_CACHE=0, gfx950): full gate 7 passed /
2 skipped / 6 xfailed (unchanged); real cos a16w4/a16w16 = 0.999995 (unchanged);
byte-equivalent ISA (gemm1 260 VGPR/0 spill, 280 buffer_load_dwordx4; gemm2 156
VGPR/0 spill; 0 scratch); s1/s2 median-of-3 within noise (tok128 s1 335/s2 1663us,
tok1024 s1 341/s2 965us vs baseline 333/1665, 342/964); style clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ilog) The moe_2stage_a16wmix extraction (678ad90) dropped the @flyc.jit decorator from _atomic_bf16_epilog. Without it the ``if token_id < i32_M`` bound check is no longer AST-rewritten into an scf.if guard; it runs as a plain Python if that is dropped at trace time, so the atomic-fadd scatter fires unconditionally on padded/OOB sorted rows. On the Kimi-K3 E896 shape this issues ~10x extra atomic HBM transactions, regressing gemm2 stage2 from ~39us to ~490us at tok16. Re-adding @flyc.jit restores the scf.if guard (byte-identical to the pre- extraction kernel). gemm2 s2 (3584x512 E896 k16, cold, median-of-3): tok16 496->38us, tok128 1669->145us, tok1024 951->176us. No VGPR spill (156, 0 spill). a16w4/a16wi4/a16w16 correctness unchanged (170+9 tests pass); a4w4/a8w4 (mxfp_moe) untouched. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…o k_wave
The tuned-CSV loader previously filtered to abf16_wfp4 only, so a16wi4 (int4)
rows were skipped and the caller fell back to the default tile config. Accept
abf16_w{fp4,int4,bf16} rows, and parse the _kb{N} grid-split-K token: since we
have no grid split-K, approximate it with our intra-block slice-K lever k_wave
(kb==2 -> kw2, kb>2 -> kw4). int4 rows (kb2/4/7/14) now resolve to k_wave 2/4;
K=7168 divides k_wave*tile_k cleanly.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…h tokens
The mxfp4 gemm1 default tile_n=256 ("fat wave") compiles to 260 VGPR, which on
gfx950 (512 VGPR/SIMD, alloc granularity 8 -> 264) allows only 1 wave/SIMD --
the occupancy cliff. Weight-load latency is then fully exposed, and because the
whole E896 grid is capped at 1 wave/SIMD regardless of M, the FlyDSL/aiter s1
gap grows with token count (~1.96x @tok8192, ~2.06x @tok16384).
tile_n=128 splits each block into 2 n-blocks with thinner waves (num_acc_n 4->2),
dropping VGPR 260->178 -> 2 waves/SIMD. rocprofv3 confirms ~2x SQ_WAVES and
~10% fewer total GPU cycles (GRBM_COUNT 185M->167M). No VGPR/AGPR spill.
Measured (3584x512 E896 k16, balanced routing, isolated gemm1 s1, median-of-5
cold, GPU4/gfx950):
tok tn256(before) tn128(after) delta
1024 320.6us 299.5us -6.6%
2048 553.5us 536.9us -3.0%
4096 806.9us 759.6us -5.9%
8192 1284.9us 1189.3us -7.4%
16384 2449.7us 2208.2us -9.9%
Also faster at small tokens (tok1 31.9->23.7us, tok64/256 ~-4%); no regression.
Scope: mxfp4 gemm1 only (D_INTER % 128 == 0). bf16 (a16w16, VGPR>=448 regardless
of tile_n) and int4 (already defaults to 128) keep _default_tile_n. gemm2 default
unchanged. The opt-in use_csv_config (aiter-compare) path is preserved -- it sets
tile_n before this default and wins. Parity: test_a16w4_moe_e2e[small,kimi512]
(mxfp4+bf16) pass cold.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…OOB safety) The gemm1 A (activation) global->LDS DMA built its buffer resource with num_records = 0xFFFFFFFF (~4GB). moe_sorting padding rows carry a sentinel sorted_token_id >= n_tokens; the A-load path decodes it (t_i32 = fused & 0x00FFFFFF -> x_row_base_div4 = t_i32 * c_k_div4) into a byte offset PAST the real [n_tokens, K] bf16 allocation. With the oversized resource the hardware OOB mask never clamps those padding-row loads, so the direct-to-LDS DMA can read unmapped memory and fault at high token counts (allocation-dependent: the read lands 1+ rows past the buffer). The epilogue already guards the STORE (token < i32_ntok) but nothing guarded this READ. Size the resource to the true allocation (i32_ntok * c_k_div4 * 4 bytes == n_tokens * K * 2). The HW mask then returns 0 for padding-row A-loads (harmless -- their store is dropped), matching aiter. Resource-descriptor only; the hot loop, VGPR (174, no spill), and s1 timing are unchanged, and valid-row output is byte-identical (verified: out.sum identical at tok 2048/8192/16384). Shared body -> covers mxfp4 / int4 / bf16 alike. Parity: test_a16w4_moe_e2e[small,kimi512] (mxfp4+bf16) pass cold; gemm1 runs clean at tok 2048/4096/8192/16384 under balanced routing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
At n_tokens>=16 the leaner TILE_K=128 (more, shorter K-tiles) plus xcd_swizzle=1 (bijective XCD/HBM-channel round-robin) is measured faster across the whole M>=16 range on both Kimi-K3 shapes (3584x384 / 3584x512): gemm1-s1 ~0.83-0.94x on the high-M tail (2048..16384), closing the residual aiter gap (3584x512 s1 was ~1.06 -1.11x aiter at tok2048..16384, now 0.96x/0.98x/1.01x/1.05x). Gated at >=16 because TILE_K=128 regresses the launch-latency-bound tok 1..8 (~+5-10%), which keep TILE_K=256. Only fires when the caller left tile_k/xcd at defaults and not use_csv_config; mxfp4 only (int4/bf16 share the body but are separately tuned). Parity held (e2e ld~1.5e-5 cold at all tokens both shapes). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…l fix) At tok 1..2 the E896 launch has only ~topk m-blocks of real work so the GPU is badly under-filled. The narrower TILE_N=64 doubles the CTA count per m-block for better latency hiding across the CUs: gemm1-s1 ~0.86-0.89x at tok 1..2 on both Kimi-K3 shapes (3584x512 tok1 s1 25.0->21.2us, tok2 27.2->24.3; total 1.44x->1.25x / 1.29x->1.21x vs aiter). Gated to tok<=2 (regresses ~+18% from tok4 up) and to the default path (not use_csv_config, needs D_INTER%64==0). Residual tok1-2 gap vs aiter (~1.4x on s1) is fixed kernel/launch latency at M=1, not empty-CTA grid-fill: a grid-size probe showed capping the launch to cover only the real work-tiles gives no speedup, so a persist/CU-capped gemm1 grid does not help. Parity held (ld~1.5e-5). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Rewrite _int4_nibble_to_bf16x8 to match aiter's tuned int4 dequant path: use v_cvt_off_f32_i4_sdwa byte_sel to read each nibble without a per-nibble shift (even nibbles 0,2,4,6 = byte_sel 0..3 on raw; odd 1,3,5,7 = byte_sel 0..3 on raw>>4 -> ONE shift total instead of seven), and pack f32->bf16 pairs with v_cvt_pk_bf16_f32 (4 packed converts instead of 8 scalar .to(bf16)). Measured on gfx950/dev4 (rocprofv3, 7168x512 E384/k8 tok4096, realistic harness memory): SQ_INSTS_VALU 470.4M -> 418.9M (-11%) and SQ_WAVES 55324 -> 45056 (now matches aiter's 45024). No VGPR spill (Scratch=0; VGPR 52->60). Parity held: cos 0.999987 (logits_diff 1.27e-5), strict kernel-mode gate PASS cold at all tokens. Same v8bf16 output layout; int4-only (mxfp4/bf16 paths untouched). Wall-clock is ~neutral at the tested MoE shapes (gemm1 is memory-bound there; the dequant VALU was already latency-hidden), but the op-count reduction matches aiter and helps compute-bound / L2-resident cases. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
int4 gemm1 stage1 is W1-load bound. With all experts active, each expert's W1 tile is re-fetched from HBM once per padded 32-row m-block. At the fill point where every expert has exactly two half-full 32-row m-blocks (avg tokens/expert in (32,64]), doubling block_m to 64 collapses them into ONE full 64-row block, halving the W1 HBM re-reads. Measured on 7168x512 E384/k8 at tok2048: TCC_MISS 3.34e7 -> 1.47e7 (2.27x), TCC_READ 5.72e7 -> 3.31e7, stage1 802 -> 749 us (-6.7%); tok3072 831 -> 769 (-7.4%); 7168x256 tok2048 -11%. No VGPR spill (92->144 VGPR, scratch 0, LDS 32->64KB). Outside that narrow band block_m=64 either wastes padding (<=1 block/expert) or halves grid parallelism with no reuse gain (>=3 blocks/expert), so it is not applied -- verified block_m stays 32 (byte-identical) at tok 512/1024/1536/ 4096/8192/16384. Add a16wi4_recommend_block_m() (mirrors aiter's estimated_m_per_expert heuristic) for the dispatcher to size routing + gemm1 block_m together (block_m sizes moe_sorting padding, so caller must couple them). int4-only; mxfp4/bf16 codegen untouched. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This reverts commit e75f0f9.
Rewrite _int4_nibble_to_bf16x8 to the VALU-lean gfx950 path: read each nibble with v_cvt_off_f32_i4_sdwa byte_sel (even nibbles from raw, odd from raw>>4 -> ONE shift total instead of seven) and pack f32->bf16 pairs with v_cvt_pk_bf16_f32 (4 packed converts instead of 8 scalar .to(bf16)). The prior attempt (e75f0f9, reverted in 712c28f) produced garbage (cos=0, output magnitude ~1e23) at the Kimi 7168xN shapes for tokens<=256. Root-caused: the byte_sel dequant was correct all along; the bug was the STATELESS rocdl.cvt_pk_bf16_f32 (has_side_effects=False). In the gemm1 hot loop the 4 identical-shaped packed converts per K-step get CSE-merged / reordered across K iterations, corrupting the v8bf16 MMA operand. Isolated-kernel tests of both primitives were correct; only the full hot-loop context miscompiled. Fix: emit the pack via a side-effecting inline-asm helper (_cvt_pk_bf16_f32_se) that pins each convert to its K-step. Bound classification (rocprofv3, gfx950/dev4, 7168x512 E384/k8): int4 gemm1 is VALU-BOUND, not memory-bound -- VALUBusy 62.7%, MfmaUtil 17.6%, MemUnitStalled 0.01%, VALU/VMEM inst ratio 43.5x, occupancy 21.7% (LDS-limited). After fix: SQ_INSTS_VALU 157.0M -> 138.2M (-12%), VALUBusy 62.7% -> 58.8%; occupancy and Scratch(=0, no spill) unchanged (VGPR 92->120, still LDS-bound). gemm1 stage1 before->after (median-of-7 cold, 7168x512 E384/k8): tok128 456.0->424.8 us (1.073x) tok1024 481.7->462.0 (1.043x) tok256 461.8->429.3 us (1.076x) tok4096 1201->1144 (1.050x) tok512 469.8->443.9 us (1.058x) tok8192 2145->2091 (1.026x) Accuracy: cos>=0.9994 at the previously-broken Kimi 7168x256 tok{16,32,64,128, 256}; output fingerprints byte-match the scalar path at every token. Official tests pass cold: test_moe_gemm_w4a16_groupwise_scale (2/2), test_moe_gemm_2stage int4_bf16 S/M/L (12/12), test_a16w4_moe_e2e mxfp4+bf16 (4/4, no regression -- int4-only change). int4-only: _int4_nibble_to_bf16x8 is called only on the _is_int4 branch; mxfp4/bf16 paths untouched. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
int4 gemm1 at the a16wi4_recommend_block_m BM=64 W1-reuse fill point (tok ~1600..3072 on 7168x512 E384/k8) was VGPR/AGPR-occupancy-crippled: TILE_N=128 compiles to VGPR=349 + AGPR=93 -> only 1 wave/SIMD (OccupancyPercent 12.3%), fully exposing the W-load + int4-dequant + A-LDS-read latency. Halving to TILE_N=64 drops the footprint to VGPR=244/AGPR=0 (no spill) -> 2 waves/SIMD (OccupancyPercent 23.0%), cutting gemm1-s1 ~23% across the whole BM=64 band (tok 1600/2048/2560/3072: 955/966/979/997us -> 734/745/764/773us), narrowing the tok2048 gap vs aiter from 1.49x to 1.15x. Scoped to w_dtype=int4 AND BM==64: the BM=32 mid band (tok 128..1024) is already 2 waves/SIMD (VGPR=234) and NOT occupancy-bound -- measured TILE_N=64 there is a net wash-to-loss (5x more SQ_WAIT_INST_LDS at equal active cycles), so it keeps TILE_N=128. mxfp4/bf16 never reach this branch (byte-identical codegen). Accuracy cos=0.999999 cold at BM=64 and tok<=256; no VGPR spill. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Not imported/exported; a measurement scratch that should not ship in the package. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
At M=1 the a16w4 gemm1 K-loop is LDS-wait-bound (~90% of SQ_BUSY cycles are SQ_WAIT_INST_LDS): the ds_read A-operand fetch cannot be hidden behind the shallow 7-tile K-loop when there is ~1 real token per 32-row block. Splitting K 4 ways (each wave computes a klen=896 slice, then an LDS-reduce) halves the per-wave K-loop depth, cutting the exposed LDS latency. Measured gemm1-s1 ~0.83x at tok 1..2 on both Kimi-K3 shapes (3584x384 / 3584x512), matching aiter's own tok1 config (t32x64x256_kw4). tok1 total 3584x512 1.27x -> 1.06x, 3584x384 1.27x -> 1.20x vs aiter. K=3584 % (4*128)==0 requires TILE_K=128 (kw4 with TILE_K=256 would need K % 1024, which 3584 fails). Regresses from tok4 up, so gated with the existing TILE_N=64 lever to tok<=2, mxfp4 only, defaults only. Parity preserved (cos ~0.99999, ld 1.5e-5); mid/high tokens and a16wi4/a16w16 untouched (byte-identical config path). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…n gfx942) The moe_2stage_a16wmix a16wi4 kernel uses the gfx950 (CDNA4) K=32 bf16 MFMA and fp4/int4 convert ops, so it aborts at launch on gfx942 (CDNA3, K=16 only) -- the CI mi325 runner hit 'Fatal Python error: Aborted' at host.py:300. Route int4_bf16 to the new kernel only when 'gfx95' in ARCH; on gfx942 fall back to the legacy moe_gemm_2stage builder (which supports gfx942 incl. groupwise, verified via FLYDSL_INT4_FORCE_LEGACY). a16w4/a16w16 tests were already gfx950-gated. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…equant
The moe_2stage_a16wmix a16wi4 (int4) and a16w16 (bf16) kernels aborted at launch
on gfx942 because they hardcoded MFMA(16,16,32,bf16) (CDNA4-only) and the int4
dequant used v_cvt_pk_bf16_f32 (gfx950-only). Arch-gate both:
* MMA: gfx950 keeps the K=32 16x16x32 bf16 MFMA (one per K-step); gfx942 uses the
K=16 16x16x16bf16_1k MFMA, splitting each v8bf16 K-step into two v4bf16 halves
and issuing TWO MFMAs into the same 16x16 f32 accumulator. Load/dequant still
produce v8bf16; only the MMA emission changes.
* int4 dequant: gfx950 keeps the VALU-lean v_cvt_pk_bf16_f32 pack; gfx942 falls
back to scalar .to(bf16) truncation per nibble (arch-agnostic).
a16w4 (mxfp4) stays gfx950-only (cvt_scalef32_pk_bf16_fp4 is CDNA4-only).
The path is selected host-side by a16wmix_use_k16() from get_rocm_arch(); the
gfx950 default (K=32) is byte-identical to before (verified: all 44 IR stages of
gemm1+gemm2 identical for mxfp4/int4/bf16). FLYDSL_A16WMIX_FORCE_K16=1 forces the
gfx942 (K=16 + scalar) codepath on any arch, to validate it on a gfx950 box.
Relax the test int4_bf16 routing gate (_A16WMIX_GFX) to include gfx94*, since the
new path now runs on gfx942.
Verified on gfx950 (GPU4): default K=32 path byte-identical + passes cold;
FORCE_K16 path correct cold for int4 + bf16 + mxfp4 incl tile_m=16/tok=64, zero
VGPR spill (only 16x16x16 MFMAs emitted). Perf: FORCE_K16 int4 is ~7-9% FASTER
than legacy int4 (moe_gemm_2stage) at 3 Kimi-like shapes on gfx950.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ly, IR byte-identical) Collapse multi-paragraph rationale/measurement-log comments to 1-2 tight lines across gemm1/gemm2/host/common/__init__, keeping the essential "why" and all load-bearing warnings terse (atomic epilog @flyc.jit, A-load OOB clamp, side-effecting bf16 convert, K16/K32 + mxfp4 arch-gates). Comment/whitespace-only: all IR stages (00_origin .. 21_final_isa) for a16w4 + int4 + bf16 gemm1+gemm2 are byte-identical to HEAD (diff -rq, cold, COMPILE_ONLY). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… on HW) The gfx942 K=16 codepath passes the gfx950 FORCE_K16 proxy but ABORTS at launch on real gfx942 (mi325 CI): the CDNA3 mfma_f32_16x16x16bf16_1k variant differs from the CDNA4 K=16 the proxy emits, so the proxy did not validate the actual gfx942 ISA. Re-gate int4_bf16 new-kernel routing to gfx950 only; gfx942 uses the legacy moe_gemm_2stage path (verified working). The arch-gated K=16 kernel code stays (dormant on gfx942) for later real-hardware validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a standalone fused 2-stage MoE kernel package
kernels/moe/moe_2stage_a16wmix/for bf16 activations × low-bit weights, aligned with aiter's
mixed_moe_gemm_2stage:cvt_scalef32_pk_bf16_fp4upconvert, e8m0 folded)The kernel is bf16 end-to-end: bf16 A staged to LDS (no A-scale), weights upconverted
to bf16 in-kernel, non-scaled
mfma_f32_16x16x32_bf16, a bf16 sorted intermediate (norequant), and a bf16 gemm2. It reuses the standard
moe_sorting+shuffle_weight+e8m0_shufflecontracts.Correctness
Strict-gated (cos ≈ 0.99999, logits_diff ≈ 1.5e-5) across the Kimi-K3 / Kimi int4 shapes and
full token sweeps (1 → 32768): a16w4 32/32 PASS, a16wi4 94/94 PASS. (For reference,
aiter's own bf16 a16w4 is numerically broken on these shapes in both SEPARATED and INTERLEAVE
gate modes — see the head-to-head notes; our kernel is the correct bf16 a16w4 there.)
Performance (gfx950, balanced routing, vs aiter's tuned kernels)
tok8→8192 + both shapes' 32768 (0.77–0.85×). Residual >1.0× only at low-M (tok1–8, LDS-wait
K-loop latency; the tok≤2 4-way slice-K lever brings tok1 to ~1.08×) and 3584×512 tok16384
(~1.04×, a small kernel-body gap on the largest tuned shape).
block_m=64 at the W1-reuse fill point, and a BM=64 occupancy fix. Residual mid-band /
high-token gap vs aiter's grid split-K (which our nonlinear-SiLU epilogue makes a large
refactor; prototyped and found not to pay with an fp32-scratch reduction).
Also included
python/flydsl/expr/rocdl/__init__.py:cvt_scalef32_pk_bf16_fp4wrapper (target-specific).tests/kernels/test_moe_gemm.py: a16w4 / a16wi4 / a16w16 e2e + variant coverage, strict gates.docs/issues/mxfp_moe_a4w4_a8w4_broken.md+tools/repro/: documents the pre-existinga4w4/a8w4 breakage in
kernels/moe/mxfp_moe(left = origin/main; not touched by this PR).Notes for review
mxfp_moe) are unchanged / byte-identical tomain(documented broken separately).FLYDSL_RUNTIME_ENABLE_CACHE=0).🤖 Generated with Claude Code