[zhiding512]

Folds in #6310 (opus fake signatures, now closed). It is kept here despite the one-backend rule because the replay below is the only code that runs those ops under FakeTensorMode, and it cannot land without the fix.

Summary

Replaces the FlyDSL MoE AOT scripts that restated the runtime's kernel selection (aiter/aot/flydsl/moe.py, mxfp4_moe.py, fhmoe.py, grouped_moe.py) with a CPU-only replay of the runtime itself. For every tuned row, moe.py builds what a caller passes fused_moe as FakeTensors and calls _fused_moe_impl (_fhmoe_impl for shared-expert rows) with COMPILE_ONLY=1, so every FlyDSL kernel on the runtime's path is compiled into the cache instead of launched. That includes the gfx1250 grouped table, which _fused_moe_impl hands to FlyDSL's grouped GEMM.

MegaMoE gets the same treatment: mega_moe.py builds the runtime's MegaMoEV2 for each rank of each deployment profile in the same dry run and calls its own preload_aot_bundles(), the method serving uses to load a profile's kernels.

Running the runtime this way also surfaced two bugs outside the AOT scripts, fixed here: the fakes of two opus a8w4 MoE ops lacked parameters, and merging tuned-config tables wrote integers as N.0, which the gfx1250 grouped MoE runtime cannot parse.

Which archs get built follows the build: each row compiles for the arch its gfx column names, and with GPU_ARCHS naming targets, as setup.py's wheel build sets it, only rows of those archs are replayed. Each arch replays in a process of its own that aiter's runtime arch queries answer for that arch, explicitly (AITER_AOT_REPLAY=1), never for the machine it runs on.

Motivation

