[FlyDSL] Add HY4 MXFP8 heterogeneous MoE fusion

Fuse HY4's FP8 shared expert into the routed MXFP8 MoE kernels to reduce launch and scheduling overhead, with AOT coverage, correctness guards, and reproducible gfx950 tuning.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

.github/scripts/split_tests.sh

@@ -56,6 +56,7 @@ if [[ "$TEST_TYPE" == "aiter" ]]; then
             printf '%s\n' \
                 "$TEST_DIR/tuning_tests/test_csv_validation.py" \
                 "$TEST_DIR/tuning_tests/test_config_shape_collision.py" \
+                "$TEST_DIR/tuning_tests/test_fhmoe_tuner.py" \
                 "$TEST_DIR/tuning_tests/test_mixed_mxfp_tuning.py"
         } | LC_ALL=C sort -u
     )
@@ -182,6 +183,7 @@ if [[ "$TEST_TYPE" == "aiter" ]]; then
     FILE_TIMES[op_tests/test_quant_mxfp6_gemm.py]=6
     FILE_TIMES[op_tests/tuning_tests/test_config_shape_collision.py]=4
     FILE_TIMES[op_tests/tuning_tests/test_csv_validation.py]=4
+    FILE_TIMES[op_tests/tuning_tests/test_fhmoe_tuner.py]=4
     FILE_TIMES[op_tests/tuning_tests/test_mixed_mxfp_tuning.py]=4
     FILE_TIMES[op_tests/test_fused_qk_rmsnorm_per_token_quant.py]=5
     FILE_TIMES[op_tests/test_groupnorm.py]=5
@@ -196,7 +198,7 @@ if [[ "$TEST_TYPE" == "aiter" ]]; then
     FILE_TIMES[op_tests/test_asm_guard.py]=4
     FILE_TIMES[op_tests/test_dsv4_rotate_quant.py]=4
     FILE_TIMES[op_tests/test_f4gemm.py]=4
