[HIP] [CK] [FlyDSL] [Kernel] Extend MXFP4 GEMM1 replacement to A4W4 - #4526
[HIP] [CK] [FlyDSL] [Kernel] Extend MXFP4 GEMM1 replacement to A4W4#4526fsx950223 wants to merge 75 commits into
Conversation
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
There was a problem hiding this comment.
🟡 Human review recommended
It changes Tier-2 MoE dispatch (aiter/fused_moe.py) and multiple FlyDSL kernel/tuning/AOT paths, so the blast radius is high despite local unit test coverage.
Pull request overview
Extends the FlyDSL MXFP4 “GEMM1 replacement” MoE pipeline beyond the prior A4W4-only port to also cover A8W4, adding cache-safe kernel naming/selection, runtime gating, AOT precompile support, and updated tuned model CSVs.
Changes:
- Add cache-safe GEMM1 kernel-name construction/parsing for new MXMOE variants (FP8 intermediates, SiTUv2 beta encoding, interleaved layouts).
- Update tuning + dispatch plumbing to couple new GEMM1 variants with mixed_moe_gemm_2stage GEMM2 kernels, plus profiling support to split stage timings.
- Expand supported shape/config coverage (aux-kernel shape set, AOT job enumeration, refreshed tuned CSVs for GLM-5 / DeepSeek-V4 / Kimi-K3).
File summaries
| File | Description |
|---|---|
| csrc/kernels/mxfp4_moe/moe_aux/codegen/gen_instances.py | Extends generated aux-kernel shape list (adds GLM-5/Kimi-K3 shapes) and minor formatting tweaks. |
| csrc/ck_gemm_moe_2stages_codegen/tune_mxfp4_flydsl.py | New CLI entrypoint for the updated FlyDSL MXFP4 GEMM1 + mixed GEMM2 coupled tuner. |
| csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py | Updates tuning logic to generate/lock candidates for A4W4 and A8W4 (FP8) replacement paths and record stage timings. |
| aiter/test_common.py | Adds optional kernel-time extraction from trace profiling for more detailed perf attribution. |
| aiter/ops/moe_mxfp4_aux.py | Adds an explicit supported-shape set + helper for MXFP4 aux-kernel coverage checks. |
| aiter/ops/flydsl/mxfp4_kname.py | Implements new cache-safe GEMM1 naming/parsing (including float-bit encoding for SiTUv2 betas) and selection helpers. |
| aiter/ops/flydsl/mxfp4_gemm1_kernels.py | Extends GEMM1 port wrapper to support dtype/activation variants, BN choices, and SiTUv2 parameters. |
| aiter/ops/flydsl/moe_kernels.py | Threads model-dim padding into reduction compilation and normalizes padded output in debug/raw-output mode. |
| aiter/ops/flydsl/kernels/mxfp4_gemm_common.py | Adds FP8/FP4 MMA atom/scaling helpers and activation math helpers used by new variants. |
| aiter/ops/flydsl/kernels/moe_gemm_2stage.py | Adds model-dim padding support to the reduction kernel path and updates comments/metadata. |
| aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py | Adds padding tags to module naming and bounds padded stores in stage2 epilog. |
| aiter/fused_moe.py | Extends runtime dispatch/gating to select MXMOE replacement kernels for both A4W4 and A8W4, including SiTUv2 beta threading and stage2 integration. |
| aiter/configs/model_configs/kimik3_a8w4_tuned_fmoe.csv | Updates tuned Kimi-K3 A8W4 configs to use new MXMOE GEMM1 kernel names. |
| aiter/configs/model_configs/glm5_fp4_newg1_mixedg2_stdsort_tuned_fmoe.csv | Adds a new tuned CSV for GLM-5 FP4 with new GEMM1 + mixed GEMM2 naming. |
| aiter/configs/model_configs/dsv4_fp8fp4_tuned_fmoe.csv | Updates DeepSeek-V4 FP8FP4 tuned rows to use MXMOE GEMM1 replacement kernels where intended. |
| aiter/aot/flydsl/mxfp4_moe.py | Expands AOT job parsing to cover fp4/fp8fp4/a8w4 tuned CSVs and generate correct compile jobs for replacement kernels. |
Review details
Suppressed comments (1)
aiter/ops/moe_mxfp4_aux.py:19
- The shape glossary line uses
? {16, 32, 64, 128}, which is unclear. Use a clear set-membership phrasing (e.g., "in {..}") so readers don't interpret it as uncertainty.
# MB = block_m (sort/gemm block size, ? {16, 32, 64, 128})
- Files reviewed: 17/17 changed files
- Comments generated: 3
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
…526-20260803 Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # aiter/fused_moe.py # aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py
Preserve the replacement and padding fixes on the MXMoE v2 APIs while applying CI formatting and review cleanups. Co-authored-by: Cursor <cursoragent@cursor.com>
Match the version used by the repository's Black action so the style check remains stable through local hooks. Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Use the tuple form so the fused MoE dispatch passes PIE810 in the pinned CI Ruff version. Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Preserve the sorted MXFP4 intermediate path across the FlyDSL MoE refactors and carry padded reduction handling into the extracted kernel. Co-authored-by: Cursor <cursoragent@cursor.com>
Stop synthesizing obsolete MXMOE plus legacy GEMM2 pairings and make GEMM1 sorted-output-only. Precompile the supported SiTUv2 beta contracts and BM16 scale layout so RUN_ONLY matches runtime. Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Keep the PR focused on the production pipeline and AOT implementation changes. Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Forward route weights unchanged to the reduction boundary, which already normalizes their dtype and layout, and remove stale dispatch commentary. Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Remove redundant historical and narrative commentary while retaining the layout, cache-key, and correctness invariants needed to maintain the kernels. Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
One thing I would like resolved before merge,
native_scale_layout defaults to False, so the tuner produces all-NaN (blocker)
See the inline comment on aiter/fused_moe.py:2713. get_2stage_cfgs binds native_scale_layout=_bm == 16, but _mxfp4_a4w4_stage1_fw defaults it to False and the tuner calls that function without the kwarg. On BM16 the two settings produce different scale layouts, and forcing it off end-to-end gives 100% NaN. Since nearly every token<=256 row in the shipped CSVs is BM16, those rows can no longer be re-tuned.
| output_aux=True, | ||
| prequant=False, | ||
| fuse_quant=_p1["out_dtype"], | ||
| output_aux=AUX_SORT_OPUS, |
There was a problem hiding this comment.
Should the tuner be updated to match this production default? _port_e2e() still passes output_aux=True, which selects the three-stage/adaptive path because _aux_uses_opus() requires the literal "opus". It also bypasses the fused-prequant path used in production, so for BM > 16 the tuner timings and candidate ranking may not reflect the default runtime pipeline.
| gfx950,256,8192,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,422.9114,flydsl_moe1_afp4_wfp4_bf16_t64x128x256_w4_bnt0_xcd4_fp4,0.1%,597.3431,flydsl_moe2_layout_afp4_wfp4_bf16_t64x256x128_reduce_sbm64_sp1601,0.0%,1020.2545,0,0,0,1060.84,3712.27, | ||
| gfx950,256,16384,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,642.4294,flydsl_moe1_afp4_wfp4_bf16_t64x128x256_w3_bnt0_xcd4_fp4,0.1%,1074.0772,flydsl_moe2_layout_afp4_wfp4_bf16_t64x256x128_reduce_sbm64_sp801,0.0%,1716.5066,0,0,0,1261.09,2257.8, | ||
| gfx950,256,32768,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,1111.7422,flydsl_moe1_afp4_wfp4_bf16_t64x128x256_w2_bnt0_xcd4_fp4,0.1%,2185.1161,flydsl_moe2_layout_afp4_wfp4_bf16_t64x256x128_reduce_sbm64_sp1601,0.0%,3296.8583,0,0,0,1313.17,1228.96, | ||
| gfx950,256,4096,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,301.2728,flydsl_mxmoe_g1_a4w4_64x256x256_situv2_xcd2,2.1%,338.8907,flydsl_moe2_layout_afp4_wfp4_bf16_t64x256x128_reduce_sbm64_sp1601,0.0%,640.1635,0,0,0,845.36,5847.59, |
There was a problem hiding this comment.
The tuner cannot reproduce 64 of the 70 GEMM1 configs changed by this PR
Mxfp4FlydslTuner._g1_kname() hard-codes A4W4 with BN/BK=256 and only varies
BM, use_nt, and inline quantization. _candidate_rows() also uses the FP4-only
_SUPPORTED set. Comparing the generated names with the 70 changed rows, only
6 are reproducible; the other 64 require BN128, XCD, K-wave, SiTUv2, or other
unsupported options.
This can overwrite optimized rows without benchmarking the existing config.
Kimi-K3 is worse: its 17 SiTUv2 rows are tuned as SiLU, and _run_candidate()
also uses a SiLU reference, so tuning passes but runtime later rejects the config
on activation mismatch and falls back.
Please use the shared kernel-name builder, expose the new tuning axes, handle
SiTUv2 explicitly, and benchmark the existing config before replacing it.
Fixing native_scale_layout alone does not resolve this coverage gap.
Fixes the round of review findings on #4526. native_scale_layout had two divergent defaults: get_2stage_cfgs bound `_bm == 16` while _mxfp4_a4w4_stage1_fw defaulted to False, so the tuner -- which calls that function directly -- evaluated every BM16 candidate against the wrong scale layout and got 100% NaN. Nearly every token<=256 row in the shipped CSVs is BM16, so those rows could not be re-tuned. The rule now lives once, in native_scale_layout_for(), and stage1 derives it from the kernel name when the caller says nothing. bias was accepted and then dropped. stage2_has_bias was unconditionally True for every MXMOE config, so bias2 landed in extra_stage2_args and _mxfp4_a4w4_stage2_fw swallowed it via **_kwargs; it now has a real parameter, forwards to _flydsl_v2_stage2_wrapper, and raises on the native GEMM2 family that cannot take it. The stage1 guard faced the wrong way and could never fire; get_2stage_cfgs now rejects, with a warning, a config whose kernel name cannot carry a bias the model actually has. Both bias predicates are decided by parse_flydsl_v2_gemm2_kernel, the same parser the stage2 dispatch uses, so config time and run time cannot disagree. Verified over all 92 shipped mxmoe rows: with no bias all 92 still dispatch MXMOE, and only the bogus stage2_has_bias flips. GEMM1 AOT jobs were keyed on the CSV's logical inter_dim while stage1 derives D_INTER from the stored weight width (w1.shape[1] // 2), which for a non-256-aligned shard is padded, with the logical width carried in w2.inter_real. The Kimi-K3 384 job therefore never matched the 512 the runtime asks for -- measured as a JIT compile on the serving path, now an AOT hit. Also drops code the review found unreachable or duplicated: the except-ValueError fallback that only deferred its own failure, the second activation check, **route_bucket_metadata in a branch where it is always empty, scale_out_uses_atomic (constant False, and it forced flydsl into import time), _effective_use_nt (made the compile key a function of runtime n_tokens, so AOT could not match it), and the unused _select_mxfp4_block_m / _make_mxfp4_g1_kname / _select_mxfp4_g1_kernel. _G1_VARIANTS and _SUPPORTED_BY_DTYPE collapse into one MXFP4_G1_VARIANTS, is_mxfp4_moe_shape_supported gains its first caller, and stage2 goes back to shape[2] to agree with the ndim assert next to it. situ_beta / situ_linear_beta / swiglu_limit stay out of the kernel symbol name: FlyDSL's cross-process cache is keyed on a hash of the launcher source and every scalar closure value it reaches, not on the symbol, and gemm1_kernel closes over all three. The comment now says so and a test pins it. They are also normalized per activation, so silu with a stray beta no longer creates a cache entry for a value the kernel never reads. Verified on gfx950 against this branch's HEAD: BM16 goes from 100% NaN to cos_diff 0.0119 on two shapes, BM32/64/128 are unchanged (BM128 bit-identical), and the AOT job for the Kimi-K3 row now hits. Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
@binding7012 all 15 threads from your review are addressed in ba12497. Replies with the details are on each thread; the short version: The blocker. You were right about the cause and the consequence. The rule now lives once — Bias. AOT
Regression check. Comparing resolved One thing worth flagging that came out of the bias work: 69 of the 92 rows (kimik2 57, glm5 8, qwen3_5_397b 4) name a native Verification caveats: the numbers above are from |
Adds `prefetch_hidden` to the a4w4 GEMM1 port. In the inline-quant path the main loop loaded the current K-tile's hidden_states immediately before quantizing them; this hoists the load one tile ahead and carries (h_v0, h_v1) across iterations, so the fetch overlaps the previous tile's MFMA. It requires inline_quant -- the cached-A path performs no hidden load -- and is rejected at every layer that can be entered independently (kernel name builder, launch-time _assert_supported, compile_gemm1_a4w4_port). The variant is selected by an `_hpf` tag in the kernel name (after `_f16in`, before `_nt`) and gets its own compiled symbol via `variant_tag += "_hpf"`, so it cannot collide with the non-prefetch build in the FlyDSL disk cache. It is threaded through the runtime launch path, the fused_moe dispatch, and the AOT job key so a tuned row is warmed under the same key inference looks up. Measured on gfx950 (MI355), kimik3 a4w4 h=3584 i=384 NE=896 topk=16, GEMM1 us from rocprofv3 --kernel-trace --stats, hpf vs the same row with `_hpf` removed: tok 2 3 4 16 32 64 256 hpf 11.29 15.62 16.67 52.60 86.68 132.96 191.99 none 13.53 15.94 17.20 52.13 87.28 134.09 194.93 gain 1.20x 1.02x 1.03x 0.99x 1.01x 1.01x 1.02x The win is concentrated at very small token counts, where the exposed load latency is the largest share of the tile; it is ~1-3% and inside noise above that. Cosine error is unchanged (~0.0114 either way) -- this is purely an instruction-scheduling change, not a numerical one. The kimik3_a4w4 tuned CSV is retuned accordingly: 12 rows change, 7 of which now name an `_hpf` GEMM1. Each retuned row was replayed through _mxfp4_a4w4_stage1_fw/_mxfp4_a4w4_stage2_fw and reproduces its recorded us1 within -3.2%/+5.9% and its recorded err1 exactly. Signed-off-by: fsx950223 <fsx950223@outlook.com>
Keep the PR focused on the production pipeline changes. Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The MXFP4 tuner could not reproduce the GEMM1 configs this PR ships. Its
_g1_kname hard-coded BN/BK=256 and varied only (BM, use_nt, inline_quant),
so of the 75 rows the PR changes to an mxmoe GEMM1 name, zero were
reachable: 51 need xcd_swizzle, 43 need BN128, 7 need k_wave, and the
Kimi-K3 rows need SiTUv2. A re-tune could therefore only overwrite those
rows with something it happened to be able to express.
_g1_variants now enumerates BN x xcd_swizzle x k_wave x num_waves x
hidden-prefetch on top of the existing variants and filters the result
through the kernel's own _assert_supported, so the tuner cannot propose a
name the runtime rejects and cannot drift from the kernel's constraints
(BN64 implies BM32 non-inline separated; num_waves==2 implies BN64;
k_wave>1 implies BM32 non-inline with num_waves*k_wave<=8). 84 candidates
per shape, every generated name round-trips through _parse_mxfp4_g1_kname,
and 71 of the 75 shipped pairs are now reachable -- the remaining 4 miss on
the GEMM2 side (the _sp801/_sp1601 spart variants are not in
get_flydsl_stage2_v2_kernels for that shape).
Two correctness fixes the expanded space depends on:
NaN no longer passes the accuracy gate. `float('nan') > errRatio` is False,
so a candidate producing garbage was accepted and, being fast, could win the
sweep -- observed on BM16 rows reporting err=nan at roughly half the runtime
of a correct candidate. The gate now rejects non-finite values explicitly.
SiTUv2 is tuned as SiTUv2. _run_candidate folded Situv2 into Silu, so the
name carried no _situv2 tag and the torch reference computed a different
activation; the row that got written was then rejected at runtime on
activation mismatch. _row_act drives both the kernel name and the reference,
and _port_e2e passes the betas run_torch_moe_stage1 defaults to.
Measured on gfx950: all 24 distinct shipped GEMM1 names now run and pass the
accuracy gate (err ~0.011-0.012), against 12 pass / 4 fail before -- and 3 of
those 12 were the NaN false positives.
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
run_config built its torch reference with run_torch_moe_stage1, whose situ_beta/situ_linear_beta default to the model values (4.0 / 25.0), but called fused_moe without beta/linear_beta, so the kernel ran the activation at 1.0. The two sides therefore computed different activations and every per_1x32 SiTUv2 shape was reported as a mismatch no matter which kernel it dispatched -- logits_diff ~0.14 with out_norm/ref_norm ~0.61, a constant ratio that is the signature of a scale disagreement rather than a bad kernel. The failure was silent in the worst way: --run_config is what tells you whether a tuned CSV is correct, and it condemned all 15 a4w4 rows of kimik3_a4w4_tuned_fmoe.csv, which are the only Situv2 rows we ship. With the betas passed, the same rows report status=ok on both this branch and main, so the verdict now tracks the kernel instead of the harness. Timing is unaffected -- beta is a compile-time scalar in the activation, not a shape or schedule -- so previously recorded us numbers remain valid. Signed-off-by: fsx950223 <fsx950223@outlook.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 23 out of 23 changed files in this pull request and generated no new comments.
Suppressed comments (1)
aiter/ops/flydsl/mxfp4_kname.py:76
- The stage-1 kernel-name parser can raise
ValueError(... expected prefix '')when_MXMOE_G1_PREFIX_REdoesn't match, becausepfxbecomes an empty string. That makes the error message unhelpful for diagnosing bad names (and the “expected prefix” string is the most actionable part here).
mode["a_dtype"] = "fp8" if prefix_match.group("a") == "8" else "fp4"
else:
pfx = _MXMOE_PREFIX[stage]
if not pfx or not kname.startswith(pfx):
raise ValueError(f"bad mxmoe kernel name: {kname!r} (expected prefix {pfx!r})")
block_m 16 has no three-stage sort instance -- codegen emits aux_sort3s_*
only for MB in {32, 64, 128} -- so it fell back to sort_quant_kernel_impl,
whose entire sort (count, cumsum, place, pad) runs under `blockIdx.x == 0`.
One CTA does O(routed_rows) work while the other 127 only zero-init the
atomic output. Opus instead spends one CTA per expert, which is exactly the
fix 61c306c already made for MB 32/64/128; block_m 16 was left behind.
The cost shows up as a prologue that grows linearly with the token count
while Opus stays flat. Measured on kimi-k3 a4w4 (NE=896, topk=16, gfx950),
prologue GPU us:
token 16 32 64 128 256 512
fused 5.43 7.04 10.31 14.53 20.31 31.36
Opus (--) (--) 9.58 9.77 11.42 13.15
main bm32 10.53 9.26 9.44 9.69 19.22 20.80
The fused sort wins while routed rows are few relative to experts and loses
linearly after, so this switches on `routed_rows >= num_experts` rather than
unconditionally: the crossover measured between token 32 and 64 and the rule
puts it at 56. token 16/32 keep the fused path and are unchanged.
block_m 16 could not simply be un-excluded, which is why the original gate
was written as a flat `block_size != 16`. _aux_uses_opus was doing double
duty: it also selected `prequant`. But block_m 16 is the only inline-quant
("_f16in") a4w4 variant, and that GEMM1 reads raw bf16 hidden_states and
ignores the A buffers, so handing it a prequantized fp4 A plus scales faults
with hipErrorIllegalAddress. The prequant decision is now keyed on block_m
directly, leaving the sort free to move.
Verified on gfx950 over all 76 tuned rows this PR changes (4 model CSVs),
main config vs PR config, same code, e2e via fused_moe: geomean 1.062,
69/76 >= 1.00, all 76 status=ok. The two rows that this PR previously
regressed both flip positive -- kimi-k3 token 256 0.991 -> 1.016 and
token 512 0.972 -> 1.036 -- with no new regression introduced; every other
sub-1.00 row reproduces as parity over three alternating repeats.
Signed-off-by: fsx950223 <fsx950223@outlook.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 23 out of 23 changed files in this pull request and generated no new comments.
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
aiter/ops/moe_mxfp4_aux.py:12
- MXFP4_MOE_SUPPORTED_SHAPES duplicates the authoritative SHAPES list in csrc/kernels/mxfp4_moe/moe_aux/codegen/gen_instances.py. Keeping these two lists manually synchronized is easy to miss during future shape additions/removals and can lead to incorrect runtime gating (accepting shapes without generated kernels, or rejecting supported ones). Consider generating this Python shape set from the same codegen source (e.g., emit a small Python module alongside the generated C++ artifacts and import it here), or otherwise enforcing the sync in the codegen pipeline.
# Keep synchronized with moe_aux/codegen/gen_instances.py::SHAPES.
MXFP4_MOE_SUPPORTED_SHAPES = frozenset(
{
(385, 7168, 512, 9),
(385, 7168, 1024, 9),
Current config validation after review fixes
gfx950 (MI355X), comparing main
f93e6cf25with PR90b634c09over all 70config rows changed by this PR. Each shape group stayed on one GPU; main,
PR+Opus, and PR+threestage were interleaved for three runs after warmup, and each
cell uses the median. Ratios below 1.000 mean the PR is faster.
All 70 cells passed accuracy under main, Opus, and threestage across all three
measured runs (630 outputs); maximum
logits_diffwas 1.944e-3 on main,1.783e-3 on Opus, and 1.759e-3 on threestage, all below the 0.01 gate. Kimi-K3
used
AITER_SITUV2_A4W4=1,AITER_FLYDSL_STAGE2_FP8=1, beta=4, andlinear_beta=25.
Opus had 42 cells faster than main by more than 2%. The only confirmed regression
above 2% after three additional interleaved repeats was GLM-5 token=4,
6144x512, E257/topk9: 44.702 us main vs 46.387 us Opus (1.0377x). The initial
Kimi-K2 token=32 signal reduced to 1.0130x over six samples.
Historical benchmark: Opus aux-sort end-to-end performance
gfx950 (MI355X), comparing main
9aa8a6b91with PR35aa68454over everymodel-config row changed by this PR: 70 cells across GLM-5, Kimi-K2, Kimi-K3, and
Qwen3.5-397B. Each arm was run three times on the same idle GPU with a warm FlyDSL
cache; every table entry uses the per-cell median. Ratios below 1.000 mean the PR is
faster.
Performance by config
geomeanis the geometric mean of per-cell PR/main ratios.sum ratiois the ratioof summed per-cell median latencies.
Overall, the Opus path is 3.60% faster geomean and 2.04% faster by summed
latency than the tested main snapshot. It is 1.03% faster geomean than the retained threestage
path. Opus has 39 cells faster by more than 2% and 5 cells slower by more than 2%.
Opus regressions above 2%
All five are Kimi-K2 rows; the other three config files have no Opus regression above
2%.
Accuracy
Strict accuracy was run for every changed row under both aux-sort paths: 140/140
cases passed, with no NaN or assertion failure.
The largest
logits_diffobserved across all repeated performance runs was 3.57e-3,below the 0.01 gate.
Test method
(
gfx,cu_num, token/model/expert dimensions, activation/dtypes, quant type,use_g1u1, anddoweight_stage1) and retaining only rows whose kernel configdiffers. This produced 5 + 57 + 4 + 4 = 70 rows.
M1 O1 T1 T2 O2 M2 M3 O3 T3to reduce clock/thermal drift.AITER_MOE_EXPERT_BALANCE=true, one idle gfx950 GPU, warm JIT caches, andthe complete fused-MoE callable timing reported as
usbyop_tests/test_moe_2stage.py.strict_accuracy=Trueseparately for Opus andthreestage. Qwen rows were JIT-compiled rather than using RUN_ONLY because the
isolated changed-row subset has no pre-generated AOT cache.
Representative command template for each shape group:
For Kimi-K3 SiTUv2 A4W4, add
AITER_SITUV2_A4W4=1andAITER_FLYDSL_STAGE2_FP8=1; the CSV strict-accuracy pass uses beta=4 andlinear_beta=25 from the model-config test path.