[Qubitium]

Summary

Makes gfx950 MXFP4 FlyDSL FMoE AOT tuning practical by moving cold kernel compilation out of the serial GPU candidate loop and replacing exhaustive per-pair measurement with an accuracy-aware two-stage scan.

This PR also addresses an important accuracy-selection issue: main can select a less accurate kernel for a negligible speed advantage. In the new --fast-scan path, final selection prefers the lowest worst observed numerical error among validated, retimed finalists within 1% of the fastest finalist. This preserves better numerical accuracy in cases where the latency difference is minor.

The tuner currently JIT-compiles each coupled G1/G2 pair inside the GPU scoring loop and rebuilds/reshuffles the full expert fixture plus torch reference for every candidate. On token=1, hidden=4096, inter=2048, experts=288, topk=8 (GLM-5.3-Flash TP=1) a cold sweep takes ~223 s, of which ~123 s is serial compilation.

Measured on MI355X (gfx950), same shape, fixed seed:

Measurement	Before	After	Speedup
Compile-only, empty cache	122.8 s	12.9 s	~9.5x
Cold full scan	223.4 s	25.4 s	~8.8x
Warm full scan	99.2 s	17.1 s	~5.8x

Accuracy-selection fix: preserve accuracy when latency is nearly equal

Main uses numerical accuracy as a pass/fail gate, then chooses the lowest measured latency. Once candidates pass --errRatio, lower error does not influence the winner. Consequently, even a tiny measured speed advantage can displace a substantially more accurate candidate; exact latency ties retain the earlier candidate regardless of accuracy.

The new --fast-scan final selection changes that behavior:

1. Find the fastest valid, retimed finalist.
2. Admit finalists whose latency is at most 1% higher.
3. Select the lowest worst observed numerical error in that set, using latency to break an accuracy tie.

The raw, unrounded error is retained across coarse warm-up, coarse measured-output, and final measured-output checks. A final retime cannot overwrite a larger previously observed error with a smaller one. Numerical rejection thresholds remain in force.

Illustrative example, assuming both candidates pass validation and reach final retiming:

Candidate	Measured latency	Observed error	Main selection	New fast-scan final selection
A	100.0 µs	0.020	Selected for lowest latency	Not selected
B	100.5 µs	0.005	Discarded despite lower error	Selected: 75% lower error for 0.5% higher latency

These are illustrative values demonstrating the selection rule, not new GPU benchmark results or a claim about downstream model accuracy. The reported token-1 benchmark below selected the same kernel pair on main and the PR.

Scope: this corrects cases of unnecessary accuracy loss within the retimed finalist set. It does not guarantee the most accurate eligible candidate across the entire coarse pool, because the funnel may exclude a candidate before final retiming. The ordinary scan path without --fast-scan remains latency-first.

Sources: main's accuracy gate, main's latency-only winner selection, and PR's 1% band and accuracy-based final selection.

What changed

• csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py — opt-in --fast-scan for the MXFP4 FlyDSL tuner:
  • legal coupled pairs are converted to deduplicated FlyDSL jobs and precompiled in parallel before GPU scoring
  • one immutable fixture and one torch reference per shape, reused across the whole sweep
  • cosine error is reduced on the device; only the final scalar is copied to the host
  • funnel: coarse scan of every legal pair -> lowest-latency performance finalists plus one distinct accuracy promotion -> longer retime with no repeated warm-up -> latency-ordered backfill when a finalist fails final validation
  • candidate-local accuracy/launch failures reject only that candidate; a shape fails only when no candidate survives
  • --fast-scan is rejected by the non-MXFP4 tuners at parse time and at dispatch
• aiter/aot/flydsl/common.py — optional start_method for the compile file pool (compile-only children can be forked from a clean fork server)
• aiter/aot/flydsl/mxfp4_moe.py — batch-compile helper and in-memory CSV parsing so tuner candidates can reuse the AOT job builder
• aiter/utility/mp_tuner.py — skip the 3 s pid stagger when mp_num == 1
• aiter/test_common.py — profile_warm_iters is now configurable instead of hard-coded to 1
• csrc/kernels/mxfp4_moe/moe_aux/codegen/gen_instances.py — add GLM-5.3 / GLM-5.3-Flash / Qwen3-Coder-Next / Inkling / MiMo-V2.5-Pro shapes
• Configs / docs / tests — glm53_fp4_tuned_fmoe.csv + glm53_fp4_untuned_fmoe.csv, a "Fast MXFP4 FlyDSL scan" section in the codegen README, and a CPU-only regression test (op_tests/tuning_tests/test_fast_tuning.py) covering the funnel, finalist selection, backfill order, accuracy parsing and CLI gating

Fast-scan arguments and usage

The following controls are implemented by Mxfp4FlydslTuner. Enable them with --mxfp4-flydsl --fast-scan; --fast-scan is rejected for unsupported tuner modes. An iteration here means one complete coupled MoE invocation (sorting, applicable quantization, G1, G2, and associated reduction), not one individual GPU kernel.

