Repository navigation
FlyDSL: port a16wi4 to new pipeline, clean old moe_gemm_2stage - #4646
Merged
Merged
Conversation
Route a16wi4 (bf16 A x int4 W) through the merged a16w-mix port
(moe_2stage_a16wmix, w_dtype="int4") instead of the bespoke int4_bf16
kernel in moe_gemm_2stage.py, which is now dead and deleted (-3590 LOC).
The kernel package is UNCHANGED: main's port already exposes w_dtype="int4"
with the standard (GGUU) N-major layout, so a16wi4 only needs host-side
weight prep + routing.
- shuffle.py: shuffle_weight_a16wi4 / shuffle_scale_a16wi4 -- FlyDSL-native
int4 prep (kpack=16 packed-nibble view + shuffle_weight, scale [E,G,N] ->
(E,N,G//2,2) bf16 pairs). aiter's CK int4 packing is kpack=8 and is
MMA-incompatible with the port's single 16x16x32 MMA.
- moe_kernels.py: fold int4 into the existing a16w-port branches
(compile_flydsl_moe_stage1/2 + _flydsl_moe_stage{1,2}_impl) by widening the
fp4 gate to b_dtype in ("fp4","int4") and passing w_dtype=b_dtype; register
_kw{n}/_bnt{n} stage1 variants (the port uses intra-block k_wave, not grid
split-K; legacy _kb{n} names stay registered so old CSVs keep parsing).
- fused_moe.py: a16wi4 get_2stage_cfgs heuristic -> port-native tiles
(token<64 -> tile_n=64/tile_k=256/kw2, <256 -> 64/128, else 128/128), no kb.
- aot/flydsl/moe.py: precompile int4 through _precompile_a16w4_to_cache
(the port's raw-data_ptr launchers), so AOT builds the runtime kernel.
- kimik2_i4_tuned_fmoe.csv retuned to port-native kernel names.
- gemm_moe_tune.py / test_moe_2stage.py: prep weights with the new helpers.
Validated on gfx950 (7168x512, E384/k8, tok16): logits_diff 1.48e-5 vs
2.18e-5 for the old kernel on main.
a16wi4 now precompiles through _precompile_a16w4_to_cache (the folded a16w-mix port), so _precompile_to_cache can no longer be reached with b_dtype="int4"; its int4 storage/shape/scale branches encoded the deleted kernel's CK layout. fhmoe (shared-expert) has no int4 path.
… 2.6x -> 1.08x)
The port has no grid split-K, so decode is workgroup-count-limited: the grid
is m_blocks * (inter_dim / tile_n). At tok1 (7168x512, E384/k8) the previous
tuning (tile_n=64/k_wave=2) yields only ~64 workgroups on 256 CUs, making
gemm1 2.58x SLOWER than the old int4 kernel, whose grid split-K filled the
machine. A narrow N-tile restores the workgroup count.
Measured us_stage1 at tok1: tile_n=64/kw2 41.6 -> tile_n=32/kw2 27.6 ->
tile_n=16/kw4 16.7 (main 16.0). Retuned CSV + heuristic:
token < 4 -> t{m}x16x256_kw4 (needs model_dim % 1024 == 0)
4..63 -> t{m}x32x256_kw2
64..255 -> t{m}x64x128 (unchanged; grid is full, wider tile wins)
>= 256 -> t{m}x128x128 (unchanged)
Interleaved kernel-only A/B vs main (min of 2 rounds, gfx950), gemm1 ratio:
tok1 2.575 -> 1.081, tok16 1.138 -> 1.051, tok64 1.008; total (g1+g2) at
tok64 0.988. tok4 stays ~1.2x -- every registered tile/b_nt/k_wave combo
lands in 39-44us there vs main's 33us, so the residual is the known
structural one (port fully-unrolled K vs main's rolled loop + grid split-K).
Correctness re-validated with the new tiles: logits_diff 1.39-1.50e-5 at
tok 1/2/4/16.
Widening tile_n to {16,32} also registered their k_wave=1 forms, where each
of the 4 N-waves would cover 4/8 cols (<16, the 16x16 MMA minimum). Derive
every name from the wave-split rule instead (kw=1 => tn>=64), and keep the
legacy _kb names only for the tile_n the old kernel actually had.
160 valid names, all 80 pre-existing names preserved.
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
…ed CSV get_2stage_cfgs' a16wi4 branch is the NO-TUNED-ROW fallback (the CSV-driven branch above it returns first whenever kernelName1/2 resolve). It had grown a per-token tile_n/tile_k/k_wave table that duplicated -- and could silently diverge from -- the tuner's output. Replace it with one shape-safe config plus the same "no tuned FlyDSL config" warning the mxfp4 fallback emits, so an untuned shape is visibly untuned instead of quietly running a guessed tile. Also raise instead of falling through when FlyDSL is unavailable: a16wi4 W1/W2 are preshuffled for the port (shuffle_weight_a16wi4, kpack=16), so a CK/ASM stage1 (CK int4 layout, kpack=8) would produce wrong numbers rather than fail. Validated: tuned shapes still dispatch the tuned tiles (tok1 _t16x16x256_kw4, tok4/16 _t16x32x256_kw2), logits_diff 1.38-1.50e-5; an untuned shape (7168x1024 tok128) warns and stays correct at 1.46e-5.
a16w4 and a16wi4 are the SAME kernel (moe_2stage_a16wmix, w_layout="standard") differing only in w_dtype, so their weight layouts must stay in lockstep. Express that by delegating: shuffle_weight_a16wi4 = nibble-pack + shuffle_weight_a16w4(..., gate_up=False), instead of separately calling shuffle_weight and hoping the two stay equal. shuffle_weight_a16w4(x, NLane, gate_up=False) is structurally identical to shuffle_weight(x, (NLane,16)) -- verified equal across E/N/K/NLane variations, not just the shapes we ship -- so this is byte-identical to the previous implementation (verified on 4 shapes) and re-validated on GPU (tok4 1.49e-5). Only the gate_up=True (GUGU) form differs; that is the fp8 stage1 layout.
The port's main K loop was fully unrolled via range_constexpr(K_TILES_TOTAL),
emitting a ~12.7K-line / 448-MFMA flat kernel body. At small M (tok4) the grid
is latency/overhead bound and this instruction footprint costs ~2x the GRBM/
SQ_BUSY cycles of the old rolled-loop kernel at an identical tile config (proven
by rocprofv3 counters + ISA diff at t16x64x128).
Convert the K loop to a manually-emitted loop-carried scf.ForOp (the DSL
range(...,init=) sugar does not fire here because _gemm1_body_a16w4 is a plain
helper, not an @flyc.kernel-decorated function, so the AST rewriter never runs).
Accumulators are carried as vector values; the prefetched B-tile is flattened/
unflattened through the loop-carried state; the final K-tile is a post-loop
epilogue. Same prefetch structure (B(kt+1) before compute(kt), A double-buffered).
ISA: 12747 -> 2174 lines, 448 -> 32 MFMA, s_cbranch 1 -> 2 (real loop).
Correctness: logits_diff <= 1.76e-5 across tok{1,2,4,16,128} x inter{256,512,1024}.
Interleaved us_stage1 (min of 3): tok4 -11%, tok16 -3.5%, tok64 -3.5%.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The port's int4 groupwise scale used an N-major (E,N,G//2,2) layout: 16 consecutive
N-lanes read scale dwords strided by G//2 (448 B at D_HIDDEN=7168), so a wave's
scale gather touched ~16 cache lines per K32 step. Switch load_b_scale_int4 (gemm1
+ gemm2) to the OLD kernel's (E,G//2,N,2) layout (shuffle_scale_for_int4), where
16 consecutive N-lanes read 16 consecutive dwords = one 64B line. The expert base
is no longer folded into the N index (separate G//2*N_OUT stride).
int4-ONLY: the whole change is under const_expr(_is_int4); the mxfp4 (a16w4) gemm1
ISA is byte-identical before/after (verified by diff of 21_final_isa.s). Host prep
in test switches shuffle_scale_a16wi4 -> shuffle_scale_for_int4(group_size=32).
Correctness: logits_diff <= 1.55e-5 across tok{1,2,4,16,128} x inter{256,512,1024}.
Interleaved us_stage1 vs old kernel (min of 3): tok4 1.275->1.038, tok16 0.994->
0.853, tok64 0.991->0.876 (combined with the rolled-loop commit).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…olled) The rolled scf.for from 703a695 applied to ALL w_dtypes, including the shipped/ tuned mxfp4 (a16w4, PR#4502) mainloop -- rolling regresses mxfp4 (+5% at tok16 on 3584x512/E896). Gate the roll behind const_expr(_is_int4): int4 (a16wi4 decode) takes the rolled scf.for; mxfp4/bf16 keep the fully-unrolled range_constexpr body. mxfp4 gemm1 ISA is now BYTE-IDENTICAL to main (diff of 21_final_isa.s at h3584_i512_ne896_bm32_tn128_situv2_xcd4: identical, 4956 lines). int4 still rolls (2181 lines, s_cbranch=2). logits_diff unchanged (<=1.45e-5). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The a16wi4 kernel now consumes the (E,G//2,N,2) int4 scale layout (commit 1b42ab4), but the 2-stage tuner (gemm_moe_tune.py) still fed the old (E,N,G//2,2) layout via shuffle_scale_a16wi4 at both prep sites -- so tuning ran against wrong numerics. Flip both to shuffle_scale_for_int4(..., group_size=32), and delete shuffle_scale_a16wi4 from aiter/ops/shuffle.py (now unreferenced). Weight prep (shuffle_weight_a16wi4) is unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
a16wi4 now consumes the replaced FlyDSL int4 kernel's weight preshuffle
(pack_int8_to_packed_int4(shuffle_weight(w.i8, (16,16))), kpack=8) in addition to
its (E,G//2,N,2) scale layout, so the caller contract is byte-identical to the
kernel this PR replaces -- zero model/serving change. shuffle_weight_a16wi4 deleted.
Kernel (int4-gated, gemm1 + gemm2):
- _int4_nibble_to_bf16x8(old_pack=True): the OLD 'interleaved-by-4' byte packing
(byte j = K_j | K_{j+4}<<4) means the SAME two cvt loads give f_lo[j]=K_j,
f_hi[j]=K_{j+4}; correct v8bf16 order is all-lows-then-all-highs -- a pure output
REORDER, zero extra instructions (verified against a host nibble-map model).
- load_b_raw_int4: the OLD layout stores K16 per klane slot (kpack=8), so a port
MFMA-K32 block c=k0*4+lane_div_16 is the concat of two adjacent klane slots
(2*(c%2), +1) 128 B apart -> TWO dwordx2 (8B) loads replace one dwordx4 (16B).
No cross-lane movement (host model: the two halves are always 128B-separated
slots of the same lane).
PERF COST (this is the trade): the dwordx4->2x dwordx2 halving is a real load-width
cost at the smallest M. Interleaved us_stage1 vs OLD kernel (min of 3):
tok1 1.099 (ABOVE old-kernel parity, ~10% slower)
tok4 1.055-1.064 (ABOVE parity, ~6%)
tok16 0.879 tok64 0.950 tok2048 0.859 (all faster than old)
vs the prior scale-only branch (kpack=16 weights) it gives back margin everywhere
(tok16 0.820->0.879, tok64 0.882->0.950) and pushes tok1/tok4 past the old-kernel
bar. int4 gemm1 B-loads: 8 dwordx4 -> 8 dwordx4(fp4 unchanged)... int4: 0 dwordx4
/ 16 dwordx2 (was 16 dwordx4/0).
Gates: logits_diff <= 1.7e-5 across tok{1,2,4,16,128} x inter{256,512,1024};
mxfp4 gemm1+gemm2 ISA BYTE-IDENTICAL to origin/main (all changes behind
const_expr(_is_int4)). Host: test + tuner (2 sites) switched to the OLD prep;
shuffle_weight_a16wi4 removed; stale docstrings fixed.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The rolled scf.for wins at decode/mid M (grid latency-bound) but COSTS ~12% at large M (tok4096, t32x128x256_bnt0 BM=32), where the kernel is HBM-adjacent and wants the unrolled body's ILP/prefetch depth. The tuned CSV uses BM==16 for every a16wi4 decode/mid row and BM>=32 only for large-M rows, so a compile-time const_expr predicate _roll_k = _is_int4 and BM<=16 cleanly selects the regime that benefits. mxfp4/bf16 and large-M int4 keep the fully-unrolled range_constexpr body (mxfp4 gemm1+gemm2 ISA byte-identical to origin/main). ISA: BM16 t16x32x256_kw2 rolled (2220 lines, s_cbranch=2); BM32 t32x128x256_bnt0 unrolled (47413 lines, s_cbranch=1). logits_diff <= 1.53e-5 (tok4 rolled, tok4096 unrolled). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Same compile-time roll gate as the weight-layout branch, applied to the scale-only tree (kpack=16 weight layout, shuffle_weight_a16wi4 retained). Roll wins at decode/mid (BM==16), unroll at large M (BM>=32) where it cost ~12% (tok4096). mxfp4 gemm1 ISA byte-identical to origin/main; logits_diff <= 1.57e-5. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ayout
The tok<16 decode rows were tuned against the kpack=16 layout (full-width dwordx4
B loads). The weight-aligned kernel does 2x dwordx2 half-width loads, moving the
tile optimum. Re-swept tile_n{16,32,64} x tile_k{128,256} x k_wave{1,2,4} x
b_nt{0,2} at tok 1/2/4/8 on the weight-aligned branch (min of 2, dispatch-confirmed):
tok1 t16x16x256_kw4 -> t16x16x256_bnt0_kw4 16.75 (was ~18.2, -8%)
tok2 t16x16x256_kw4 -> t16x16x128_bnt0_kw4 23.34
tok4 t16x32x256_kw2 -> t16x32x128_bnt0_kw4 34.61 (was ~36.4, -5%)
tok8 t16x32x256_kw2 -> t16x32x256_bnt0_kw2 59.75
Pattern: b_nt=0 (L2-cached B) now wins at small M -- the 2x request count of the
half-width loads benefits from L2 reuse, which streaming (b_nt=2) gave up. tok16+
unchanged. CSV-only; no kernel change. logits_diff <= 1.50e-5 on all retuned rows.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Found while retuning: the rolled scf.for produces WRONG results (logits_diff ~6e-5) for the narrow combo k_wave==1 (num_n_waves==4) + num_acc_n==1 + TILE_K>=256 (e.g. t16x64x256), while the UNROLLED path is correct there (1.5e-5) and every other rolled combo is correct (kw2/kw4-tk256, kw1-tk128, kw1-tk256 num_acc_n==2). Root cause is in the rolled loop's iter-arg/A-DMA interaction at this shape. No production CSV row hits it (kw1-tk256 in prod is only BM=32, which is unrolled), but the BM<=16 roll gate would roll it if a future retune picked such a decode row. Exclude the combo from the roll (fall back to the correct unrolled body) -- a correctness guard, not a perf choice. mxfp4 gemm1 ISA still byte-identical to main. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ack)
The tuned CSV stopped at token=8192, so a 16k-token shape fell through to
get_2stage_cfgs' untuned fallback (t{m}x128x128, nt) and only logged a
warning. Sweep at 7168 x {256,512,1024}, E384/k8, gfx950:
t64x128x128_bnt0 2742.9 <- chosen (inter=512)
t64x128x128 2941.7 (= what the fallback picked; -6.1% by tuning)
t64x64x128_bnt0 3557.2
t64x64x256_bnt0 3938.8
t64x128x256_bnt0 5660.1
BM=128 variants 3927-4922 (all worse; BM=64 is right at 16k)
One row per shape (no duplicate-shape collisions), us1/us2 are measured.
Verified: dispatches t64x128x128_bnt0 with no fallback warning,
logits_diff 1.48e-5.
NOTE - this exposes a large-batch deficit vs the kernel being replaced that
is NOT a tuning artifact: at tok16384 the best available config is 1.37x
(inter=512) / 1.53x (inter=1024) slower than the old kernel, though inter=256
is at parity (0.98x). Compare tok4096, which is 1.02x. Tracked separately;
the tuned rows are still a strict improvement over the fallback.
…ack)
The tuned CSV stopped at token=8192, so a 16k-token shape fell through to
get_2stage_cfgs' untuned fallback (t{m}x128x128, nt) and only logged a
warning. Sweep at 7168 x {256,512,1024}, E384/k8, gfx950:
t64x128x128_bnt0 2742.9 <- chosen (inter=512)
t64x128x128 2941.7 (= what the fallback picked; -6.1% by tuning)
t64x64x128_bnt0 3557.2
t64x64x256_bnt0 3938.8
t64x128x256_bnt0 5660.1
BM=128 variants 3927-4922 (all worse; BM=64 is right at 16k)
(For future large-M sweeps, b_nt=2 and tile_n=64 need not be tested: at large
M the weights are reused across many M-blocks so L2-cached b_nt=0 wins, and
tile_n=128 amortises the A-side LDS traffic. Both lost here, as expected.)
One row per shape (no duplicate-shape collisions), us1/us2 are measured.
Verified: dispatches t64x128x128_bnt0 with no fallback warning,
logits_diff 1.48e-5.
At this shape gemm1 is 2743us vs the old kernel's 3970us (0.69x, i.e. the
port is ~31% faster); gemm2 is at parity (1988 vs 2005).
coderfeli
force-pushed
the
flydsl-a16wi4-replace-int4-moe
branch
from
August 10, 2026 03:36
9981990 to
2b71b5c
Compare
token=32768 fell through to the untuned fallback for the same reason 16384 did. Reused the 16k winner rather than re-sweeping, after confirming the choice still holds at 2x the tokens (7168x512, E384/k8, gfx950, us_stage1): t64x128x128_bnt0 5040 <- chosen (same as 16k) t128x128x128_bnt0 8415 t64x128x256_bnt0 10388 Only b_nt=0 x tile_n=128 candidates were tried: at large M the weights are reused across many M-blocks so L2-cached beats streaming, and tile_n=128 amortises the A-side LDS traffic, so nt and tile_n=64 are not worth machine time here. One row per shape, us1/us2 measured. Verified: dispatches t64x128x128_bnt0 with no fallback warning, logits_diff 1.48e-5. vs the old kernel at this shape: gemm1 5031 vs 7368 (0.68x), gemm2 3780 vs 3848 (0.98x).
# Conflicts: # aiter/configs/model_configs/kimik2_i4_tuned_fmoe.csv # aiter/ops/flydsl/kernels/moe_2stage_a16wmix/gemm1.py
The roll/no-roll reasoning was written out in full three times (at the _roll_k predicate and again at both mainloop branches). Keep the canonical explanation at the predicate, make the branches point at it, and tighten the guard comment. Comment-only: mxfp4 gemm1+gemm2 ISA re-verified byte-identical to origin/main, int4 logits_diff 1.48-1.51e-5 at tok 1/4/16384. -10 lines.
57 -> 29 added comment lines; keeps each fact (layout strides, the K32 = two klane slots relation, the scale dword formula, the _roll_bad guard) in one line instead of a paragraph. mxfp4 gemm1+gemm2 ISA byte-identical to origin/main, int4 1.49-1.51e-5.
_bf16_frag / _bf16_frag4 / _mma were byte-identical in gemm1.py and gemm2.py (21 lines duplicated). Moved to utils.mma_bf16 (+ two private frag builders); each kernel now binds it with functools.partial(mma_bf16, mma_atom, use_k16), so the mainloop call sites are unchanged. gemm1 1159->1141, gemm2 820->802, utils 179->204: net -13 lines and one copy of the MMA fragment logic instead of two. Codegen-neutral: mxfp4 gemm1+gemm2 ISA byte-identical to origin/main, a16wfp4 pytest 12/12, int4 logits_diff 1.41-1.57e-5 at tok 1/4/16384.
…lace accs
Replace the hand-emitted scf.ForOp / scf.YieldOp / as_ir_value / as_dsl_value
marshalling with the DSL range(..., init=) loop-carried sugar. My earlier claim that
the sugar 'is unavailable in a plain helper' was half-right: it needs the AST
rewriter, which does NOT process a plain helper -- but DOES process a local
@flyc.jit body. Wrapping the loop in a local @flyc.jit (_run_rolled) fires the sugar
(verified: emits scf.for with iter_args and launches correctly).
Also drop the accumulator round-trip entirely: rmem tensors are memref-backed, so
mutating acc_gate/acc_up in place across scf.for iterations is a side effect, not SSA
dataflow -- the fly-promote-regmem-to-vectorssa pass lifts them to loop iter_args by
itself (verified in the promoted IR). This deletes _acc_to_vals / _vals_to_acc /
_fresh_accs and compute_tile's accs_gate/accs_up params; only the prefetched B-tile
(genuine SSA) is still carried, via _flatten_b_tile/_unflatten_b_tile.
Net -51 lines in gemm1; imports ir/scf/as_dsl_value/as_ir_value removed. Cleanup
only -- decode is perf-neutral (tok1 1.00, tok4 1.06, tok16 0.88 vs old, matching
the pre-refactor curve) and mxfp4 gemm1 ISA is byte-identical to origin/main (the
rolled branch is int4/BM<=16-only; mxfp4 never enters it). logits_diff <= 1.66e-5
across tok{1,2,4,16,128} x inter{256,512,1024}. The BM<=16 gate and _roll_bad guard
STAY: rolling still loses ~3% at BM=32 (tok4096) AND fails correctness at BM=64
(logits 2.1-2.4e-5 > 2e-5) even with in-place accs, so the large-M unroll and both
guards remain necessary (see report).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
CI failed with RuntimeError: FLYDSL_RUNTIME_RUN_ONLY=1 but no usable AOT cache for launch_gemm2 ... aiter.ops.flydsl.kernels.moe_2stage_a16wmix.gemm2 get_flydsl_stage2_kernels_int4_bf16 never registered `b_nt`, so its two readers each fell back to their own default: _flydsl_stage2_wrapper to 0, _precompile_a16w4_to_cache to 2. b_nt is baked into the gemm2 kernel name, so AOT wrote `..._tn128` while the runtime looked up `..._tn128_bcm0` -- the artifact was on disk under a key nobody queried. get_flydsl_stage2_kernels does register b_nt for fp4, which is why only the new int4 path tripped it; stage1's int4 entries already register it, matching the CI log naming only launch_gemm2. No runtime behaviour change: b_nt=0 is what stage2 already ran and what the tuned CSV was measured with. Only the AOT artifact moves. Verified: reproduced the exact RuntimeError by AOT-compiling one int4 stage2 job and re-running it under FLYDSL_RUNTIME_RUN_ONLY=1; after the fix the same sequence compiles `..._bcm0` and loads it. A harness diffing the gemm compile kwargs of both paths over all 88 CSV rows x 2 stages reports 88 divergences before (every stage2 row, b_cache_mod 0 vs 2) and 0 after. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
gemm1 and gemm2 contract the same way -- N-major preshuffled W, per-lane scalar
scale gather, upconvert to v8bf16 MMA fragments, A staged through LDS and read
back on the CK sub-lane map -- and differ only in what N and K mean (stage1:
N=2*inter_dim gate|up, K=model_dim; stage2: N=model_dim, K=inter_dim).
Both files carried their own copy. Compared with comments stripped, these were
already character-identical: load_b_raw / load_b_raw_int4 / load_b_raw_bf16,
load_b_scale / load_b_scale_int4, _a_col_bytes_for_ku. upconvert_b differed only
by a local name (p vs pp), and the W/W-scale layout + buffer-resource setup
feeding them was duplicated too. The copies had ALREADY drifted in their
comments, which is the real cost: a reader can't tell whether they've also
drifted in behaviour.
Extract into utils:
make_b_loader() -> _BLoader W layouts, buffer resources, copy atoms and the
raw/scale/upconvert closures, keyed on w_dtype
(fp4 mxfp4 / int4 / bf16)
make_a_loader() -> _ALoader A-LDS view + v8bf16 fragment reader; `tiles` is
exposed because the global->LDS DMA writes
through the same view
The genuine stage differences become named parameters rather than divergent
copies: make_a_loader takes `swizzle` (gemm2 XOR-swizzles A-LDS, gemm1 is
linear) and `k_grp_base_bytes`/`A_SLOT_BYTES` (gemm1's k_wave group +
double-buffer slot; None for gemm2's single unslotted region).
load_b_scale_int4 now takes scale_expert_base as an argument instead of reaching
for a late-bound enclosing local defined 100 lines below its definition.
Also rename mma_bf16 -> _mma_bf16: every other helper in utils is underscore-
private, including its own _bf16_frag8/_bf16_frag4, and there is no __all__ or
external consumer.
Verified BYTE-EXACT: 168 configs (stage1+stage2 x gfx950/gfx942 x fp4/int4/bf16
x tile_m 16/64 x tile_n x k_wave 1/2/4 x b_nt 0/2 x persist) compiled cold into
a scratch cache before and after; all 118 that produce a kernel have identical
gpu.binary ELF sha256, and the 30 LDS-overflow aborts + 20 gfx942-fp4 compile
rejects are unchanged too (so the refactor can't have silently turned an error
into a kernel). Re-run after each extraction step.
op_tests/test_moe_2stage.py --csv-filter wint4 abf16_wfp4: 88/88, err
1.42e-05..1.62e-05 (threshold 0.01).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Review pass over the non-kernel parts of this PR, keeping only what the a16wi4
port actually needs:
fused_moe.py Restore the original `and is_flydsl_available()` condition
instead of restructuring it into an early raise -- the
no-FlyDSL fall-through is pre-existing behaviour and not
this PR's business. Condense the fallback comment to the
load-bearing facts (tiles come from the tuned CSV; ksplit
is 0 because the port has no grid split-K; block_m must
equal tile_m).
test_moe_2stage.py Revert entirely -- the only change was reflowing the
a16wi4 weight prep into chained .view() calls. The file is
now untouched by this PR, which is a stronger statement of
the drop-in claim: the new port consumes byte-identically
what the deleted kernel did.
moe_reduce.py Drop the "Extracted from moe_gemm_2stage.py" provenance
rather than annotating it as deleted; a pointer to a file
that no longer exists helps nobody, and git history has it.
Note on ksplit (removed from the fallback in an earlier commit): that is
required, not cosmetic. The deleted kernel implemented grid split-K; the
a16w-mix port raises on k_batch != 1 and uses intra-block k_wave instead, so
emitting a `_kb{n}` name and a nonzero MOEMetadata.ksplit would advertise
partials the kernel never writes.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
coderfeli
force-pushed
the
flydsl-a16wi4-replace-int4-moe
branch
from
August 11, 2026 07:24
d757e94 to
d0b22c2
Compare
The int4 registry still emitted the deleted kernel's `_kb{n}` grid-split-K names
"so pre-existing tuned CSVs keep parsing". That is the wrong trade: the a16w-mix
port has no grid split-K (it uses intra-block k_wave), so such a row parsed
fine and then ran WITHOUT the split-K its tuned timing was measured with -- a
silent substitution, and the CSV's own ksplit column would still tell the
framework to expect partials the kernel never writes. A stale row must fail
loudly instead.
- drop the `_kb{n}` names (and the now-always-1 `k_batch` key). A stale name
hits "Invalid FlyDSL kernel name" in _flydsl_stage{1,2}_wrapper.
- forward k_batch from _flydsl_moe_stage1_impl's a16w branch to the port, so
its existing `k_batch != 1` guard actually fires instead of the request
being dropped on the floor:
NotImplementedError: a16w4 gemm1 only supports k_batch=1, got 4
- drop the int4 stage2 `_persist` names for the same reason: the a16w branch
of _flydsl_moe_stage2_impl never forwards `persist` to the port, so they
would silently run non-persist.
No shipped CSV names any of these (checked all of aiter/configs), and every row
of kimik2_i4_tuned_fmoe.csv still resolves; its ksplit column is already 0 for
all 88 rows. Tuning cannot regenerate them either, since the port rejects
k_batch>1.
Pre-existing and NOT touched here: get_flydsl_stage2_kernels registers the same
`_persist` names for fp4, and bf16-A x fp4-W routes through that same a16w
branch, so `flydsl_moe2_abf16_wfp4_..._persist` is silently non-persist on main
too. No CSV uses it.
Verified: parity harness 176/176 compile keys agree; test_moe_2stage.py
--csv-filter wint4 abf16_wfp4 88/88, err 1.44e-05..1.65e-05.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The kernel indexes the scale as shuffle_scale_for_int4 emits it, (E, G//2, N, 2) -- dword = e*(G//2*N) + (group//2)*N + n -- but five comments still described the superseded (E, N, G//2, 2), including the flydsl_a16w4_gemm1 docstring, which is the caller-facing contract and spelled out the replaced formula dword = n*(G//2)+group//2. No code change; the layout itself is byte-identical to the deleted kernel's (verified: the deleted kernel is bit-exact vs the torch ref at E=384/topk=8, and the port fed those same tensors reproduces it to cos=0.999996 / 0.23% L1, i.e. bf16 accumulation-order noise only). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
get_flydsl_stage1_kernels_int4_bf16 enumerated k_wave for every tile_m, so 20 of the 96 registered names request more LDS than a workgroup has: t64x16x256_kw4 wants 4*2*64*256*2 = 256 KB and dies at compile with "local memory (262144) exceeds limit (163840)". Such a name is not merely slow -- the tuner can never time it and the AOT precompile silently drops it. Add _gemm1_lds_bytes(), mirroring the lds_bytes computation in moe_2stage_a16wmix.gemm1 (per-k-wave double-buffered BM x TILE_K bf16 A tile, overlaid with the k_wave>1 slice-K reduce scratch), and skip any config over the budget. 96 -> 76 names; every kernel named by the tuned CSV survives. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Tuning a16wi4 produced no row for any shape -- "stage1 and stage2 should
be valid together" -> "no valid candidate found". Three separate causes,
all a16wi4-only:
1. gen_flydsl_i4_2stages_task mapped bf16 -> "fp16" in _a_dtype_map, but
a_dtype_str is what selects the sorted bf16 intermediate
("a2_a16w_sorted") for the stage2 candidates. Stage2 got the packed
a2_qt instead and every candidate raised
"requires D_INTER (K) % 64 == 0, got D_INTER=8".
gen_flydsl_2stages_task already carries a comment warning about
exactly this mapping.
2. a2_a16w_sorted was only built for fp4x2, and a redundant earlier
i4x2 branch shadowed the a16w-mix branch, so it stayed None. Fold
i4x2 into the a16w-mix branch (same a2_qt/a2_scale handling, plus
the sorted build).
3. Stage1 was scored with the unsorted run_torch_moe_stage1 while the
port returns a sorted [max_sorted, inter] intermediate, so every
candidate scored ~84% err. Use run_a16w_stage1_sorted_ref +
_a16w_sorted_cos, as the mxfp4 a16w4 arm already does. (The deleted
int4 kernel returned unsorted output, which is why the old reference
was correct for it and became wrong at the port.)
After this, all 88 kimi a16wi4 shapes tune with err1 0.0%.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The committed metrics described the deleted kernel, not the kernels the rows name: 67 rows had kernelName1 rewritten while us1/us/err1/tflops/bw stayed byte-identical to the pre-port values (tok=1 inter=512 kept us1=18.1998 across a _kb7 -> _kw4 swap), and the six rows added for token 16384/32768 carried tflops/bw copied verbatim from their 8192 siblings -- impossible at 2x and 4x the tokens. The cause was not carelessness: the tuner could not emit an a16wi4 row at all until the previous commit. Re-tuned all 88 shapes on gfx950 (8x MI355, --errRatio 0.5), 0 failures, and merged keep-better so a shape is only rewritten when the fresh run is faster. 87 rows updated, geomean 0.861 (-14%), best cells -20..-31%. The six fabricated rows are rewritten unconditionally -- they had no valid baseline to keep. The hand-transcribed names were often not the best config either: most low-M rows move off t16x32x256_kw2 onto t16x32x128_kw4 / t16x64x128. One shape is a genuine regression and keeps its old row: tok=1 inter=256 E384 measures 27.37 us on the port vs 25.95 for the split-K kernel it replaces (+5.4%). Rows sorted by inter_dim, expert, topk, token. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 task done
…tch) Review follow-ups on the shared a16w4/a16wi4/a16w16 operand paths: * A path: dma_x_tile_to_lds (gemm1) and dma_a_tile_to_lds (gemm2) were the same kernel twice -- same per-thread tile decomposition, same 16 B copy, same gfx942 VGPR-staging fallback. Folded into make_a_loader as store_tile(), so the A loader owns both halves of a layout the two must agree on. The stages differ only via row_base_dwords / a_num_bytes / swizzle / slot params. Drops gemm1's dead _a16w4_swizzle_xor16(enable=False) call, which read as if stage1 swizzled. * B path: _BLoader's 7 dtype-shaped fields (two of them never called directly, two None-by-dtype) collapse to col/col_pair/load_raw/load_scale/upconvert, all specialized for w_dtype inside the loader. The N addressing both stages open-coded identically -- row//16, row%16, ng//32, (ng//16)%2 and the readfirstlane'd int4 scale base -- moves in behind a _BCol descriptor. gemm1's load_b_tile keeps only what is stage1's business (pairing gate|up into one value for the software pipeline); the dtype dispatch it used to repeat is gone, as is gemm2's copy of it. * act: gemm1 no longer branches on the act string or threads five loose f32 situv2 scalars through four layers; act.py gains SituParams/situ_params and gate_up_act, and its helpers are public now that they are shared. Net -392/+349 with no weight-dtype branches left in either stage body. Codegen: verified by disassembling the compiled artifact (llvm-objdump) for 38 tile configs across fp4/int4/bf16 x silu/situv2 x k_wave 1/2/4 x b_nt x xcd_swizzle, vs the same matrix built from the parent commit. Every tile shape the tuned CSVs use is unchanged to within +-1 scalar setup instruction (most are bit-identical; the rest differ only in SGPR numbering). Two synthetic BM=64 sweep tiles that are already register-starved move by 0.1-0.2% in accvgpr/scratch traffic. MFMA, ds_read, buffer_load and cvt counts are identical everywhere. Tests: op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py (6/6), and a16wi4 token=1/128 e2e at 7168x256/E384 -- logits_diff 1.6e-5, 76.5->76.8 us and 254.3->254.1 us (run-to-run noise; the harness does not seed its inputs). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Emitting the readfirstlane'd e*(G//2*N_OUT) term up front in make_b_loader, rather than where the stages used to compute it, was enough to shift downstream scalar scheduling: stage2's int4 kernel came out with the same instruction mix in a slightly different order. Deferring it to first use restores gemm2 int4 to a bit-identical disassembly (t16x128x128: 0/970 opcode lines differ, was 12). It stays loop-invariant either way -- one readfirstlane, CSE'd across every scale gather in the kernel. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
6 tasks done
Contributor
|
Validated the replacement int4 path on gfx942 with vLLM#51274. The strict production-shape comparison passed at E=112, top-k=15, 3584x3072, and tokens 1, 4, 7, 16, and 32. Target-only TP8 with expert parallelism, graph capture, and chat sanity passed. Full five-shot GSM8K scored 0.960576 on all 1,319 questions. Can this receive final review and be included in the next AITER release? vLLM#51274 now depends on this replacement entry point. |
3 of 13 tasks
Collaborator
Author
Merged. @maeehart It will be included. |
This was referenced Aug 18, 2026
samremes
pushed a commit
that referenced
this pull request
Aug 27, 2026
…#5017) * [FlyDSL] gfx942 a16wi4: pack f32->bf16 with lshr-16 instead of scalar truncf v_cvt_pk_bf16_f32 is gfx950-only. After #4646 the gfx942 int4 fallback used f32.to(bf16)/truncf, which is much more VALU than the old moe_gemm_2stage shift-pack. Same nibble order; gfx950 packed convert and MXFP4 are unchanged. * [FlyDSL] Clarify gfx942 a16wi4 upconvert comments
JiaoliangYu
added a commit
that referenced
this pull request
Aug 28, 2026
…error check (#5075) * [FlyDSL] gfx942 a16wi4: pack f32->bf16 with lshr-16 instead of scalar (#5017) * [FlyDSL] gfx942 a16wi4: pack f32->bf16 with lshr-16 instead of scalar truncf v_cvt_pk_bf16_f32 is gfx950-only. After #4646 the gfx942 int4 fallback used f32.to(bf16)/truncf, which is much more VALU than the old moe_gemm_2stage shift-pack. Same nibble order; gfx950 packed convert and MXFP4 are unchanged. * [FlyDSL] Clarify gfx942 a16wi4 upconvert comments * ci: allow multigpu label to trigger tests (#5008) * [HIP] [CK] [MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE (#4620) * [MoE] Added Gelu with tanh approx for CK XDL 2-stage MoE * applied copilot's comment for str2ActivationType * Dropping cross-activation CK configs * Applied Ying comment * Added block for run_1stage for unsuported activations * fixed test * [Triton] Move attention configs to nested layout and unify their resolution (#5019) Relocate 14 attention config files from the flat arch-prefixed layout to configs/<arch>/triton/attention/<d_type>/DEFAULT.json - mha, extend_attention, mla_decode_rope, hstu_attn_fwd and hstu_attn_bwd - retiring configs/hstu_attn/. The redundant -DEFAULT suffix is dropped from directory names (the file is already DEFAULT.json), matching the chunk_delta_attn precedent. The six reader modules resolve through the shared resolve_config_dir("attention", ...) probe instead of hand-built paths. LEANATTN is not migrated: upstream removed the lean_atten kernel and its config. * [Triton] Migrate the GMM tuned configs to the nested layout (#5020) Move configs/<arch>-GMM.json (gfx942, gfx950, gfx1250) to configs/<arch>/triton/gmm/gmm/DEFAULT.json and point the reader at it. GMM gets its own op directory instead of folding under gemm/. The doubled gmm/gmm is just the <op>/<d_type> layout: the op is "gmm" and the family's config name is "GMM", so _dtype_dir() yields "gmm" too. _triton_kernels/gmm.py now resolves the directory through the shared resolve_config_dir("gmm", "GMM", backend="triton") probe and loads DEFAULT.json from it; arch_info and AITER_TRITON_CONFIGS_PATH are dead there and are dropped. No legacy_dir is passed - the files move and the loader flips in this one commit, so every revision resolves. * [Triton] Move MOE tuned configs to the nested layout (#5022) Move the three remaining MOE tuned configs from the flat configs/moe/ directory into configs/<arch>/<backend>/<op>/<d_type>/: moe/gfx950-A8W4.json -> gfx950/triton/moe/a8w4/DEFAULT.json moe/gfx1250-A8W4.json -> gfx1250/gluon/moe/a8w4/DEFAULT.json moe/gfx1250-A4W4.json -> gfx1250/gluon/moe/a4w4/DEFAULT.json The backend directory follows the dispatch path the table actually feeds, not the arch: gfx950's a8w4 table is keyed bm<block_m>_n<N>_k<K> and is read by the Triton path, while both gfx1250 tables are bucket-keyed and read by the Gluon path. So the a8w4 family spans backends and a4w4 is Gluon-only. These three are all that is left of configs/moe/: PR #4833 removed the rest of the legacy MOE stack (utils/moe_config_utils.py, the fused sigmoid-top1 routing kernel, the moe_op/moe_op_e2e/mxfp4 variants and every configs/moe/*-MOE-*.json), so this completes the directory. The two surviving loaders are rewired onto the shared probe in the same commit. _get_a8w4_dispatch() and _get_a4w4_dispatch() now resolve their directory with resolve_config_dir("moe", "<A8W4|A4W4>") and read DEFAULT.json from it, instead of hand-building an arch-prefixed path under configs/moe/. Neither call passes backend=: because the backend differs per arch for the same family, pinning one would make the other arch's file unreachable. The documented probe order -- nested triton, then nested gluon -- picks whichever directory the running arch ships. a4w4 also moves off its private os.path.exists + json.load pair onto load_config_json(..., required=False), matching a8w4; both still return {} when no tuned file is shipped for the arch, so the safe-default fallback paths are unchanged. resolve_config_dir() lives in utils/gemm_config_utils.py and is added by the config-unification branch -- merge that one first. * [Doc][Skill] port udpated flydsl kernel code cleanup skill (#5051) * [Triton/Gluon] MOE a8w4 cudagraph updates (#5037) * [Triton] Remove legacy flat-layout support from config resolution (#4948) * [Triton/Gluon] Move gluon gemm_a8w8 kernel into _gluon_kernels/gfx950 (#4866) * [Triton] Migrate conv configs to the nested arch/backend layout (#5018) Move all 59 flat configs/conv/<arch>-<CONFIG_NAME>.json files to configs/<arch>/triton/conv/<d_type>/DEFAULT.json, the layout GEMM already uses, and point _conv_config_path() at the shared resolve_config_dir() probe instead of building the legacy path by hand. This picks up the ten tables #4869 added (CONV-PREPACK on all seven arches, CONV-3X3-NCHW on gfx1100/gfx1151/gfx1201) alongside the original 49. The renames and the loader flip land in one commit so every revision resolves conv configs from exactly one layout: no legacy_dir fallback is needed and bisect stays clean. File contents are untouched (pure renames). _conv_config_path() is the single choke point, so get_conv_config(), has_conv_config(), conv_config_uses_exact_routes() and has_exact_conv_config() all pick up the nested path; the variant-aware four-tier walk, STANDARD_M_BOUNDS and the lru_caches are untouched. * Tune MoE GEMM A8W8 blockscale (#5028) * [Triton] Migrate MHC configs to the nested arch/backend layout (#5021) Move all 15 flat configs/<arch>-MHC_*.json files to configs/<arch>/triton/mhc/<d_type>/, keeping the C=<n> specialized file stems and naming each family default DEFAULT.json, and rewire mhc_config_utils onto the shared resolve_config_dir() probe. The documented gfx942 fallback retry resolves through the probe's arch= override (added by the legacy-removal PR, which merges first); the C-bucket walk and _FALLBACK_DEV semantics are unchanged. The renames and the loader flip land in one commit so every revision resolves MHC configs from exactly one layout, and this branch touches no shared resolver code. * [Gluon] add bench for mxfp8 GEMM (#5029) * [FlyDSL] One-stage split-K for the a8w8 preshuffle GEMM (#5007) * [FlyDSL] one-stage split-K for the a8w8 preshuffle GEMM Fold the split-K reduction into the GEMM launch: every split publishes an fp32 partial, and the last one to arrive at the tile's semaphore reduces and converts in the same kernel, so split-K costs one launch rather than two. The partials cross CTAs that may sit on different XCDs, each with its own L2, so they have to reach a common point. Doing that with an agent-scope fence costs a whole-L2 buffer_wbl2 per CTA plus a buffer_inv on the reader, which also evicts the A/B tiles every other in-flight CTA is still reading -- measured 2-3x slower than not splitting at all. Marking just these accesses sc0|sc1 writes them through and leaves L2 alone. That turns split-K from a 18-180% regression into a 9-50% win over k_split=1. k_split == 1 is untouched: same kernel, same cached stores, and its output is bitwise identical to main across the shapes checked. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] drop the k_split == 2 special case It had its own epilogue: whichever split arrived first published its fp32 fragment, and the second spun on a ready flag, kept its own fragment in VGPRs and wrote the final tile -- saving one workspace plane and one round trip. It does not pay for itself. Spinning is slower than just going through the generic path: 4.7 vs 5.8 us at 1x576, 18-27% across the six shapes measured. Removing it also drops two fragments, two copy atoms, the doubled semaphore, and the split-plane special cases in the launcher and the AOT pre-compile. k_split == 2 now takes the same path as every other split count, which also fixes the per-split workspace offset: it was guarded on split_k > 2, so a k_split == 2 launch routed through the generic path would have had every split write the same plane. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] fix the split-K keyword the tuner passes to the launcher The tuner called the preshuffle launcher with k_split=, but that launcher names the argument split_k= (matching the hgemm split-K path it sits next to). Every flydsl preshuffle candidate therefore raised TypeError. The tuner records a raising candidate as rejected rather than as an error, so the run completed, kept only the 8wave candidates, and picked ck or cktile for four shapes that flydsl had previously won -- a result indistinguishable from a legitimate tuning outcome. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] trim the split-K diff to what split-K needs Three things that were not split-K: out_dtype grew an fp32 branch and a raise. Nothing needs it -- fp32 is the type of the *partial*, which "Float32 if split_k > 1" already covers, and the final output is still bf16/fp16. The bias element-type change existed only to feed that branch. Both are back to main's two-case form. The K-tile index has to gain a bid_z offset, which is genuine, but the name k_tile_base pushed several one-line fx.copy calls past the line limit and a trailing comma pinned others open, so a one-token change read as +5 -1. Renaming to k_off and dropping the magic trailing commas keeps them one-liners. The copy atom for the output no longer branches on out_elem_bytes; it picks the op from split_k directly. Kernel diff: +201 -22 -> +168 -21, with no behaviour change. k_split == 1 still compiles to bytes identical to main. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [Config] retune Kimi-K3 a8w8 M<=32 with one-stage split-K 29 rows, all M <= 32: 21 stay on flydsl with a better config and 8 move from ck to flydsl. 20 of them use split-K, mostly k_split=7. The tuner proposed 37 rows. Each changed row was then re-measured old config against new on an idle GPU, and the 8 that were actually slower there were kept at main's value -- the tuner picks its winner while four GPUs are saturated, and for shapes where several configs sit within noise of each other that choice does not survive on an idle card. Nearly all of them were k_split=2 at N=6400, which lost 7-13%. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] hoist the partial store out of the split_k branch Both arms opened with the same fx.copy; only what follows it differs. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] address review: semaphore dtype, buffer lifetime, reduce guard The k_split == 1 path passed an empty bf16/fp16 tensor in the semaphore slot while the AOT pre-compile passed an empty int32 one. dtype is part of FlyDSL's executable cache signature, so every non-split-K preshuffle kernel missed its AOT entry and JIT-compiled at first call -- a regression across all existing tuned configs, not just split-K. Both sides now pass int32. The split-K buffers were cached per (m, n, tile, k_split). m is in the key, so a server sweeping batch sizes grows the cache without bound, and an eviction frees memory whose address a captured CUDA graph still holds. k_split_candidates only proposes split-K while the tile grid is under one CTA per CU and caps k_split * tile_count at four per CU, which bounds tile_count below CU_NUM and the workspace at 4 * CU_NUM * tile_m * tile_n floats -- so the buffers are now fixed-size and keyed on (device, stream) only, the way _get_split_k_tensors already does it, with a capacity check for anything that would exceed the bound. The reduce derives its vector count as tile_n // 4 and would have dropped the tail columns for a tile_n that is not a multiple of 4; the comment claimed the invariant but nothing enforced it. Now rejected at compile time. The semaphore reset was a plain cached store while every other cross-CTA access in that block carries sc0|sc1. It is the same cross-XCD handoff, one launch later, so it writes through too. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] guard M against the layout bound the kernel assumes The kernel views A and C through layouts with a hardcoded 65536 rows, in three places, with nothing on the host stopping a larger M from indexing past them. Named the bound, used it at all three sites, and rejected an out-of-range M in the launcher with a message that says why. gemm_kernels keeps its own copy of the literal because that module has to import without FlyDSL present; a test asserts the two agree. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] cut the comments back to what the code cannot say 57 added comment lines down to 26. Dropped the ones restating the line below them -- what k_off is, which path split_k > 1 takes, that the partial store publishes a partial -- and shortened the rest. What is left is the reasoning that is not recoverable from the code: why the partials cannot use an agent-scope fence, why the buffers are fixed size rather than shape-keyed, why the semaphore dtype has to match the AOT side, why _REDUCE_VEC is 4, and why the k_split candidates are enumerated per shape. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * Extract the flydsl split-K reduce into a reusable copy-atom epilogue Move the one-stage split-K reduction out of preshuffle_gemm into splitk_epilogue.splitk_reduce_epilogue, with the output element class as its only dtype knob so other GEMMs can reuse it. The reduce now goes through copy atoms and a buffer-tensor descriptor instead of raw buffer_ops: make_layout_tv gives each thread 4 contiguous columns, so the loads stay dwordx4 and the stores dwordx2, and the descriptor bounds cover the ragged-M tail. Resetting the semaphore with atomic_add(-split_k) rather than a plain store also drops a next-launch increment race. Verified on gfx950: rel_err matches k_split=1 for M in {1,8,64} x k_split in {1,2,7,14}, CUDA-graph replays clean, and the k_split=1 output hashes are identical to origin/main. Over the 20 tuned Kimi-K3 split-K shapes the reduce is 0-3% faster than the buffer_ops version. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * Let --splitK gate the FlyDSL candidates as it does the other backends The ck, cktile and asm task builders all take useSplitK and collapse the split-K dimension to a single splitK=0 candidate when it is off. The FlyDSL builder never received the flag, so it generated k_split candidates unconditionally: on the Kimi-K3 shape set that is 98488 extra candidates on top of 99136, roughly double overall and 2.5-3.1x over M in 1..128. It also made the flag useless as a switch. Split-K wins often enough at small M that 20 of the 176 tuned rows are split-K winners, so a run without --splitK still produced split-K configs. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> * [FlyDSL] Retune GLM5.2 mxfp4 MoE and fix a scale-view cache leak (#5045) * [FlyDSL] Retune GLM5.2 mxfp4 MoE and fix a scale-view cache leak Retune all 64 GLM5.2 shapes (model_dim=6144, inter_dim 256..2048, E=257, topk=9) for gfx950. 27 shapes move to the coupled flydsl_mxmoe port, which the previous config only reached on 5 rows. Measured through the production fused_moe path, each shape timed on one GPU under both configs: median +6.7%, mean +8.6%; 44/64 faster by >1%, 6 slower (worst -3.5%). Small batches gain most (token<=64 median +11.4%). Six shapes (2048/256, 4/1024, and 2/16/64/128 at 2048) are left on their existing main entries rather than retuned. Two fixes fell out of the tuning runs: _mxfp4_scale_u8 was wrapped in lru_cache(maxsize=2048). Its body is a bare .view(), so the memo buys nothing, but tensors hash by identity: every per-call intermediate scale misses and is then pinned by the cache. A tuning sweep leaked ~0.75 GiB per timed iteration and exhausted a 288 GiB card. v2_stage1_dequant_cosine_err looped per sorted row, costing one .item() sync each -- ~295k syncs per timed candidate at token=32768/topk=9. Now batched in chunks, which bounds the int64 gather in mxfp4_to_f32 while keeping the equal-weight average over rows. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * [FlyDSL] Default FMoE tuning to FlyDSL v2 and update GLM5 FP4 layout configs * fix black test * Emit the non-f4out AOT job for _f4out mxmoe stage-2 rows An `_f4out` GEMM2 kernel only really runs the mxfp4-out path when both gates are open: AITER_MXFP4_INTERMEDIATE, and the shape check in fused_moe (`D_HIDDEN == 7168`). Otherwise `_f4out` is stripped from the kernel name and the plain kernel launches instead. The AOT generator skipped such rows outright, so it never pre-compiled the kernel that actually launches. GLM5 is D_HIDDEN=6144, so every `_f4out` row there falls back -- and the config only survived because an unrelated row happened to name the plain kernel and seed the same cache entry. Retuning that row to `_f4out` removed the last such seed and CI hit `FLYDSL_RUNTIME_RUN_ONLY=1 but no usable AOT cache for launch_gemm2` on token=16384, inter_dim=512. Emit the fallback job unconditionally, plus the f4out one when AITER_MXFP4_INTERMEDIATE is set. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: charlieguo1106 <cguo@amd.com> * [UT] Support a4w4 in test_mega_moe (#5052) * support a4w4 in test_mega_moe_gfx1250 * support 64K * [Triton/Gluon] [ASM] [HIP] add mla v4 prefill asm kernel (#4926) * Add MLA v4 sparse prefill asm support Integrate the gfx1250 MLA implementation and consolidate sparse prefill correctness and performance tests. Co-authored-by: Cursor <cursoragent@cursor.com> * Update op_tests/test_pa_sparse_prefill.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * Apply black formatting to test_pa_sparse_prefill Pure reformat, no behaviour change. Fixes the failing black CI job (black[colorama]==26.5.1, default line length). * Fix asm candidate reference in test_pa_sparse_prefill The asm candidate passed split["ref"] -- the raw input dict -- where checkAllclose expects the reference tensor, so the first asm comparison died with: TypeError: isclose(): argument 'other' (position 2) must be Tensor, not dict meaning the asm path could never run. Compute the fp8 reference the same way the opus fp8 candidate above it does. * Report per-row nnz and default the CLI to the asm comparison sweep Two test-driver changes: * nnz_prefix/nnz_extend columns now report per-row nnz instead of the pool-wide total, so they match the --nnz-prefix/--nnz-extend asked for rather than scaling with N. total_nnz still carries the full count -- the TFLOPS/TB-s figures need the real work done. * CLI defaults now describe the three-way opus/triton/asm comparison out of the box: N in [512, 1024, 2048, 4096] x nnz_prefix in [256, 1024, 4096, 8192, 16384] x nnz_extend 128, at H_Q=128 fp8 (the only shape the asm candidate registers for). --mode/--total_pages default empty so the unrelated mode sweep stays off unless asked for. Every flag still overrides. Pytest coverage is unaffected: it reads _PYTEST_SHAPES/_PYTEST_MODES, not argparse. * Accept an over-allocated CSR indptr in mla_sparse_prefill check_csr required indptr->numel() == T+1 exactly. Decode reuses this kernel with the extend region empty and sizes its CSR row-pointer buffers once at [max_batch+1], launching with the live batch, so numel > T+1 is the normal case there rather than a mistake -- and the exact test rejected it outright. The kernel reads indptr[0..T] and nothing past it, so the extra tail is inert: verified bit-identical output against the exactly-sized call. An undersized indptr is still rejected. Trade-off: an indptr built for a different T is no longer caught here. Separating that from the legitimate case needs device data (indptr[T] against the indices length), i.e. a sync per call. Callers that can slice to [:T+1] should. * Fix int32 overflow in sparse prefill query offset `_sparse_attn_prefill_kernel` derived `query_idx` from `tl.program_id(0)`, which Triton types as int32. The q/out addresses are computed as `query_idx * q_stride_t` and `query_idx * out_stride_t`, and in the V4 layout that stride is `num_heads * head_dim` = 128 * 512 = 65536. The product therefore leaves the int32 positive range at `query_idx >= 32768` and wraps to a negative offset, so the kernel reads and writes outside the q/out allocations. Observed as NaNs followed by a hard GPU page fault: Memory access fault by GPU node-2 ... Reason: Page not present Verified on gfx1250 with a fixed-pattern sparse prefill case (H=128, D=512, pool=16384, nnz_prefix=256, nnz_extend=128): N=32768 before: clean (largest size that still fits int32) N=32769 before: fault after: nan=0 inf=0 N=65536 before: fault after: nan=0 inf=0 Promoting `query_idx` to int64 moves both offsets to 64-bit address arithmetic. This mirrors the existing `slot_off` cast a few lines below, which already handles the same class of overflow on the pool index; the difference is that the wrapped pool offset stays inside the allocation and reads silently, while this one faults. --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> * [HIP] FIX MLA the nhead fold error for cp round robin (#4964) * fix the nhead fold error for cp round robin * fix the split test * support varlen * gqa96 qseqlen<6 not fold * [HIP] [Bugfix] Fix DSV4 FP4 KV-cache scattered row writes (#5034) * fix(dsv4): scatter FP4 KV cache by row-local offset Signed-off-by: AMD-yanfeiwang <yanfei.wang@amd.com> * test(dsv4): remove specialized KV-cache regression Keep the bug fix focused without carrying a narrow special-case test. --------- Signed-off-by: AMD-yanfeiwang <yanfei.wang@amd.com> * [ASM] [HIP] 1x32 mxfp4 asm kernel (#4890) * 1x32 mxfp4 asm kernel * Update tuned config * Upate 1x32 kernel to embedd X quant * Drop the standalone MXFP4 X quant pre-pass plumbing The FLAT MXFP4 kernels dynamic-quantize X in-kernel, so the host-side pre-pass entrypoint, its Python binding and the test helper have no caller left. Removing them also restores the per_1x32 scale-sorting condition, which still tested a pre-pass flag that no longer exists. Co-authored-by: Cursor <cursoragent@cursor.com> * Tune 1x32 kernel * Fix 1x32 race condition for O buffer clearning * SImplify zero protocall and bind it to TG0 always * Fix lm_eval utter failure with 1x32 kernel --------- Co-authored-by: Sergey Solo <ssolovye@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> * [HIP] update MHA CPP reademe (#4874) * update the supported arguments configuration * update the perf data * update the image * fix the log * fix * benchmark_fwd support opus kernel * add opus perf data * perf data * fix the comment * fix * Test FFM bringup on MI250 build runner (#5071) * perf(gfx1250): drop a16w16's 4 GiB pre-check, fail on wrong results The bench skipped any a16w16 shape whose largest operand passed 4 GiB, on the stated grounds that "the heuristic refuses" it. That reads the guard too broadly. opus_dispatch_a16w16_gfx1250 (opus_gemm_arch_gfx1250.cuh:150-183) searches the tuned table first and returns on a hit. check_shape_4g runs only after that misses, on the way to the split-K heuristic kid, and it is that kid's launcher that builds the 32-bit gmem descriptors. A tuned 4wave_wl_co winner never reaches the check: per gen_instances_gfx1250.py:770-778 the pipeline "builds no gmem descriptor at all" and clamps every dimension through TDM descriptors instead. So 4 GiB bounds one fallback path, not a16w16, and a pre-check in Python keeps skipping shapes that tuning has already made runnable. Removed; the kernel raises if it must, and the exception is recorded as a row. The 20260827 sweep shows what the fallback costs. At N=129280 the tuned 4wave_wl_co kid does M=512 in 449us (2112 TFLOPS); M=1024 has no tuned winner, drops to split-K, and takes 3343us (568 TFLOPS) -- 7.4x slower for 2x the work. 11 of 60 shapes hit a 4wave_wl_co kid; the rest are split-K, so most of the low numbers in this table measure tuning coverage rather than the hardware. Widening that coverage is a job for csrc/gemm_a16w16/gemm_a16w16_tune.py --libtype opus, not for this file. Worse, split-K is not just slow at the top of the range: all four M=65536 shapes came back err=0.96-0.99, an unrelated result, while every other row was 0 or ~1e-5. None of them trip the 4 GiB guard (M*K*2 = 896 MB, M*N*2 <= 256 MB), and the UT neither raises nor warns -- it returns the ratio and prints a number. The sweep reported them as data. a16w16 now checks the returned ratio against _A16W16_MAX_ERR and calls _note_failure, so a silent miscompare shows up in the failed-op list. a16w16 also gets its own M list. The global sweep jumps 2048 -> 65536, so the prefill chunk sizes were never measured on the BF16 linears; _A16W16_MS adds 4096/8192/16384 and AITER_BENCH_TOKENS still overrides it. The lm_head cap stays. It is a statement about what DSv4 runs -- one row per sequence -- not about what the kernel can do, and its comment no longer leans on the 4 GiB number. Separately, put a8w8_blockscale back in --dsv4 and correct its note. The note blamed #4773's gluon tuning rows for the make_llir crash. The real cause is the UT's extra "ck strided x_scale" check (test_gemm_a8w8_blockscale.py:120), added by #4406 and gated on ck_preshuffle alone. The mxfp8_128 path declares its layout with is_x_scale_transposed=True and never reads the stride, so a strided x_scale tests nothing there and only gives triton a specialization that fails to compile. A/B with that line as the only variable, over a 162-case matrix (27 default M x six (n,k)): case 2 before it dies, case 160 after -- M=16 and M=64 included, which is what #4773 covers. Fixing it properly is upstream's call; meanwhile _A8W8_BLOCKSCALE_TOKENS starts at 1024, clear of the M that reach those rows. Verified on gfx1250-atom--20260827-ubench: 36/36 cases, err=0, 2207-7003 TFLOPS. --------- Signed-off-by: AMD-yanfeiwang <yanfei.wang@amd.com> Co-authored-by: msaffari-amd <msaffari@amd.com> Co-authored-by: Xin Huang <Xin.Huang@amd.com> Co-authored-by: Alexandra Sidorova <asidorov@amd.com> Co-authored-by: Satya Nikhil Kodukula <nikhil.kodukula@gmail.com> Co-authored-by: Felix Li <felix.li@amd.com> Co-authored-by: Lukasz Burzawa <lukasz.burzawa@amd.com> Co-authored-by: Vinayak Gokhale <vinayak.gokhale@amd.com> Co-authored-by: Nidal Danial <81209936+nidal567@users.noreply.github.com> Co-authored-by: Shao-Chun Lee <Shao-Chun.Lee@amd.com> Co-authored-by: XiaobingZhang <xiaobingzhangupc@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: charlieguo1106 <cguo@amd.com> Co-authored-by: yanboshao <yashao@amd.com> Co-authored-by: junxiaguo <JunXia.Guo@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: minmengdie <memin@amd.com> Co-authored-by: AMD-yanfeiwang <yanfei.wang@amd.com> Co-authored-by: Sergey Solovyev <sergey.solovyev@amd.com> Co-authored-by: Sergey Solo <ssolovye@amd.com> Co-authored-by: Yu <jiaolyu@amd.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
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.
FlyDSL: port a16wi4 to new pipeline, clean old moe_gemm_2stage