[OPUS] [FlyDSL] [JIT] [Feature] Build the MoE and MegaMoE AOT caches by replaying the runtime on CPU

aiter/aot/flydsl/README.md

@@ -7,14 +7,96 @@ compiles them up front so the runtime JIT path hits the cache.
 
 | Module | OpKind | Description |
 | --- | --- | --- |
-| `moe.py` | `MOE` | MoE / Mixed-MoE kernels (stage1 + stage2) |
+| `moe.py` | `MOE` | Every FlyDSL kernel the tuned fused-MoE / FHMoE / gfx1250 grouped rows reach, by replaying `fused_moe` itself (see below) |
 | `gemm.py` | `GEMM` | GEMM kernels |
 | `conv.py` | `CONV` | conv3d implicit-GEMM kernels, plus the NCDHW->NDHWC pre-transpose each NCDHW caller reaches first |
-| `grouped_moe.py` | `GROUPED_MOE` | gfx1250 grouped MoE GEMM kernels |
 | `chunk_gdn_h.py` | `CHUNK_GDN_H` | chunk-gdn-h opt (K5) kernels |
-| `mega_moe.py` | `MEGA_MOE` | MegaMoE A8W4 profile bundles for MTPR 8192/16384/32768 |
+| `mega_moe.py` | `MEGA_MOE` | MegaMoE A8W4 and A4W4 bundles of every deployment profile, by preloading `MegaMoEV2` itself (see below) |
 | `fmha_fp8.py` | `FMHA_FP8` | gfx950 FP8 flash-attention forward, per head shape in `DEFAULT_SHAPES` |
+| `mxfp4_moe.py` | — | Not part of the AOT build; the MXFP4 tuner's `--fast-scan` precompiles its candidates with it |
 | `common.py` | — | Shared job collection, the deadlock-free fork pool, and cache-hit checking logic |