Argument	Type	Default	Behavior and limits
--fast-scan	Boolean flag	Off	Enable parallel candidate precompilation, one shared fixture/reference per shape, and the coarse/final measurement funnel. Omit the flag to retain the ordinary scan path.
--fast-scan-iters N	Integer	10	Measured invocations per candidate in the initial round. Must be at least 2; smaller values are rejected at argument parsing and runtime.
--fast-scan-final-iters N	Integer	100	Measured invocations per finalist in the final round. Must be at least 2. This is a configurable measurement budget, not a kernel correctness requirement.
--fast-scan-finalists N	Integer	5	Number of lowest-latency coarse survivors selected for final retiming. Runtime clamps the requested count to at least 1; the available survivor pool limits the actual count. One additional distinct accuracy candidate is promoted when available, so the default target is up to 6 successful retimes.
--fast-scan-warmup-accuracy-checks N	Integer	1	Untimed invocations whose outputs are validated before each candidate's coarse measurement. Runtime clamps negative values to 0. Setting 0 skips these premeasurement checks; the final output of the measured block is still checked.

In fast-scan mode, the dedicated controls above supply the measurement and warm-up counts instead of the ordinary --iters/--warmup settings. The existing --errRatio accuracy gate still applies. Candidate enumeration is separate: the existing --mxfp4-search-mode defaults to prune; use --mxfp4-search-mode full to enumerate all statically supported GEMM1 variants. The reported 174 candidate pairs describe the benchmark's configured search space, not every possible shape/search mode.

Default-budget example

Run from the repository root:

python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py \
  -i aiter/configs/model_configs/glm53_fp4_untuned_fmoe.csv \
  -o /tmp/glm53_fp4_fast_scan_tuned.csv \
  --mxfp4-flydsl --fast-scan --mp 1 --batch 1 \
  --fast-scan-iters 10 \
  --fast-scan-final-iters 100 \
  --fast-scan-finalists 5 \
  --fast-scan-warmup-accuracy-checks 1

The four numeric settings are shown explicitly for discoverability; omitting them uses these same defaults. For the reported 174-candidate case with no failures, this means 174 untimed warm-up/check invocations, 1,740 coarse measured invocations, and 600 final measured invocations.

Validate the generated rows through production dispatch:

python3 csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py \
  --mxfp4-flydsl \
  --run_config /tmp/glm53_fp4_fast_scan_tuned.csv

Adjusting the measurement budget

For an experimental shorter final round, use the same tuning command with --fast-scan-final-iters 32 and a separate output path. Keep --fast-scan-iters 10 to leave initial screening unchanged. This is an example of the configurable budget, not a claim that 32 iterations has been validated to retain every winner or reproduce the published benchmark. Increase the initial-round count when coarse rankings are unstable; adding final samples cannot recover a candidate already excluded from the finalist set.

Finalists do not receive another untimed warm-up: they have already been compiled and exercised during the initial round (normally 1 untimed + 10 measured invocations). Repeating initialization warm-up is intentionally avoided. This does not assert that cache contents or GPU clock conditions remain identical between rounds. There is currently no separate final-round warm-up argument.

Timing, filtering, and selection semantics

• Both rounds use the profiler-derived mean of summed active GPU execution durations, excluding host dispatch and idle gaps. The fast path uses profile_warm_iters=0, so it does not automatically discard the first profiled invocation.
• Outlier filtering is enabled independently for a measured block only when its requested iteration count is greater than 30. It removes whole-iteration totals outside the Q1 - 1.5*IQR / Q3 + 1.5*IQR fences and averages the retained samples. Therefore the default 10-iteration coarse round is unfiltered, whereas the 100-iteration final round is filtered. Changing 100 to 30 disables this filtering; changing it to 32 retains it. The threshold is currently hardcoded.
• Warm-up outputs and the final output of each measured block are accuracy-checked outside timing; every measured invocation's output is not independently checked. Worst observed error is retained across coarse and final checks. Non-finite errors, excessive error, or launch failures reject that candidate; failed finalists are backfilled in coarse-latency order.
• Final selection admits retimed candidates within 1% of the fastest retimed finalist, then chooses the lowest worst-case error, using latency to break an accuracy tie. This 1% band is currently hardcoded and is relative to measured finalists, not a guarantee against an unmeasured global optimum.
• A fixed 100-sample final round is not adaptive: it neither stops early on a clear winner nor establishes a confidence interval. Fewer iterations reduce measurement work; winner stability should be checked on the target GPU and shapes before adopting a reduced default.

Implementation references at the reviewed head: arguments, candidate execution and two-stage scan, and timing/filtering.

Testing

• python3 -m py_compile on all changed files
• python3 -m unittest op_tests.tuning_tests.test_fast_tuning (CPU-only; no GPU required)
• GPU: --mxfp4-flydsl --fast-scan sweep and the --run_config verification pass on MI355X / gfx950

