[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,69 @@ the JIT path hits the cache instead of compiling again.
 
 | 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 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` |
 | `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 where the
+  quant family supports it (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, and the default and model
+  SiTUv2 betas. Combinations the runtime rejects are reported as `[SKIP]`.
+- 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.
+- aiter fixes some arch-dependent values at import, so the replay runs in one
+  process per target arch with the GPUs hidden and `GPU_ARCHS` set before
+  aiter is imported. `python -m aiter.aot.flydsl.moe` starts those processes
+  side by side and splits `AITER_FLYDSL_AOT_WORKERS` between them by job count.
+- 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.
 
 ---
 
@@ -48,9 +103,6 @@ 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
 
@@ -75,16 +127,20 @@ 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.
+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
@@ -105,7 +161,9 @@ python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv
 > 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, edit the `cu_num`
-> column in the CSV.
+> column in the CSV. `moe.py` additionally runs each target arch in its own
+> process with `GPU_ARCHS` pinned and the GPUs hidden, so it builds 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"
@@ -148,16 +146,10 @@ def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]:
         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:
@@ -170,16 +162,10 @@ def _compile_one_config_for(kind: OpKind) -> Callable[..., dict[str, Any]]:
         from .mega_moe import compile_one_config
     elif 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:
@@ -432,12 +418,27 @@ def run_aot(cache_dir: str) -> None:
     os.makedirs(cache_dir, exist_ok=True)
     os.environ["FLYDSL_RUNTIME_CACHE_DIR"] = cache_dir
 
+    # MoE replays the runtime itself, which needs aiter fully imported for each
+    # target arch, so it runs in per-arch processes of its own.
+    from .moe import DEFAULT_CSVS as MOE_CSVS
+    from .moe import run as run_moe_replay
+
+    moe_failed_archs = run_moe_replay([p for p in MOE_CSVS if os.path.isfile(p)])
+    moe_errors = (
+        [f"FlyDSL MOE replay failed for {moe_failed_archs} arch(es)"]
+        if moe_failed_archs
+        else []
+    )
+
+    pooled_kinds = [kind for kind in OpKind if kind is not OpKind.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 moe_errors:
+            raise AssertionError(f"[aiter] FlyDSL AOT failures: {moe_errors[0]}")
         print("[aiter] FlyDSL AOT: no kernels to compile, skipping")
         return
 
@@ -452,7 +453,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})"
     )
 
@@ -475,9 +476,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(moe_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:
@@ -491,7 +492,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"
@@ -505,7 +506,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
             )

aiter/aot/flydsl/dry_run.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Run aiter's real host code on CPU, compiling every FlyDSL launch it reaches.
+
+Inside :func:`dry_run`, tensors are FakeTensors on a logical ``cuda:0``, aiter's
+device custom ops dispatch to their fakes, and ``COMPILE_ONLY=1`` makes each
+``@flyc.jit`` call compile and persist its artifact instead of launching. The
+host code that picks kernels, sizes buffers and derives compile-time
+parameters is the runtime's own, so the cache keys it produces are the ones
+the runtime will look up. Any attempt to reach the GPU runtime raises.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from collections.abc import Iterator
+from contextlib import ExitStack, contextmanager
+from types import SimpleNamespace
+from unittest import mock
+
+
+class _Stream:
+    """Stand-in for a ``torch.cuda`` stream; launches never happen."""
+
+    cuda_stream = 0
+
+    def __init__(self, device=None, priority=0):
+        self.device = device
+
+
+class _Event:
+    """Stand-in for ``torch.cuda.Event``; nothing is ever recorded."""
+
+    def __init__(self, *_args, **_kwargs):
+        pass
+
+
+def _forbid(*_args, **_kwargs):
+    raise RuntimeError("FlyDSL AOT dry run reached the GPU runtime")
+
+
+def _torch_guard_modules() -> list:
+    # jit/core.py imports torch_guard through its own sys.path entry, so the
+    # module is loaded twice and each copy decorates its own ops.
+    modules = []
+    for name in ("torch_guard", "aiter.jit.utils.torch_guard"):
+        module = sys.modules.get(name)
+        if module is not None and module not in modules:
+            modules.append(module)
+    return modules
+
+
+def _register_stream_stub() -> None:
+    from flydsl.compiler.jit_argument import JitArgumentRegistry
+    from flydsl.expr.typing import Stream
+
+    if _Stream not in JitArgumentRegistry.registry:
+        JitArgumentRegistry.register(_Stream)(Stream)
+
+
+@contextmanager
+def dry_run(arch: str, cu_num: int, name: str | None = None) -> Iterator[None]:
+    """Make the current process look like an ``arch`` GPU with ``cu_num`` CUs.
+    ``name`` is the device name host code reads, e.g. to pick a per-GPU-model
+    tuning table; it defaults to ``arch``."""
+    import torch
+    from flydsl.compiler.jit_executor import CompiledArtifact
+    from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
+
+    from aiter.jit import core as jit_core
+    from aiter.jit.utils import chip_info
+
+    _register_stream_stub()
+    stream = _Stream(torch.device("cuda", 0))
+    properties = SimpleNamespace(
+        name=name or arch, gcnArchName=arch, multi_processor_count=cu_num
+    )
+    real_data_ptr = torch.Tensor.data_ptr
+
+    def data_ptr(tensor):
+        if not isinstance(tensor, FakeTensor):
+            return real_data_ptr(tensor)
+        if tensor.numel() == 0:
+            return 0
+        # Storage identity keeps aliasing visible to host code comparing
+        # pointers; the address is never dereferenced.
+        return (
+            tensor.untyped_storage()._cdata
+            + tensor.storage_offset() * tensor.element_size()
+        )
+
+    convert, tensor_cls, _, _ = jit_core._pybind_develop_hooks()
+    cached = (
+        chip_info.get_gfx,
+        chip_info.get_gfx_custom_op_core,
+        chip_info.get_gfx_runtime,
+        chip_info.get_cu_num,
+    )
+    env = {
+        "COMPILE_ONLY": "1",
+        "FLYDSL_GPU_ARCH": arch,
+        "ARCH": arch,
+        "GPU_ARCHS": arch,
+        "CU_NUM": str(cu_num),
+    }
+    with ExitStack() as stack:
+        stack.enter_context(mock.patch.dict(os.environ, env))
+        for function in cached:
+            function.cache_clear()
+            stack.callback(function.cache_clear)
+        for owner, attribute, value in (
+            (torch.cuda, "_lazy_init", _forbid),
+            (CompiledArtifact, "_ensure_engine", _forbid),
+            (torch.cuda, "current_stream", lambda *_a, **_k: stream),
+            (torch.cuda, "Stream", _Stream),
+            (torch.cuda, "Event", _Event),
+            (torch.cuda, "current_device", lambda: 0),
+            (torch.cuda, "get_device_properties", lambda *_a, **_k: properties),
+            (torch.cuda, "synchronize", lambda *_a, **_k: None),
+            (torch.Tensor, "data_ptr", data_ptr),
+            (chip_info, "_detect_native", lambda: [arch]),
+            (
+                jit_core,
+                "_pybind_develop_hooks_cache",
+                (convert, tensor_cls, lambda _index: 0, lambda: 0),
+            ),
+        ):
+            stack.enter_context(mock.patch.object(owner, attribute, value))
+        for module in _torch_guard_modules():
+            previous = module.run_host_scalar_ops_for_real(True)
+            stack.callback(module.run_host_scalar_ops_for_real, previous)
+        stack.enter_context(FakeTensorMode(allow_non_fake_inputs=True))
+        yield