• The old scripts mirrored kernel names, launch parameters and host-side decisions, so every kernel or wrapper change needed a matching AOT change. Recent examples: Kimi-K3 MegaMoE ([FlyDSL] [CI] megamoe tune kimi-k3 #6259) changed the MegaMoE AOT by +121/−11 lines, A4W4 MegaMoE ([FlyDSL] [MegaMoE] A4W4 mega_moe operator #4980) by +17/−7 to teach it the FP4 quant mode and tensor shapes, and [FlyDSL] fix(flydsl): support kmk3 ep moe & opt for ep moe gemm2 #5804 had to teach mxfp4_moe.py its new GEMM2 compile parameter. With the replay, [FlyDSL] [MegaMoE] A4W4 mega_moe operator #4980 needs only "a4w4" added to the profile list.
• They drifted:
  • Against a full build of the old MoE scripts, the runtime needs 249 kernels they never compiled, which JIT on the first inference call. Most are the heuristic kernels the runtime switches to when a tuned config cannot take an expert bias; others have compile-time parameters that follow the actual M.
  • grouped_moe.py compiled only the auxiliary kernels around the gfx1250 grouped GEMM. Since the runtime moved that GEMM to the TDM batched kernel, the GEMM itself compiled on first use, and most of the auxiliary variants it did compile were for the removed non-TDM path.
  • The old mega_moe.py compiled only the mask_topk_ids=False combine kernels, so forward(mask_invalid_slots=True) compiles online, and fails under FLYDSL_RUNTIME_RUN_ONLY=1.
• AOT must run without a GPU.

Changes

• aiter/aot/flydsl/dry_run.py (new): context manager for the replay. It sets COMPILE_ONLY=1, pins the target arch, CU count and optionally the device name (MegaMoE's combine geometry is tuned per GPU model), enters FakeTensorMode, and stubs or forbids the GPU runtime (CUDA init, streams, events and device queries, the FlyDSL execution engine). It also turns off the runtime's rebuild of a C++ module whose device code lacks the target arch: only host code runs there, and the per-arch replays share one build.
• aiter/aot/flydsl/moe.py (rewritten, 1186 to 726 lines):
  • Parses the tuned FMoE, FHMoE and gfx1250 grouped tables with strings only, so the parent process stays on the light AITER_AOT_IMPORT import.
  • Enumerates per row what the row does not fix. On the caller side: expert bias for Swiglu + MXFP4 rows, the GPT-OSS path, as fused_moe enables it (both stages, or stage 1 only), both gate/up layouts for MXMOE stage 1, SiTUv2 betas (unset, 1.0, model defaults), num_local_tokens (passed by callers behind an EP dispatch), and for FHMoE rows with MXFP8 routed weights both values of clamp_shared ([FlyDSL] [CI] Add HY4 MXFP8 heterogeneous MoE fusion #6050), which the runtime compiles into stage 1. For M: both values get_padded_M maps to the row (token // 2 + 1 and token).
  • Where a variant's two token counts request different FlyDSL kernels, identified by FlyDSL's own cache key, it bisects the bucket down to every M at which they change, because a row can switch more than once in its bucket: for 7168x768, E384, topk 6 at token 2048, the gfx1250 stage-1 quant uses K split 7, then 4, then the compact quant. These replays are reported as [BISECT].
  • Replays gfx1250 grouped rows with the gate/up-interleaved weights and n32k4 scales gfx1250 callers pass. Grouped rows tuned for EP only (topk −1, or ep_fused) are skipped: the runtime picks them only under a caller's expert mask or MegaMoE's stage-2 scatter, and the table records neither the global expert count nor the real topk.
  • Replays each job. Caller-side combinations the runtime itself rejects (NotImplementedError) are reported as [SKIP]. A row's default caller must be served: token counts at which the runtime refuses it while serving it at others (an M-dependent limit, such as the 32-bit route-reduce workspace that op_tests/test_moe_2stage.py also skips) are reported as [REFUSED]. A row refused at every token count, or any other error, fails the build.
  • Replays the archs GPU_ARCHS names, intersected with the tables' gfx column (every arch with rows when it is unset or native), so CI's gfx942;gfx950 wheel build skips the gfx1250 rows.
  • Runs one process per arch, started with GPU_ARCHS naming it and AITER_AOT_REPLAY=1 and with the GPUs hidden, since aiter fixes arch-dependent values such as dtypes.fp8 at import. Each process checks at entry that it was started that way and that dtypes.fp8 is its arch's type. aiter's C++ core module, which has host code only, is built once before they start: not every toolchain builds every arch (ROCm 7.2.4 cannot build gfx1250's cluster launch).
  • Schedules the jobs:
    • A worker replays up to 12 jobs of one row in turn, so the kernels they share compile once instead of concurrent variants waiting on each other's compile locks.
    • Equal-sized batches are shuffled stably, so rows of one model do not start together.
    • The arch processes run side by side and split the AITER_MAX_JOBS worker budget ([HIP] [OPUS] [FlyDSL] Simplify AITER compiler worker fan-out to smooth-out regressions #5164's shared policy) by job count.
  • --check compiles nothing and fails on any kernel the replay needs but the cache lacks (FLYDSL_RUNTIME_RUN_ONLY=1).
• aiter/aot/flydsl/mega_moe.py (rewritten, 590 to 305 lines): one job per rank of each deployment profile (DeepSeek-V4-Pro r0/r32/r64, Kimi-K3; MTPR 8192/16384/32768) and quant (A8W4, and A4W4 since [FlyDSL] [MegaMoE] A4W4 mega_moe operator #4980). Each job builds MegaMoEV2 with FakeTensor weights and calls preload_aot_bundles() and quantize(), which the preload leaves out but fixed-slot forward and the per-stage benchmarks use. mori shmem, through which MegaMoEV2 allocates symmetric buffers and exchanges peer pointers, hands out FakeTensors and fake addresses; a build host without mori gets a stand-in module. One job covers both stages, so the old script's 384 jobs become 192. The module keeps only the profiles; CLI options are unchanged. Like the MoE replay, it runs in a process of its own, started with GPU_ARCHS=gfx950 and AITER_AOT_REPLAY=1, and it is skipped when GPU_ARCHS names targets without gfx950.
• aiter/aot/flydsl/common.py: run_aot runs the MoE and MegaMoE replays first, in processes of their own (replay_env), then the pooled kinds. OpKind.MXFP4_MOE and OpKind.GROUPED_MOE are removed.
• Removed aiter/aot/flydsl/fhmoe.py and grouped_moe.py; the replay covers their rows. quant_layout_supported() in moe_fused_route_quant_scatter.py, which [FlyDSL] Kimi-K3 TP8 gfx1250 grouped-MoE tuning (experts split 8 ways) #6768 added for grouped_moe.py alone, goes with it. mxfp4_moe.py leaves the AOT build too, but stays: the MXFP4 tuner's --fast-scan ([HIP] [CK] [FlyDSL] Accelerate MXFP4 FMoE AOT tuning ~9.5x on gfx950 #6405) precompiles its candidates with it.
• aiter/jit/core.py: update_config_files reads the tables it merges as text, so the merged file reproduces every cell as written. Parsed, a column with a blank cell turned float and its integers were written back as N.0: 6 of the 15 merged tables had such cells, and since [Config] [Tuning] Add tuned grouped MoE config for Qwen3.5 FP4 on gfx1250 #6090 added a second grouped table, the gfx1250 grouped MoE runtime raised ValueError on every row of the merged one. Shapes are still compared on numeric values, so 256 and 256.0 stay one shape, and the fastest duplicate is still picked on a numeric us. Since [JIT] Check single-file tuned configs for duplicate shapes #6365 that check lives in _raise_on_duplicate_shapes, shared with the single-file path, which now reads its table as text too.
• Runtime-side support:
  • aiter/jit/utils/torch_guard.py: run_host_scalar_ops_for_real() lets scalar-returning custom ops run their implementation under FakeTensorMode, since their fakes carry no value. Off by default.
  • aiter/jit/utils/chip_info.py: in an AOT replay process (AITER_AOT_REPLAY=1, set only by replay_env), get_gfx_runtime() returns the one arch GPU_ARCHS names and does not probe the device, so a visible GPU of another arch cannot leak in. Elsewhere it behaves as on main, which never lets GPU_ARCHS override it.
  • aiter/utility/dtypes.py: since [FlyDSL] [JIT] Fix Windows build arguments and optional imports #6589, get_dtype_fp8() returns uint8 when torch sees no device. An AOT replay process, which hides the GPUs, now takes its arch's fp8 type from get_gfx_runtime() instead; every other process keeps main's behavior. Without it, every fp8 path of the replays took the wrong branch: fp8 activations were quantized to fp4, MXFP8 rows found no input recipe, and gfx1250 grouped rows compiled other kernels than the runtime requests without any error. op_tests/test_dtype_parser.py (main's tests unchanged) and op_tests/test_build_targets_architecture.py cover the flag.
  • aiter/fhmoe.py: fhmoe_ is split into a plain _fhmoe_impl and the decorated op (same signature, including [FlyDSL] [CI] Add HY4 MXFP8 heterogeneous MoE fusion #6050's clamp_shared), so the replay can call the implementation under FakeTensorMode.
  • aiter/ops/opus/moe_stage1_a8w4.py, aiter/ops/opus/moe_stage2_a8w4.py: the fakes registered for _opus_moe_stage1_a8w4_fwd_raw and _opus_moe_stage2_a8w4_decode_fwd_raw lacked the ops' scalar parameters (topk, block_m, kernelName, inter_dim_pad, activation, swiglu_limit, situ_beta, situ_linear_beta and token_num, topk, block_m, kernel_id, inter_dim_pad). torch_compile_guard calls the fake with every argument the op takes, so any call under FakeTensorMode raised TypeError. The fakes still return out.
• op_tests/test_fhmoe.py: the tests of the removed script's internals are replaced by five tests of the replay: row parsing (DSV4 i384 and HY4 MXFP8 rows, clamp_shared), import under AITER_AOT_IMPORT=1, padding columns, the caller choices handed to _fhmoe_impl, and, on gfx950, replaying HY4 rows into a scratch cache and loading their kernels under FLYDSL_RUNTIME_RUN_ONLY=1.
• aiter/aot/flydsl/README.md: documents how both replays work and --check.

Runtime behavior is unchanged apart from the config merge: the torch_guard switch is off outside the replay, the chip_info and dtypes changes only apply under AITER_AOT_REPLAY=1, and eager execution never calls the opus fakes. Merged tables now carry their sources' cells unchanged; no value differs from before, and the gfx1250 grouped lookup that raised now parses.

Performance (AOT build time, 256-CPU host, 96 workers)

MoE kinds, measured before the last merge of main:

Build	Old scripts (pooled run_aot path)	This PR
From scratch	170 s for MoE + MXFP4 MoE (3713 jobs, 2595 artifacts), plus 10 s for grouped MoE	166 s (22350 jobs plus 1588 bisection replays, 2353 artifacts)
Rebuild, cache already complete	not measured	30 s
--check of a complete cache	n/a	44 s

Run-to-run variation is about ±6%. Without the scheduling above (one forked worker per job, archs one after another), the from-scratch build took 537 s. That build was bound by the parent's fork rate and by concurrent variants of a row waiting on each other's compile locks. The artifact count follows the runtime: it fell from 2751 to 2353 when #6376 limited FlyDSL expert bias to Swiglu rows, rose to 3463 (22698 jobs) after #6460 reverted that and #6405 added the glm53 fp4 table, and is 2406 (22992 jobs) since #6478 limits it to Swiglu again.

On the current head with #6700 (FlyDSL 0.4.0, AITER_MAX_JOBS=160), the MoE build takes 124.5 s from scratch for every arch (9440 jobs plus 660 bisection replays; 8948 of the jobs remain under CI's GPU_ARCHS=gfx942;gfx950). Before 2f85842 limited expert bias to Swiglu rows, it took 182 s at AITER_MAX_JOBS=96 (22992 jobs plus 1696 bisection replays). MegaMoE, A8W4 and A4W4, takes 415 s (192 jobs).

The wheel build's run_aot now runs MegaMoE as a phase of its own instead of inside the shared pool (the 180 s below).

MegaMoE, same host and workers, A8W4 only (measured before #4980 added A4W4, which builds about as much again):

Build	Old script	This PR
From scratch	222 s (192 jobs, 3018 artifacts)	180 s (96 jobs, 3162 artifacts)

Testing

Unless marked as measured on the current head with #6700, the numbers below were measured on 61fa528 or earlier, before 2f85842 limited expert bias to Swiglu rows.

• Full replay on the current head with [JIT] Replay biased interleaved MoE callers again and fix #6362's CI failures #6700 (main through [MHA] Fix -DONLY_FAV3=0 quoting that breaks the backward JIT build #6616, FlyDSL 0.4.0; all tuned FMoE, FHMoE and gfx1250 grouped rows; gfx942, gfx950, gfx1250): 1551 rows, 9440 jobs and 660 bisection replays (gfx1250 422, gfx950 238), 0 failures, 2191 artifacts, and --check of that cache passes. Caller-side combinations the runtime itself rejects are reported as [SKIP] (153, all on gfx950): no safe MXFP4 backend at inter_dim 384/640/1152, which includes MiniMax-M3 a4w4 with interleaved weights and a bias, and the 32-bit route-reduce limit. One token count is [REFUSED] (see Known gaps). Without [JIT] Replay biased interleaved MoE callers again and fix #6362's CI failures #6700, head 2f85842 replays 9116 jobs into 2156 artifacts and lacks 35 the runtime requests: GPT-OSS a8w4's interleaved, biased kernels on gfx950 and the biased gfx1250 callers'.

• Target archs: with GPU_ARCHS=gfx942;gfx950, as CI's wheel build sets it, the replay skips the 1428 gfx1250 jobs and leaves exactly the gfx942 and gfx950 artifacts of the full replay (84 and 2057, same keys); with GPU_ARCHS unset, all 2406. MegaMoE leaves the same 6236 artifacts as before and is skipped under GPU_ARCHS=gfx942.

• Without the arch reaching the replay processes (each inheriting CI's GPU_ARCHS=gfx942;gfx950, so dtypes.fp8 is uint8 since [FlyDSL] [JIT] Fix Windows build arguments and optional imports #6589), the gfx950 replay fails 530 jobs, misses 503 of its 2057 artifacts and builds 294 wrong ones; the gfx1250 replay reports no failure while 185 of its 265 artifacts are missing, all 130 fp8 TDM grouped GEMMs among them. gfx942 is unaffected only because none of its rows has fp8 activations. With the flag, a gfx942 replay process on this gfx950 host gets float8_e4m3fnuz and gfx942 from get_gfx_runtime() even with a GPU visible.

• Artifact count across the merge: replaying the pre-merge head and the current head with the same FlyDSL 0.4.0 gives 3463 and 2406 artifacts. The difference is [HIP] [CK] [MoE] Fix missing SwiGLU non-split no-bias CK-Tile MXFP4 dispatch (follow-up to #6376 revert) #6478's limit of FlyDSL expert bias to Swiglu: all 405 MXMOE gemm2 kernels that disappear are bias variants whose bias-free twin is still built, the FlyDSL a8w4/a4w4 GEMMs drop their bias specialization (per kernel name, exactly half the artifacts where the bias variant applied), and a bias=both replay of a Silu row requests one extra stage-1 kernel before the merge and none after. The only new kernel names come from tables main added.

• Wheel-build emulation: the AOT step run the way CI's setup.py bdist_wheel runs it (aiter/jit/core.py loaded a second time as jit.core, aiter imported with AITER_AOT_IMPORT=1, every JIT .so removed first, GPU_ARCHS=gfx942;gfx950). Before the fix, CI's wheel build failed on both replays: the gfx950 MoE replay found module_moe_mxfp4_aux built for gfx942 only, by the gfx942 replay running alongside, and died; all 96 MegaMoE jobs raised KeyError: 'module_aiter_core'. Now the MoE replay (all three archs), MegaMoE (96/96) and the full run_aot (GEMM 6787, CONV 133, CHUNK_GDN_H 288, FMHA_FP8 92) pass, and the two replays leave exactly the artifacts of the standalone builds (5515 = 2353 + 3162). A later CI run failed again when the gfx1250 replay happened to build aiter's C++ core module first: gfx1250 builds turn on cluster launch, which the HIP of ROCm 7.2.4 does not declare. That module has host code only and is now built once before the replays start; in the emulation, only that step builds it. On the current head the emulation replays gfx942 and gfx950 only, as GPU_ARCHS names, with 0 failures. CI's prebuild wheel job passed on the head before the GPU_ARCHS filter, which still replayed gfx1250.

• Dense M check: replaying every bucket at about 16 sampled M (every M up to token 64) under FLYDSL_RUNTIME_RUN_ONLY=1 finds no missing kernel: 132834 jobs on gfx950, 16716 on gfx942, 11160 on gfx1250. Before the bisection, 138 of the gfx1250 probes missed the stage-1 quant kernel of an intermediate K split; gfx942/gfx950 had no such case, and bisection compiles nothing new there.

• The scheduled build leaves exactly the same artifact set as one worker per job, and --check passes on it.

• Opus fakes: called under FakeTensorMode with every argument the ops take, both raise TypeError on main (takes 10 positional arguments but 18 were given, and 9 of 14 for stage 2) and return out here. Without the fix, 43 replay jobs fail.

• MoE artifacts vs. a full build of the old scripts (same FlyDSL 0.3.4.1, on the base before the merges):

  • 2502 shared artifacts, all byte-identical ELF.
  • 93 only in the old build, none reachable by the runtime:
    • 78 are bias variants of tuned configs that the runtime discards when a bias is present; this is confirmed per row from the runtime's own log.
    • 13 are rows with token=3, which get_padded_M never returns.
    • 2 are CK-Tile epilogues for rows that either have no split-K or are tagged flydsl_fallback, which the runtime drops at load.
  • 249 only in this build: kernels the runtime uses that the old scripts missed.

• [FlyDSL] fix(flydsl): support kmk3 ep moe & opt for ep moe gemm2 #5804's use_valid_token_count kernels: of the 442 the old mxfp4_moe.py builds for it, 215 are in the replay's cache, byte-identical. The other 227 are never requested by the runtime: bias variants of Silu/SiTUv2 rows, inline-sort rows whose tuned config the runtime drops under num_local_tokens (it then uses heuristic FlyDSL kernels, which the replay builds and the old script did not), and a token=3 row.

• gfx1250 grouped artifacts vs. grouped_moe.py (its own table, 117 rows): 21 shared, byte-identical. 162 only in this build, including 112 TDM grouped GEMM kernels the old script never built. 149 only in the old build: 109 auxiliary variants for the removed non-TDM path, and 40 built with topk −1 for the EP-only rows, which no call matches. No gfx1250 GPU was available, so this path is validated on CPU only.

• Kimi-K3 TP8 grouped table ([FlyDSL] Kimi-K3 TP8 gfx1250 grouped-MoE tuning (experts split 8 ways) #6768, inter_dim 384): [FlyDSL] Kimi-K3 TP8 gfx1250 grouped-MoE tuning (experts split 8 ways) #6768 made grouped_moe.py skip the stage-2 activation quant, whose layout needs inter_dim % 256 == 0. The replay needs no such rule: the runtime launches that quant only when gemm1 cannot requantize its own output (a gemm1 bias, or AITER_FLYDSL_DISABLE_GEMM1_REQUANT=1), and SiTUv2 callers pass no bias. The table's 15 rows replay as 168 jobs and 78 bisection replays, 0 failures, into 31 artifacts with no inter_dim-384 quant kernel; every gemm1 requantizes in its epilogue. With AITER_FLYDSL_DISABLE_GEMM1_REQUANT=1, every job fails on the assertion [FlyDSL] Kimi-K3 TP8 gfx1250 grouped-MoE tuning (experts split 8 ways) #6768 worked around, as the runtime itself would. Of the 11 artifacts main's grouped_moe.py builds for the table, 8 are in the replay's cache, byte-identical, and 3 are never requested: K-split route quant for wmma_rep 2 and 8, whose rows serve only 1025+ tokens, where the runtime turns K-split off, and a dense route-quant-scatter kernel the runtime no longer calls. The replay adds the table's 23 TDM GEMMs.

• Config merge: regenerating all 15 merged tables changes no value, and each now has exactly as many N.0 cells as its sources. A duplicate written as 256 in one table and 256.0 in another is still caught and resolved to the faster row. op_tests/tuning_tests/test_config_shape_collision.py, test_tuner_infra.py and test_fast_tuning.py pass (72 tests).

• Strict cache-hit test on MI355X (gfx950): op_tests/test_moe_2stage.py --no-legacy under fail_on_aot_cache_miss, with a cache built by this PR. On the current head (main through [MHA] Fix -DONLY_FAV3=0 quoting that breaks the backward JIT build #6616, FlyDSL 0.4.0), every model CSV with FlyDSL rows on gfx950 passes: 24 CSVs, 1198 cases, no cache miss. Every CSV has the same case count as on the previous head (main through [FlyDSL] fix(comm-fused-moe): key winners by GPU model #6541, FlyDSL 0.3.4.1, 1199 cases) except qwen3_8_2400b_fp4, where the test now skips the token-32768 row itself. These runs predate 2f85842. On the current head with [JIT] Replay biased interleaved MoE callers again and fix #6362's CI failures #6700, the 8 gptoss_fp8fp4 cases (interleaved weights with stage-1 and stage-2 bias) pass with no cache miss against a cache built from that table; without [JIT] Replay biased interleaved MoE callers again and fix #6362's CI failures #6700 they fail as CI's MI35X shard 7 did. On the base before the merges, 662 cases across 12 model CSVs all passed with no cache miss:

  Model CSV	Cases
  dsv4_fp8fp4	182
  dsv3_fp4	89
  kimik2_i4	88
  kimik2_fp4	80
  dsv41_fp8	77
  kimik3_a8w4	34
  glm5_mxfp8, minimax_m3_mxfp8	32 each
  gptoss_fp4	24
  gptoss_fp8fp4, minimax_m3_a16w4, minimax_m3_mxfp4	8 each

  kimik3_a16w4 has no cases because the test skips SiTUv2 a16w4. Re-run after the merges with the full default build: glm5_mxfp8 (32) and dsv3_fp4 (89) on idle GPUs, kimik2_i4 and dsv4_fp8fp4 (59 and 85 cases; the test skipped the rest because other jobs left too little free GPU memory for its reference) all pass with no cache miss.

• AOT vs. on-GPU JIT: glm5_mxfp8 and dsv3_fp4 run again with an empty cache and the AOT check off. Every kernel the runtime JIT-compiles on the GPU (41 and 83) is byte-identical to the replay's. logits_diff varies from run to run either way: two runs on the same AOT cache differ in more cases than AOT vs. JIT.

• MegaMoE full build: on the current head (FlyDSL 0.4.0), A8W4 and A4W4 each pass all 96 jobs.

• MegaMoE A4W4 artifacts vs. the old script on current main (same FlyDSL 0.3.4.1): all 3218 old artifacts are present; 3082 are byte-identical and the other 136 differ only in the same v_add3_u32 vs. v_add_u32 + v_sub_u32 selection (2-7 asm lines). The 144 only in this build are again the mask_topk_ids=True combine kernels.

• MegaMoE artifacts vs. a full build of the old script (same FlyDSL 0.3.4.1):

  • All 3018 old artifacts are present; 2921 are byte-identical.
  • The other 97 differ only in how one address sum is selected (v_add3_u32 vs. v_add_u32 + v_sub_u32, 2-7 asm lines). Building the old script twice gives 109 such differences, so this is FlyDSL/LLVM run-to-run variation, not the replay.
  • 144 only in this build: the mask_topk_ids=True combine kernels (combine kernels 144 -> 288).

• MegaMoE on 8x MI355X under FLYDSL_RUNTIME_RUN_ONLY=1, with the cache built by this PR: op_tests/multigpu_tests/test_mega_moe_v2.py --max-tok-per-rank 8192 --bs-list 1,128,1024,8192 --pad-tail-rows 4 passes for v4_pro and kimi_k3_route, A8W4 and A4W4, on the previous head (main through [FlyDSL] fix(comm-fused-moe): key winners by GPU model #6541, FlyDSL 0.3.4.1; not yet re-run on this one) (relL2 0.059 and 0.057 for A8W4 against rtol 0.10, 0.212 and 0.219 for A4W4 against rtol 0.25; CUDA-graph replay unchanged, masked rows zero). With the old script's cache, the --pad-tail-rows run fails with no usable AOT cache for combine_launch.

• The MegaMoE replay also runs on a host without mori, including twice in one process.

• op_tests/test_fhmoe.py, the AOT tests plus the clamp_shared and op-schema tests, on gfx950: 12 passed, including the run-only load of replayed HY4 kernels.

• op_tests/test_flydsl_moe_aux.py (script, gfx950): every case matches exactly, including gather-reduce with f32, bf16 and f16 route weights. op_tests/test_dtype_parser.py and op_tests/test_build_targets_architecture.py pass with GPU_ARCHS unset, gfx950 and gfx942;gfx950; under the earlier GPU_ARCHS fallback, main's test_linux_runtime_errors_are_not_fallbacks failed with GPU_ARCHS=gfx950.

• black and ruff are clean on the changed files.

To reproduce:

python -m aiter.aot.flydsl.moe            # build, CPU only
python -m aiter.aot.flydsl.moe --check    # verify an existing cache without compiling
python -m aiter.aot.flydsl.mega_moe       # MegaMoE profiles, CPU only

Known gaps (follow-up)

• EP-only rows of the gfx1250 grouped table are not replayed: 30 of its 153 rows (7168x3072 with 96 local experts, a8w4 and a4w4, two of them ep_fused; 7168x2048 with 64 and with 65 local experts). The runtime picks them only when the caller passes an expert_mask (it then looks up topk −1) or uses MegaMoE's stage-2 scatter (ep_fused), and the table records neither the global expert count nor the real topk the kernels are specialized on. Covering them needs an EP deployment profile per shape (global experts, EP size, topk), like mega_moe.py's profiles, plus a Stage2ScatterContext for the two ep_fused rows. Until then, those kernels compile on first use. The old grouped_moe.py did not cover them either: it built them with topk −1, which no call matches.
• With the old merge, the two ep_fused rows also come out of the merged table as 1.0, which the runtime's string match never selects; the merge fix here keeps them 1.
• Token counts above a row's tuned token are not replayed: get_padded_M serves M up to 131071 with the row tuned at 32768 and every M from 131072 up with the row tuned at 131072, while the replay covers M from token // 2 + 1 to token. Probing the extra range under FLYDSL_RUNTIME_RUN_ONLY=1 finds nothing missing for the 32768 rows or on gfx1250, but on gfx950, 42 of 4020 probes miss a stage-1 or stage-2 kernel at M >= 262144, on all three shapes tuned at 131072 (for a8w4 7168x3072 with 48 local experts, from M = 262144). Those kernels compile on first use; the old scripts did not cover them either. Covering them means replaying each tier's upper end and bisecting up to a chosen M cap.
• The Qwen3 8x2400B fp4 row tuned at token 32768 cannot be served on main: its reduce-mode stage 2 needs a 5.4 GB route-reduce workspace, beyond the 32-bit buffer addressing [FlyDSL] [CI] Add HY4 MXFP8 heterogeneous MoE fusion #6050 guards, so the runtime raises NotImplementedError for every M from 26215 to 131071 that maps to the row, and test_moe_2stage.py skips it. The replay reports that token count as [REFUSED] and builds the row's kernels from the token counts it does serve. The fix belongs in the table (retuning that token count), not in the replay.
• CI's wheel no longer ships gfx1250 MoE kernels: its GPU_ARCHS is gfx942;gfx950, and listing gfx1250 there would also build the C++ modules for gfx1250, which the HIP of ROCm 7.2.4 cannot (hipLaunchAttributeClusterDimension is undeclared). The old scripts, and this PR before this change, built them regardless of GPU_ARCHS. They build on their own with GPU_ARCHS=gfx1250 python -m aiter.aot.flydsl.moe (README §1): an installed wheel already ships the host-only C++ core module that build imports, and on ROCm 7.2.4 from source, GPU_ARCHS=gfx950 python -c "import aiter" builds it first. From a fresh JIT directory that sequence replays the 1428 gfx1250 jobs and leaves the same 265 artifacts as the full replay (counts measured before 2f85842).
• gfx1250 is validated on CPU only so far; a run on a gfx1250 machine is pending: a strict cache-hit sweep of every grouped row (op_tests/flydsl_tests/test_flydsl_grouped_gemm.py --scenario csv, which runs under FLYDSL_RUNTIME_RUN_ONLY=1 by default), token counts inside the buckets the bisection covers (--scenario verify --tokens ...), and a byte comparison of the kernels the runtime JIT-compiles there (--no-check-aot-cache) against the AOT cache.

Documentation

• aiter/aot/flydsl/README.md updated

Dependencies

• No new third-party dependencies

Breaking Changes

• python -m aiter.aot.flydsl.grouped_moe and the aiter.aot.flydsl.fhmoe module are removed; python -m aiter.aot.flydsl.moe covers their rows. Nothing in the repo or CI calls them. mxfp4_moe.py stays, outside the AOT build.
• OpKind.MXFP4_MOE and OpKind.GROUPED_MOE are removed.
• python -m aiter.aot.flydsl.moe keeps --csv and adds --row-tokens-only and --check.
• When GPU_ARCHS names targets, the MoE and MegaMoE AOT (standalone and in run_aot) build only those archs; unset or native, every arch the tables have, as before.
• Merged tuned-config tables (<tmpdir>/aiter_configs/*.csv) now keep each cell as written in its source, e.g. 64 rather than 64.0.

[Copilot]

🟡 Changes recommended

Capped token buckets miss reachable runtime variants, temporary Mori modules leak globally, and check mode can still compile a host extension.

3 open findings

• Scope fake Mori modules to the context manager · New
• Enumerate all M values mapped to capped padding buckets · New
• Prevent --check from compiling host helpers · New

What changed in this PR

Replaces hand-maintained FlyDSL MoE AOT compilation with CPU-only runtime replay using FakeTensors, while extending the same approach to MegaMoE.

Changes:

• Adds dry-run infrastructure and runtime replay for MoE/FHMoE and MegaMoE.
• Adds CPU-only architecture detection and FakeTensor support.
• Consolidates AOT modules and updates tests/documentation.

File	Description
aiter/​aot/​flydsl/​dry_run.py	Adds GPU-free replay context.
aiter/​aot/​flydsl/​moe.py	Rewrites MoE AOT around runtime replay.
aiter/​aot/​flydsl/​mega_moe.py	Replays MegaMoE bundle loading.
aiter/​aot/​flydsl/​common.py	Runs MoE replay separately from pooled jobs.
aiter/​aot/​flydsl/​fhmoe.py	Removes obsolete FHMoE adapter.
aiter/​aot/​flydsl/​mxfp4_moe.py	Removes obsolete MXFP4 compiler.
aiter/​aot/​flydsl/​README.md	Documents replay and cache checking.
aiter/​fhmoe.py	Extracts the undecorated FHMoE implementation.
aiter/​jit/​utils/​torch_guard.py	Allows real scalar operations during replay.
aiter/​jit/​utils/​chip_info.py	Adds explicit CPU-only architecture fallback.
aiter/​ops/​opus/​moe_stage1_a8w4.py	Completes the stage-one fake signature.
aiter/​ops/​opus/​moe_stage2_a8w4.py	Completes the stage-two fake signature.
op_tests/​test_fhmoe.py	Updates FHMoE AOT parsing coverage.

🧠 Review effort: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Replay omits part of saturated token buckets, and runtime architecture detection can fail open to an incompatible target.

3 open findings

• Do not mask GPU detection failures with architecture fallback · New
• Prevent --check from compiling host helpers
• Scope fake Mori modules to the context manager

1 resolved since last review

• Enumerate all M values mapped to capped padding buckets

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

The broad cross-architecture replay and JIT integration needs human verification, particularly with gfx1250 GPU validation still pending.

8 open issues

• Use host-build fallback when no two-stage jobs exist
• Enforce total AOT worker cap across architectures
• Bias interleave layouts are incorrectly filtered from replay
• Accept FlyDSL implementation rows in stage-one filtering
• Add cache-key bisection tests for transitions and failures
• Test CPU replay execution and global patch restoration
• Add regressions for blank cells and numeric key deduplication
• Prevent --check from compiling host helpers

1 resolved since last review

• Pin the fp8 alias when emulating gfx950

_{Review effort: Balanced}

[zhiding512]

AOT cache diff: this branch vs main

Built the FlyDSL AOT cache from empty on both sides and compared the resulting artifacts.
Setup: main @ c2a682572, this branch @ 3563493b1, same box, FlyDSL 0.4.0, ROCm 7.2.4,
GPU_ARCHS unset (all archs), separate FLYDSL_RUNTIME_CACHE_DIR per side.

Modules run — main: moe + mxfp4_moe + grouped_moe + mega_moe; this branch: moe + mega_moe
(matching what run_aot dispatches on each side).

Both sides build clean

	kernels	result
main moe / mxfp4_moe / grouped_moe / mega_moe	2710 / 1783 / 356 / 384	all 0 failed
this branch, replay	gfx942 1218, gfx950 7576, gfx1250 492	all 0 failed

Plus 153 [SKIP] caller-side variants and 1 [REFUSED] token count on gfx950
(the 32-bit route-reduce workspace limit), and 660 [BISECT] probes, 0 failed.

Artifact diff

	main	this branch
artifacts	9507 (166 MB)	8427 (155 MB)
common	7972	7972
only on this side	1535	455

Split by module (MegaMoE isolated into its own cache dirs to attribute it exactly):

	main	branch	common	main-only	branch-only
MoE family	3415	2191	1944	1471	247
MegaMoE	6092	6236	6028	64	208

Which cache does the runtime actually want

--check replays the runtime and fails on anything the cache lacks, so it answers this directly:

--check run against	result
this branch's cache	0 failed on all three archs ✅
main's cache	1312 misses ❌ — gfx1250 0 ok / 492 failed, gfx942 252 failed, gfx950 568 failed

Missing kernels against main's cache:

493  launch_mixed_moe_gemm1
492  launch_topids_to_rows
252  launch
 45  launch_gemm1
 28  launch_silu_and_mul_fq
  2  launch_mixed_moe_gemm2

So the asymmetry is real: main's extra 1535 artifacts are largely combinations no caller reaches,
while main misses 1312 that are reached.

The cleanest instance is cache-key drift rather than a missing compile — launch_topids_to_rows:

main:   launch_topids_to_rows_193a5a5fbae511fbef597d5fc98e9e98
branch: launch_topids_to_rows_dbe460ae9a7818cb0e35d0afd200bbd7   <- what the runtime looks up

grouped_moe.py:252 does build it, but under a manager_key no caller ever requests. Every
gfx1250 grouped MoE row therefore JITs online today, while the AOT log reports 356 ok, 0 failed.

Four kernel names are absent from main's cache entirely — launch_gemm_a8w4_tdm (94),
launch_token_multidest (17), launch_compact_quant (9), launch_scatter_preshuffle (5) —
and MegaMoE combine_launch coverage doubles (144 -> 288).

Two caveats

1. The replay needs up-to-date C++ modules. It executes aiter's real host code, so stale
   prebuilt aiter/jit/*.so break it where the old enumeration never noticed. On my checkout all
   9440 jobs failed with
   AttributeError: type object 'aiter.jit.module_aiter_core.ActivationType' has no attribute 'Relu2',
   and the _load_host_helpers prebuild failed with
   module 'aiter.jit.module_moe_mxfp4_aux' has no attribute '_mxfp4_moe_sort_internal_is_supported',
   until I deleted the stale .so and let them rebuild. Both symbols exist in csrc
   (aiter_enum.h:15, rocm_ops.hpp:2871). Not a bug in this PR, but it is a new environment
   requirement worth a line in the README next to the gfx1250 note.

2. --check only proves coverage of what the replay enumerates. The caller-side dimensions are
   fixed (bias, gate/up layout, SiTUv2 betas, num_local_tokens, clamp_shared); a real caller
   outside that set would be missed here too. The difference is that such a gap now shows up in
   --check, whereas main's key drift was silent. I have no independent evidence that all 1471
   main-only artifacts are dead — only that the replay never requests them.

[Copilot]

🔵 Needs a closer look

The rewrite spans runtime dispatch, cache generation, and process isolation, while gfx1250 GPU validation remains pending.

7 open issues

• Use host-build fallback when no two-stage jobs exist
• Enforce total AOT worker cap across architectures
• Accept FlyDSL implementation rows in stage-one filtering
• Add cache-key bisection tests for transitions and failures
• Test CPU replay execution and global patch restoration
• Add regressions for blank cells and numeric key deduplication
• Prevent --check from compiling host helpers

1 resolved since last review

• Bias interleave layouts are incorrectly filtered from replay

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Cross-architecture runtime replay and shared JIT changes warrant final human review, with gfx1250 hardware validation still pending.

7 open issues

• Use host-build fallback when no two-stage jobs exist
• Enforce total AOT worker cap across architectures
• Accept FlyDSL implementation rows in stage-one filtering
• Add cache-key bisection tests for transitions and failures
• Test CPU replay execution and global patch restoration
• Add regressions for blank cells and numeric key deduplication
• Prevent --check from compiling host helpers

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The replay filter excludes registered whole-graph FlyDSL implementations, allowing cache verification to succeed without checking their kernels.

7 open issues

• Use host-build fallback when no two-stage jobs exist
• Enforce total AOT worker cap across architectures
• Accept FlyDSL implementation rows in stage-one filtering
• Add cache-key bisection tests for transitions and failures
• Test CPU replay execution and global patch restoration
• Add regressions for blank cells and numeric key deduplication
• Prevent --check from compiling host helpers

1 previously missed issue

In code that hasn't changed since last review

Include FlyDSL implementations and No-input replay recipes

aiter/​aot/​flydsl/​moe.py:151

This filter excludes registered whole-graph FlyDSL names such as impl__flydsl_gfx942__16_16_16_False:default:0, which fused_moe.py resolves through fused_moe_registry. It rejects all 27 FlyDSL rows in qwen38_flash_next_bf16_tuned_fmoe.csv, leaving that table with zero replay jobs, so --csv <that table> --check succeeds without checking those kernels. Include these implementations and add their QuantType.No input recipe in _caller_inputs, with a regression test; accepting the prefix alone would hit the unsupported-recipe branch.

_{Review effort: Balanced}