MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel (replaces old a16w4) - #4502
Open
coderfeli wants to merge 42 commits into
Open
MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel (replaces old a16w4)#4502coderfeli wants to merge 42 commits into
coderfeli wants to merge 42 commits into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
Replace the a16w4 SiTUv2 SEPARATED fused-MoE path with a numerically-correct
ported FlyDSL 2-stage kernel (aiter/ops/flydsl/kernels/moe_2stage_a16wmix).
aiter's previous compile_mixed_moe_gemm{1,2}_a16w4 fails aiter's own strict
accuracy gate on this path (logits_diff ~1.0); the port passes at ~1.5e-5.
- Vendored kernel package: gemm1 (gate+up GEMM + SiTUv2 -> sorted bf16
intermediate) and gemm2 (down-proj + routing-weighted atomic bf16 scatter),
non-scaled MFMA(16,16,32,bf16), CDNA (gfx942/gfx950).
- fused_moe_ a16w4-SiTUv2-SEPARATED branch: converts the incoming guinterleave
(a8w4/mxfp8-shared) weight/scale layout to the ported kernel's standard
shuffle_weight/e8m0 layout (pure byte permutations, verified round-trip) and
routes to the port. Gated + additive; a8w4/mxfp8/INTERLEAVE paths untouched.
- Focused test: op_tests/flydsl_tests/test_a16w4_flydsl_situv2.py.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Step 1 (perf, CSV-driven, --kernel isolated, balanced routing, median-of-3, gfx950) confirmed the ported kernel is on-par-or-better than aiter's old a16w4 across the throughput band: OURS wins the entire M>=16 range on both Kimi shapes (0.83-0.99x), with the old kernel ahead only on the tok<=8 fixed-launch-overhead points and the single 3584x512 tok16384 cell. Given that (and that the old kernel is numerically broken, 0/N on the strict gate), remove aiter's a16w4 path entirely so the port is the sole a16w4 kernel: - mixed_moe_gemm_2stage.py: drop compile_mixed_moe_gemm1_a16w4, compile_mixed_moe_gemm2_a16w4, the compat wrapper compile_a16w4_moe_gemm2, and the a16w4-only _decode_e8m0_byte_to_f32 helper (~4076 lines). Generic compile_mixed_moe_gemm1/2 and all other dtypes (a8w4/int4/mxfp8/a4w4) are untouched; now-unused imports pruned. - moe_kernels.py: remove the a16w4 branch of compile_flydsl_moe_stage1/2; it now raises NotImplementedError (the port is dispatched directly from fused_moe_). - fused_moe.py: broaden the a16w4 dispatch to serve bf16-a16w4 SiTUv2 in both gate modes (the kernel is gate-mode agnostic; only activation-dtype routing decides bf16 vs fp8), and remove the dead a16w4 sub-clause / _a_type bf16-fp16 branches from get_2stage_cfgs' use_mxfp4_flydsl path. - Fallback handling: features the port does not implement (non-default SiTUv2 beta/linear_beta, per-expert bias, expert-parallel masking) now raise a clear NotImplementedError instead of routing to the numerically-broken old kernel (explicit-unsupported over silently-wrong). - Remove op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py (tested the removed low-level a16w4 kernel API); the production path is covered by test_a16w4_flydsl_situv2.py. Correctness gate re-run (aiter test_fmoe, a16w4 SiTUv2 SEPARATED, strict, gfx950, no-bias supported path): PASS at logits_diff ~1.5e-5 across the token sweep on both shapes. a8w4 FlyDSL tests still pass (shared mixed_moe path intact). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…riginal test - Port FlyDSL PR #948 consolidation (commit 3d1e03d6) into the vendored package: common.py + host.py removed; shared helpers folded into gemm1.py, launch glue + CSV resolvers + a16wi4_recommend_block_m into __init__.py. Carried over the buffer_ops modernization (scale-load -> make_buffer_tensor/BufferCopy32b). aiter adaptations preserved (kernels.common.* -> aiter.ops.flydsl.kernels.*, the CDNA3-encoded s_waitcnt shim). Also dropped the unused vendored _gemm_common_utils.py (never imported). Public API (flydsl_a16w4_gemm1/2) unchanged. - Condense moe_2stage_a16w4_dispatch.py to the essential dispatch (guinterleave-> standard-shuffle inversion, CSV gemm2 tiles); keep load-bearing notes. - gemm2 now takes aiter's tuned CSV tiles via resolve_a16w4_gemm2_config (on-par-to-faster: 0.96-1.02x vs kernel default). gemm1 keeps its own tuned default (tile_n=128 + M>=16 TILE_K=128/xcd lever): perf-gate showed the CSV gemm1 tiles regress mid/high-M by 1.09x-2.02x, so they are NOT taken. - Route aiter's ORIGINAL a16w4 test through the new kernel: _effective_gate_mode returns SEPARATED for bf16-a16w4 (was INTERLEAVE, which downgraded abf16_wfp4 rows to a8w4/fp8 and never exercised the a16w4 kernel), and the a16w4 SiTUv2 path is bias-free (the port has no per-expert bias). Removed the bespoke test_a16w4_flydsl_situv2.py; the original test_moe_2stage.py abf16_wfp4 rows now cover the new kernel. Gate: aiter test_fmoe a16w4 SiTUv2 SEPARATED strict (gfx950): PASS at logits_diff ~1.5e-5 across the token sweep, both shapes. a8w4 FlyDSL tests still pass (25/25). black-88 + ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
coderfeli
force-pushed
the
flydsl-a16w4-mxfp-moe
branch
from
August 1, 2026 11:23
330290c to
06bc8a1
Compare
- Sync the FlyDSL CSV refactor (a16wmix_tuned.csv + resolve_a16wmix_gemm1/2_config + _default_tiles_fallback) into the vendored package. The default path of flydsl_a16w4_gemm1/2 now resolves tiles from our tuned CSV (data-vs-code; provably perf-neutral: CSV values equal the removed inline heuristics, IR byte-identical). CSV auto-located in the package dir via os.path.dirname(__file__). - Reroute op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py to test the NEW a16w-mix a16w4 kernel through the production fused_moe dispatch (bf16 A x mxfp4 W, SiTUv2 / SEPARATED) with a strict cos/logits_diff gate -- explicit FlyDSL-side coverage complementing the routed test_moe_2stage.py rows. Does NOT use the removed low-level compile_mixed_moe_gemm1_a16w4 API. Gate: a16w4 SiTUv2 strict (gfx950): PASS at logits_diff ~1.5e-5 (unchanged -- CSV == old defaults). Rerouted test 6/6 pass. black-88 + ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
coderfeli
added a commit
that referenced
this pull request
Aug 1, 2026
… resolvers) Sync the FlyDSL CSV refactor into the vendored package to keep it identical with PR #4502: add a16wmix_tuned.csv (our tuned tiles, incl. int4 rows) and the resolve_a16wmix_gemm1/2_config + _default_tiles_fallback resolvers; the flydsl_a16w4_gemm1/2 default path is now CSV-driven (data-vs-code, perf-neutral: CSV values equal the removed inline heuristics). CSV auto-located in the package dir. The a16wi4 dispatch keeps reading aiter's tuned int4 CSV (kimik2_i4_tuned_fmoe.csv, the head-to-head-validated on-par config) via resolve_a16w4_gemm1/2_config; the package sync does not change its behavior. Gate: a16wi4 strict (gfx950): PASS at logits_diff ~1.5e-5 across the Kimi int4 shapes. black-88 + ruff clean. Package byte-identical with #4502. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…age1/2, drop bespoke layer
Rewire the a16w4 (bf16 A x MXFP4 W) SiTUv2 MoE path so the ported FlyDSL kernel is a
genuine drop-in, removing the per-forward weight relayout and the redundant tuned CSV.
Native layouts (no relayout). The ported gemm1 gains a w_layout="guinterleave" mode that
consumes aiter's native GUGU stage1 W1+scale (shuffle_weight_a16w4 / shuffle_scale_a16w4)
directly. gemm2's native gate_up=False W2+scale is byte-identical to the standard layout
it already reads (E*model_dim % 256 == 0), so no mode is needed. This deletes the four
per-forward _unshuffle_guinterleave -> standard-shuffle passes (hundreds of MB for Kimi
E896) the old dispatch ran on every call.
Route through compile_flydsl_moe_stage1/2. The a16w4 branches of
compile_flydsl_moe_stage{1,2} now build the ported kernel (was: raise). The a16w4 run
path (run_flydsl_a16w4_moe in moe_kernels.py, replacing the deleted
moe_2stage_a16w4_dispatch.py) sorts, resolves gemm2 tiles from aiter's tuned kimik3 CSV
(resolve_a16w4_gemm2_config), and launches both stages; fused_moe_'s a16w4 case calls it.
Drop the redundant tuned config. Delete a16wmix_tuned.csv (321 rows) and the
_load_ours_tuned_csv / _resolve_ours_tuned stack; resolve_a16wmix_gemm{1,2}_config now
return the documented heuristic directly. gemm1 keeps that heuristic (aiter's kimik3
gemm1 tiles regress the ported body); gemm2 uses the kimik3 CSV.
Verified cold on gfx950 (MI355X, cu=256):
- op_tests/test_moe_2stage.py --csv-filter abf16_wfp4: 32/32 kimik3 a16w4 rows pass
strict accuracy (logits_diff <= 1.6e-5); perf 54us@tok1 .. 3779us@tok16384 (inter512),
strictly faster than the relayout dispatch.
- Native-vs-standard kernel parity bit-identical (max|d|=0.0).
- op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py: 6/6 passed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment-only cleanup of the a16w4 shared-kernel package: collapse ASCII banner blocks to one-line section headers and drop per-line narration that restated adjacent code (the k_wave=1 else-branch note, the BufferCopy128b restatement, verbose route-through and use_csv_config comments). No code change: the Python AST is byte-identical on all three files, so codegen/ISA and perf are unchanged. All load-bearing "why" comments are kept -- the @flyc.jit LOAD-BEARING / ~13x-regression warning above the atomic epilog, the bf16-W num_records i64-fold ABI note, offset-unit gotchas, and the preshuffle layout/stride math. black-88 + ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Trim the verbose a16w4 SiTUv2 dispatch/guard comments in _fused_moe_impl. No code change (token-identical to parent); a16w4 + a8w4 rows still pass the strict test_moe_2stage accuracy gate. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…16w4_moe Reuse moe_sorting's already-zeroed moe_buf as the gemm2 atomic-scatter target (matches the a8w4/mxfp4 sibling path) instead of allocating+zeroing a separate out_buf; make inter_sorted torch.empty (gemm1 masked-store + gemm2 token_id<M gate mean padding never reaches the output); and drop the two unconditional torch.cuda.synchronize() calls (g1->g2 order is guaranteed on one stream). Removes 2 host syncs + 2 device memsets per call. End-to-end fused_moe latency (gfx950, kimik3 3584x512 a16w4): tok1 -26%, tok32 -7%, tok4096 neutral. Accuracy unchanged (logits_diff ~1.5e-5). The syncs also blocked CUDA-graph capture. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
a16wmix_use_k16 now takes the target arch as a required param (callers pass get_rocm_arch()) instead of reading it internally, and the FLYDSL_A16WMIX_FORCE_K16 test-hook env is removed. Behavior on gfx950 is unchanged (use_k16=False); validated logits_diff ~1.5e-5 (a16w4) / 2e-5 (a8w4). Removes the only os use in gemm1.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move the built-in tile heuristic (resolve_a16wmix_gemm{1,2}_config,
_default_tiles_fallback, _default_tile_n) and the aiter tuned-CSV resolvers
(resolve_a16w4_gemm{1,2}_config, pick_a16w4_config, kernelName decoders, CSV
loader) out of the package __init__ into csv_dispatch.py. __init__ drops from
810 to ~505 lines and keeps only the compile dispatch + launch glue. Pure code
move (re-imported into __init__); tile selection and the compiled kernels are
unchanged -- validated logits_diff ~1.5e-5 (a16w4) / 2e-5 (a8w4).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Extract _silu_mul_batch/_sigmoid_f32/_tanh_f32/_situ_mul_batch (and LOG2E) from gemm1.py into a leaf act.py that uses tensor_shim._to_raw (equivalent to gemm1's local _raw for these fx inputs, so no circular import). gemm1 imports the two batch helpers it calls. Pure move; validated logits_diff ~1.5e-5 (a16w4) / 2e-5 (a8w4). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace gemm1's local _raw IR-unwrap shim with tensor_shim._to_raw (aliased as _raw), so the whole package (gemm1, gemm2 via import, act) shares the one tensor_shim helper. Verified equivalent at every call site (all pass fx values): accuracy smoke unchanged at logits_diff ~1.5e-5 (a16w4) / 2e-5 (a8w4), pyflakes clean. Remaining raw llvm/rocdl (MFMA, BufferCopy, masked buffer_store, ptr/gep shims, side-effecting cvt_pk) are load-bearing or have no tensor_shim equivalent and are intentionally kept. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…inside)
compile_gemm{1,2}_a16w4_port now take use_k16 as a required keyword arg instead
of calling a16wmix_use_k16(get_rocm_arch()) internally. The dispatch callers
(compile_flydsl_moe_stage{1,2} and _get_compiled_gemm{1,2}_a16w4) resolve it once
and pass it in; the arch query no longer lives in the kernel-building path.
Drops the now-unused get_rocm_arch import from gemm1/gemm2. Validated
logits_diff ~1.4e-5 (a16w4) / 1.9e-5 (a8w4); pyflakes clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Remove the use_csv_config/csv_path opt-in from flydsl_a16w4_gemm{1,2}: no caller
ever set it True (production uses the heuristic for gemm1 and passes CSV-resolved
tiles to gemm2 explicitly), so the branch was dead. The heuristic
resolve_a16wmix_gemm{1,2}_config path is now unconditional.
- Drop resolve_a16w4_gemm1_config and _kw_tile_k_for (only reachable via the dead
use_csv_config branch). resolve_a16w4_gemm2_config stays (run_flydsl_a16w4_moe
calls it directly).
- Drop _kwave_from_kbatch + the _kb kernelName tag decode: a16w4 tuned rows never
carry _kb, so k_wave always came from _kw (or 1).
Dead-code removal; tile selection and compiled kernels unchanged (CPU-verified:
heuristic configs + kernelName decode identical). pyflakes clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…sl_a16w4_moe
Route a16w4 (bf16 A x MXFP4 W, SiTUv2) through get_2stage_cfgs -> fused_moe_2stages
like a4w4/a8w4, instead of the bespoke pre-dispatch intercept:
- get_2stage_cfgs: re-add a16w4 to use_mxfp4_flydsl (bf16 A, mxfp4 W, Situv2) with
_a_type="bf16" (kernelName abf16_wfp4, no _gui) and fuse_quant="" -- reusing the
existing skip-quant branch (a2_scale=None, bf16 intermediate, no inter-stage quant).
- _flydsl_moe_stage{1,2}_impl: replace the stale a16w4 handling (which built the
standard _s1/_s2 args -- wrong ABI for the port kernel) with early branches that
call the ported flydsl_a16w4_gemm{1,2} launchers. stage1 returns the bf16
[sorted_size, inter_dim] sorted intermediate (threaded to stage2 unchanged);
stage2 atomic-scatters into the moe_sorting-zeroed moe_buf. gemm1 keeps the tuned
heuristic tiles (tile_n=None); gemm2 uses the CSV tiles (resolve_a16w4_gemm2_config).
- _fused_moe_impl: drop the run_flydsl_a16w4_moe intercept; keep the NotImplementedError
guards for features the port can't serve (beta/linear_beta, bias, expert_mask, gfx).
- stage2_uses_route_reduce: treat any a16w4 (bf16 A) stage2 as atomic, so a "reduce"-
named CSV config still gets accumulate=True sorting (moe_buf zeroed) -- the ported
gemm2 only supports atomic scatter. Fixes a NaN at small-M reduce-named rows.
- Remove run_flydsl_a16w4_moe (its per-stage logic now lives in the impls).
Validated (gfx950, kimik3): logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; dedicated pytest
6/6; e2e perf parity with the bespoke path (tok4/8/256/16384 within ~2%); a4w4/a8w4
metadata unchanged. pyflakes clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… guard raises
- _flydsl_moe_stage1_impl a16w4 branch: allocate inter_sorted with torch.zeros when
inter_dim_pad > 0 (torch.empty otherwise) -- restores the pre-fold pad-zeroing that
was dropped when the alloc moved into the early branch. No-op for unpadded shapes.
- Trim the verbose fold-branch comments in _flydsl_moe_stage{1,2}_impl.
- Collapse the five repetitive a16w4 NotImplementedError guards in _fused_moe_impl
into one (cond, why) loop.
Behavior-preserving for the tuned (unpadded) a16w4 path; pyflakes clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… kernels dir #2 (dead CSV code): after the fold, a16w4 gemm2 tiles are already carried by the get_2stage_cfgs kernelName2 (parsed into tile_n/tile_k/b_nt/waves_per_eu/xcd_swizzle by _flydsl_stage2_wrapper), so the stage2 impl re-reading the CSV via resolve_a16w4_gemm2_config was redundant. Switch the a16w4 stage2 branch to the passed tiles (same as a4w4/a8w4) and delete the a16w4 CSV family (resolve_a16w4_gemm2_config, pick_a16w4_config, _load_a16w4_csv, _decode_a16w4_kname, _default_a16w4_csv_path, kernelName regexes, FLYDSL_A16W4_TUNED_CSV env). csv_dispatch.py is now heuristic-only -> renamed tile_heuristic.py. #4: move act.py (SiLU/SiTUv2 helpers) out of moe_2stage_a16wmix/ to the shared aiter/ops/flydsl/kernels/act.py so other FlyDSL MoE kernels can reuse it. Validated (gfx950 kimik3): logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; dedicated pytest 6/6; a16w4 e2e perf parity (tok4/8/256/16384 within noise; get_padded_M is identity at those tokens so kernelName tiles == the old CSV-resolved tiles). pyflakes clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…branches)
The flydsl_a16w4_gemm{1,2} launchers (the only callers of resolve_a16wmix_*) are
mxfp4-only; a16wi4 (int4) dispatches via its own get_2stage_cfgs block + kimik2_i4
tuned CSV, never through this heuristic. So the int4/bf16 branches in
_default_tile_n/_default_tiles_fallback (and the w_dtype param) were dead. Remove
them and the w_dtype plumbing from the resolver calls. mxfp4 heuristic outputs are
byte-identical (CPU-verified: gemm1 tok2/8/256 tiles unchanged); the load-bearing
gemm1 heuristic stays (CSV gemm1 tiles regress). pyflakes clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…pret Replace the fx.Float32(_raw(x).bitcast(T.f32)) boundary in the int4/mxfp4 scale loads (gemm1 load_b_scale_int4 + _e8m0_byte_to_f32, gemm2 load_b_scale_int4) with the native .bitcast(fx.Float32) on the already-typed fx.Int32 expr (same emitted arith.bitcast; drops the _raw unwrap + raw T.f32 + redundant fx.Float32 re-wrap), and drop one redundant fx.Float32 re-wrap of `eff` in _int4_nibble_to_bf16x8. Load-bearing raw ops kept: _udiv/_umod (fx.Int32 is signed -> // / % emit floordivsi/remsi != divui/remui), the masked buffer_store output scatter (fx.copy can't express runtime-num_records masked scalar scatter), _buffer_i32_scalar_read (already on the layout API). Validated: pytest 6/6, logits_diff a16w4 ~1.5e-5 / a8w4 ~1.9e-5, perf parity. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… a16wi4 helpers #6: _get_compiled_gemm{1,2}_a16w4 is @functools.cache'd, so the builder runs once and the returned @flyc.jit launcher is dispatched via _run_compiled. Drop the compile_flydsl_moe_stage{1,2} round-trip in _get_compiled and call compile_gemm{1,2}_a16w4_port directly -- the kernel name/cache is set by the builder's name_suffix, so this yields the identical cached kernel (AOT, which still uses compile_flydsl_moe_stage{1,2}, matches). One fewer hop. Also remove a16wi4_recommend_block_m and a16wi4_scale_to_kernel_layout: dead on this branch (only in __all__; a16wi4 dispatches via its own get_2stage_cfgs block + kimik2_i4 CSV, not these helpers). Validated: pytest 6/6, logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…istic
The gemm1 tile heuristic (tile_heuristic.py resolve_a16wmix_gemm1_config) existed
only because it bypassed the kernelName registry to use tile_k=128 / xcd=1 tiles
that weren't registered. Register those variants and drive gemm1 from the CSV like
gemm2 / a4w4 / a8w4, then delete the heuristic:
- get_flydsl_stage1_kernels: for a16w4 (bf16 A x fp4 W) only, register tile_k in
{128,256} and xcd_swizzle in {0,1,4} (other dtypes unchanged).
- kimik3_fp4_tuned_fmoe.csv: rewrite the 32 gfx950 kernelName1 to the tuned tiles
(tok1-2 t32x64x128_bnt0_kw4; tok3-8 t32x128x256_bnt0; tok16-1024 t32x128x128_xcd1;
tok2048+ t32x128x128_bnt0_xcd1) -- both shapes.
- _flydsl_moe_stage1_impl a16w4 branch: pass the wrapper-parsed CSV tiles
(tile_n/tile_k/k_wave/b_nt/xcd_swizzle) to flydsl_a16w4_gemm1 instead of tile_n=None;
force waves_per_eu=None (kernel default) since a no-_w name parses to wpe=1.
- flydsl_a16w4_gemm1/2: drop the resolve_a16wmix_* / _default_tile_n fallback; use the
caller tiles. Delete tile_heuristic.py + its imports/__all__.
Perf-neutral (CSV tiles == the old heuristic, wpe=None matched): a16w4 e2e 3584x512
tok8=83 / tok256=458 / tok16384=3758 us (ref 84/451/3749, within ~3%). Correctness
logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; dedicated pytest 6/6 (incl. token=1 kw4);
pyflakes clean; a8w4/a4w4/other-dtype registration unchanged.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- CI: black-reformat moe_kernels.py; wrap the implicit f-string concat in the a16w4 guard (ruff ISC004). - Align the "fp4"/"mxfp4" dtype naming to "fp4": b_dtype is only ever "fp4" at runtime (parsed from the abf16_wfp4 kernelName), so the "mxfp4" alias in the b_dtype checks was dead -> drop it; rename the port w_dtype "mxfp4"->"fp4" (byte-exact: name_suffix maps both to ""). Covers moe_kernels, fused_moe, aot, and the moe_2stage_a16wmix package. No CSV change needed (kernelName is wfp4). - Trim the verbose a16w4 comment blocks in moe_kernels (compile entries + stage impl branches). Behavior-preserving: dedicated pytest 6/6, logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; black + ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Merge the byte-identical _gep1/_gep3 into one _gep. The GEP op is polymorphic in the base ptr's address space (the ptr already carries ptr<1>/ptr<3>), so the two names were pure duplication. - Drop the duplicated fx.PointerType.get(i32, Global) construction: move _global_i32_ptr above _global_i32_buffer_view and have the buffer view reuse it. Pure refactor, byte-behavior unchanged: dedicated pytest 6/6, logits_diff a16w4 ~1.5e-5 / a8w4 ~2.1e-5; black + ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… gap) The decode-latency tok1/tok2 cells used t32x64x128_kw4 / t32x128x256, which lost to the (numerically-broken) old a16w4 kernel by ~10-20% fixed launch overhead. A tile_n=32 gemm1 tile (t32x32x128_kw4) is 25-30% faster there: 3584x384 tok1 gemm1 17.4 -> 12.2 us (total ~24.6 -> ~19.4, now beats old ~22) 3584x512 tok1 gemm1 17.9 -> 13.5 us (total ~24.8 -> ~20.1, now beats old ~24) 3584x384 tok2 gemm1 19.6 -> 17.3 us tok2/512 and tok3+ keep their current tiles (tile_n=32 regresses there as M grows and the N tile starves). t32x32x128_kw4 is a registry-valid kernel (enum 4*tn<=tk holds: 128<=128), so it stays AOT-cacheable. Correctness unchanged: logits_diff ~1.4-1.8e-5 at every changed cell; dedicated pytest 6/6. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The AOT precompile drove the a16w4 (bf16 A x fp4 W) port through the generic MX
gemm launch-arg builder (_s1_args_fp4), but the ported moe_2stage_a16wmix kernels
have a different launch ABI (raw fx.Int64 device pointers) -> "too many positional
arguments" at flyc.compile -> the artifact never persisted -> every a16w4 kernel
reported "produced no kernel" (128 MOE AOT failures, CI break).
Add a dedicated _precompile_a16w4_to_cache that drives the SAME runtime launchers
(flydsl_a16w4_gemm{1,2}) the fused-MoE op uses, with real CPU dummies under
COMPILE_ONLY (outside FakeTensorMode, since the port takes raw .data_ptr()). This
guarantees the AOT cache key is identical to runtime by construction:
waves_per_eu=None (both stages), persist=False, w_layout="guinterleave", and the
runtime g2 tile downgrade are all applied inside the launcher.
Verified: kimik3_fp4 AOT now 128 ok / 0 failed; run-only (FLYDSL_RUNTIME_RUN_ONLY=1)
loads every a16w4 artifact for all 32 rows x both shapes (0 cache misses ->
keys match runtime, no JIT fallback). Non-a16w4 dtypes unchanged (gated on
a_dtype=="bf16" and b_dtype=="fp4").
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The artifact is keyed only on the port kernel's constexpr params (shapes / tiles / topk / act), never on the launch pointers, grid, or n_tokens — and COMPILE_ONLY never dereferences them. So the per-buffer dummy tensors (correctly-shaped a / w1 / w1_scale / inter / w2 / ... plus the routing-size math) are unnecessary: a single 1-elem placeholder covers every pointer arg, and n_tokens/M_logical/max_sorted can be 1. Shared kwargs are factored into one `common` dict across the two stages. Dropped the now-dead CU_NUM/get_cu_num plumbing (a16w4 persist=False never reads cu_num). Behavior unchanged: kimik3_fp4 AOT 128 ok / 0 failed, run-only 32/32 with 0 cache misses (keys still match runtime). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…cache Behavior-identical; keeps CI ruff green. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The a16w4 port baked situ_beta/linear_beta at 1.0 (epilog called _situ_mul_batch with defaults), so fused_moe hard-blocked non-default beta with NotImplementedError. Thread beta/linear_beta as runtime f32 scalars (host passes value + reciprocal, mirroring swiglu_limit); the compile cache key stays beta-free so one kernel serves any beta (incl. AOT). Unblocks kimi k3 (beta=4.0, linear_beta=25.0). _situ_mul_batch clamp_limit now defaults to +inf (no clamp), matching the situv2() reference and the a8w4 situv2 path (swiglu_limit=inf). The old hardcoded 7.0 was invisible at beta=1.0 (tanh/sigmoid saturate) but diverged badly for large linear_beta (tanh(u/25) is near-linear). Rename _is_a16w4 -> _is_a16w4_situv2 and document why the SiTUv2 gate is load-bearing: gpt-oss is also bf16 A x MXFP4 W but Swiglu and takes swiglu_mxfp4_bf16_cktile; the activation is what disambiguates them. Validated gfx950: k3 4.0/25.0 cos_diff 1.5e-5; beta=1.0 unchanged; N betas -> 1 compile; gptoss still routes to cktile. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
coderfeli
force-pushed
the
flydsl-a16w4-mxfp-moe
branch
from
August 3, 2026 11:03
c044d33 to
05bd519
Compare
# Conflicts: # aiter/fused_moe.py
# Conflicts: # aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py # aiter/ops/flydsl/moe_kernels.py
… a8w4) The prior runtime-beta commit changed the a16w4 SiTUv2 clamp to +inf (no clamp), which was wrong: the model (and the a8w4 mixed_moe path) clamp SiTUv2 at swiglu_limit. At large linear_beta (k3 uses 25) tanh(u/25) is near-linear, so no-clamp vs clamp-at-7 diverges massively — a16w4 vs a8w4 cos_diff was 0.46. (It passed earlier only because the torch reference used for validation also didn't clamp.) Thread swiglu_limit into the a16w4 port as a runtime f32 scalar (same pattern as beta) and clamp via maximumf(-limit), identical to the a8w4/mixed_moe situv2 clamp: +inf -> no clamp, finite -> clamp. Host passes runtime_swiglu_limit(swiglu_limit, act) so a16w4 == a8w4 for any limit; the scalar is out of the compile cache key (one kernel serves any limit/beta). Validated gfx950: a16w4(lim=7,b=4,lb=25) vs a8w4 cos_diff 0.46 -> 5.5e-3 (residual is a8w4 fp8 quant); no-limit default unchanged (1.5e-5); a16wfp4 pytest 6/6. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
# Conflicts: # aiter/fused_moe.py
9 tasks
Align the a16w-mix a16w4 (bf16 A x mxfp4 W, SiTUv2) kernel to main's GGUU
(standard, separated gate/up) W1 layout, and add real tuner support so the
kimik3_fp4 CSV is produced by gemm_moe_tune (was hand-written).
- Layout: a16w4 dispatch uses w_layout="standard" (GGUU) in the runtime
(moe_kernels) and AOT (aot/flydsl/moe) paths; tuner W1 prep and both tests
set to GGUU. a8w4 (fp8) stays gate/up-interleaved (GUGU).
- gemm_moe_tune a16w4 support: skip the unsupported CK b16xfp4 enumeration;
pass situ_beta to ref+kernel; map the SORTED stage1 output <-> unsorted ref
by expert-match (_a16w_sorted_row_map) for both the stage1 compare and the
stage2 sorted intermediate input; guard out invalid/degenerate candidates
(stage2 sbm, tile_n=256 LDS overflow, tile_k not dividing inter_dim, and
stage1 num_acc_n<1 e.g. tile_n=32/kw1 which emits NaN).
- Retune kimik3_fp4_tuned_fmoe.csv (all 32 shapes).
Validated on gfx950: a16w4 pytest 6/6 (tok 1/16/128 x inter {512,384}) and
tok 1024/4096 x {512,384} 4/4; CSV has no degenerate/invalid tile picks.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Merge the three separate a16w-mix stage2 skip checks (sbm tile_m, tile_n=256 LDS overflow, tile_k not dividing inter_dim) into a single guard. Logically identical (same candidates skipped) -> no retune needed; a16w4 pytest 6/6. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This was referenced Aug 6, 2026
The stage1 gather / stage2 scatter I added ran inside the tuner's timed region (run_perftest times run_func), so the recorded us was kernel + a full re-sort over the sorted buffer (~28k rows) -- 2-3x inflated, and for small tokens the re-sort dominated, corrupting kernel selection too. Make run_func kernel-only; move the sort into the untimed ref/data-prep: - stage1: run_flydsl_stage1_out returns the kernel's SORTED output as-is; accuracy compares in sorted space against run_a16w_stage1_sorted_ref (computes the torch ref then gathers to sorted order) via _a16w_sorted_cos. - stage2: the sorted intermediate input (a2_a16w_sorted) is built once in generate_data and fed via a conditional run-arg; run_flydsl_stage2_out no longer re-sorts (just a2_scale=None for bf16). - drop the timed-path _a16w_sorted_row_map / topk_ids plumbing. Retune kimik3_fp4_tuned_fmoe.csv: us now tracks e2e (e.g. 512/1 333->27, 384/4096 1893->944), not inflated. Validated: a16w4 pytest 6/6 (tok 1/16/128) and tok 1024/4096 4/4; no degenerate/invalid tile picks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…cnt (ISA byte-exact)
Cleanup of the a16w4/a16wi4 package leaf helpers (no codegen change; gemm1+gemm2
AMDGCN verified byte-exact A/B via FLYDSL_DUMP_IR):
- Raw ptr casts -> flydsl llvm-ptr API: _lds_ptr3/_global_base_ptr1 now go through
a single _int_to_llvm_ptr() using fx.to_llvm_ptr(fx.inttoptr(..., AS.Shared/Global)),
dropping the hardcoded ir.Type.parse("!llvm.ptr<3|1>") (backend resolves the AS).
- Drop the hand-encoded CDNA3 s_waitcnt shim: flydsl 0.3.0 has the arch-dispatched
keyword form fx.rocdl.s_waitcnt(lgkmcnt=0); same encoding (ISA identical).
- b_cache_mod: use the CSV-driven b_nt directly (constant default when unset); remove
the dead/inconsistent per-token U-shape heuristic (moe_kernels always passes b_nt).
- Extract the 18 shared leaf helpers gemm2 imported from gemm1 into a package-local
utils.py; gemm1 shrinks ~185 LOC. gemm2/__init__/moe_kernels repoint to .utils.
Validated: a16w4 SiTUv2 pytest 6/6 cold (cache cleared), pyflakes + black clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
CI Black check flagged gemm_moe_tune.py; apply black to the a16w4 stage2 ternary run-arg lists. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
valarLip
reviewed
Aug 8, 2026
|
|
||
| _act = "situv2" if act in ("situv2", "situ") else act | ||
| sorted_size = int(sorted_expert_ids.shape[0]) * int(tile_m) | ||
| _alloc = torch.zeros if inter_dim_pad > 0 else torch.empty |
Collaborator
Author
There was a problem hiding this comment.
ported from old code, https://github.com/ROCm/aiter/pull/4502/changes#diff-d5d52b4d0fede962b29252c6940ba9aa70479833fc6b767058574c607e140004L1496. seems for bug fixes of gptoss. No sure if needed for other models. will not be trigger for inter_dim_pad=0
| _alloc = torch.zeros if inter_dim_pad > 0 else torch.empty | ||
| inter_sorted = _alloc(sorted_size, inter_dim, dtype=torch.bfloat16, device=dev) | ||
| flydsl_a16w4_gemm1( | ||
| a_bf16=a.to(torch.bfloat16).contiguous(), |
Collaborator
Author
There was a problem hiding this comment.
no real call. just make sure no error. can be deleted.
…a!=1 - gemm1/gemm2 builders: assert TILE_N//(4//k_wave) >= 16 (gemm1) and TILE_N//4 >= 16 (gemm2). num_acc_n==0 (e.g. tile_n=32,k_wave=1) empties every accumulate/store loop -> silent all-zero output that times fast. The shipped CSV avoids it and the tuner skips it; this closes the sharp edge for future tuning. - fused_moe_2stages: skip the throwaway (token_num, topk, inter_dim) `a2` allocation for the a16w4 port. Its stage1 allocates its own sorted [sorted_size, inter_dim] bf16 intermediate and ignores `out`, so the pre-alloc (up to tens of MB at prefill) was pure waste. Tightly gated to the a16w4 flydsl path; a8w4/a4w4/mxfp8/non-flydsl unaffected. - test: parametrize situ_beta/linear_beta with (4,25) (kimi-k3 betas) to cover the runtime SiTUv2 beta path directly (beta is not a compile key). Validated gfx950: a16w4 pytest 12/12 (incl. beta!=1), a8w4 15/15. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli
force-pushed
the
flydsl-a16w4-mxfp-moe
branch
from
August 8, 2026 05:02
d6df09f to
e2fb8da
Compare
… types
Cleanup only; ISA byte-exact (gemm1+gemm2 21_final_isa.s identical A/B),
a16w4 pytest 12/12 cold.
- Remove a16wmix_use_k16(arch); inline `"gfx95" not in str(get_rocm_arch())`
at the 4 call sites (2 __init__ launchers, 2 moe_kernels compile entries).
- Remove kmchunks_for/lds_acc_bytes_for from the package utils; inline
BM//16 and BM*TILE_N*4 (the shared mxfp4_gemm_common copies are untouched).
- Drop the _get_compiled_gemm{1,2}_a16w4 @functools.cache wrappers from
__init__; move @functools.cache directly onto compile_gemm{1,2}_a16w4_port
and call them all-keyword from the launchers, matching moe_kernels.py's
direct calls so the compile cache key is shared (fixes the previously
uncached moe_kernels compile-entry path).
- ir dialect -> fx types: ir.Type.parse("vector<2xbf16>") -> T.vec(2, T.bf16)
(gemm1/gemm2), ir.IntegerType.get_signless(32) -> T.i32 (utils cvt_pk);
drop the now-unused `from flydsl._mlir import ir` imports. The load-bearing
raw llvm.load/StoreOp/AtomicRMWOp/inline_asm (raw !llvm.ptr consumers,
side-effecting cvt_pk) are kept.
Co-Authored-By: Claude Opus 4.8 (1M context) <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
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
Replaces aiter's a16w4 (bf16 activation x MXFP4 weight) SiTUv2 fused-MoE
kernel with a numerically-correct, faster ported FlyDSL 2-stage kernel
(
aiter/ops/flydsl/kernels/moe_2stage_a16wmix), and removes the oldcompile_mixed_moe_gemm{1,2}_a16w4path entirely.aiter's previous a16w4 SiTUv2 kernel fails aiter's own strict accuracy gate
(
op_tests/test_moe_2stage.py, bf16 SiTUv2 reference):logits_diff ~1.0(essentially wrong output). The port passes the same gate at
logits_diff ~1.5e-5across the full token sweep on both tuned Kimi shapes (3584x384,3584x512, E896, topk16).
Reproduction (independent re-measurement, gfx950)
Re-measured the shipped port against its own tuned CSV
(
kimik3_fp4_tuned_fmoe.csv,abf16_wfp4) on gfx950 (cu=256) viatest_moe_2stage.py --kernel --no-legacy --csv-filter abf16_wfp4, both Kimishapes. The kernel runs correctly and meets or beats its tuned latency at
every token (0.80–1.06×), monotonic across M. Ratio = measured / tuned-CSV
(lower = faster than the recorded tune).
Correctness:
op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py6/6 pass(2 shapes × tok 1/16/128) vs the authoritative torch SiTUv2 reference.
3584x512
3584x384
What changed
aiter/ops/flydsl/kernels/moe_2stage_a16wmix/(3 files:
gemm1.py,gemm2.py,__init__.py): gemm1 (gate+up GEMM + SiTUv2-> sorted bf16 intermediate), gemm2 (down-proj + routing-weighted atomic bf16
scatter). Non-scaled
MFMA(16,16,32,bf16), CDNA (gfx942/gfx950). Consolidatedlayout (shared helpers in gemm1, launch glue + CSV tile resolvers in
__init__); scale loads use themake_buffer_tensor/BufferCopy32blayout API.fused_moe_dispatch (aiter/ops/flydsl/moe_2stage_a16w4_dispatch.py):serves bf16-a16w4 SiTUv2 (the kernel is gate-mode agnostic). Converts the
incoming guinterleave (a8w4/mxfp8-shared) weight/scale layout to the port's
standard
shuffle_weight/e8m0_shufflelayout via pure byte permutations(verified round-trip). gemm2 tiles come from aiter's tuned CSV
(
resolve_a16w4_gemm2_config, on-par-to-faster); gemm1 keeps its own tuneddefault — a perf-gate showed the CSV gemm1 tiles regress mid/high-M by
1.1x–2.0x, so they are not taken.
compile_mixed_moe_gemm{1,2}_a16w4, the compatwrapper
compile_a16w4_moe_gemm2, and the a16w4-only_decode_e8m0_byte_to_f32helper (~4076 lines) from
mixed_moe_gemm_2stage.py; the a16w4 branches ofcompile_flydsl_moe_stage1/2; and the dead a16w4 sub-clause /_a_typebf16-fp16 branches in
get_2stage_cfgs. Genericcompile_mixed_moe_gemm1/2and all other dtypes (a8w4 / int4 / mxfp8 / a4w4) are untouched.
beta/linear_beta, per-expert bias, expert-parallel masking) raise a clear
NotImplementedErrorrather than routing to the numerically-broken old kernel.op_tests/test_moe_2stage.pya16w4 (abf16_wfp4)SiTUv2 rows now exercise the new kernel through the production dispatch —
_effective_gate_modereturns SEPARATED for bf16-a16w4 (previously INTERLEAVE,which silently downgraded those rows to a8w4/fp8), and the a16w4 SiTUv2 path is
bias-free (the port has no per-expert bias). No bespoke test file.