The fast scan still only checks kernel-level correctness. Selected rows should be validated through the production operator (for example via --run_config) before the tuned table is committed, because that also covers runtime dispatch and integration.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs (added automatically when gfx942 configs change); main branch always runs both MI35X and MI300X
ci:triton-355	Run the full Triton test suite on MI35X, not only the tests the change affects
ci:multi-gpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: DeepSeek-V4-Pro, MiniMax-M3-MXFP4
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6405 --add-label <label>

One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.

[Qubitium]

Note to AMD devs, this code/pr was originally created in early Sept but I currently do not have access to MI350 to re-validate. Please triggee CI and validate this again on MI350.

[Qubitium]

Added a Fast-scan arguments and usage section to the PR description, with the complete argument table, explicit tuning and production-validation commands, and measurement semantics checked against head 53f99ce.

Control	Default	Meaning
--fast-scan	Off	Enable the two-stage path; requires --mxfp4-flydsl.
--fast-scan-iters	10	Initial measured coupled MoE invocations per candidate; minimum 2.
--fast-scan-final-iters	100	Final measured invocations per finalist; minimum 2.
--fast-scan-finalists	5	Performance finalists, plus one distinct accuracy candidate when available. Values below 1 are clamped to 1.
--fast-scan-warmup-accuracy-checks	1	Initial untimed, accuracy-checked invocations. Values below 0 are clamped to 0.

Warm-up clarification: the first round already compiles and exercises every surviving candidate. A finalist normally reaches retiming after 1 untimed accuracy-check invocation and 10 measured invocations. The final round intentionally adds no redundant initialization warm-up. Cache/clock conditions may evolve between rounds, but that is separate from needing to compile or initialize the candidate again.

100 final iterations is adjustable, not a correctness requirement. The description includes a 32-iteration experimental alternative without claiming it has the same validated winner retention. The current collector enables 1.5×IQR outlier filtering only above 30 measured iterations: 10 and 30 are unfiltered; 32 and 100 are filtered. The fast path does not automatically discard the first profiled sample. These settings replace the ordinary iteration/warm-up counts in fast-scan mode.

The new documentation also explains accuracy checks, finalist backfill, the fixed 1% final selection band, and the distinction between candidate search-space selection and measurement budgets. Published benchmark results remain attributed to their original configuration. This update changes documentation only; no iteration defaults or kernel behavior were changed.

[Qubitium]

In my private AITER fork, I already run AOT compilation and tuning automatically during model loading when a model encounters a previously unseen configuration within a supported module/kernel path. The kernel implementations are prevalidated; new configurations are compiled, tuned, and numerically checked before use.

Once this PR is validated, I believe we can integrate the faster tuning flow into model loading upstream. That would make selecting and preparing kernels for new models largely automatic wherever existing kernel support applies, without requiring a separate offline tuning sweep.

[zhiding512]

IGTM.

[zufayu]

One area that may deserve a closer look before merging is the runtime performance impact. The tuning wall-time improvement is clear, but glm53_fp4_tuned_fmoe.csv is also used by the production dispatch path.

Could the reviewers please confirm that the committed configs do not regress performance versus the current main fallback? Ideally, it would be helpful to see fast scan vs full scan kernel selections and --run_config latency for the committed GLM-5.3 shapes on the same gfx950 system. The 1% threshold is applied only after the coarse scan narrows the candidate set, so the comparison would help close this risk.

[Qubitium]

One area that may deserve a closer look before merging is the runtime performance impact. The tuning wall-time improvement is clear, but glm53_fp4_tuned_fmoe.csv is also used by the production dispatch path.

Could the reviewers please confirm that the committed configs do not regress performance versus the current main fallback? Ideally, it would be helpful to see fast scan vs full scan kernel selections and --run_config latency for the committed GLM-5.3 shapes on the same gfx950 system. The 1% threshold is applied only after the coarse scan narrows the candidate set, so the comparison would help close this risk.

@zufayu @zhiding512 I no longer have access to MI355X with FP4 support. I only just got AMD dev cloud credits today for MI300 so I need help from AMD staff regarding validation. I just don't have the hardware.

[zufayu]

@zhiding512 Could you or another reviewer with access to gfx950/MI350 please help run the performance validation discussed above? The author no longer has access to the hardware.

The main remaining gap is to compare fast scan vs full scan for the committed GLM-5.3 shapes, including the selected kernel pairs and --run_config latency, and to compare the committed PR configs against the current main fallback. This would help us close the runtime regression risk before merging.

[zufayu]

One CI failure does need a PR-side fix: op_tests/tuning_tests/test_fast_tuning.py is executed directly by the standard test runner. It currently cannot import csrc, and the GLM53 config lookup resolves through the installed aiter_meta path. Please derive the repo root from Path(__file__).resolve().parents[2], add it to sys.path before the csrc imports, and resolve the config from that root.

The separate HSTU dtype failure appears unrelated to this PR; it looks like global default dtype leakage from another test.