+| `dry_run.py` | — | Runs aiter's real host code on CPU (FakeTensors, `COMPILE_ONLY=1`, GPU runtime forbidden) |
+
+### How `moe.py` works
+
+`moe.py` does not restate any MoE launch logic. For every tuned row it builds
+what a caller passes `fused_moe` (bf16 activations, routing, and weights
+quantized and preshuffled the way `op_tests/test_moe_2stage.py` prepares them)
+as FakeTensors and calls the runtime's own implementation (`_fused_moe_impl`,
+or `_fhmoe_impl` for shared-expert rows) inside `dry_run`. The runtime looks the
+row up, picks and parameterizes its kernels, sizes its buffers and launches its
+auxiliary kernels as at inference; every FlyDSL launch on that path is
+compiled into the cache instead of run. Kernel or wrapper changes therefore
+need no AOT change.
+
+- Each row is replayed at its tuned token count and at the smallest M that
+  `get_padded_M` maps to it, because some compile-time parameters follow the
+  actual M. Where the two request different FlyDSL kernels, the replay bisects
+  the M between them down to every M at which the runtime switches kernels,
+  since a row can switch more than once in its bucket (e.g. the gfx1250
+  stage-1 quant's K split). These extra replays are reported as `[BISECT]`.
+  `--row-tokens-only` replays the tuned count only.
+- Caller-side choices a row does not fix are enumerated: expert bias for
+  Swiglu + MXFP4 rows only (the GPT-OSS path; on both stages, and on stage 1
+  alone, since a stage-2 bias makes the runtime re-select configs whose stage 2
+  cannot take one), both gate/up layouts for MXMOE
+  stage-1 rows, the default and model SiTUv2 betas, a received token count
+  (`num_local_tokens`, as callers behind an EP dispatch pass), and for FHMoE
+  rows with MXFP8 routed weights both values of `clamp_shared`, which the
+  runtime compiles into stage 1 (MXFP4 rows keep the shared expert clamped).
+  Combinations the runtime rejects (`NotImplementedError`) are reported as
+  `[SKIP]`.
+- The default caller of a row must be served. If the runtime refuses it at
+  only some of the row's token counts, an M-dependent runtime limit applies
+  (e.g. the FlyDSL route-reduce workspace beyond 32-bit addressing, which
+  `op_tests/test_moe_2stage.py` skips as well); those token counts compile
+  nothing and are reported as `[REFUSED]`. Refused at every token count, or
+  failing with any other error, the row fails the build.
+- Rows of the gfx1250 grouped table replay with the gate/up-interleaved
+  weights and n32k4 scales gfx1250 callers pass; `_fused_moe_impl` hands them
+  to FlyDSL's grouped GEMM. 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.
+- Each row compiles for the arch its `gfx` column names. When `GPU_ARCHS`
+  names targets, as `setup.py`'s wheel build has it, only rows of those archs
+  are replayed; unset or `native`, rows of every arch are.
+- aiter fixes some arch-dependent values at import (e.g. `dtypes.fp8`), so the
+  replay runs in one process per arch, started with `GPU_ARCHS` naming that
+  arch and `AITER_AOT_REPLAY=1`. Under that flag `get_gfx_runtime()` and
+  `dtypes.fp8` answer the `GPU_ARCHS` arch instead of probing a device, so a
+  visible GPU of another arch cannot leak in; the GPUs stay hidden all the
+  same. `python -m aiter.aot.flydsl.moe` starts those processes side by side
+  and splits the worker budget (`AITER_MAX_JOBS`, below) between them by job
+  count. Beforehand, it builds aiter's C++ core module and the one
+  behind the two-stage host-side query once, for the arch with the most
+  two-stage jobs and with the whole budget: the replays run only their host
+  code, and not every toolchain builds every arch (ROCm 7.2 cannot build the
+  cluster launch gfx1250 builds turn on).
+- A worker replays up to 12 jobs of one row in turn, so the kernels they share
+  compile once. `AITER_FLYDSL_AOT_TIMEOUT` applies to the worker's whole batch.
+- `--check` compiles nothing and fails on any kernel the replay would need but
+  the cache lacks; it is the CPU-side counterpart of the step-2 test below.
+
+### How `mega_moe.py` works
+
+`mega_moe.py` does not restate MegaMoE's kernel selection either. For each rank
+of each deployment profile it builds the runtime's own `MegaMoEV2` inside
+`dry_run`, with FakeTensor expert weights, and calls `preload_aot_bundles()`:
+the method serving uses to load every Stage1/prepare, Stage2 and combine
+variant of the profile's bundle plan.
+
+- mori shmem, through which `MegaMoEV2` allocates its symmetric buffers and
+  exchanges peer pointers, hands out FakeTensors and fake addresses, so no GPU,
+  process group or mori install is needed.
+- The device reports itself as an MI355X, because the combine launch geometry
+  is tuned per GPU model.
+- `quantize()` is replayed too: `preload_aot_bundles()` leaves it out, but
+  fixed-slot `forward` and the per-stage benchmarks call it.
+- Like the MoE replay, it runs in a process of its own, started with
+  `GPU_ARCHS=gfx950` and `AITER_AOT_REPLAY=1`; both
+  `python -m aiter.aot.flydsl.mega_moe` and `setup.py` start it that way. It
+  is skipped when `GPU_ARCHS` names targets without gfx950.
 
 ---
 
@@ -42,15 +124,15 @@ into CI.
 # MoE / Mixed-MoE (default CSVs)
 python -m aiter.aot.flydsl.moe
 
+# MoE rows of gfx1250 only, which CI's gfx942;gfx950 wheel build leaves out
+GPU_ARCHS=gfx1250 python -m aiter.aot.flydsl.moe
+
 # GEMM
 python -m aiter.aot.flydsl.gemm
 
 # conv3d
 python -m aiter.aot.flydsl.conv
 
-# grouped MoE (gfx1250)
-python -m aiter.aot.flydsl.grouped_moe
-
 # chunk-gdn-h
 python -m aiter.aot.flydsl.chunk_gdn_h
 
@@ -61,11 +143,6 @@ python -m aiter.aot.flydsl.mega_moe
 python -m aiter.aot.flydsl.mega_moe --experts-per-rank 48
 ```
 
-FHMoE stage-1 kernels use distinct cache identities for clamped and unclamped
-shared SiLU. MXFP8/FP8 shared-expert rows automatically emit both
-`clamp_shared=True` and `False` jobs; existing MXFP4/FP8 DSV4 rows remain
-clamped-only.
-
 The FP8 flash-attention kernel picks its tile, rescale threshold and split-K
 factor per call, so a serving process JIT-compiles each new combination on the
 request path. `fmha_fp8` enumerates every combination the wrapper's heuristics
@@ -80,16 +157,30 @@ python -m aiter.aot.flydsl.fmha_fp8
 python -m aiter.aot.flydsl.fmha_fp8 --shape 12:12:192:128
 ```
 
-MegaMoE defaults to all three DeepSeek-V4-Pro deployment profiles: r0/r32/r64
-(`experts_per_rank=48/52/56`). The expert count is part of the compiled ABI, so
-one profile cannot safely reuse another profile's bundle.
+The gfx1250 MoE build imports aiter, so it needs aiter's C++ core module. An
+installed wheel ships it. Built for gfx1250, the module needs a HIP that
+declares the cluster launch API, which ROCm 7.2.4's does not; there, build it
+once for a CDNA arch first (it has host code only):
+
+```bash
+GPU_ARCHS=gfx950 python -c "import aiter"
+GPU_ARCHS=gfx1250 python -m aiter.aot.flydsl.moe
+```
+
+MegaMoE defaults to every profile the service ships: DeepSeek-V4-Pro r0/r32/r64
+(`experts_per_rank=48/52/56`) and Kimi-K3 (`experts_per_rank=112`). The expert
+count is part of the compiled ABI, so one profile cannot safely reuse another
+profile's bundle.
 
 ### Common arguments
 
 ```bash
 # Custom CSV(s) — every module supports --csv and accepts multiple paths
 python -m aiter.aot.flydsl.moe --csv /path/to/config1.csv /path/to/config2.csv
 python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv
+
+# MoE: verify an existing cache on CPU instead of compiling
+python -m aiter.aot.flydsl.moe --csv /path/to/config.csv --check
 ```
 
 ### Environment variables
@@ -102,15 +193,17 @@ python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv
 | `AITER_FLYDSL_AOT_TIMEOUT` | Per-kernel wall-clock cap (seconds). A worker stuck *alive* past this is killed (and retried); `0` disables. | `1200` |
 | `AITER_FLYDSL_AOT_MAX_RETRIES` | Retries for a worker that **died abnormally** (OOM-kill / segfault / timeout-kill). A clean compile error is never retried. `0` disables. | `2` |
 | `AITER_CONFIGS` | Resolves the default CSV lookup path (same as the runtime JIT) | repo built-in |
-| `ARCH` / `GPU_ARCHS` | Selects which jobs to build, not what arch a job compiles *for* (that comes from the CSV's `gfx`, with legacy `cu_num` fallback). `conv.py` applies it inside `parse_csv`, so both `python -m` and the `setup.py` path (`run_aot`) honour it. `gemm.py` still filters in `main()` only, so `run_aot` builds all of its archs. | auto-detect |
+| `ARCH` / `GPU_ARCHS` | Selects which jobs to build, not what arch a job compiles *for* (that comes from the CSV's `gfx`, with legacy `cu_num` fallback). `conv.py` applies it inside `parse_csv`, and `moe.py` and `mega_moe.py` (`GPU_ARCHS` only) before they start their replays, so both `python -m` and the `setup.py` path (`run_aot`) honour it. `gemm.py` still filters in `main()` only, so `run_aot` builds all of its archs. | auto-detect |
 
 > **About the compile target arch.** The arch each kernel is actually compiled
 > for is taken from the CSV's `gfx` column, with `cu_num_to_arch(...)` retained
 > as a fallback for legacy rows, and applied internally via `FLYDSL_GPU_ARCH`.
 > That internal var is overwritten
 > for every job, so setting `ARCH` / `GPU_ARCHS` / `FLYDSL_GPU_ARCH` in your shell
 > does **not** change what gets built. To cross-compile, use the target `gfx` and
-> `cu_num` values in the CSV.
+> `cu_num` values in the CSV. `moe.py` and `mega_moe.py` additionally run each
+> target arch in its own process with `GPU_ARCHS` pinned and `AITER_AOT_REPLAY=1`,
+> so they build the same kernels on a machine with or without GPUs.
 
 Example:
 

aiter/aot/flydsl/common.py

@@ -30,10 +30,8 @@ class OpKind(enum.Enum):
     construction errors instead of silently routing to the wrong code path."""
 
     MOE = "moe"
-    MXFP4_MOE = "mxfp4_moe"
     GEMM = "gemm"
     CONV = "conv"
-    GROUPED_MOE = "grouped_moe"
     CHUNK_GDN_H = "chunk_gdn_h"
     MEGA_MOE = "mega_moe"
     FMHA_FP8 = "fmha_fp8"
@@ -133,31 +131,68 @@ def override_env(var_name: str, value: str | None) -> Iterator[None]:
             os.environ[var_name] = prev
 
 
+def aot_target_archs() -> set[str] | None:
+    """The archs ``GPU_ARCHS`` names, as the wheel build sets it; ``None``
+    when it is unset or ``native``, meaning every arch the tables have."""
+    value = os.environ.get("GPU_ARCHS", "").strip()
+    if not value or value.lower() == "native":
+        return None
+    return {arch.strip() for arch in value.split(";") if arch.strip()}
+
+
+def replay_env(arch: str, workers: int | None = None) -> dict[str, str]:
+    """Environment of a process that replays the runtime for ``arch``. aiter
+    fixes some arch-dependent values at import (e.g. ``dtypes.fp8``), so the
+    process starts with ``GPU_ARCHS`` naming ``arch`` and with
+    ``AITER_AOT_REPLAY=1``, under which aiter's runtime arch queries answer
+    ``arch`` instead of probing a device. The GPUs stay hidden, and aiter is
+    imported in full rather than in AOT mode."""
+    env = dict(os.environ)
+    env.pop("AITER_AOT_IMPORT", None)
+    env.update(
+        GPU_ARCHS=arch,
+        AITER_AOT_REPLAY="1",
+        HIP_VISIBLE_DEVICES="-1",
+        ROCR_VISIBLE_DEVICES="-1",
+        CUDA_VISIBLE_DEVICES="-1",
+    )
+    if workers is not None:
+        env["AITER_MAX_JOBS"] = str(workers)
+    return env
+
+
+def check_replay_process(arch: str) -> None:
+    """Raise unless this process started with ``replay_env(arch)``, so that
+    what aiter fixed at import belongs to ``arch``."""
+    from aiter.utility import dtypes
+
+    if os.environ.get("AITER_AOT_REPLAY") != "1" or os.environ.get("GPU_ARCHS") != arch:
+        raise RuntimeError(
+            f"replaying {arch} needs AITER_AOT_REPLAY=1 and GPU_ARCHS={arch} set "
+            "before aiter is imported (see replay_env)"
+        )
+    expected = dtypes.defaultDtypes[arch]["fp8"]
+    if dtypes.fp8 is not expected:
+        raise RuntimeError(
+            f"aiter was imported for another arch: dtypes.fp8 is {dtypes.fp8}, "
+            f"{arch} needs {expected}"
+        )
+
+
 def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]:
     """Load DEFAULT_CSVS + parse_csv for the named kind and return its
     job list. Note: importing .gemm / .moe / .chunk_gdn_h here also
     runs their module-level imports, which pull in FlyDSL (e.g.
     ``flydsl.expr``). Job collection is therefore not free in the
     parent process, just shifted once out of every child."""
-    if kind is OpKind.MEGA_MOE:
-        from .mega_moe import production_jobs
-
-        # Ships both V4-Pro and Kimi-K3 so neither JITs online at serving time.
-        return production_jobs()
     if kind is OpKind.FMHA_FP8:
         from .fmha_fp8 import default_jobs
 
         return default_jobs()
-    if kind is OpKind.MOE:
-        from .moe import DEFAULT_CSVS, parse_csv
-    elif kind is OpKind.MXFP4_MOE:
-        from .mxfp4_moe import DEFAULT_CSVS, parse_csv
-    elif kind is OpKind.GEMM:
+    if kind is OpKind.GEMM:
         from .gemm import DEFAULT_CSVS, parse_csv
     elif kind is OpKind.CONV:
         from .conv import DEFAULT_CSVS, parse_csv
-    elif kind is OpKind.GROUPED_MOE:
-        from .grouped_moe import DEFAULT_CSVS, parse_csv
     elif kind is OpKind.CHUNK_GDN_H:
         from .chunk_gdn_h import DEFAULT_CSVS, parse_csv
     else:
@@ -166,20 +201,12 @@ def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]:
 
 
 def _compile_one_config_for(kind: OpKind) -> Callable[..., dict[str, Any]]:
-    if kind is OpKind.MEGA_MOE:
-        from .mega_moe import compile_one_config
-    elif kind is OpKind.FMHA_FP8:
+    if kind is OpKind.FMHA_FP8:
         from .fmha_fp8 import compile_one_config
-    elif kind is OpKind.MOE:
-        from .moe import compile_one_config
-    elif kind is OpKind.MXFP4_MOE:
-        from .mxfp4_moe import compile_one_config
     elif kind is OpKind.GEMM:
         from .gemm import compile_one_config
     elif kind is OpKind.CONV:
         from .conv import compile_one_config
-    elif kind is OpKind.GROUPED_MOE:
-        from .grouped_moe import compile_one_config
     elif kind is OpKind.CHUNK_GDN_H:
         from .chunk_gdn_h import compile_one_config
     else:
@@ -399,12 +426,34 @@ def run_aot(cache_dir: str) -> None:
     os.makedirs(cache_dir, exist_ok=True)
     os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = cache_dir
 
+    # MoE and MegaMoE replay the runtime itself, which needs aiter fully
+    # imported for the target arch, so they run in processes of their own.
+    from .mega_moe import run as run_mega_moe_replay
+    from .moe import DEFAULT_CSVS as MOE_CSVS
+    from .moe import run as run_moe_replay
+
+    replay_errors = []
+    moe_failed_archs = run_moe_replay([p for p in MOE_CSVS if os.path.isfile(p)])
+    if moe_failed_archs:
+        replay_errors.append(
+            f"FlyDSL MOE replay failed for {moe_failed_archs} arch(es)"
+        )
+    if run_mega_moe_replay():
+        replay_errors.append("FlyDSL MEGA_MOE replay failed")
+
+    pooled_kinds = [
+        kind for kind in OpKind if kind not in (OpKind.MOE, OpKind.MEGA_MOE)
+    ]
     all_jobs: list[tuple[OpKind, dict[str, Any]]] = []
-    for kind in OpKind:
+    for kind in pooled_kinds:
         for job in _collect_aot_jobs_for(kind):
             all_jobs.append((kind, job))
 
     if not all_jobs:
+        if replay_errors:
+            raise AssertionError(
+                f"[aiter] FlyDSL AOT failures: {'; '.join(replay_errors)}"
+            )
         print("[aiter] FlyDSL AOT: no kernels to compile, skipping")
         return
 
@@ -419,7 +468,7 @@ def run_aot(cache_dir: str) -> None:
 
     print(
         f"[aiter] FlyDSL AOT: {len(all_jobs)} kernels "
-        f"({'+'.join(k.name for k in OpKind)}), "
+        f"({'+'.join(k.name for k in pooled_kinds)}), "
         f"{max_workers} worker processes (cache: {cache_dir})"
     )
 
@@ -442,9 +491,9 @@ def run_aot(cache_dir: str) -> None:
             result_dir,
         )
 
-        ok_by_kind: dict[OpKind, int] = {k: 0 for k in OpKind}
-        fail_by_kind: dict[OpKind, int] = {k: 0 for k in OpKind}
-        errors: list[str] = []
+        ok_by_kind: dict[OpKind, int] = {k: 0 for k in pooled_kinds}
+        fail_by_kind: dict[OpKind, int] = {k: 0 for k in pooled_kinds}
+        errors: list[str] = list(replay_errors)
         for (kind, _job), result, spec in zip(all_jobs, raw, specs):
             label = spec[2]
             if result is not None and result.get("compile_time") is not None:
@@ -458,7 +507,7 @@ def run_aot(cache_dir: str) -> None:
                 fail_by_kind[kind] += 1
                 errors.append(f"FlyDSL {label} produced no kernel")
 
-        for kind in OpKind:
+        for kind in pooled_kinds:
             print(
                 f"[aiter] FlyDSL {kind.name} AOT: "
                 f"compiled {ok_by_kind[kind]} ok, {fail_by_kind[kind]} failed"
@@ -472,7 +521,9 @@ def run_aot(cache_dir: str) -> None:
                 suffix = (
                     f"; ... ({len(unique_errors) - _MAX_ERRORS_IN_MSG} more unique)"
                 )
-            tally = ", ".join(f"{k.name}: {fail_by_kind[k]} failed" for k in OpKind)
+            tally = ", ".join(
+                f"{k.name}: {fail_by_kind[k]} failed" for k in pooled_kinds
+            )
             raise AssertionError(
                 f"[aiter] FlyDSL AOT failures ({tally}): " + "; ".join(head) + suffix
             )