-    FILE_TIMES[op_tests/test_fhmoe.py]=4
+    FILE_TIMES[op_tests/test_fhmoe.py]=40
     FILE_TIMES[op_tests/test_flydsl_batched_gemm.py]=4
     FILE_TIMES[op_tests/test_flydsl_fmha.py]=4
     FILE_TIMES[op_tests/test_fmha_fwd_mxfp8_asm.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_a6w6_tuning op_tests.tuning_tests.test_mixed_mxfp_tuning -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_a6w6_tuning op_tests.tuning_tests.test_mixed_mxfp_tuning -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/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
@@ -127,11 +125,25 @@ def build_stage2_args(
         )
 
     def compile_stage1(self, **kwargs):
+        from aiter.fhmoe import _is_hy4_mxfp8_fhmoe_contract
         from aiter.ops.flydsl.fhmoe import compile_flydsl_fhmoe_stage1
 
+        is_hy4 = _is_hy4_mxfp8_fhmoe_contract(
+            model_dim=kwargs["model_dim"],
+            inter_dim=kwargs["inter_dim"],
+            experts=kwargs["experts"],
+            topk=kwargs["topk"],
+            routed_mxfp8=kwargs["b_dtype"] == "fp8",
+            hidden_pad=kwargs.get("model_dim_pad", 0),
+            intermediate_pad=kwargs.get("inter_dim_pad", 0),
+            gate_interleaved=kwargs.get("gate_mode", "separated") == "interleave",
+            doweight_stage1=kwargs.get("doweight_stage1", False),
+            shared_expert_id=self.shared_expert_id,
+        )
         return compile_flydsl_fhmoe_stage1(
             **kwargs,
             shared_expert_id=self.shared_expert_id,
+            clamp_shared=not is_hy4,
         )
 
     def compile_stage2(self, **kwargs):
@@ -145,6 +157,7 @@ def compile_stage2(self, **kwargs):
 
 def precompile_fhmoe_to_cache(
     *,
+    stage: int,
     experts: int,
     shared_expert_id: int,
     a_dtype: str = "fp8",
@@ -160,21 +173,30 @@ def precompile_fhmoe_to_cache(
             "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 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=}")
+    if cu_num not in (0, 256):
+        raise ValueError(
+            f"FHMoE AOT supports only gfx950 or the default CU sentinel, got {cu_num=}"
+        )
 
     from aiter.aot.flydsl.moe import _precompile_to_cache
 
     return _precompile_to_cache(
+        stage=stage,
         experts=experts,
         a_dtype=a_dtype,
         b_dtype=b_dtype,

aiter/aot/flydsl/moe.py

@@ -65,6 +65,11 @@
 MOE_AOT_ARCH_DEFAULT = "gfx950"
 
 
+def _mx_w2_scale_numel(experts: int, model_dim: int, inter_dim: int) -> int:
+    scale_k = (inter_dim + 255) // 256 * 256
+    return experts * model_dim * (scale_k // 32)
+
+
 def parse_csv(csv_path: str):
     """Parse the CSV and return a list of unique compile jobs.
 
@@ -90,6 +95,8 @@ def parse_csv(csv_path: str):
             cu_num = int(row.get("cu_num", "0"))
             block_m = int(row.get("block_m", "0") or "0")
             shared_expert_id = int(row.get("shared_expert_id", "-1") or "-1")
+            model_dim_pad = int(row.get("hidden_pad", "0") or "0")
+            inter_dim_pad = int(row.get("intermediate_pad", "0") or "0")
             act_type = row.get("act_type", "")
             act_name = act_type.strip().split(".")[-1].lower()
             act = act_name if act_name in ("swiglu", "situv2") else "silu"
@@ -181,6 +188,8 @@ def parse_csv(csv_path: str):
                         "enable_bias": enable_bias,
                         "token_num": token,
                         "block_m": block_m,
+                        "model_dim_pad": model_dim_pad,
+                        "inter_dim_pad": inter_dim_pad,
                     }
                     if shared_expert_id >= 0:
                         job["shared_expert_id"] = shared_expert_id
@@ -239,6 +248,8 @@ def _precompile_to_cache(
     cu_num: int = 0,
     token_num: int = 0,
     block_m: int = 0,
+    model_dim_pad: int = 0,
+    inter_dim_pad: int = 0,
     a_scale_one: bool = False,
     xcd_swizzle: int = 0,
     enable_bias: bool = False,
@@ -324,7 +335,7 @@ def _make_sorted_weights(doweight: bool):
         return None
 
     def _make_a1_scale():
-        """Mirror fused_moe_2stages a1_scale construction (per_1x32 + fp4-weight path)."""
+        """Mirror per-1x32 MXFP4/MXFP8 runtime a1_scale construction."""
         if not use_mx_gemm:
             return None
         if a_dtype == "fp8":
@@ -400,7 +411,7 @@ def _make_a2_scale_for_stage2():
         return None
 
     def _make_w_scale(scale_storage_numel: int):
-        # mxfp4 e8m0 scale -- viewed as uint8 by _view_safe before kernel launch.
+        # MX E8M0 scale storage is viewed as uint8 before kernel launch.
         return torch.zeros(scale_storage_numel, dtype=torch.uint8, device=dev)
 
     def _make_a_user(a_dtype_user_shape):
@@ -619,6 +630,8 @@ def _make_a_user(a_dtype_user_shape):
                 waves_per_eu=waves_per_eu,
                 b_nt=b_nt,
                 gate_mode=gate_mode,
+                model_dim_pad=model_dim_pad,
+                inter_dim_pad=inter_dim_pad,
                 enable_bias=(kernel_bias is not None),
                 a_scale_one=a_scale_one,
                 xcd_swizzle=xcd_swizzle,
@@ -682,7 +695,7 @@ def _make_a_user(a_dtype_user_shape):
 
             a2_scale = _make_a2_scale_for_stage2()
             if use_mx_gemm:
-                w2_scale = _make_w_scale(E * model_dim * (inter_dim // 32))
+                w2_scale = _make_w_scale(_mx_w2_scale_numel(E, model_dim, inter_dim))
             else:
                 w2_scale = torch.zeros(1, device=dev, dtype=torch.float32)
 
@@ -810,6 +823,8 @@ def _make_a_user(a_dtype_user_shape):
                 use_global_a=requires_flydsl_stage2_global_a(a),
                 cu_num_mul=cu_num_mul,
                 b_nt=b_nt,
+                model_dim_pad=model_dim_pad,
+                inter_dim_pad=inter_dim_pad,
                 xcd_swizzle=xcd_swizzle,
                 enable_bias=enable_bias,
             )

aiter/configs/tuned_fhmoe.csv

@@ -11,3 +11,20 @@ gfx950,256,256,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float
 gfx950,256,512,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8,flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_bnt2_persist
 gfx950,256,1024,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8,flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_bnt2_persist
 gfx950,256,2048,7168,384,385,7,384,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,64,0,flydsl_moe1_afp8_wfp4_bf16_t64x128x256_w3_bnt0_gui,flydsl_moe2_afp8_wfp4_bf16_t64x128x128_atomic
+gfx950,256,1,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w4_gui_kw2_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,2,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_bnt0_gui_kw2_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,4,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w4_bnt0_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,8,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w3_gui_xcd4_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,16,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x128x256_w3_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,32,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w3_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x256_atomic
+gfx950,256,64,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x128x256_w3_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,128,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w2_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,256,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w3_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,512,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,32,0,flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w3_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic
+gfx950,256,1024,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,64,0,flydsl_moe1_afp8_wfp8_bf16_t64x128x256_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t64x128x256_atomic
+gfx950,256,2048,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,128,0,flydsl_moe1_afp8_wfp8_bf16_t128x128x256_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t64x128x128_atomic_persist_sbm128
+gfx950,256,4096,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,64,0,flydsl_moe1_afp8_wfp8_bf16_t64x128x256_w4_bnt0_gui_xcd4_fp8,flydsl_moe2_afp8_wfp8_bf16_t64x128x256_reduce
+gfx950,256,8192,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,128,0,flydsl_moe1_afp8_wfp8_bf16_t128x128x256_w2_bnt0_gui_xcd4_fp8,flydsl_moe2_afp8_wfp8_bf16_t128x128x128_reduce_persist
+gfx950,256,16384,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,128,0,flydsl_moe1_afp8_wfp8_bf16_t128x128x256_bnt0_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t128x128x128_reduce_persist
+gfx950,256,32768,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,128,0,flydsl_moe1_afp8_wfp8_bf16_t128x128x256_w2_bnt0_gui_xcd4_fp8,flydsl_moe2_afp8_wfp8_bf16_t64x128x128_atomic_persist_sbm128
+gfx950,256,131072,6144,256,257,9,256,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float8_e4m3fn,QuantType.per_1x32,1,0,0,0,GateMode.INTERLEAVE,128,0,flydsl_moe1_afp8_wfp8_bf16_t128x256x256_w4_bnt0_gui_fp8,flydsl_moe2_afp8_wfp8_bf16_t128x128x128_reduce_bnt2_xcd4