[FlyDSL] [CI] Add HY4 MXFP8 heterogeneous MoE fusion

.github/scripts/aiter_test.sh

@@ -85,6 +85,11 @@ for file in "${sharded_files[@]}"; do
     # batch gate so they exercise the persistent kernel at every batch size.
     test_cmd=(timeout 60m python3 "$file")
     case "$file" in
+        op_tests/test_fhmoe.py|op_tests/tuning_tests/test_fhmoe_tuner.py)
+            # These files contain pytest-style test functions and intentionally
+            # have no script entry point; direct execution would only import them.
+            test_cmd=(timeout 60m python3 -m pytest -q "$file")
+            ;;
         op_tests/multigpu_tests/bench_mega_moe.py)
             {
                 echo "Running MegaMoE fused-scatter accuracy on 8 GPUs when supported"

.github/scripts/split_tests.sh

@@ -60,6 +60,7 @@ if [[ "$TEST_TYPE" == "aiter" ]]; then
                 "$TEST_DIR/tuning_tests/test_csv_validation.py" \
                 "$TEST_DIR/tuning_tests/test_config_shape_collision.py" \
                 "$TEST_DIR/tuning_tests/test_fast_tuning.py" \
+                "$TEST_DIR/tuning_tests/test_fhmoe_tuner.py" \
                 "$TEST_DIR/tuning_tests/test_mixed_mxfp_tuning.py"
         } | LC_ALL=C sort -u
     )

.github/split-test-times/aiter.tsv

@@ -14,7 +14,7 @@ op_tests/test_concat_cache_mla.py	268
 op_tests/test_deepgemm.py	10
 op_tests/test_dsv4_rotate_quant.py	4
 op_tests/test_f4gemm.py	4
-op_tests/test_fhmoe.py	4
+op_tests/test_fhmoe.py	300
 op_tests/test_flydsl_batched_gemm.py	10
 op_tests/test_flydsl_causal_conv1d_update.py	90
 op_tests/test_flydsl_compress_attn.py	62
@@ -166,4 +166,5 @@ op_tests/test_topk_softmax_shared_gate.py	26
 op_tests/test_vsa_sparse_attention.py	4
 op_tests/tuning_tests/test_config_shape_collision.py	4
 op_tests/tuning_tests/test_csv_validation.py	1
+op_tests/tuning_tests/test_fhmoe_tuner.py	12
 op_tests/tuning_tests/test_mixed_mxfp_tuning.py	4

.github/workflows/tuning-tests.yaml

@@ -93,7 +93,7 @@ jobs:
             -e PYTHONUNBUFFERED=1 \
             -w /workspace \
             tuning_tests_level01 \
-            bash -lc "set -euo pipefail; python3 -m unittest op_tests.tuning_tests.test_csv_validation op_tests.tuning_tests.test_tuner_infra op_tests.tuning_tests.test_mp_tuner_logic op_tests.tuning_tests.test_config_shape_collision op_tests.tuning_tests.test_fast_tuning op_tests.tuning_tests.test_a6w6_tuning op_tests.tuning_tests.test_mixed_mxfp_tuning op_tests.tuning_tests.test_a16w16_tuned_config_lookup -v 2>&1 | tee /workspace/${LOG_DIR}/level01.log"
+            bash -lc "set -euo pipefail; python3 -m unittest op_tests.tuning_tests.test_csv_validation op_tests.tuning_tests.test_tuner_infra op_tests.tuning_tests.test_mp_tuner_logic op_tests.tuning_tests.test_config_shape_collision op_tests.tuning_tests.test_fast_tuning op_tests.tuning_tests.test_a6w6_tuning op_tests.tuning_tests.test_mixed_mxfp_tuning op_tests.tuning_tests.test_a16w16_tuned_config_lookup -v 2>&1 | tee /workspace/${LOG_DIR}/level01.log && python3 -m pytest op_tests/tuning_tests/test_fhmoe_tuner.py -v 2>&1 | tee -a /workspace/${LOG_DIR}/level01.log"
 
       - name: Upload level 0+1 artifacts
         if: always()

aiter/aot/flydsl/README.md

@@ -1,9 +1,9 @@
 # FlyDSL AOT Pre-compilation & Tests
 
 This directory holds the **AOT (Ahead-Of-Time) pre-compilation entry points** for
-FlyDSL kernels. Each module extracts every unique FlyDSL kernel name from aiter's
-tuned CSV configs and compiles them into the cache up front, so that at runtime
-the JIT path hits the cache instead of compiling again.
+FlyDSL kernels. Each module derives the runtime-reachable jobs from aiter's
+tuned CSV configs—including ABI/cache variants that share a kernel name—and
+compiles them up front so the runtime JIT path hits the cache.
 
 | Module | OpKind | Description |
 | --- | --- | --- |
@@ -61,6 +61,11 @@ 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
@@ -98,14 +103,15 @@ 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 `cu_num`). `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`, 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 derived per-job from the CSV's `cu_num` column (`cu_num_to_arch(...)`)
-> and applied internally via `FLYDSL_GPU_ARCH`. That internal var is overwritten
+> 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, edit the `cu_num`
-> column in the CSV.
+> does **not** change what gets built. To cross-compile, use the target `gfx` and
+> `cu_num` values in the CSV.
 
 Example:
 

aiter/aot/flydsl/fhmoe.py

@@ -9,8 +9,6 @@
 
 from dataclasses import dataclass
 
-from aiter.aot.flydsl.common import cu_num_to_arch
-
 
 def _shared_weight(device, n_in: int, k_in: int):
     import torch
@@ -129,9 +127,11 @@ def build_stage2_args(
     def compile_stage1(self, **kwargs):
         from aiter.ops.flydsl.fhmoe import compile_flydsl_fhmoe_stage1
 
+        clamp_shared = kwargs.pop("clamp_shared", True)
         return compile_flydsl_fhmoe_stage1(
             **kwargs,
             shared_expert_id=self.shared_expert_id,
+            clamp_shared=clamp_shared,
         )
 
     def compile_stage2(self, **kwargs):
@@ -145,42 +145,66 @@ def compile_stage2(self, **kwargs):
 
 def precompile_fhmoe_to_cache(
     *,
+    stage: int,
     experts: int,
     shared_expert_id: int,
     a_dtype: str = "fp8",
     b_dtype: str = "fp4",
     act: str = "silu",
     cu_num: int = 0,
+    gfx: str = "gfx950",
     enable_bias: bool = False,
     **kwargs,
 ):
     """Precompile one heterogeneous MoE job through the shared AOT harness."""
+    if stage not in (1, 2):
+        raise ValueError(f"FHMoE AOT stage must be 1 or 2, got {stage}")
     if shared_expert_id != experts - 1:
         raise ValueError(
             "FHMoE AOT expects the shared expert to be the final logical expert; "
             f"got {shared_expert_id=} for {experts=}"
         )
-    if a_dtype != "fp8" or b_dtype != "fp4":
+    if a_dtype != "fp8" or b_dtype not in ("fp4", "fp8"):
         raise ValueError(
-            "FHMoE AOT supports routed FP8 activations and MXFP4 weights; "
+            "FHMoE AOT supports routed FP8 activations with MXFP4 or MXFP8 weights; "
             f"got {a_dtype=} and {b_dtype=}"
         )
+    gate_mode = kwargs.get("gate_mode", "separated")
+    if stage == 1 and b_dtype == "fp8" and gate_mode != "interleave":
+        raise ValueError(
+            "FHMoE AOT requires interleaved gate/up layout for MXFP8 routed "
+            f"weights, got {gate_mode=}"
+        )
+    if stage == 1 and not kwargs.get("clamp_shared", True) and b_dtype != "fp8":
+        raise ValueError(
+            "clamp_shared=False requires MXFP8 routed weights; "
+            "MXFP4/FP8 FHMoE keeps the shared expert clamped"
+        )
     if enable_bias:
         raise ValueError("FHMoE AOT does not support expert bias")
     if act != "silu":
         raise ValueError(f"FHMoE AOT supports only SiLU, got {act=}")
-    if cu_num_to_arch(cu_num) != "gfx950":
-        raise ValueError(f"FHMoE AOT supports only gfx950, got {cu_num=}")
+    gfx = gfx or "gfx950"
+    if cu_num <= 0:
+        cu_num = 256
+    if gfx != "gfx950" or cu_num != 256:
+        raise ValueError(
+            "FHMoE AOT requires target gfx950 with 256 CUs, "
+            f"got {gfx=} and {cu_num=}"
+        )
 
+    from aiter.aot.flydsl.common import override_env
     from aiter.aot.flydsl.moe import _precompile_to_cache
 
-    return _precompile_to_cache(
-        experts=experts,
-        a_dtype=a_dtype,
-        b_dtype=b_dtype,
-        act=act,
-        cu_num=cu_num,
-        enable_bias=enable_bias,
-        _aot_backend=_FHMoEAOTBackend(shared_expert_id),
-        **kwargs,
-    )
+    with override_env("FLYDSL_GPU_ARCH", gfx):
+        return _precompile_to_cache(
+            stage=stage,
+            experts=experts,
+            a_dtype=a_dtype,
+            b_dtype=b_dtype,
+            act=act,
+            cu_num=cu_num,
+            enable_bias=enable_bias,
+            _aot_backend=_FHMoEAOTBackend(shared_expert_id),
+            **kwargs,
+        )