From 68a134b80e03b5b92e6de4e9c9a38f48b8b3c15b Mon Sep 17 00:00:00 2001 From: root Date: Thu, 1 Oct 2026 12:03:09 +0000 Subject: [PATCH] [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 --- .github/scripts/aiter_test.sh | 5 + .github/scripts/split_tests.sh | 1 + .github/split-test-times/aiter.tsv | 3 +- .github/workflows/tuning-tests.yaml | 2 +- aiter/aot/flydsl/README.md | 22 +- aiter/aot/flydsl/fhmoe.py | 56 +- aiter/aot/flydsl/moe.py | 113 +- aiter/aot/flydsl/mxfp4_moe.py | 7 +- aiter/configs/tuned_fhmoe.csv | 16 + aiter/fhmoe.py | 524 +++- aiter/fhmoe_contract.py | 37 + aiter/fused_moe.py | 225 +- aiter/ops/flydsl/fhmoe.py | 47 +- aiter/ops/flydsl/kernels/fhmoe.py | 8 +- .../flydsl/kernels/mixed_moe_gemm_2stage.py | 2 +- .../kernels/mixed_moe_gemm_2stage_common.py | 49 +- aiter/ops/flydsl/moe_kernels.py | 50 +- aiter/utility/fp4_utils.py | 13 + op_tests/test_fhmoe.py | 2290 +++++++++++++++-- op_tests/test_moe_2stage.py | 89 + op_tests/tuners/tune_fhmoe.py | 1291 ++++++++++ op_tests/tuning_tests/README.md | 72 +- .../test_config_shape_collision.py | 65 +- op_tests/tuning_tests/test_fhmoe_tuner.py | 1202 +++++++++ op_tests/tuning_tests/test_tune_pipeline.py | 82 + 25 files changed, 5927 insertions(+), 344 deletions(-) create mode 100644 aiter/fhmoe_contract.py create mode 100644 op_tests/tuners/tune_fhmoe.py create mode 100644 op_tests/tuning_tests/test_fhmoe_tuner.py diff --git a/.github/scripts/aiter_test.sh b/.github/scripts/aiter_test.sh index 5f814f1be1a..71f1cb171d2 100755 --- a/.github/scripts/aiter_test.sh +++ b/.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" diff --git a/.github/scripts/split_tests.sh b/.github/scripts/split_tests.sh index 08dff013071..66ee7adca88 100755 --- a/.github/scripts/split_tests.sh +++ b/.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 ) diff --git a/.github/split-test-times/aiter.tsv b/.github/split-test-times/aiter.tsv index 5e6b110b2f4..fda0dc4f688 100644 --- a/.github/split-test-times/aiter.tsv +++ b/.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 diff --git a/.github/workflows/tuning-tests.yaml b/.github/workflows/tuning-tests.yaml index 121e8160411..26868e2ee6a 100644 --- a/.github/workflows/tuning-tests.yaml +++ b/.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() diff --git a/aiter/aot/flydsl/README.md b/aiter/aot/flydsl/README.md index 02352b38eb8..cd2429797b4 100644 --- a/aiter/aot/flydsl/README.md +++ b/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: diff --git a/aiter/aot/flydsl/fhmoe.py b/aiter/aot/flydsl/fhmoe.py index 16d983d7c48..bae562d243c 100755 --- a/aiter/aot/flydsl/fhmoe.py +++ b/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, + ) diff --git a/aiter/aot/flydsl/moe.py b/aiter/aot/flydsl/moe.py index 63418ca0354..9c9521b2ba4 100755 --- a/aiter/aot/flydsl/moe.py +++ b/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. @@ -88,8 +93,11 @@ def parse_csv(csv_path: str): topk = int(row["topk"]) doweight_stage1 = bool(int(row.get("doweight_stage1", "0"))) cu_num = int(row.get("cu_num", "0")) + gfx = (row.get("gfx") or "").strip() 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" @@ -140,8 +148,9 @@ def parse_csv(csv_path: str): "inter_dim": inter_dim, "topk": topk, "cu_num": cu_num, + "gfx": gfx, # Not used by the epilogue compile; zeroed so dedup keys on - # (act, inter_dim, topk, cu_num) only. + # (act, inter_dim, topk, gfx, cu_num) only. "model_dim": 0, "experts": 0, } @@ -177,10 +186,13 @@ def parse_csv(csv_path: str): "topk": topk, "doweight_stage1": doweight_stage1, "cu_num": cu_num, + "gfx": gfx, "act": act, "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 @@ -198,12 +210,23 @@ def parse_csv(csv_path: str): full_job["v2_output_layout"] = True if stage2_v2_params is not None: full_job["out_dtype"] = stage2_v2_params["a_dtype"] - key = job_identity(full_job) - if key in seen: - continue - seen.add(key) - - jobs.append(full_job) + is_fhmoe_stage1 = ( + params["stage"] == 1 + and shared_expert_id >= 0 + and params["b_dtype"] == "fp8" + ) + compile_jobs = [full_job] + if is_fhmoe_stage1: + compile_jobs = [ + {**full_job, "clamp_shared": True}, + {**full_job, "clamp_shared": False}, + ] + for compile_job in compile_jobs: + key = job_identity(compile_job) + if key in seen: + continue + seen.add(key) + jobs.append(compile_job) return jobs @@ -239,12 +262,15 @@ 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, stage1_fuse_quant=None, k_wave: int = 1, v2_output_layout: bool = False, + clamp_shared: bool = True, # Stage2-only kernel tuning knobs (registered by the production-variant # entries in `get_flydsl_stage2_kernels`). Forwarded into # `compile_flydsl_moe_stage2` for stage 2 AOT compilation. @@ -262,8 +288,8 @@ def _precompile_to_cache( topk, a_dtype, b_dtype, ...)`` combination, then dispatches into the same runtime entry points used by the fused-MoE op. ``COMPILE_ONLY=1`` causes the executor to compile and persist the artifact without launching a - kernel. This guarantees that the cache key written here equals the cache - key the runtime will look up at inference time. + kernel. Tests compare representative AOT/runtime identities and run-only + retrieval so cache-key drift is detected. """ import torch @@ -324,7 +350,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 +426,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): @@ -601,30 +627,35 @@ def _make_a_user(a_dtype_user_shape): if _aot_backend is None else _aot_backend.compile_stage1 ) - exe = compile_stage1( - model_dim=model_dim, - inter_dim=inter_dim, - experts=E, - topk=topk, - tile_m=tile_m, - tile_n=tile_n, - tile_k=tile_k, - doweight_stage1=(sw is not None), - a_dtype=a_dtype, - b_dtype=b_dtype, - out_dtype=_gemm_out_dtype, - act=act, - use_async_copy=True, - k_batch=k_batch, - waves_per_eu=waves_per_eu, - b_nt=b_nt, - gate_mode=gate_mode, - enable_bias=(kernel_bias is not None), - a_scale_one=a_scale_one, - xcd_swizzle=xcd_swizzle, - k_wave=k_wave, - v2_output_layout=_v2_output_layout, - ) + stage1_compile_kwargs = { + "model_dim": model_dim, + "inter_dim": inter_dim, + "experts": E, + "topk": topk, + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "doweight_stage1": (sw is not None), + "a_dtype": a_dtype, + "b_dtype": b_dtype, + "out_dtype": _gemm_out_dtype, + "act": act, + "use_async_copy": True, + "k_batch": k_batch, + "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, + "k_wave": k_wave, + "v2_output_layout": _v2_output_layout, + } + if _aot_backend is not None: + stage1_compile_kwargs["clamp_shared"] = clamp_shared + exe = compile_stage1(**stage1_compile_kwargs) _run_compiled(exe, args) if _gui_sk_fused or _gui_sk or _splitk_fp4: @@ -682,7 +713,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) @@ -737,6 +768,8 @@ def _make_a_user(a_dtype_user_shape): else: _persist_m = -1 if m_blocks > 256 else 1 if a_dtype == "fp8": + # Mirror runtime normalization for legacy tuned names that + # retain `_persist` as part of their identity. _persist_m = resolve_flydsl_grid_y_persist_m(m_blocks) _n_in = model_dim @@ -810,6 +843,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, ) @@ -1015,6 +1050,7 @@ def compile_one_config( experts: int, topk: int, cu_num: int = 0, + gfx: str = "", **kwargs, ) -> dict: """Compile one MoE kernel configuration and save to cache. @@ -1024,7 +1060,7 @@ def compile_one_config( Returns a dict with timing info. """ - aot_arch = cu_num_to_arch(cu_num, default=MOE_AOT_ARCH_DEFAULT) + aot_arch = gfx or cu_num_to_arch(cu_num, default=MOE_AOT_ARCH_DEFAULT) is_epilogue = kwargs.get("stage") == "epilogue" shape_str = ( f"{kernel_name} inter_dim={inter_dim} topk={topk}" @@ -1086,6 +1122,7 @@ def compile_one_config( precompile = precompile_fhmoe_to_cache kwargs["shared_expert_id"] = shared_expert_id + kwargs["gfx"] = aot_arch else: precompile = _precompile_to_cache precompile( @@ -1143,7 +1180,7 @@ def main(): print(f" Stage2 jobs: {len(stage2_jobs)}") print(f" Epilogue jobs: {len(epilogue_jobs)}") print(f" Total jobs: {len(all_jobs)}") - print(" Compile arch: (from cu_num)") + print(" Compile arch: (from gfx, with legacy cu_num fallback)") print(f" Cache dir: {cache_dir}") print(f" Target arch: {arch}") print("=" * 72) diff --git a/aiter/aot/flydsl/mxfp4_moe.py b/aiter/aot/flydsl/mxfp4_moe.py index 4d4934e3b89..b45ad3446e0 100644 --- a/aiter/aot/flydsl/mxfp4_moe.py +++ b/aiter/aot/flydsl/mxfp4_moe.py @@ -41,7 +41,6 @@ # Mirror the runtime gate so the default build skips the opt-in mxfp4-out path. _MXFP4_INTERMEDIATE = os.environ.get("AITER_MXFP4_INTERMEDIATE", "0") not in ("0", "") -_STAGE2_FP8_ROUTE_OUT = os.environ.get("AITER_FLYDSL_STAGE2_FP8", "0") == "1" def _job_key(job: dict) -> tuple: @@ -215,11 +214,7 @@ def _add(job): if v2_g2 is not None: bm = v2_g2["tile_m"] - out_dtype = ( - "fp8" - if v2_g2["epilog"] == "reduce" and _STAGE2_FP8_ROUTE_OUT - else "bf16" - ) + out_dtype = "bf16" bias_supported = ( row.get("q_type", "").strip().split(".")[-1] == "per_1x32" and row.get("dtype", "") in ("torch.bfloat16", "torch.float16") diff --git a/aiter/configs/tuned_fhmoe.csv b/aiter/configs/tuned_fhmoe.csv index 4172b32a28f..0d537d90c28 100644 --- a/aiter/configs/tuned_fhmoe.csv +++ b/aiter/configs/tuned_fhmoe.csv @@ -11,3 +11,19 @@ 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_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 +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 +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_t128x128x128_reduce diff --git a/aiter/fhmoe.py b/aiter/fhmoe.py index b5c88b756c2..88f016e5554 100644 --- a/aiter/fhmoe.py +++ b/aiter/fhmoe.py @@ -1,8 +1,26 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. -"""Fused heterogeneous MoE (FHMoE) MXFP4/FP8 dispatch.""" +"""Fused heterogeneous MoE (FHMoE) MXFP4/FP8 and MXFP8/FP8 dispatch. +The notation names routed/shared storage respectively. Shared ``FP8`` means +E4M3 payloads with per-32 E8M0 block scales, stored separately from routed +weights. + +FHMoE requires gfx950, SiLU, per-1x32 quantization, contiguous preshuffled +weights/scales, all four shared tensors, and a dummy final routed row identified +by ``shared_expert_id == E - 1``. Expert bias and expert-mask dispatch are not +supported. + +MXFP8/FP8 rows in ``tuned_fhmoe.csv`` are selected from request geometry, +dtypes, padding, and gate layout; model identity is not inferred. +``clamp_shared`` is an explicit caller-provided semantic. The existing DSV4 +I384 MXFP4/FP8 compatibility path is retained unchanged. Unknown MXFP8 +geometries may use generic fallback, while partial matches to dedicated +geometries fail closed on contract mismatches. +""" + +import csv import functools import os from dataclasses import replace @@ -10,7 +28,10 @@ import torch from aiter import ActivationType, QuantType, dtypes -from aiter.jit.utils.chip_info import get_gfx +from aiter.fhmoe_contract import ( + FHMOE_BUFFER_MAX_BYTES, +) +from aiter.jit.utils.chip_info import get_gfx_runtime from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.ops.flydsl.moe_common import GateMode @@ -48,28 +69,32 @@ def _validate_fhmoe_contract( E = w1.shape[0] model_dim = w2.shape[1] inter_dim = w2.shape[2] * (model_dim // w1.shape[-1]) - if get_gfx() != "gfx950": + if get_gfx_runtime() != "gfx950": raise NotImplementedError( - "Heterogeneous MXFP4/FP8 experts currently require gfx950" + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts currently require gfx950" ) if quant_type != QuantType.per_1x32: raise ValueError( - "Heterogeneous MXFP4/FP8 experts require per_1x32 quantization" + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts require " + "per_1x32 quantization" ) if activation != ActivationType.Silu: - raise ValueError("Heterogeneous MXFP4/FP8 experts currently require SiLU") + raise ValueError( + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts currently require SiLU" + ) if gate_mode not in (GateMode.INTERLEAVE, GateMode.SEPARATED): raise ValueError( - "Heterogeneous MXFP4/FP8 experts require interleaved or " - "separated gate/up weights" + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts require " + "interleaved or separated gate/up weights" ) if expert_mask is not None: raise NotImplementedError( - "Heterogeneous MXFP4/FP8 experts do not yet support expert masks" + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts do not yet " + "support expert masks" ) if bias1 is not None or bias2 is not None: raise NotImplementedError( - "Heterogeneous MXFP4/FP8 experts do not support expert biases" + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts do not support expert biases" ) if shared_expert_id != E - 1: raise ValueError( @@ -77,11 +102,21 @@ def _validate_fhmoe_contract( f"weight row and shared_expert_id == E - 1; got {shared_expert_id=} " f"and E={E}" ) - if w1.dtype != dtypes.fp4x2 or w2.dtype != dtypes.fp4x2: - raise ValueError("Heterogeneous routed weights must use MXFP4") + if w1.dtype != w2.dtype: + raise ValueError( + "Heterogeneous routed W1/W2 must use the same dtype, got " + f"{w1.dtype} and {w2.dtype}" + ) + if w1.dtype not in (dtypes.fp4x2, dtypes.fp8): + raise ValueError("Heterogeneous routed weights must use MXFP4 or MXFP8") + if w1.dtype == dtypes.fp8 and gate_mode != GateMode.INTERLEAVE: + raise ValueError( + "Heterogeneous MXFP8 routed weights require interleaved gate/up layout" + ) if w1.shape[1] != 2 * inter_dim: raise ValueError( - "Heterogeneous MXFP4/FP8 experts require gate and up projections" + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts require " + "gate and up projections" ) assert shared_w1 is not None and shared_w2 is not None assert shared_w1_scale is not None and shared_w2_scale is not None @@ -129,6 +164,12 @@ def _validate_fhmoe_contract( assert w1_scale is not None and w2_scale is not None expected_w1_scale_numel = E * 2 * inter_dim * (model_dim // 32) expected_w2_scale_numel = E * model_dim * (stage2_scale_k // 32) + if w1_scale.ndim not in (1, 2) or w2_scale.ndim not in (1, 2): + raise ValueError( + "Preshuffled routed scales must be flattened or 2D; raw " + f"[E, N, K/32] layouts are unsupported, got {w1_scale.ndim}D/" + f"{w2_scale.ndim}D" + ) if w1_scale.numel() != expected_w1_scale_numel: raise ValueError( "Expected preshuffled routed w1_scale to contain " @@ -179,12 +220,14 @@ def _flydsl_fhmoe_stage1_wrapper( out_scale_sorted=None, bias1=None, topk_ids=None, + block_m=None, swiglu_limit: float | None = None, inter_dim_pad: int = 0, model_dim_pad: int = 0, shared_w1=None, shared_w1_scale=None, shared_expert_id: int = -1, + clamp_shared: bool = True, v2_output_layout: bool = False, **_kwargs, ): @@ -230,6 +273,7 @@ def _flydsl_fhmoe_stage1_wrapper( k_wave=parsed.get("k_wave", 1), v2_output_layout=v2_output_layout, shared_expert_id=shared_expert_id, + clamp_shared=clamp_shared, ) @@ -247,6 +291,7 @@ def _flydsl_fhmoe_stage2_wrapper( a2_scale=None, sorted_weights=None, bias2=None, + block_m=None, inter_dim_pad: int = 0, model_dim_pad: int = 0, expert_mask=None, @@ -303,7 +348,6 @@ def _flydsl_fhmoe_stage2_wrapper( def _use_fhmoe_wrappers(metadata): from aiter.fused_moe import _flydsl_stage1_wrapper, _flydsl_stage2_wrapper - from aiter.ops.flydsl.moe_kernels import get_flydsl_kernel_params stage1_func = getattr(metadata.stage1, "func", metadata.stage1) stage2_func = getattr(metadata.stage2, "func", metadata.stage2) @@ -313,16 +357,10 @@ def _use_fhmoe_wrappers(metadata): or stage2_func is not _flydsl_stage2_wrapper ): raise NotImplementedError( - "Heterogeneous MXFP4/FP8 experts require the two-stage FlyDSL path" - ) - kernel_names = ( - metadata.stage1.keywords.get("kernelName", ""), - metadata.stage2.keywords.get("kernelName", ""), - ) - if any(get_flydsl_kernel_params(name) is None for name in kernel_names): - raise NotImplementedError( - "Heterogeneous MXFP4/FP8 experts require valid FlyDSL kernels" + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts require the " + "two-stage FlyDSL path" ) + _validate_fhmoe_metadata_consistency(metadata) stage1 = functools.partial( _flydsl_fhmoe_stage1_wrapper, *metadata.stage1.args, @@ -336,32 +374,204 @@ def _use_fhmoe_wrappers(metadata): return replace(metadata, stage1=stage1, stage2=stage2) -def _dsv4_i384_fhmoe_config_file() -> str: +def _validate_fhmoe_metadata_consistency(metadata): + """Validate shared stage identity and sorting-block metadata.""" + from aiter.ops.flydsl.moe_kernels import get_flydsl_kernel_params + + stage1 = get_flydsl_kernel_params(metadata.stage1.keywords.get("kernelName", "")) + stage2 = get_flydsl_kernel_params(metadata.stage2.keywords.get("kernelName", "")) + if stage1 is None or stage2 is None: + raise NotImplementedError( + "Heterogeneous MXFP4/FP8 and MXFP8/FP8 experts require valid FlyDSL kernels" + ) + stage2_sort_block = int(stage2.get("sort_block_m", 0) or stage2["tile_m"]) + if ( + stage1["stage"] != 1 + or stage2["stage"] != 2 + or int(stage1["tile_m"]) != int(metadata.block_m) + or stage2_sort_block != int(metadata.block_m) + ): + raise NotImplementedError( + "FHMoE metadata block_m must match stage1 tile_m and stage2 sort_block_m" + ) + return stage1, stage2 + + +def _dtype_name(dtype: torch.dtype) -> str: + if dtype == dtypes.fp8: + return "fp8" + if dtype == dtypes.fp4x2: + return "fp4" + if dtype == dtypes.bf16: + return "bf16" + if dtype == dtypes.fp16: + return "f16" + return str(dtype) + + +def _is_fhmoe_metadata_compatible( + metadata, + *, + model_dim: int, + dtype: torch.dtype, + q_dtype_a: torch.dtype, + q_dtype_w: torch.dtype, + gate_mode: GateMode, +) -> bool: + """Validate a dedicated row against model-agnostic FHMoE request fields.""" + stage1, stage2 = _validate_fhmoe_metadata_consistency(metadata) + a_dtype = _dtype_name(q_dtype_a) + b_dtype = _dtype_name(q_dtype_w) + out_dtype = _dtype_name(dtype) + stage2_a_dtype = metadata.fuse_quant or a_dtype + return ( + stage1["a_dtype"] == a_dtype + and stage1["b_dtype"] == stage2["b_dtype"] == b_dtype + and stage2["a_dtype"] == stage2_a_dtype + and stage2["out_dtype"] == out_dtype + and stage1.get("gate_mode") == gate_mode.value + and _valid_fhmoe_kernel_geometry(stage1, stage2, model_dim=model_dim) + ) + + +def _is_dsv4_i384_metadata(metadata) -> bool: + """Validate that a dedicated DSV4 row names compatible FlyDSL kernels.""" + stage1, stage2 = _validate_fhmoe_metadata_consistency(metadata) + return ( + stage1["a_dtype"] == stage2["a_dtype"] == "fp8" + and stage1["b_dtype"] == stage2["b_dtype"] == "fp4" + and stage2["out_dtype"] == "bf16" + and stage1.get("gate_mode") == "interleave" + and _valid_fhmoe_kernel_geometry(stage1, stage2, model_dim=7168) + ) + + +def _valid_fhmoe_kernel_geometry(stage1, stage2, *, model_dim: int) -> bool: + """Return whether split-K/wave and stage-2 sorting geometry can compile.""" + stage1_k_wave = int(stage1.get("k_wave", 1)) + stage1_k_batch = int(stage1.get("k_batch", 1)) + if stage1_k_batch != 1 or stage1_k_wave < 1: + return False + stage1_k_dim = model_dim // stage1_k_batch + stage2_sort_block = int(stage2.get("sort_block_m", 0) or stage2["tile_m"]) + return ( + stage1_k_dim % stage1_k_wave == 0 + and (stage1_k_dim // stage1_k_wave) % int(stage1["tile_k"]) == 0 + and stage2_sort_block % int(stage2["tile_m"]) == 0 + ) + + +def _fhmoe_config_file() -> str: from aiter.jit.core import AITER_CONFIGS return AITER_CONFIGS.AITER_CONFIG_FHMOE_FILE @functools.cache -def _supports_dsv4_i384_fhmoe_config(max_tokens: int, config_file: str) -> bool: +def _config_has_mxfp8_fhmoe_geometry( + config_file: str, + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, +) -> bool: + """Return whether the dedicated CSV declares this MXFP8/FP8 geometry.""" + with open(config_file, newline="") as stream: + rows = csv.DictReader(stream) + return any( + int(row["model_dim"]) == model_dim + and int(row["inter_dim"]) == inter_dim + and int(row["expert"]) == experts + and int(row["topk"]) == topk + and int(row["shared_expert_id"]) == experts - 1 + and row["q_dtype_w"] == str(dtypes.fp8) + for row in rows + ) + + +@functools.cache +def _config_has_mxfp8_fhmoe_contract( + config_file: str, + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + dtype: torch.dtype, + hidden_pad: int, + intermediate_pad: int, + gate_mode: GateMode, + doweight_stage1: bool, +) -> bool: + """Return whether the dedicated CSV declares this full MXFP8 contract.""" + if not _config_has_mxfp8_fhmoe_geometry( + config_file, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + ): + return False + with open(config_file, newline="") as stream: + rows = csv.DictReader(stream) + return any( + int(row["model_dim"]) == model_dim + and int(row["inter_dim"]) == inter_dim + and int(row["expert"]) == experts + and int(row["topk"]) == topk + and int(row["shared_expert_id"]) == experts - 1 + and row["act_type"] == str(ActivationType.Silu) + and row["dtype"] == str(dtype) + and row["q_dtype_a"] == str(dtypes.fp8) + and row["q_dtype_w"] == str(dtypes.fp8) + and row["q_type"] == str(QuantType.per_1x32) + and int(row["use_g1u1"]) == 1 + and int(row["doweight_stage1"]) == int(doweight_stage1) + and int(row["hidden_pad"]) == hidden_pad + and int(row["intermediate_pad"]) == intermediate_pad + and row["gate_mode"] == str(gate_mode) + for row in rows + ) + + +def _dsv4_i384_fhmoe_config_file() -> str: + """Compatibility wrapper for existing DSV4 capability tests.""" + return _fhmoe_config_file() + + +def _default_fhmoe_bucket(tokens: int, get_padded_m) -> int: + return get_padded_m(tokens) + + +def _supports_fhmoe_config( + max_tokens: int, + config_file: str, + *, + shape: tuple[int, int, int, int], + q_dtype_w: torch.dtype, + bucket_fn, + validate_fn, +) -> bool: try: from aiter.fused_moe import get_2stage_cfgs, get_padded_M - # Config keys are padded tiers; tensors retain the original M. required_tokens = { - get_padded_M(1 << exponent) for exponent in range(max_tokens.bit_length()) + bucket_fn(1 << exponent, get_padded_M) + for exponent in range(max_tokens.bit_length()) } - required_tokens.add(get_padded_M(max_tokens)) + required_tokens.add(bucket_fn(max_tokens, get_padded_M)) + model_dim, inter_dim, experts, topk = shape for token in required_tokens: metadata = get_2stage_cfgs( token, - 7168, - 384, - 385, - 7, + model_dim, + inter_dim, + experts, + topk, torch.bfloat16, dtypes.fp8, - dtypes.fp4x2, + q_dtype_w, QuantType.per_1x32, True, ActivationType.Silu, @@ -372,7 +582,9 @@ def _supports_dsv4_i384_fhmoe_config(max_tokens: int, config_file: str) -> bool: GateMode.INTERLEAVE, config_file=config_file, ) - _use_fhmoe_wrappers(metadata) + metadata = _use_fhmoe_wrappers(metadata) + if not validate_fn(metadata): + return False except ( ImportError, KeyError, @@ -385,8 +597,20 @@ def _supports_dsv4_i384_fhmoe_config(max_tokens: int, config_file: str) -> bool: return True +@functools.cache +def _supports_dsv4_i384_fhmoe_config(max_tokens: int, config_file: str) -> bool: + return _supports_fhmoe_config( + max_tokens, + config_file, + shape=(7168, 384, 385, 7), + q_dtype_w=dtypes.fp4x2, + bucket_fn=_default_fhmoe_bucket, + validate_fn=_is_dsv4_i384_metadata, + ) + + def supports_dsv4_i384_fhmoe(max_tokens: int) -> bool: - """Return whether the active FHMoE CSV covers every M through the limit.""" + """Check complete tuned DSV4 I384 CSV coverage through ``max_tokens``.""" if type(max_tokens) is not int or max_tokens <= 0: return False @@ -399,6 +623,205 @@ def supports_dsv4_i384_fhmoe(max_tokens: int) -> bool: return _supports_dsv4_i384_fhmoe_config(max_tokens, config_file) +@functools.cache +def _supports_mxfp8_fhmoe_config( + max_tokens: int, + config_file: str, + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + dtype: torch.dtype, + hidden_pad: int, + intermediate_pad: int, + gate_mode: GateMode, + doweight_stage1: bool, + force_reduce: bool, + stage2_fp8: bool, +) -> bool: + try: + from aiter.fused_moe import ( + _flydsl_route_output_nbytes, + _stage2_runtime_uses_route_reduce, + get_2stage_cfgs, + get_padded_M, + ) + + max_bucket = get_padded_M(max_tokens) + max_metadata = None + required_tokens = { + get_padded_M(1 << exponent) for exponent in range(max_tokens.bit_length()) + } + required_tokens.add(max_bucket) + for token in required_tokens: + metadata = get_2stage_cfgs( + token, + model_dim, + inter_dim, + experts, + topk, + dtype, + dtypes.fp8, + dtypes.fp8, + QuantType.per_1x32, + True, + ActivationType.Silu, + doweight_stage1, + hidden_pad, + intermediate_pad, + True, + gate_mode, + config_file=config_file, + ) + metadata = _use_fhmoe_wrappers(metadata) + if not _is_fhmoe_metadata_compatible( + metadata, + model_dim=model_dim, + dtype=dtype, + q_dtype_a=dtypes.fp8, + q_dtype_w=dtypes.fp8, + gate_mode=gate_mode, + ): + return False + if token == max_bucket: + max_metadata = metadata + if max_metadata is None: + return False + if _stage2_runtime_uses_route_reduce( + max_metadata.stage2, + max_tokens, + model_dim, + force_reduce=force_reduce, + stage2_fp8=stage2_fp8, + ): + output_element_size = torch.empty((), dtype=dtype).element_size() + if ( + _flydsl_route_output_nbytes( + max_metadata.stage2, + max_tokens, + topk, + model_dim, + output_element_size, + stage2_fp8=stage2_fp8, + ) + > FHMOE_BUFFER_MAX_BYTES + ): + return False + except ( + ImportError, + KeyError, + OSError, + RuntimeError, + TypeError, + ValueError, + ): + return False + return True + + +def supports_mxfp8_fhmoe_config( + max_tokens: int, + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + dtype: torch.dtype, + hidden_pad: int, + intermediate_pad: int, + gate_mode: GateMode, + doweight_stage1: bool, +) -> bool: + """Return whether dedicated MXFP8/FP8 coverage is executable through M.""" + if type(max_tokens) is not int or max_tokens <= 0: + return False + try: + if get_gfx_runtime() != "gfx950": + return False + if int(os.environ.get("AITER_BYPASS_TUNE_CONFIG", "0")) != 0: + return False + config_file = _fhmoe_config_file() + except (ImportError, KeyError, OSError, RuntimeError, TypeError, ValueError): + return False + return _supports_mxfp8_fhmoe_config( + max_tokens, + config_file, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + dtype=dtype, + hidden_pad=hidden_pad, + intermediate_pad=intermediate_pad, + gate_mode=gate_mode, + doweight_stage1=doweight_stage1, + force_reduce=os.environ.get("AITER_FLYDSL_FORCE_REDUCE", "0") == "1", + stage2_fp8=os.environ.get("AITER_FLYDSL_STAGE2_FP8", "0") == "1", + ) + + +def _resolve_mxfp8_fhmoe_config_file( + max_tokens: int, + *, + model_dim: int, + inter_dim: int, + experts: int, + topk: int, + dtype: torch.dtype, + hidden_pad: int, + intermediate_pad: int, + gate_mode: GateMode, + doweight_stage1: bool, +) -> str | None: + """Resolve a dedicated MXFP8 config, failing closed on partial matches.""" + config_file = _fhmoe_config_file() + has_geometry = _config_has_mxfp8_fhmoe_geometry( + config_file, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + ) + if not has_geometry: + return None + has_contract = _config_has_mxfp8_fhmoe_contract( + config_file, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + dtype=dtype, + hidden_pad=hidden_pad, + intermediate_pad=intermediate_pad, + gate_mode=gate_mode, + doweight_stage1=doweight_stage1, + ) + if not has_contract: + raise ValueError( + "MXFP8/FP8 FHMoE request partially matches a dedicated CSV " + "contract but differs in dtype, padding, gate mode, or " + "doweight_stage1" + ) + if not supports_mxfp8_fhmoe_config( + max_tokens, + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + dtype=dtype, + hidden_pad=hidden_pad, + intermediate_pad=intermediate_pad, + gate_mode=gate_mode, + doweight_stage1=doweight_stage1, + ): + raise NotImplementedError( + "The dedicated MXFP8/FP8 FHMoE config does not cover " + f"this request at M={max_tokens}" + ) + return config_file + + def _is_dsv4_i384_fhmoe_contract( model_dim: int, inter_dim: int, @@ -408,9 +831,11 @@ def _is_dsv4_i384_fhmoe_contract( intermediate_pad: int, gate_mode: GateMode, doweight_stage1: bool, + routed_mxfp4: bool = True, ) -> bool: return ( (model_dim, inter_dim, experts, topk) == (7168, 384, 385, 7) + and routed_mxfp4 and hidden_pad == 0 and intermediate_pad == 0 and gate_mode == GateMode.INTERLEAVE @@ -428,6 +853,7 @@ def _uses_dsv4_fhmoe_config( intermediate_pad: int, gate_mode: GateMode, doweight_stage1: bool, + routed_mxfp4: bool = True, ) -> bool: """Return whether dispatch should use the dedicated DSV4 FHMoE table.""" return _is_dsv4_i384_fhmoe_contract( @@ -439,6 +865,7 @@ def _uses_dsv4_fhmoe_config( intermediate_pad, gate_mode, doweight_stage1, + routed_mxfp4, ) and supports_dsv4_i384_fhmoe(token_num) @@ -472,6 +899,8 @@ def fhmoe_fake( shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, output: torch.Tensor | None = None, + # Model semantic forwarded by fused_moe; False leaves shared SiLU unclamped. + clamp_shared: bool = True, ) -> torch.Tensor: del ( w1, @@ -498,6 +927,7 @@ def fhmoe_fake( shared_w1_scale, shared_w2_scale, shared_expert_id, + clamp_shared, ) if output is not None: return output @@ -540,6 +970,8 @@ def fhmoe_( shared_w2_scale: torch.Tensor | None = None, shared_expert_id: int = -1, output: torch.Tensor | None = None, + # Keep this schema default aligned with fused_moe's public API. + clamp_shared: bool = True, ) -> torch.Tensor: from aiter.fused_moe import _fused_moe_impl @@ -565,6 +997,11 @@ def fhmoe_( bias2, ) q_dtype_a = dtypes.fp8 if gate_mode_enum == GateMode.INTERLEAVE else dtypes.fp4x2 + if not clamp_shared and w1.dtype != dtypes.fp8: + raise ValueError( + "clamp_shared=False requires MXFP8 routed weights; " + "MXFP4/FP8 FHMoE keeps the shared expert clamped" + ) experts = w1.shape[0] model_dim = w2.shape[1] inter_dim = w2.shape[2] * (model_dim // w1.shape[-1]) @@ -579,6 +1016,7 @@ def fhmoe_( intermediate_pad, gate_mode_enum, doweight_stage1, + routed_mxfp4=w1.dtype == dtypes.fp4x2, ): if not supports_dsv4_i384_fhmoe(hidden_states.shape[0]): raise NotImplementedError( @@ -588,6 +1026,20 @@ def fhmoe_( from aiter.jit.core import AITER_CONFIGS metadata_config_file = AITER_CONFIGS.AITER_CONFIG_FHMOE_FILE + elif w1.dtype == dtypes.fp8: + request_dtype = hidden_states.dtype if dtype is None else dtype + metadata_config_file = _resolve_mxfp8_fhmoe_config_file( + hidden_states.shape[0], + model_dim=model_dim, + inter_dim=inter_dim, + experts=experts, + topk=topk, + dtype=request_dtype, + hidden_pad=hidden_pad, + intermediate_pad=intermediate_pad, + gate_mode=gate_mode_enum, + doweight_stage1=doweight_stage1, + ) return _fused_moe_impl( hidden_states=hidden_states, w1=w1, @@ -616,11 +1068,13 @@ def fhmoe_( _q_dtype_a=q_dtype_a, _metadata_transform=_use_fhmoe_wrappers, _metadata_config_file=metadata_config_file, + _enforce_metadata_block_size=True, _stage1_extra_args={ "shared_w1": shared_w1, "shared_w1_scale": shared_w1_scale, "shared_expert_id": shared_expert_id, "swiglu_limit": swiglu_limit, + "clamp_shared": clamp_shared, }, _stage2_extra_args={ "shared_w2": shared_w2, diff --git a/aiter/fhmoe_contract.py b/aiter/fhmoe_contract.py new file mode 100644 index 00000000000..020de73ae67 --- /dev/null +++ b/aiter/fhmoe_contract.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Lightweight FHMoE safety helpers shared by runtime, AOT, and tuners. + +This module must remain free of imports from the top-level ``aiter`` namespace +so it is safe when ``AITER_AOT_IMPORT=1`` exposes only the lightweight JIT core. +""" + +FHMOE_BUFFER_MAX_BYTES = (1 << 32) - 1 + + +def fhmoe_reduce_output_nbytes( + tokens: int, + topk: int, + model_dim: int, + element_size: int, +) -> int: + """Return bytes in the reduce-mode ``[M, topk, model_dim]`` temporary.""" + return int(tokens) * int(topk) * int(model_dim) * int(element_size) + + +def max_fhmoe_reduce_tokens( + topk: int, + model_dim: int, + element_size: int, +) -> int: + """Largest M addressable by the current 32-bit reduce-output resource.""" + bytes_per_token = fhmoe_reduce_output_nbytes( + 1, + topk, + model_dim, + element_size, + ) + if bytes_per_token <= 0: + raise ValueError("FHMoE reduce-output dimensions must be positive") + return FHMOE_BUFFER_MAX_BYTES // bytes_per_token diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index e964795b01c..d595d4dbaa6 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -22,6 +22,10 @@ mxfp4_moe_sort_fwd, ) from aiter import get_hip_quant as get_quant +from aiter.fhmoe_contract import ( + FHMOE_BUFFER_MAX_BYTES, + fhmoe_reduce_output_nbytes, +) from aiter.fused_moe_registry import ( BoundFusedMoeImpl, FusedMoeImplResolutionError, @@ -837,11 +841,34 @@ def get_topk_valid_mask( return expert_mask[topk_ids] +def _is_flydsl_stage2(stage2: Callable) -> bool: + func = getattr(stage2, "func", stage2) + return func is _flydsl_stage2_wrapper or getattr(func, "_is_flydsl_stage2", False) + + +def _stage2_honors_runtime_reduce_override(stage2: Callable) -> bool: + """Return whether stage2 reaches common runtime reduce resolution.""" + if not _is_flydsl_stage2(stage2): + return False + kernel_name = getattr(stage2, "keywords", {}).get("kernelName", "") + parsed = _get_flydsl_moe_kernels().get_flydsl_kernel_params(kernel_name) + if parsed is None: + raise ValueError( + "Runtime reduce overrides require a recognized FlyDSL stage2 " + f"kernel, got {kernel_name!r}" + ) + # The a16w-mix launcher returns before the environment override. MXFP4 + # always atomically scatters; INT4 follows only its configured kernel mode. + return not ( + parsed.get("a_dtype") == "bf16" and parsed.get("b_dtype") in ("fp4", "int4") + ) + + def stage2_uses_route_reduce(stage2: Callable) -> bool: """Return True when stage2 writes per-slot route output then reduces it.""" func = getattr(stage2, "func", stage2) kernel_name = getattr(stage2, "keywords", {}).get("kernelName", "") - if func is _flydsl_stage2_wrapper or getattr(func, "_is_flydsl_stage2", False): + if _is_flydsl_stage2(func): parsed = _get_flydsl_moe_kernels().get_flydsl_kernel_params(kernel_name) if parsed is None: return False @@ -860,6 +887,95 @@ def stage2_uses_route_reduce(stage2: Callable) -> bool: return False +def _stage2_runtime_uses_route_reduce( + stage2: Callable, + token_num: int, + model_dim: int, + *, + force_reduce: bool | None = None, + stage2_fp8: bool | None = None, +) -> bool: + """Return the effective route-output mode used by stage2 at runtime.""" + if stage2_uses_route_reduce(stage2): + return True + if not _stage2_honors_runtime_reduce_override(stage2): + return False + kernel_name = getattr(stage2, "keywords", {}).get("kernelName", "") + parsed = _get_flydsl_moe_kernels().get_flydsl_kernel_params(kernel_name) + assert parsed is not None + mode = _get_flydsl_moe_kernels().resolve_flydsl_stage2_mode( + parsed.get("mode", "atomic"), + token_num, + model_dim, + 2, + parsed.get("b_dtype", ""), + False, + force_reduce=force_reduce, + stage2_fp8=stage2_fp8, + ) + return mode == "reduce" + + +def _stage2_uses_fp8_route_output( + stage2: Callable, + *, + stage2_fp8: bool | None = None, + return_per_slot: bool = False, +) -> bool: + """Return whether runtime reduce output uses packed FP8 value/scale rows.""" + if stage2_fp8 is None: + stage2_fp8 = os.environ.get("AITER_FLYDSL_STAGE2_FP8", "0") == "1" + if not stage2_fp8 or return_per_slot: + return False + if not _stage2_honors_runtime_reduce_override(stage2): + return False + kernel_name = getattr(stage2, "keywords", {}).get("kernelName", "") + parsed = _get_flydsl_moe_kernels().get_flydsl_kernel_params(kernel_name) + assert parsed is not None + return parsed.get("b_dtype") in ("fp4", "fp8") + + +def _flydsl_route_output_nbytes( + stage2: Callable, + token_num: int, + topk: int, + model_dim: int, + output_element_size: int, + *, + stage2_fp8: bool | None = None, + return_per_slot: bool = False, +) -> int: + """Return bytes allocated by the effective common stage2 reduce output.""" + if _stage2_uses_fp8_route_output( + stage2, + stage2_fp8=stage2_fp8, + return_per_slot=return_per_slot, + ): + from aiter.ops.flydsl.kernels.mxfp4_gemm_common import fp8out_row_bytes + + kernels = _get_flydsl_moe_kernels() + row_bytes = fp8out_row_bytes( + model_dim, + scale_blk=kernels._S2_LEGACY_FP8_SCALE_BLK, + pitch_align=kernels._S2_LEGACY_FP8_PITCH_ALIGN, + ) + return int(token_num) * int(topk) * int(row_bytes) + return fhmoe_reduce_output_nbytes( + token_num, + topk, + model_dim, + output_element_size, + ) + + +def _validate_flydsl_route_output_nbytes(route_output_nbytes: int) -> None: + if route_output_nbytes > FHMOE_BUFFER_MAX_BYTES: + raise NotImplementedError( + "FlyDSL stage2 reduce output exceeds the current 32-bit " + f"buffer-addressing limit: {route_output_nbytes} bytes" + ) + + # Lru cache will using hash to create key, which makes error when w1,w2 shape is symint. # We can use torch.compile(dynamic=False) to avoid @functools.lru_cache(maxsize=2048) @@ -1005,14 +1121,67 @@ def fused_moe( quant_type_a: QuantType | None = None, quant_dtype_a: torch.dtype | None = None, quant_dtype_a2: torch.dtype | None = None, + # False explicitly selects unclamped shared-expert SiLU semantics. + clamp_shared: bool = True, ): - if ( + """Run fused MoE, dispatching to FHMoE when shared-expert tensors are present. + + Supplying any shared tensor (or ``shared_expert_id``) requires all four + shared weight/scale tensors and dispatches through the gfx950 FHMoE path. + Inputs must satisfy the contracts documented in :mod:`aiter.fhmoe`. + + ``clamp_shared`` is an FHMoE model-semantic input. The default ``True`` + applies the configured SiLU clamp to routed and shared experts. ``False`` + changes only the shared expert to unclamped SiLU; routed experts still use + ``swiglu_limit``. Clamped and unclamped stage-1 kernels have distinct AOT + cache identities, so deployments changing this setting must rebuild their + FlyDSL AOT cache. + + """ + has_shared_expert = ( any( tensor is not None for tensor in (shared_w1, shared_w2, shared_w1_scale, shared_w2_scale) ) or shared_expert_id != -1 - ): + ) + if not clamp_shared and not has_shared_expert: + raise ValueError("clamp_shared=False requires FHMoE shared-expert dispatch") + if has_shared_expert: + routed_mxfp8 = w1 is not None and w1.dtype == dtypes.fp8 + unsupported = [] + if stage2_scatter is not None: + unsupported.append("stage2_scatter") + if quant_type_a is not None: + unsupported.append("quant_type_a") + fhmoe_quant_dtype_a = ( + dtypes.fp8 if GateMode(gate_mode) == GateMode.INTERLEAVE else dtypes.fp4x2 + ) + if quant_dtype_a is not None and quant_dtype_a != fhmoe_quant_dtype_a: + unsupported.append("quant_dtype_a") + if quant_dtype_a2 is not None: + unsupported.append("quant_dtype_a2") + # The MXFP8 path does not consume caller-prequantized activations and + # requires an explicit preshuffle marker. Preserve the existing DSV4 + # MXFP4 path, which accepts activation scales and trusts its layout. + if routed_mxfp8 and a1_scale is not None: + unsupported.append("a1_scale") + if routed_mxfp8 and a2_scale is not None: + unsupported.append("a2_scale") + if unsupported: + raise NotImplementedError( + "FHMoE shared-expert dispatch does not support " + + ", ".join(unsupported) + ) + weights = (w1, w2, shared_w1, shared_w2) + if routed_mxfp8 and any( + weight is not None and not getattr(weight, "is_shuffled", False) + for weight in weights + ): + raise ValueError( + "FHMoE weights must be AITER-preshuffled and marked " + "is_shuffled=True; re-run the shuffle after tensor copies" + ) from aiter.fhmoe import _fhmoe return _fhmoe( @@ -1044,6 +1213,7 @@ def fused_moe( shared_w1_scale=shared_w1_scale, shared_w2_scale=shared_w2_scale, shared_expert_id=shared_expert_id, + clamp_shared=clamp_shared, output=output, ) if not block_size_M: @@ -1282,6 +1452,7 @@ def _fused_moe_impl( _stage1_extra_args: dict | None = None, _stage2_extra_args: dict | None = None, _stage2_override: Callable | None = None, + _enforce_metadata_block_size: bool = False, ) -> torch.Tensor: # We do such convert since custom_op schema restriction on block_size_M, and Enum type activation = ActivationType(activation) @@ -1437,8 +1608,7 @@ def _fused_moe_impl( ): if _bad: raise NotImplementedError( - f"a16w4 (bf16 A x MXFP4 W) {_a16w4_why} is not supported: " - f"{_why}." + f"a16w4 (bf16 A x MXFP4 W) {_a16w4_why} is not supported: {_why}." ) config_situ_beta, config_situ_linear_beta, _ = _normalize_mxfp4_activation_params( @@ -1447,7 +1617,7 @@ def _fused_moe_impl( def _resolve_metadata(disable_inline_sort=False): metadata = get_2stage_cfgs( - get_padded_M(M), # consider token_num > 1024 as prefill + get_padded_M(M), model_dim, inter_dim, E, @@ -1529,6 +1699,15 @@ def _resolve_metadata(disable_inline_sort=False): "FlyDSL fallback cannot accept this invocation." ) + if ( + _enforce_metadata_block_size + and block_size_M is not None + and int(block_size_M) != int(metadata.block_m) + ): + raise ValueError( + f"FHMoE block_size_M={block_size_M} does not match the selected " + f"kernel block size {metadata.block_m}" + ) block_size_M = metadata.block_m if block_size_M is None else block_size_M if metadata.full_impl is not None: full_output = metadata.full_impl( @@ -1585,6 +1764,25 @@ def _resolve_metadata(disable_inline_sort=False): sort_m_indices = None sort_reverse_sorted = None + runtime_stage2_reduce = _stage2_runtime_uses_route_reduce( + metadata.stage2, + M, + model_dim, + ) + uses_common_route_reduce = ( + _is_flydsl_stage2(metadata.stage2) and runtime_stage2_reduce + ) + if uses_common_route_reduce: + output_element_size = torch.empty((), dtype=dtype).element_size() + _validate_flydsl_route_output_nbytes( + _flydsl_route_output_nbytes( + metadata.stage2, + M, + topk, + model_dim, + output_element_size, + ) + ) prequantized_a1 = None prequantized_a1_scale = None # The a4w4 aux-sort families already fuse activation quantization into GEMM1, @@ -1639,7 +1837,7 @@ def _resolve_metadata(disable_inline_sort=False): expert_mask, num_local_tokens, return_local_topk_ids=need_local_topk_ids, - accumulate=_atomic, + accumulate=_atomic and not runtime_stage2_reduce, output_aux=metadata.output_aux, output=output, ) @@ -1670,7 +1868,7 @@ def _resolve_metadata(disable_inline_sort=False): global_experts=global_E, block_m=block_size_M, moebuf_dtype=dtype, - accumulate=not stage2_uses_route_reduce(metadata.stage2), + accumulate=not runtime_stage2_reduce, compact_scale=fused_compact_scale, ) local_topk_ids = None @@ -1686,7 +1884,7 @@ def _resolve_metadata(disable_inline_sort=False): num_local_tokens, moe_sorting_dispatch_policy, return_local_topk_ids=need_local_topk_ids, - accumulate=not stage2_uses_route_reduce(metadata.stage2), + accumulate=not runtime_stage2_reduce, flat=metadata.flat, output=None if metadata.flat else output, ) @@ -2067,6 +2265,12 @@ def nextPow2(n): def get_padded_M(M): + """Map runtime M to a tuning bucket. + + Values below 32768 use the next power of two, 32768..131071 use the 32768 + tier, and values from 131072 upward use the 131072 tier. Individual model + contracts may impose a stricter maximum before lookup. + """ if M < _PADDED_M_TIERS[0]: return nextPow2(M) for tier in reversed(_PADDED_M_TIERS): @@ -3446,8 +3650,7 @@ def _lookup_cfg(c2s): reject_reason = f"no MXMOE kernel for activation {activation!r}" elif configured_act != expected_act: reject_reason = ( - f"activation {configured_act!r} does not match runtime " - f"{expected_act!r}" + f"activation {configured_act!r} does not match runtime {expected_act!r}" ) elif swiglu_limit and expected_act not in ("silu", "swiglu"): reject_reason = ( diff --git a/aiter/ops/flydsl/fhmoe.py b/aiter/ops/flydsl/fhmoe.py index 8896e68b3d0..004e8f0e407 100644 --- a/aiter/ops/flydsl/fhmoe.py +++ b/aiter/ops/flydsl/fhmoe.py @@ -4,6 +4,7 @@ """FlyDSL runtime support for fused heterogeneous MoE (FHMoE).""" import functools +import math import torch @@ -12,6 +13,7 @@ from .moe_kernels import ( _flydsl_moe_stage1_impl, _flydsl_moe_stage2_impl, + runtime_swiglu_limit, ) @@ -42,13 +44,21 @@ def compile_flydsl_fhmoe_stage1( k_wave: int = 1, v2_output_layout: bool = False, shared_expert_id: int = -1, + clamp_shared: bool = True, ): """Compile the heterogeneous stage1 kernel.""" from .kernels.fhmoe import compile_mixed_fhmoe_gemm1 from .moe_common import GateMode - if b_dtype != "fp4": - raise ValueError(f"FHMoE stage1 requires routed MXFP4 weights, got {b_dtype}") + if b_dtype not in ("fp4", "fp8"): + raise ValueError( + f"FHMoE stage1 requires routed MXFP4 or MXFP8 weights, got {b_dtype}" + ) + normalized_gate_mode = GateMode(gate_mode) + if b_dtype == "fp8" and normalized_gate_mode != GateMode.INTERLEAVE: + raise ValueError( + "FHMoE MXFP8 routed weights require interleaved gate/up layout" + ) return compile_mixed_fhmoe_gemm1( model_dim=model_dim, inter_dim=inter_dim, @@ -67,7 +77,7 @@ def compile_flydsl_fhmoe_stage1( k_batch=k_batch, waves_per_eu=waves_per_eu, b_nt=b_nt, - gate_mode=GateMode(gate_mode), + gate_mode=normalized_gate_mode, model_dim_pad=model_dim_pad, inter_dim_pad=inter_dim_pad, enable_bias=enable_bias, @@ -76,6 +86,7 @@ def compile_flydsl_fhmoe_stage1( k_wave=k_wave, v2_output_layout=v2_output_layout, shared_expert_id=shared_expert_id, + clamp_shared=clamp_shared, ) @@ -108,8 +119,10 @@ def compile_flydsl_fhmoe_stage2( """Compile the heterogeneous stage2 kernel.""" from .kernels.fhmoe import compile_mixed_fhmoe_gemm2 - if b_dtype != "fp4": - raise ValueError(f"FHMoE stage2 requires routed MXFP4 weights, got {b_dtype}") + if b_dtype not in ("fp4", "fp8"): + raise ValueError( + f"FHMoE stage2 requires routed MXFP4 or MXFP8 weights, got {b_dtype}" + ) return compile_mixed_fhmoe_gemm2( model_dim=model_dim, inter_dim=inter_dim, @@ -291,11 +304,31 @@ def flydsl_fhmoe_stage1( shared_w1: torch.Tensor, shared_w1_scale: torch.Tensor, shared_expert_id: int, + clamp_shared: bool = True, ): - """Run stage1 with MXFP4 routed experts and one FP8 shared expert.""" + """Run stage1 with MXFP4/MXFP8 routed and separately stored FP8 shared weights.""" + if b_dtype == "fp8" and gate_mode != "interleave": + raise ValueError( + "FHMoE MXFP8 routed weights require interleaved gate/up layout" + ) + if not clamp_shared and act != "silu": + raise ValueError("clamp_shared=False is supported only with act='silu'") + activation_limit = runtime_swiglu_limit(swiglu_limit, act) + if ( + not clamp_shared + and k_batch > 1 + and act == "silu" + and math.isfinite(activation_limit) + ): + raise NotImplementedError( + "Split-K FHMoE postprocessing cannot clamp routed SiLU while " + "leaving the shared expert unclamped; use k_batch=1 when " + "clamp_shared=False and swiglu_limit is finite" + ) compile_kernel = functools.partial( compile_flydsl_fhmoe_stage1, shared_expert_id=shared_expert_id, + clamp_shared=clamp_shared, ) build_mx_args = functools.partial( _s1_args_fhmoe, @@ -376,7 +409,7 @@ def flydsl_fhmoe_stage2( shared_w2_scale: torch.Tensor, shared_expert_id: int, ) -> torch.Tensor: - """Run stage2 with MXFP4 routed experts and one FP8 shared expert.""" + """Run stage2 with MXFP4/MXFP8 routed and separately stored FP8 shared weights.""" compile_kernel = functools.partial( compile_flydsl_fhmoe_stage2, shared_expert_id=shared_expert_id, diff --git a/aiter/ops/flydsl/kernels/fhmoe.py b/aiter/ops/flydsl/kernels/fhmoe.py index 73ffe7b7207..5c247e18531 100644 --- a/aiter/ops/flydsl/kernels/fhmoe.py +++ b/aiter/ops/flydsl/kernels/fhmoe.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Heterogeneous MoE facades for the shared MXFP4/FP8 kernel builders.""" +"""Heterogeneous MoE facades for MXFP4/MXFP8 routed and FP8 shared weights.""" import functools @@ -43,9 +43,10 @@ def compile_mixed_fhmoe_gemm1( k_wave: int = 1, v2_output_layout: bool = False, shared_expert_id: int, + clamp_shared: bool = True, ): """Compile a stage1 kernel with an FP8 shared expert.""" - if shared_expert_id is None: + if shared_expert_id != experts - 1: raise ValueError( "FHMoE stage1 requires shared_expert_id == experts - 1; " f"got {shared_expert_id=} and {experts=}" @@ -78,6 +79,7 @@ def compile_mixed_fhmoe_gemm1( k_wave=k_wave, v2_output_layout=v2_output_layout, shared_expert_id=shared_expert_id, + clamp_shared=clamp_shared, ) @@ -111,7 +113,7 @@ def compile_mixed_fhmoe_gemm2( use_global_a: bool = True, ): """Compile a stage2 kernel with an FP8 shared expert.""" - if shared_expert_id is None: + if shared_expert_id != experts - 1: raise ValueError( "FHMoE stage2 requires shared_expert_id == experts - 1; " f"got {shared_expert_id=} and {experts=}" diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index ea0257e4ae4..97567674fb7 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (C) 2025-2026 FlyDSL Project Contributors -"""Ordinary MoE facades for the shared MXFP4/FP8 kernel builders.""" +"""Ordinary MoE facades for the shared MXFP4/MXFP8 kernel builders.""" import functools diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py index a00c486396f..83e08d75427 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (C) 2025-2026 FlyDSL Project Contributors -"""Shared MXFP4/FP8 MoE and heterogeneous MoE kernel builders.""" +"""Shared MXFP4/MXFP8 MoE and heterogeneous MoE kernel builders.""" """MoE GEMM stage1/stage2 kernel implementations (FlyDSL MFMA FP8/FP16/FP4). @@ -138,6 +138,7 @@ def compile_mixed_moe_gemm1_common( xcd_swizzle: int = 0, k_wave: int = 1, shared_expert_id: int | None = None, + clamp_shared: bool = True, v2_output_layout: bool = False, ): """Compile stage1 kernel: act(X @ W_gate.T, X @ W_up.T) -> [tokens*topk, inter_dim].""" @@ -147,6 +148,8 @@ def compile_mixed_moe_gemm1_common( "FHMoE stage1 requires shared_expert_id == experts - 1; " f"got {shared_expert_id=} and {experts=}" ) + if heterogeneous_b and not clamp_shared and act != "silu": + raise ValueError("clamp_shared=False is supported only with act='silu'") gpu_arch = get_hip_arch() def _al(x, a): @@ -160,9 +163,6 @@ def _al(x, a): is_f4_a = a_dtype == "fp4" is_f4_b = b_dtype == "fp4" is_f8_b = b_dtype == "fp8" - if heterogeneous_b and not is_f4_b: - raise ValueError("Heterogeneous B requires MXFP4 routed weights") - sort_block_m = tile_m num_waves = min(4, tile_n // 32) # K-wave groups reduce accumulators in LDS before the epilogue. @@ -272,12 +272,17 @@ def load_bias_scalar(bias_rsrc, offset): # therefore must not be part of the on-disk symbol/cache identity. act_tag = "" if act == "silu" else f"_{act}" heterogeneous_tag = f"_shared_fp8_e{shared_expert_id}" if heterogeneous_b else "" - # ABI v33 adds four runtime SiTUv2 beta scalars; heterogeneous ABI tracks one - # version ahead of the ordinary kernel. - kernel_version = 34 if heterogeneous_b else 33 + # Keep the existing MXFP4/FP8 DSV4 cache identity unchanged. The no-clamp + # tag is HY4-specific; heterogeneous MXFP8 kernels use the v35 revision. + shared_clamp_tag = "_shared_noclamp" if heterogeneous_b and not clamp_shared else "" + # ABI v33 adds four runtime SiTUv2 beta scalars. Heterogeneous MXFP4 keeps + # v34; v35 identifies the MXFP8 routed/shared K-offset implementation. + kernel_version = ( + 35 if heterogeneous_b and is_f8_b else (34 if heterogeneous_b else 33) + ) module_name = ( f"mfma_moe1_silu_mul_a{a_dtype}_w{b_dtype}_{out_s}" - f"_t{tile_m}x{tile_n}x{tile_k}_pm{persist_m}{fp4q_tag}{fp8q_tag}{sort_tag}{async_tag}{sk_tag}{kw_tag}{go_tag}{gui_tag}{as1_tag}{xcd_tag}{act_tag}{v2out_tag}{heterogeneous_tag}_v{kernel_version}" + f"_t{tile_m}x{tile_n}x{tile_k}_pm{persist_m}{fp4q_tag}{fp8q_tag}{sort_tag}{async_tag}{sk_tag}{kw_tag}{go_tag}{gui_tag}{as1_tag}{xcd_tag}{act_tag}{v2out_tag}{heterogeneous_tag}{shared_clamp_tag}_v{kernel_version}" ).replace("-", "_") cshuffle_elem_bytes = 4 if need_quant else (4 if out_is_f32 else 2) @@ -928,7 +933,10 @@ def load_cell(rsrc, b_layout_arg, elem_type, k0): ku * b_cells_per_ku, index=True ) if const_expr(shared_b): - shared_k0_base = (base_k * arith.index(2)) // c64 + # ``base_k`` is expressed in routed-weight storage + # elements. Convert it back to the logical K offset + # before indexing the always-FP8 shared weights. + shared_k0_base = (base_k * arith.index(b_byte_div)) // c64 shared_k0_base += arith.constant(ku * 2, index=True) s0, s1 = load_cell( shared_w_rsrc, @@ -1951,7 +1959,15 @@ def _tid_dispatch(): def activate_pair(g, u): if const_expr(act == "silu"): - g, u = clamp_gate_up(g, u, swiglu_neg_limit) + # Some models leave the dense shared expert unclamped. + # ``shared_b`` specializes the whole tile, so + # disabling the clamp here adds no per-lane divergence. + activation_neg_limit = ( + fx.Float32(float("-inf")) + if const_expr(shared_b and not clamp_shared) + else swiglu_neg_limit + ) + g, u = clamp_gate_up(g, u, activation_neg_limit) return silu_mul_batch([g], [u])[0] return gate_up_act(act, [g], [u], activation_params)[0] @@ -3189,8 +3205,6 @@ def compile_mixed_moe_gemm2_common( is_f4_a = a_dtype == "fp4" is_f4_b = b_dtype == "fp4" is_f8_b = b_dtype == "fp8" - if heterogeneous_b and not is_f4_b: - raise ValueError("Heterogeneous B requires MXFP4 routed weights") serial_shared_n = heterogeneous_b and tile_n == 256 scale_pack_m = 2 @@ -3332,7 +3346,12 @@ def load_bias_scalar(bias_rsrc, offset): xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" heterogeneous_tag = f"_shared_fp8_e{shared_expert_id}" if heterogeneous_b else "" serial_n_tag = "_serialn128" if serial_shared_n else "" - if heterogeneous_b: + if heterogeneous_b and is_f8_b: + variant_tags = ( + f"_vscale_fix3_fp4opt_v3{pm_tag}{sbm_tag}{wpe_tag}{async_tag}" + f"{cumul_tag}{acc_tag}{xcd_tag}{heterogeneous_tag}{serial_n_tag}" + ) + elif heterogeneous_b: variant_tags = ( f"_vscale_fix3_fp4opt_v2{pm_tag}{sbm_tag}{wpe_tag}{async_tag}" f"{cumul_tag}{acc_tag}{xcd_tag}{heterogeneous_tag}{serial_n_tag}" @@ -3886,7 +3905,9 @@ def load_cell(rsrc, expert_base, stride_n0, elem_type, k0): ) if const_expr(shared_b): - shared_k0_base = _div_pow2(base_k * arith.index(2), 64) + # ``base_k`` uses the routed-weight storage width; + # recover the logical K offset for FP8 shared weights. + shared_k0_base = _div_pow2(base_k * arith.index(b_byte_div), 64) shared_k0_base += arith.constant(ku * 2, index=True) s0, s1 = load_cell( shared_w_rsrc, diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 26431abdd73..acc327faaed 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -111,6 +111,34 @@ def requires_flydsl_stage2_reduce( return int(token_num) * int(model_dim) * int(element_size) > 0xFFFFFFFF +def resolve_flydsl_stage2_mode( + mode: str, + token_num: int, + model_dim: int, + element_size: int, + b_dtype: str, + return_per_slot: bool, + *, + force_reduce: bool | None = None, + stage2_fp8: bool | None = None, +) -> str: + """Resolve runtime reduce overrides shared by sorting and stage 2.""" + if force_reduce is None: + force_reduce = os.environ.get("AITER_FLYDSL_FORCE_REDUCE", "0") == "1" + if stage2_fp8 is None: + stage2_fp8 = os.environ.get("AITER_FLYDSL_STAGE2_FP8", "0") == "1" + fp8_reduce = stage2_fp8 and b_dtype in ("fp4", "fp8") and not return_per_slot + if force_reduce or fp8_reduce: + return "reduce" + if ( + mode != "reduce" + and not return_per_slot + and requires_flydsl_stage2_reduce(token_num, model_dim, element_size) + ): + return "reduce" + return mode + + def requires_flydsl_stage2_global_a(a: torch.Tensor) -> bool: """Return whether A's stored bytes reach the 4 GiB buffer descriptor limit.""" return a.numel() * a.element_size() >= (1 << 32) @@ -2112,17 +2140,14 @@ def _flydsl_moe_stage2_impl( model_dim = w2.shape[1] inter_dim = inter_states.shape[2] - # Debug: force stage2 to use the masked reduce epilogue instead of atomic - # accumulate. Enabled by default; set AITER_FLYDSL_FORCE_REDUCE=0 to opt out. - if os.environ.get("AITER_FLYDSL_FORCE_REDUCE", "0") == "1": - mode = "reduce" - elif ( - mode != "reduce" - and not return_per_slot - and requires_flydsl_stage2_reduce(token_num, model_dim, 2) - ): - # Buffer atomics use 32-bit offsets; reduce outputs larger than 4 GiB. - mode = "reduce" + mode = resolve_flydsl_stage2_mode( + mode, + token_num, + model_dim, + 2, + b_dtype, + return_per_slot, + ) accumulate = mode != "reduce" and not return_per_slot @@ -2181,7 +2206,8 @@ def _flydsl_moe_stage2_impl( _persist_m = -1 if m_blocks > 256 else 1 if a_dtype == "fp8": - # FP8 uses non-persistent scheduling, so cap grid.y via persist_m. + # Normalize FP8-A grid scheduling from runtime occupancy, even for + # legacy tuned names that retain `_persist` as part of their identity. _persist_m = resolve_flydsl_grid_y_persist_m(m_blocks) if bias is not None and bias.dtype != torch.float32: diff --git a/aiter/utility/fp4_utils.py b/aiter/utility/fp4_utils.py index e15bceee500..539d43bb003 100644 --- a/aiter/utility/fp4_utils.py +++ b/aiter/utility/fp4_utils.py @@ -158,6 +158,19 @@ def f32_to_mx_e8m0_scale( ) +def f32_to_fused_moe_mxfp8_scale(amax: Tensor) -> Tensor: + """Match the FlyDSL fused-MoE FP8 epilogue's E8M0 scale rounding. + + The epilogue rounds ``amax`` to the nearest power of two with ties upward, + then reserves eight exponent steps (a factor of 2^8) of FP8 headroom. This + differs from the project-wide MXFP8 ``RoundUp`` default. + """ + bits = amax.float().contiguous().view(torch.int32) + rounded = (bits + 0x400000) & -0x800000 + exponent = ((rounded >> 23) - 8).clamp_(0, 255).to(torch.uint8) + return exponent.view(dtypes.fp8_e8m0) + + def fp4_f32_to_e8m0_scale(amax: Tensor) -> Tensor: """Default MXFP4 E8M0 block scale: NV ROUND_UP / RCEIL with FP4 E2M1. diff --git a/op_tests/test_fhmoe.py b/op_tests/test_fhmoe.py index 6062c50110a..9f75b1697d5 100644 --- a/op_tests/test_fhmoe.py +++ b/op_tests/test_fhmoe.py @@ -3,11 +3,14 @@ """Correctness oracle for fused heterogeneous MoE (FHMoE). -The default case is deliberately small. Set AITER_HETERO_MOE_DSV4=1 to use -the exact DeepSeek-V4-Pro TP8 dimensions; that profile allocates several GiB. +The generic MXFP4 case is deliberately small. Dedicated MXFP8 tests use the +production HY4 TP8 dimensions (6144x256, E=257) and allocate several GiB. +Set AITER_HETERO_MOE_DSV4=1 to run the generic path with exact +DeepSeek-V4-Pro TP8 dimensions, or AITER_HETERO_MOE_HY4=1 to run that sweep +with HY4 MXFP8 routed weights, unclamped shared semantics, and tuned dispatch. Set AITER_HETERO_MOE_FULL_SWEEP=1 to cover M=1,4,8,16,24,32,40,64 and -AITER_HETERO_MOE_STRESS_REPEATS=500 to stress both stage-2 epilogues and -bound the default atomic path's run-to-run variation. +AITER_HETERO_MOE_STRESS_REPEATS=500 to stress forced-reduce and default +stage-2 execution and bound atomic run-to-run variation where applicable. """ from __future__ import annotations @@ -26,6 +29,7 @@ from aiter.fused_moe import fused_moe from aiter.jit.utils.chip_info import get_gfx from aiter.ops.flydsl.moe_common import GateMode +from aiter.ops.flydsl.moe_kernels import runtime_swiglu_limit from aiter.ops.quant import per_1x32_f4_quant from aiter.ops.shuffle import ( shuffle_scale, @@ -33,6 +37,7 @@ shuffle_weight, shuffle_weight_a16w4, ) +from aiter.ops.triton.quant import dynamic_mxfp8_quant from aiter.utility import fp4_utils from aiter.utility.mx_types import MxDtypeInt @@ -64,6 +69,7 @@ class _Weights: raw_routed_w2: torch.Tensor raw_routed_s1: torch.Tensor raw_routed_s2: torch.Tensor + routed_expert_ids: tuple[int, ...] shared_w1: torch.Tensor shared_w2: torch.Tensor shared_s1: torch.Tensor @@ -79,11 +85,22 @@ class _Weights: def _profile() -> _Profile: + if _use_hy4_profile(): + return _Profile(6144, 256, 256, 257, 8) if os.environ.get("AITER_HETERO_MOE_DSV4", "0") == "1": return _Profile(7168, 384, 384, 385, 6) return _Profile(256, 128, 128, 9, 2) +def _use_hy4_profile() -> bool: + return os.environ.get("AITER_HETERO_MOE_HY4", "0") == "1" + + +def _mxfp8_profile() -> _Profile: + """The production contract for MXFP8 routed and FP8 shared experts.""" + return _Profile(6144, 256, 256, 257, 8) + + def _m_values() -> list[int]: if os.environ.get("AITER_HETERO_MOE_FULL_SWEEP", "0") == "1": return [1, 4, 8, 16, 24, 32, 40, 64] @@ -94,6 +111,10 @@ def _swiglu_limit(profile: _Profile) -> float: return 10.0 +def _shared_swiglu_limit(swiglu_limit: float, *, clamp_shared: bool) -> float: + return swiglu_limit if clamp_shared else float("inf") + + def _rel_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: delta = actual.float() - expected.float() denominator = torch.linalg.vector_norm(expected.float()).clamp_min(1e-12) @@ -275,11 +296,68 @@ def _make_routed_weights( w2 = w2_u8.view(dtypes.fp4x2) s1 = s1_u8.view(dtypes.fp8_e8m0) s2 = s2_u8.view(dtypes.fp8_e8m0) - active = slice(0, profile.routed_topk) - raw_w1 = w1[active].clone() - raw_w2 = w2[active].clone() - raw_s1 = s1[active].clone() - raw_s2 = s2[active].clone() + routed_expert_ids = tuple(range(profile.routed_topk)) + raw_w1 = w1[: profile.routed_topk].clone() + raw_w2 = w2[: profile.routed_topk].clone() + raw_s1 = s1[: profile.routed_topk].clone() + raw_s2 = s2[: profile.routed_topk].clone() + return ( + *_shuffle_w1(w1, s1, interleave), + *_shuffle_w2(w2, s2, interleave), + raw_w1, + raw_w2, + raw_s1, + raw_s2, + routed_expert_ids, + ) + + +def _make_mxfp8_routed_weights( + profile: _Profile, + device: torch.device, + interleave: bool = True, +) -> tuple: + h = profile.hidden + i = profile.inter + e = profile.experts + generator = torch.Generator(device=device).manual_seed(29) + + w1 = torch.zeros((e, 2 * i, h), dtype=dtypes.fp8, device=device) + w2 = torch.zeros((e, h, i), dtype=dtypes.fp8, device=device) + s1_u8 = torch.full((e, 2 * i, h // 32), 0x7F, dtype=torch.uint8, device=device) + s2_u8 = torch.full((e, h, i // 32), 0x7F, dtype=torch.uint8, device=device) + + routed_expert_ids = ( + (0, 1, 7, 8, 127, 128, 254, 255) + if (profile.experts, profile.routed_topk) == (257, 8) + else tuple(range(profile.routed_topk)) + ) + for slot, expert_id in enumerate(routed_expert_ids): + scale = 0.02 + 0.005 * slot + dense_w1 = ( + torch.randn((1, 2 * i, h), generator=generator, device=device) * scale + ).to(dtypes.bf16) + dense_w2 = ( + torch.randn((1, h, i), generator=generator, device=device) * scale + ).to(dtypes.bf16) + if profile.intermediate_pad: + dense_w1[:, profile.logical_inter : i] = 0 + dense_w1[:, i + profile.logical_inter :] = 0 + dense_w2[:, :, profile.logical_inter :] = 0 + quant_w1, quant_s1 = dynamic_mxfp8_quant(dense_w1, quant_dtype=dtypes.fp8) + quant_w2, quant_s2 = dynamic_mxfp8_quant(dense_w2, quant_dtype=dtypes.fp8) + w1[expert_id].copy_(quant_w1[0]) + w2[expert_id].copy_(quant_w2[0]) + s1_u8[expert_id].copy_(quant_s1[0]) + s2_u8[expert_id].copy_(quant_s2[0]) + + s1 = s1_u8.view(dtypes.fp8_e8m0) + s2 = s2_u8.view(dtypes.fp8_e8m0) + active = torch.tensor(routed_expert_ids, dtype=torch.long, device=device) + raw_w1 = w1.index_select(0, active).clone() + raw_w2 = w2.index_select(0, active).clone() + raw_s1 = s1.index_select(0, active).clone() + raw_s2 = s2.index_select(0, active).clone() return ( *_shuffle_w1(w1, s1, interleave), *_shuffle_w2(w2, s2, interleave), @@ -287,12 +365,15 @@ def _make_routed_weights( raw_w2, raw_s1, raw_s2, + routed_expert_ids, ) -def _build_weights(profile: _Profile, interleave: bool) -> _Weights: +def _build_weights( + profile: _Profile, interleave: bool, *, routed_mxfp8: bool = False +) -> _Weights: if get_gfx() != "gfx950": - pytest.skip("heterogeneous MXFP4/FP8 MoE requires gfx950") + pytest.skip("heterogeneous MXFP4/FP8 or MXFP8/FP8 MoE requires gfx950") if "shared_w1" not in inspect.signature(fused_moe).parameters: pytest.fail("AITER fused_moe is missing the heterogeneous shared arguments") @@ -306,7 +387,12 @@ def _build_weights(profile: _Profile, interleave: bool) -> _Weights: raw_routed_w2, raw_routed_s1, raw_routed_s2, - ) = _make_routed_weights(profile, device, interleave) + routed_expert_ids, + ) = ( + _make_mxfp8_routed_weights(profile, device, interleave) + if routed_mxfp8 + else _make_routed_weights(profile, device, interleave) + ) ( shared_w1, shared_w2, @@ -330,6 +416,7 @@ def _build_weights(profile: _Profile, interleave: bool) -> _Weights: raw_routed_w2, raw_routed_s1, raw_routed_s2, + routed_expert_ids, shared_w1, shared_w2, shared_s1, @@ -347,7 +434,16 @@ def _build_weights(profile: _Profile, interleave: bool) -> _Weights: @pytest.fixture(scope="module") def weights() -> _Weights: - return _build_weights(_profile(), interleave=True) + return _build_weights( + _profile(), + interleave=True, + routed_mxfp8=_use_hy4_profile(), + ) + + +@pytest.fixture(scope="module") +def generic_weights() -> _Weights: + return _build_weights(_Profile(256, 128, 128, 9, 2), interleave=True) @pytest.fixture(scope="module") @@ -356,7 +452,22 @@ def a4_weights() -> _Weights: return _build_weights(profile, interleave=False) -def _route_inputs(profile: _Profile, m: int, device: torch.device) -> tuple: +@pytest.fixture(scope="module") +def mxfp8_weights() -> _Weights: + return _build_weights(_mxfp8_profile(), interleave=True, routed_mxfp8=True) + + +@pytest.fixture(scope="module") +def dsv4_weights() -> _Weights: + return _build_weights(_Profile(7168, 384, 384, 385, 6), interleave=True) + + +def _route_inputs( + profile: _Profile, + m: int, + device: torch.device, + routed_expert_ids: tuple[int, ...] | None = None, +) -> tuple: hidden_generator = torch.Generator(device=device).manual_seed(41 + m) hidden = torch.randn( (m, profile.hidden), @@ -368,7 +479,13 @@ def _route_inputs(profile: _Profile, m: int, device: torch.device) -> tuple: slot = torch.arange(profile.routed_topk, device=device, dtype=dtypes.i32).unsqueeze( 0 ) - routed_ids = (row + slot) % profile.routed_topk + active_ids = torch.tensor( + routed_expert_ids or tuple(range(profile.routed_topk)), + device=device, + dtype=dtypes.i32, + ) + assert active_ids.numel() == profile.routed_topk + routed_ids = active_ids[(row + slot) % profile.routed_topk] raw_weights = torch.arange( 1, profile.routed_topk + 1, @@ -413,17 +530,25 @@ def _run_composed_oracle( profile: _Profile, weights: _Weights, m: int, + *, + clamp_shared: bool, + swiglu_limit: float = 10.0, ) -> tuple[torch.Tensor, torch.Tensor, dict]: hidden, routed_weight, routed_ids, all_weight, all_ids = _route_inputs( - profile, m, weights.routed_w1.device + profile, + m, + weights.routed_w1.device, + weights.routed_expert_ids, ) hetero_kwargs = _common_kwargs(profile, weights.routed_s1, weights.routed_s2) + hetero_kwargs["swiglu_limit"] = swiglu_limit hetero_kwargs.update( shared_w1=weights.shared_w1, shared_w2=weights.shared_w2, shared_w1_scale=weights.shared_s1, shared_w2_scale=weights.shared_s2, shared_expert_id=profile.shared_id, + clamp_shared=clamp_shared, ) actual = fused_moe( hidden, @@ -439,13 +564,15 @@ def _run_composed_oracle( routed_w2 = _mark_shuffled(weights.routed_w2[:routed_e]) routed_s1 = weights.routed_s1[: routed_e * 2 * profile.inter] routed_s2 = weights.routed_s2[: routed_e * profile.hidden] + routed_kwargs = _common_kwargs(profile, routed_s1, routed_s2) + routed_kwargs["swiglu_limit"] = swiglu_limit routed_out = fused_moe( hidden, routed_w1, routed_w2, routed_weight, routed_ids, - **_common_kwargs(profile, routed_s1, routed_s2), + **routed_kwargs, ) shared_ids = torch.zeros((m, 1), dtype=dtypes.i32, device=hidden.device) @@ -457,13 +584,17 @@ def _run_composed_oracle( 1, 1, ) + shared_kwargs = _common_kwargs(native_profile, weights.native_s1, weights.native_s2) + shared_kwargs["swiglu_limit"] = _shared_swiglu_limit( + swiglu_limit, clamp_shared=clamp_shared + ) shared_out = fused_moe( hidden, weights.native_w1, weights.native_w2, shared_weight, shared_ids, - **_common_kwargs(native_profile, weights.native_s1, weights.native_s2), + **shared_kwargs, ) return ( actual, @@ -477,18 +608,25 @@ def _run_composed_oracle( "hetero_kwargs": hetero_kwargs, "routed_out": routed_out, "shared_out": shared_out, + "clamp_shared": clamp_shared, + "swiglu_limit": swiglu_limit, }, ) -def _fp8_group_quant_dequant(x: torch.Tensor, group_size: int) -> torch.Tensor: +def _fp8_group_quant_dequant( + x: torch.Tensor, group_size: int, *, fused_intermediate: bool = False +) -> torch.Tensor: shape = x.shape blocks = x.float().view(-1, group_size) - scale = fp4_utils.f32_to_mx_e8m0_scale( - blocks.abs().amax(dim=1), dtype=MxDtypeInt.FP8_E4M3 + amax = blocks.abs().amax(dim=1) + scale = ( + fp4_utils.f32_to_fused_moe_mxfp8_scale(amax) + if fused_intermediate + else fp4_utils.f32_to_mx_e8m0_scale(amax, dtype=MxDtypeInt.FP8_E4M3) ) scale_f32 = fp4_utils.e8m0_to_f32(scale).view(-1, 1) - quant = (blocks / scale_f32).to(dtypes.fp8) + quant = (blocks / scale_f32).to(torch.float8_e4m3fn) return (quant.float() * scale_f32).view(shape) @@ -498,11 +636,13 @@ def _fp4_group_quant_dequant(x: torch.Tensor) -> torch.Tensor: return fp4_utils.mxfp4_to_f32(quant) * scale_f32 -def _activation_quant_dequant(x: torch.Tensor, quantization: str) -> torch.Tensor: +def _activation_quant_dequant( + x: torch.Tensor, quantization: str, *, fused_intermediate: bool = False +) -> torch.Tensor: if quantization == "fp4": return _fp4_group_quant_dequant(x) if quantization == "fp8": - return _fp8_group_quant_dequant(x, 32) + return _fp8_group_quant_dequant(x, 32, fused_intermediate=fused_intermediate) raise ValueError(f"Unsupported activation quantization: {quantization}") @@ -525,10 +665,16 @@ def _torch_routed_reference( weights: _Weights, profile: _Profile, activation_quantization: str = "fp8", + swiglu_limit: float = 10.0, ) -> torch.Tensor: - """Dequantized FP32 reference for the routed A8W4 experts.""" - w1 = _dequant_fp4_weight(weights.raw_routed_w1, weights.raw_routed_s1) - w2 = _dequant_fp4_weight(weights.raw_routed_w2, weights.raw_routed_s2) + """Dequantized FP32 reference for routed MXFP4 or MXFP8 experts.""" + dequant_weight = ( + _dequant_fp8_weight + if weights.raw_routed_w1.dtype == dtypes.fp8 + else _dequant_fp4_weight + ) + w1 = dequant_weight(weights.raw_routed_w1, weights.raw_routed_s1) + w2 = dequant_weight(weights.raw_routed_w2, weights.raw_routed_s2) x = _activation_quant_dequant(hidden.float(), activation_quantization) expanded_x = x[:, None, :].expand(-1, profile.routed_topk, -1) slot_out = torch.zeros( @@ -536,21 +682,23 @@ def _torch_routed_reference( dtype=dtypes.fp32, device=hidden.device, ) + activation_limit = runtime_swiglu_limit(swiglu_limit, "silu") - for expert_id in range(profile.routed_topk): + for weight_index, expert_id in enumerate(weights.routed_expert_ids): mask = routed_ids == expert_id if not mask.any(): continue - gate_up = F.linear(expanded_x[mask], w1[expert_id]) + gate_up = F.linear(expanded_x[mask], w1[weight_index]) gate, up = gate_up.chunk(2, dim=-1) - limit = _swiglu_limit(profile) - gate = gate.clamp(max=limit) - up = up.clamp(min=-limit, max=limit) + gate = gate.clamp(max=activation_limit) + up = up.clamp(min=-activation_limit, max=activation_limit) inter = F.silu(gate) * up - if activation_quantization == "fp8": - inter = inter.to(dtypes.bf16).float() - inter = _activation_quant_dequant(inter, activation_quantization) - slot_out[mask] = F.linear(inter, w2[expert_id]) + inter = _activation_quant_dequant( + inter, + activation_quantization, + fused_intermediate=activation_quantization == "fp8", + ) + slot_out[mask] = F.linear(inter, w2[weight_index]) return (slot_out * routed_weight[..., None]).sum(dim=1) @@ -572,14 +720,15 @@ def _torch_shared_reference( if group_size is not None: gate_up = gate_up.to(dtypes.bf16).float() gate, up = gate_up.chunk(2, dim=-1) - gate = gate.clamp(max=swiglu_limit) - up = up.clamp(min=-swiglu_limit, max=swiglu_limit) + activation_limit = runtime_swiglu_limit(swiglu_limit, "silu") + gate = gate.clamp(max=activation_limit) + up = up.clamp(min=-activation_limit, max=activation_limit) inter = F.silu(gate) * up if group_size == "fp4": inter = _fp4_group_quant_dequant(inter) elif group_size is not None: inter = inter.to(dtypes.bf16).float() - inter = _fp8_group_quant_dequant(inter, group_size) + inter = _fp8_group_quant_dequant(inter, group_size, fused_intermediate=True) return F.linear(inter, w2[0]) @@ -594,9 +743,15 @@ def _torch_heterogeneous_reference( context["routed_ids"], weights, profile, + swiglu_limit=context["swiglu_limit"], ) shared = _torch_shared_reference( - context["hidden"], weights, 32, _swiglu_limit(profile) + context["hidden"], + weights, + 32, + _shared_swiglu_limit( + context["swiglu_limit"], clamp_shared=context["clamp_shared"] + ), ) return routed + shared, routed, shared @@ -610,7 +765,12 @@ def test_heterogeneous_moe_matches_precision_oracles( profile = _profile() monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") - forced, unfused, context = _run_composed_oracle(profile, weights, m) + forced, unfused, context = _run_composed_oracle( + profile, + weights, + m, + clamp_shared=not _use_hy4_profile(), + ) stress_repeats = int(os.environ.get("AITER_HETERO_MOE_STRESS_REPEATS", "1")) assert stress_repeats > 0 for repeat in range(stress_repeats): @@ -727,10 +887,12 @@ def test_heterogeneous_moe_matches_precision_oracles( atomic_cosine <= 1e-3 ), f"default atomic FP32 cosine distance: {atomic_cosine:.3e}" - limit = _swiglu_limit(profile) - golden = _torch_shared_reference(context["hidden"], weights, None, limit) - per_32 = _torch_shared_reference(context["hidden"], weights, 32, limit) - per_128 = _torch_shared_reference(context["hidden"], weights, 128, limit) + shared_limit = _shared_swiglu_limit( + context["swiglu_limit"], clamp_shared=context["clamp_shared"] + ) + golden = _torch_shared_reference(context["hidden"], weights, None, shared_limit) + per_32 = _torch_shared_reference(context["hidden"], weights, 32, shared_limit) + per_128 = _torch_shared_reference(context["hidden"], weights, 128, shared_limit) error_32 = _rel_l2(per_32, golden) error_128 = _rel_l2(per_128, golden) assert error_32 <= error_128 + 5e-4, ( @@ -739,119 +901,503 @@ def test_heterogeneous_moe_matches_precision_oracles( ) -def test_no_shared_explicit_defaults_preserve_old_api_output( +def test_mxfp8_routed_fp8_shared_heterogeneous_path( monkeypatch: pytest.MonkeyPatch, - weights: _Weights, + mxfp8_weights: _Weights, ): - profile = _profile() + """Exercise HY4 MXFP8 routed and FP8 shared experts across the ID range.""" + profile = _mxfp8_profile() + assert mxfp8_weights.routed_expert_ids == (0, 1, 7, 8, 127, 128, 254, 255) monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") - hidden, routed_weight, routed_ids, _, _ = _route_inputs( - profile, 4, weights.routed_w1.device - ) - routed_e = profile.experts - 1 - routed_w1 = _mark_shuffled(weights.routed_w1[:routed_e]) - routed_w2 = _mark_shuffled(weights.routed_w2[:routed_e]) - routed_s1 = weights.routed_s1[: routed_e * 2 * profile.inter] - routed_s2 = weights.routed_s2[: routed_e * profile.hidden] - kwargs = _common_kwargs(profile, routed_s1, routed_s2) - omitted = fused_moe( - hidden, - routed_w1, - routed_w2, - routed_weight, - routed_ids, - **kwargs, + fused, composed, context = _run_composed_oracle( + profile, mxfp8_weights, 4, clamp_shared=False ) - explicit = fused_moe( - hidden, - routed_w1, - routed_w2, - routed_weight, - routed_ids, - shared_w1=None, - shared_w2=None, - shared_w1_scale=None, - shared_w2_scale=None, - shared_expert_id=-1, - **kwargs, + high_precision, _, _ = _torch_heterogeneous_reference( + profile, mxfp8_weights, context ) - assert torch.equal(omitted, explicit) + + assert torch.isfinite(fused).all() + fused_error = _rel_l2(fused, high_precision) + composed_error = _rel_l2(composed, high_precision) + cosine = _cosine_distance(fused, high_precision) + assert fused_error <= 4e-2, f"MXFP8 FHMoE FP32 error: {fused_error:.3e}" + # The fused and composed paths can select different tiles and reduction + # orders. Keep a comparative guard, but allow normal MXFP8 accumulation + # variance while the independent FP32 and cosine limits enforce accuracy. + assert fused_error <= composed_error + 2e-2, ( + "MXFP8 FHMoE error regressed against separately scheduled routed and " + f"shared kernels: {fused_error=:.3e}, {composed_error=:.3e}" + ) + assert cosine <= 1e-3, f"MXFP8 FHMoE cosine distance: {cosine:.3e}" + + # The final routed row is a dummy placeholder for the separately supplied + # shared expert. Corrupt it to prove shared tiles never read routed storage. + dummy_w1 = mxfp8_weights.routed_w1[-1].clone() + dummy_w2 = mxfp8_weights.routed_w2[-1].clone() + try: + mxfp8_weights.routed_w1[-1].fill_(1) + mxfp8_weights.routed_w2[-1].fill_(-1) + corrupted_dummy_output = fused_moe( + context["hidden"], + mxfp8_weights.routed_w1, + mxfp8_weights.routed_w2, + context["all_weight"], + context["all_ids"], + **context["hetero_kwargs"], + ) + finally: + mxfp8_weights.routed_w1[-1].copy_(dummy_w1) + mxfp8_weights.routed_w2[-1].copy_(dummy_w2) + assert torch.equal(fused, corrupted_dummy_output) -def test_heterogeneous_moe_uses_a_separate_custom_op_schema(): - from aiter.fhmoe import fhmoe_ - from aiter.ops.flydsl.fhmoe import ( - flydsl_fhmoe_stage1, - flydsl_fhmoe_stage2, +def test_mxfp8_hy4_shared_silu_is_unclamped( + monkeypatch: pytest.MonkeyPatch, + mxfp8_weights: _Weights, +): + profile = _mxfp8_profile() + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") + fused, composed, context = _run_composed_oracle( + profile, + mxfp8_weights, + 4, + clamp_shared=False, + swiglu_limit=0.25, ) - from aiter.ops.flydsl.kernels.fhmoe import ( - compile_mixed_fhmoe_gemm1, - compile_mixed_fhmoe_gemm2, + + half_shared_weight = context["all_weight"].clone() + half_shared_weight[:, profile.routed_topk] = 0.5 + half_shared = fused_moe( + context["hidden"], + mxfp8_weights.routed_w1, + mxfp8_weights.routed_w2, + half_shared_weight, + context["all_ids"], + **context["hetero_kwargs"], ) - from aiter.ops.flydsl.kernels.mixed_moe_gemm_2stage import ( - compile_mixed_moe_gemm1, - compile_mixed_moe_gemm2, + observed_shared = 2 * (fused.float() - half_shared.float()) + unclamped = _torch_shared_reference( + context["hidden"], mxfp8_weights, 32, float("inf") ) - from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2 + clamped = _torch_shared_reference(context["hidden"], mxfp8_weights, 32, 0.25) - assert callable(fhmoe_) - legacy_schema = torch.ops.aiter.fused_moe_.default._schema - fhmoe_schema = torch.ops.aiter.fhmoe_.default._schema - fhmoe_fields = { - "shared_w1", - "shared_w2", - "shared_w1_scale", - "shared_w2_scale", - "shared_expert_id", - } - legacy_schema_fields = {argument.name for argument in legacy_schema.arguments} - fhmoe_schema_fields = {argument.name for argument in fhmoe_schema.arguments} + assert _rel_l2(fused, composed) <= 4e-2 + assert _rel_l2(unclamped, clamped) >= 1e-2 + assert _rel_l2(observed_shared, unclamped) < _rel_l2(observed_shared, clamped) - assert fhmoe_fields.isdisjoint(legacy_schema_fields) - assert fhmoe_fields <= fhmoe_schema_fields - ordinary_apis = ( - flydsl_moe_stage1, - flydsl_moe_stage2, - compile_mixed_moe_gemm1, - compile_mixed_moe_gemm2, +def test_mxfp8_hy4_compatible_geometry_supports_explicit_shared_clamp( + monkeypatch: pytest.MonkeyPatch, + mxfp8_weights: _Weights, +): + profile = _mxfp8_profile() + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") + fused, _, context = _run_composed_oracle( + profile, + mxfp8_weights, + 4, + clamp_shared=True, + swiglu_limit=0.25, ) - assert all( - fhmoe_fields.isdisjoint(inspect.signature(api).parameters) - for api in ordinary_apis + expected, _, _ = _torch_heterogeneous_reference(profile, mxfp8_weights, context) + + assert torch.isfinite(fused).all() + assert _rel_l2(fused, expected) <= 4e-2 + + +def test_generic_shared_silu_is_clamped( + monkeypatch: pytest.MonkeyPatch, + generic_weights: _Weights, +): + profile = _Profile(256, 128, 128, 9, 2) + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") + fused, composed, context = _run_composed_oracle( + profile, + generic_weights, + 4, + clamp_shared=True, + swiglu_limit=0.25, ) - assert {"situ_beta", "situ_linear_beta", "k_batch_intra_block"} <= set( - inspect.signature(flydsl_moe_stage1).parameters + + half_shared_weight = context["all_weight"].clone() + half_shared_weight[:, profile.routed_topk] = 0.5 + half_shared = fused_moe( + context["hidden"], + generic_weights.routed_w1, + generic_weights.routed_w2, + half_shared_weight, + context["all_ids"], + **context["hetero_kwargs"], ) - assert {"shared_w1", "shared_w1_scale", "shared_expert_id"} <= set( - inspect.signature(flydsl_fhmoe_stage1).parameters + observed_shared = 2 * (fused.float() - half_shared.float()) + clamped = _torch_shared_reference(context["hidden"], generic_weights, 32, 0.25) + unclamped = _torch_shared_reference( + context["hidden"], generic_weights, 32, float("inf") ) - assert {"shared_w2", "shared_w2_scale", "shared_expert_id"} <= set( - inspect.signature(flydsl_fhmoe_stage2).parameters + + assert _rel_l2(fused, composed) <= 4e-2 + assert _rel_l2(clamped, unclamped) >= 1e-2 + assert _rel_l2(observed_shared, clamped) < _rel_l2(observed_shared, unclamped) + + +@pytest.mark.parametrize( + ("m", "stage1_fragment", "stage2_fragment"), + ( + (1, "_kw2", "_atomic"), + (2, "_kw2", "_atomic"), + (4, "_t32", "_atomic"), + (8, "_xcd4", "_atomic"), + (16, "_t32x128", "_atomic"), + (32, "_t32x64", "_t32x128x256_atomic"), + (64, "_t32x128", "_atomic"), + (128, "_w2", "_atomic"), + (256, "_t32x64", "_atomic"), + (512, "_t32x64", "_atomic"), + (1024, "_t64", "_t64x128x256_atomic"), + (2048, "_t128", "_atomic_sbm128"), + (4096, "_t64", "_reduce"), + (8192, "_t128", "_reduce"), + (16384, "_t128", "_reduce"), + (32768, "_t128", "_reduce"), + (38836, "_t128", "_reduce"), + ), +) +def test_mxfp8_production_decode_and_prefill_match_sampled_fp32( + monkeypatch: pytest.MonkeyPatch, + mxfp8_weights: _Weights, + m: int, + stage1_fragment: str, + stage2_fragment: str, +): + """Exercise every tuned HY4 bucket and its representative runtime shapes.""" + import importlib + + fused_moe_module = importlib.import_module("aiter.fused_moe") + from aiter.jit.core import AITER_CONFIGS + + profile = _mxfp8_profile() + config_path = Path(AITER_CONFIGS.AITER_CONFIG_FHMOE_FILE) + metadata = fused_moe_module.get_2stage_cfgs( + fused_moe_module.get_padded_M(m), + profile.hidden, + profile.inter, + profile.experts, + profile.routed_topk + 1, + torch.bfloat16, + dtypes.fp8, + dtypes.fp8, + aiter.QuantType.per_1x32, + True, + aiter.ActivationType.Silu, + False, + 0, + 0, + True, + GateMode.INTERLEAVE, + config_file=str(config_path), ) - assert all( - "shared_expert_id" in inspect.signature(api).parameters - for api in (compile_mixed_fhmoe_gemm1, compile_mixed_fhmoe_gemm2) + assert stage1_fragment in metadata.stage1.keywords["kernelName"] + assert stage2_fragment in metadata.stage2.keywords["kernelName"] + + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "0") + hidden, routed_weight, routed_ids, all_weight, all_ids = _route_inputs( + profile, + m, + mxfp8_weights.routed_w1.device, + mxfp8_weights.routed_expert_ids, ) - assert all( - "v2_output_layout" in inspect.signature(api).parameters - for api in (flydsl_fhmoe_stage1, compile_mixed_fhmoe_gemm1) + kwargs = _common_kwargs(profile, mxfp8_weights.routed_s1, mxfp8_weights.routed_s2) + kwargs.update( + shared_w1=mxfp8_weights.shared_w1, + shared_w2=mxfp8_weights.shared_w2, + shared_w1_scale=mxfp8_weights.shared_s1, + shared_w2_scale=mxfp8_weights.shared_s2, + shared_expert_id=profile.shared_id, + clamp_shared=False, ) - fhmoe_apis = ( - flydsl_fhmoe_stage1, - flydsl_fhmoe_stage2, - compile_mixed_fhmoe_gemm1, - compile_mixed_fhmoe_gemm2, + actual = fused_moe( + hidden, + mxfp8_weights.routed_w1, + mxfp8_weights.routed_w2, + all_weight, + all_ids, + **kwargs, ) - assert all("xcd_swizzle" in inspect.signature(api).parameters for api in fhmoe_apis) + if stage2_fragment == "_atomic": + out_buf = torch.full_like(actual, -7) + buffered = fused_moe( + hidden, + mxfp8_weights.routed_w1, + mxfp8_weights.routed_w2, + all_weight, + all_ids, + output=out_buf, + **kwargs, + ) + assert buffered is out_buf + assert torch.isfinite(buffered).all() + assert _rel_l2(buffered, actual) <= 8e-3 + sample_rows = torch.tensor( + sorted({0, m // 2, m - 1}), dtype=torch.long, device=hidden.device + ) + sample_context = { + "hidden": hidden.index_select(0, sample_rows), + "routed_weight": routed_weight.index_select(0, sample_rows), + "routed_ids": routed_ids.index_select(0, sample_rows), + "clamp_shared": False, + "swiglu_limit": _swiglu_limit(profile), + } + expected, _, _ = _torch_heterogeneous_reference( + profile, mxfp8_weights, sample_context + ) + actual_sample = actual.index_select(0, sample_rows) -def test_fhmoe_runtime_compile_bridge_forwards_xcd(monkeypatch: pytest.MonkeyPatch): - from aiter.ops.flydsl import fhmoe + assert torch.isfinite(actual_sample).all() + assert _rel_l2(actual_sample, expected) <= 4e-2 + assert _cosine_distance(actual_sample, expected) <= 1e-3 - tensor = torch.empty(0) + +def test_dsv4_production_reduce_matches_fp32( + monkeypatch: pytest.MonkeyPatch, + dsv4_weights: _Weights, +): + profile = _Profile(7168, 384, 384, 385, 6) + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "0") + fused, composed, context = _run_composed_oracle( + profile, dsv4_weights, 4, clamp_shared=True + ) + expected, _, _ = _torch_heterogeneous_reference(profile, dsv4_weights, context) + + assert torch.isfinite(fused).all() + assert _rel_l2(fused, expected) <= 5e-2 + assert _rel_l2(fused, expected) <= _rel_l2(composed, expected) + 1e-2 + + +def test_mxfp8_shared_k_offsets_cover_multiple_tiles( + monkeypatch: pytest.MonkeyPatch, + mxfp8_weights: _Weights, +): + """Exercise multiple K tiles in both stages with the production HY4 shape. + + M remains small intentionally: stage-1 H=6144 and stage-2 I=256, rather + than token count, are what force shared-weight loads beyond the first tile. + """ + profile = _mxfp8_profile() + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") + + fused, composed, _ = _run_composed_oracle( + profile, mxfp8_weights, 4, clamp_shared=False + ) + + assert torch.isfinite(fused).all() + assert _rel_l2(fused, composed) <= 4e-2 + + +def test_mxfp8_heterogeneous_path_requires_interleaved_gate( + mxfp8_weights: _Weights, +): + profile = _mxfp8_profile() + hidden, _, _, all_weight, all_ids = _route_inputs( + profile, + 4, + mxfp8_weights.routed_w1.device, + mxfp8_weights.routed_expert_ids, + ) + kwargs = _common_kwargs( + profile, + mxfp8_weights.routed_s1, + mxfp8_weights.routed_s2, + gate_mode=GateMode.SEPARATED, + ) + kwargs.update( + shared_w1=mxfp8_weights.shared_w1, + shared_w2=mxfp8_weights.shared_w2, + shared_w1_scale=mxfp8_weights.shared_s1, + shared_w2_scale=mxfp8_weights.shared_s2, + shared_expert_id=profile.shared_id, + clamp_shared=False, + ) + + with pytest.raises(ValueError, match="require interleaved gate/up layout"): + fused_moe( + hidden, + mxfp8_weights.routed_w1, + mxfp8_weights.routed_w2, + all_weight, + all_ids, + **kwargs, + ) + + +def test_mxfp8_heterogeneous_path_supports_non_hy4_contract(): + profile = _Profile(256, 128, 128, 9, 2) + weights = _build_weights(profile, interleave=True, routed_mxfp8=True) + fused, composed, _ = _run_composed_oracle(profile, weights, 4, clamp_shared=True) + + assert torch.isfinite(fused).all() + assert _rel_l2(fused, composed) <= 4e-2 + + +def test_zero_swiglu_limit_references_are_unclamped(generic_weights: _Weights): + profile = _Profile(256, 128, 128, 9, 2) + hidden, routed_weight, routed_ids, _, _ = _route_inputs( + profile, + 4, + generic_weights.routed_w1.device, + generic_weights.routed_expert_ids, + ) + routed_zero = _torch_routed_reference( + hidden, + routed_weight, + routed_ids, + generic_weights, + profile, + swiglu_limit=0, + ) + routed_unclamped = _torch_routed_reference( + hidden, + routed_weight, + routed_ids, + generic_weights, + profile, + swiglu_limit=float("inf"), + ) + shared_zero = _torch_shared_reference(hidden, generic_weights, 32, 0) + shared_unclamped = _torch_shared_reference( + hidden, generic_weights, 32, float("inf") + ) + + torch.testing.assert_close(routed_zero, routed_unclamped) + torch.testing.assert_close(shared_zero, shared_unclamped) + + +def test_no_shared_explicit_defaults_preserve_old_api_output( + monkeypatch: pytest.MonkeyPatch, + weights: _Weights, +): + profile = _profile() + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") + hidden, routed_weight, routed_ids, _, _ = _route_inputs( + profile, 4, weights.routed_w1.device + ) + routed_e = profile.experts - 1 + routed_w1 = _mark_shuffled(weights.routed_w1[:routed_e]) + routed_w2 = _mark_shuffled(weights.routed_w2[:routed_e]) + routed_s1 = weights.routed_s1[: routed_e * 2 * profile.inter] + routed_s2 = weights.routed_s2[: routed_e * profile.hidden] + kwargs = _common_kwargs(profile, routed_s1, routed_s2) + + omitted = fused_moe( + hidden, + routed_w1, + routed_w2, + routed_weight, + routed_ids, + **kwargs, + ) + explicit = fused_moe( + hidden, + routed_w1, + routed_w2, + routed_weight, + routed_ids, + shared_w1=None, + shared_w2=None, + shared_w1_scale=None, + shared_w2_scale=None, + shared_expert_id=-1, + **kwargs, + ) + assert torch.equal(omitted, explicit) + + +def test_heterogeneous_moe_uses_a_separate_custom_op_schema(): + from aiter.fhmoe import fhmoe_ + from aiter.ops.flydsl.fhmoe import ( + flydsl_fhmoe_stage1, + flydsl_fhmoe_stage2, + ) + from aiter.ops.flydsl.kernels.fhmoe import ( + compile_mixed_fhmoe_gemm1, + compile_mixed_fhmoe_gemm2, + ) + from aiter.ops.flydsl.kernels.mixed_moe_gemm_2stage import ( + compile_mixed_moe_gemm1, + compile_mixed_moe_gemm2, + ) + from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2 + + assert callable(fhmoe_) + legacy_schema = torch.ops.aiter.fused_moe_.default._schema + fhmoe_schema = torch.ops.aiter.fhmoe_.default._schema + fhmoe_fields = { + "shared_w1", + "shared_w2", + "shared_w1_scale", + "shared_w2_scale", + "shared_expert_id", + "clamp_shared", + } + legacy_schema_fields = {argument.name for argument in legacy_schema.arguments} + fhmoe_schema_fields = {argument.name for argument in fhmoe_schema.arguments} + + assert fhmoe_fields.isdisjoint(legacy_schema_fields) + assert fhmoe_fields <= fhmoe_schema_fields + + ordinary_apis = ( + flydsl_moe_stage1, + flydsl_moe_stage2, + compile_mixed_moe_gemm1, + compile_mixed_moe_gemm2, + ) + assert all( + fhmoe_fields.isdisjoint(inspect.signature(api).parameters) + for api in ordinary_apis + ) + assert {"situ_beta", "situ_linear_beta", "k_batch_intra_block"} <= set( + inspect.signature(flydsl_moe_stage1).parameters + ) + assert {"shared_w1", "shared_w1_scale", "shared_expert_id"} <= set( + inspect.signature(flydsl_fhmoe_stage1).parameters + ) + assert {"shared_w2", "shared_w2_scale", "shared_expert_id"} <= set( + inspect.signature(flydsl_fhmoe_stage2).parameters + ) + assert all( + "shared_expert_id" in inspect.signature(api).parameters + for api in (compile_mixed_fhmoe_gemm1, compile_mixed_fhmoe_gemm2) + ) + assert all( + "v2_output_layout" in inspect.signature(api).parameters + for api in (flydsl_fhmoe_stage1, compile_mixed_fhmoe_gemm1) + ) + assert all( + "clamp_shared" in inspect.signature(api).parameters + for api in (flydsl_fhmoe_stage1, compile_mixed_fhmoe_gemm1) + ) + fhmoe_apis = ( + flydsl_fhmoe_stage1, + flydsl_fhmoe_stage2, + compile_mixed_fhmoe_gemm1, + compile_mixed_fhmoe_gemm2, + ) + assert all("xcd_swizzle" in inspect.signature(api).parameters for api in fhmoe_apis) + + +def test_fhmoe_runtime_compile_bridge_forwards_xcd(monkeypatch: pytest.MonkeyPatch): + from aiter.ops.flydsl import fhmoe + + tensor = torch.empty(0) compile_calls = [] def invoke_compiler(**kwargs): @@ -883,6 +1429,7 @@ def compile_stage2(**kwargs): shared_w1=tensor, shared_w1_scale=tensor, shared_expert_id=8, + clamp_shared=False, xcd_swizzle=4, v2_output_layout=True, ) @@ -907,6 +1454,7 @@ def compile_stage2(**kwargs): 1, { "shared_expert_id": 8, + "clamp_shared": False, "xcd_swizzle": 4, "v2_output_layout": True, }, @@ -915,18 +1463,443 @@ def compile_stage2(**kwargs): ] -def _mock_dsv4_i384_fhmoe_metadata(monkeypatch: pytest.MonkeyPatch): +def test_fhmoe_rejects_unclamped_splitk_with_finite_limit(): + from aiter.ops.flydsl.fhmoe import flydsl_fhmoe_stage1 + + tensor = torch.empty(0) + with pytest.raises( + NotImplementedError, + match="cannot clamp routed SiLU while leaving the shared expert unclamped", + ): + flydsl_fhmoe_stage1( + tensor, + tensor, + tensor, + tensor, + tensor, + out_dtype="fp8", + act="silu", + k_batch=2, + gate_mode="interleave", + swiglu_limit=10.0, + shared_w1=tensor, + shared_w1_scale=tensor, + shared_expert_id=8, + clamp_shared=False, + ) + + +@pytest.mark.parametrize("act", ("swiglu", "situv2")) +def test_fhmoe_rejects_unclamped_shared_expert_for_non_silu(act: str): + from aiter.ops.flydsl.fhmoe import ( + compile_flydsl_fhmoe_stage1, + flydsl_fhmoe_stage1, + ) + + error = "clamp_shared=False is supported only with act='silu'" + with pytest.raises(ValueError, match=error): + compile_flydsl_fhmoe_stage1( + model_dim=256, + inter_dim=128, + experts=9, + topk=3, + tile_m=32, + tile_n=128, + tile_k=256, + doweight_stage1=False, + a_dtype="fp8", + b_dtype="fp4", + out_dtype="fp8", + act=act, + shared_expert_id=8, + clamp_shared=False, + ) + + tensor = torch.empty(0) + with pytest.raises(ValueError, match=error): + flydsl_fhmoe_stage1( + tensor, + tensor, + tensor, + tensor, + tensor, + act=act, + shared_w1=tensor, + shared_w1_scale=tensor, + shared_expert_id=8, + clamp_shared=False, + ) + + +def test_low_level_fhmoe_rejects_separated_mxfp8(): + from aiter.ops.flydsl.fhmoe import ( + compile_flydsl_fhmoe_stage1, + flydsl_fhmoe_stage1, + ) + + with pytest.raises(ValueError, match="interleaved gate/up"): + compile_flydsl_fhmoe_stage1( + model_dim=256, + inter_dim=128, + experts=9, + topk=3, + tile_m=32, + tile_n=128, + tile_k=256, + doweight_stage1=False, + a_dtype="fp8", + b_dtype="fp8", + out_dtype="fp8", + gate_mode="separated", + shared_expert_id=8, + ) + + tensor = torch.empty(0) + with pytest.raises(ValueError, match="interleaved gate/up"): + flydsl_fhmoe_stage1( + tensor, + tensor, + tensor, + tensor, + tensor, + a_dtype="fp8", + b_dtype="fp8", + out_dtype="fp8", + gate_mode="separated", + shared_w1=tensor, + shared_w1_scale=tensor, + shared_expert_id=8, + ) + + +@pytest.mark.parametrize( + ("argument", "value"), + ( + ("stage2_scatter", object()), + ("quant_type_a", aiter.QuantType.per_1x32), + ("quant_dtype_a", dtypes.fp8), + ("quant_dtype_a2", dtypes.fp8), + ), +) +def test_fhmoe_rejects_dropped_public_controls(argument: str, value): + tensor = torch.empty(0) + kwargs = { + "shared_w1": tensor, + "shared_w2": tensor, + "shared_w1_scale": tensor, + "shared_w2_scale": tensor, + "shared_expert_id": 0, + argument: value, + } + with pytest.raises(NotImplementedError, match=argument): + fused_moe(tensor, tensor, tensor, tensor, tensor, **kwargs) + + +@pytest.mark.parametrize("argument", ("a1_scale", "a2_scale")) +def test_mxfp8_fhmoe_rejects_external_activation_scales(argument: str): + weight = torch.empty(0, dtype=dtypes.fp8) + tensor = torch.empty(0) + kwargs = { + "shared_w1": weight, + "shared_w2": weight, + "shared_w1_scale": tensor, + "shared_w2_scale": tensor, + "shared_expert_id": 0, + argument: torch.ones(1), + } + with pytest.raises(NotImplementedError, match=argument): + fused_moe(tensor, weight, weight, tensor, tensor, **kwargs) + + +def test_mxfp4_fhmoe_preserves_activation_scales_and_unmarked_weights( + monkeypatch: pytest.MonkeyPatch, +): + import aiter.fhmoe as fhmoe_module + + captured = {} + + def fake_fhmoe(**kwargs): + captured.update(kwargs) + return torch.empty(0) + + monkeypatch.setattr(fhmoe_module, "_fhmoe", fake_fhmoe) + tensor = torch.empty(0) + routed_weight = torch.empty(0, dtype=dtypes.fp4x2) + shared_weight = torch.empty(0, dtype=dtypes.fp8) + a1_scale = torch.ones(1) + a2_scale = torch.ones(1) + + fused_moe( + tensor, + routed_weight, + routed_weight, + tensor, + tensor, + shared_w1=shared_weight, + shared_w2=shared_weight, + shared_w1_scale=tensor, + shared_w2_scale=tensor, + shared_expert_id=0, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + + assert captured["a1_scale"] is a1_scale + assert captured["a2_scale"] is a2_scale + + +@pytest.mark.parametrize( + ("gate_mode", "quant_dtype_a"), + ( + (GateMode.SEPARATED.value, dtypes.fp4x2), + (GateMode.INTERLEAVE.value, dtypes.fp8), + ), +) +def test_fhmoe_accepts_matching_public_activation_dtype( + monkeypatch: pytest.MonkeyPatch, + gate_mode: str, + quant_dtype_a: torch.dtype, +): + import aiter.fhmoe as fhmoe_module + + captured = {} + + def fake_fhmoe(**kwargs): + captured.update(kwargs) + return torch.empty(0) + + monkeypatch.setattr(fhmoe_module, "_fhmoe", fake_fhmoe) + tensor = torch.empty(0) + tensor.is_shuffled = True + + fused_moe( + tensor, + tensor, + tensor, + tensor, + tensor, + gate_mode=gate_mode, + shared_w1=tensor, + shared_w2=tensor, + shared_w1_scale=tensor, + shared_w2_scale=tensor, + shared_expert_id=0, + quant_dtype_a=quant_dtype_a, + ) + + assert captured["gate_mode"] == gate_mode + assert captured["clamp_shared"] is True + + +def test_common_fhmoe_has_no_model_identity_predicates(): import importlib + contract_module = importlib.import_module("aiter.fhmoe_contract") + assert not hasattr(contract_module, "_is_hy4_mxfp8_fhmoe_contract") + + +def test_clamp_shared_preserves_existing_positional_argument_order(): + import aiter.fhmoe as fhmoe_module + + parameters = list(inspect.signature(fused_moe).parameters) + assert parameters[-4:] == [ + "quant_type_a", + "quant_dtype_a", + "quant_dtype_a2", + "clamp_shared", + ] + assert list(inspect.signature(fhmoe_module.fhmoe_fake).parameters)[-2:] == [ + "output", + "clamp_shared", + ] + + +def test_unclamped_request_requires_shared_expert_dispatch(): + with pytest.raises(ValueError, match="requires FHMoE"): + fused_moe(None, None, None, None, None, clamp_shared=False) + + +def test_flydsl_stage2_classification_covers_plain_and_fhmoe_wrappers( + monkeypatch: pytest.MonkeyPatch, +): + import functools + import importlib + + fhmoe_module = importlib.import_module("aiter.fhmoe") fused_moe_module = importlib.import_module("aiter.fused_moe") + + assert fused_moe_module._is_flydsl_stage2(fused_moe_module._flydsl_stage2_wrapper) + assert fused_moe_module._is_flydsl_stage2(fhmoe_module._flydsl_fhmoe_stage2_wrapper) + assert not fused_moe_module._is_flydsl_stage2(lambda: None) + + fp8_stage2 = functools.partial( + fused_moe_module._flydsl_stage2_wrapper, + kernelName="flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic", + ) + fhmoe_fp8_stage2 = functools.partial( + fhmoe_module._flydsl_fhmoe_stage2_wrapper, + kernelName="flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic", + ) + a16w4_stage2 = functools.partial( + fused_moe_module._flydsl_stage2_wrapper, + kernelName="flydsl_moe2_abf16_wfp4_bf16_t32x128x128_atomic", + ) + a16wi4_stage2 = functools.partial( + fused_moe_module._flydsl_stage2_wrapper, + kernelName="flydsl_moe2_abf16_wint4_bf16_t16x128x128_reduce", + ) + + assert fused_moe_module._stage2_honors_runtime_reduce_override(fp8_stage2) + assert fused_moe_module._stage2_honors_runtime_reduce_override(fhmoe_fp8_stage2) + assert not fused_moe_module._stage2_honors_runtime_reduce_override(a16w4_stage2) + assert not fused_moe_module._stage2_honors_runtime_reduce_override(a16wi4_stage2) + unknown_stage2 = functools.partial( + fused_moe_module._flydsl_stage2_wrapper, + kernelName="flydsl_moe2_unknown", + ) + with pytest.raises(ValueError, match="require a recognized FlyDSL stage2"): + fused_moe_module._stage2_honors_runtime_reduce_override(unknown_stage2) + + monkeypatch.delenv("AITER_FLYDSL_FORCE_REDUCE", raising=False) + monkeypatch.delenv("AITER_FLYDSL_STAGE2_FP8", raising=False) + assert not fused_moe_module._stage2_runtime_uses_route_reduce(fp8_stage2, 1, 6144) + monkeypatch.setenv("AITER_FLYDSL_STAGE2_FP8", "1") + assert fused_moe_module._stage2_runtime_uses_route_reduce(fp8_stage2, 1, 6144) + bf16_route_bytes = 50000 * 9 * 6144 * 2 + fp8_route_bytes = fused_moe_module._flydsl_route_output_nbytes( + fp8_stage2, + 50000, + 9, + 6144, + 2, + ) + assert bf16_route_bytes > (1 << 32) - 1 + assert fp8_route_bytes == 50000 * 9 * (6144 + 6144 // 8) + assert fp8_route_bytes < (1 << 32) - 1 + fused_moe_module._validate_flydsl_route_output_nbytes(fp8_route_bytes) + fp8_bytes_per_token = 9 * (6144 + 6144 // 8) + fp8_max_tokens = ((1 << 32) - 1) // fp8_bytes_per_token + fused_moe_module._validate_flydsl_route_output_nbytes( + fp8_max_tokens * fp8_bytes_per_token + ) + with pytest.raises(NotImplementedError, match="32-bit buffer-addressing"): + fused_moe_module._validate_flydsl_route_output_nbytes( + (fp8_max_tokens + 1) * fp8_bytes_per_token + ) + assert not fused_moe_module._stage2_runtime_uses_route_reduce(a16w4_stage2, 1, 6144) + monkeypatch.delenv("AITER_FLYDSL_STAGE2_FP8") + monkeypatch.setenv("AITER_FLYDSL_FORCE_REDUCE", "1") + assert fused_moe_module._stage2_runtime_uses_route_reduce(fp8_stage2, 1, 6144) + monkeypatch.delenv("AITER_FLYDSL_FORCE_REDUCE") + auto_reduce_m = 0xFFFFFFFF // (6144 * 2) + 1 + assert fused_moe_module._stage2_runtime_uses_route_reduce( + fp8_stage2, auto_reduce_m, 6144 + ) + + +def test_fhmoe_sort_block_override_must_match_metadata( + monkeypatch: pytest.MonkeyPatch, + mxfp8_weights: _Weights, +): + profile = _mxfp8_profile() + monkeypatch.setenv("AITER_BF16_FP8_MOE_BOUND", "0") + hidden, _, _, all_weight, all_ids = _route_inputs( + profile, 1, mxfp8_weights.routed_w1.device, mxfp8_weights.routed_expert_ids + ) + kwargs = _common_kwargs(profile, mxfp8_weights.routed_s1, mxfp8_weights.routed_s2) + kwargs.update( + shared_w1=mxfp8_weights.shared_w1, + shared_w2=mxfp8_weights.shared_w2, + shared_w1_scale=mxfp8_weights.shared_s1, + shared_w2_scale=mxfp8_weights.shared_s2, + shared_expert_id=profile.shared_id, + clamp_shared=False, + ) + + matched = fused_moe( + hidden, + mxfp8_weights.routed_w1, + mxfp8_weights.routed_w2, + all_weight, + all_ids, + block_size_M=32, + **kwargs, + ) + assert torch.isfinite(matched).all() + with pytest.raises(ValueError, match="does not match"): + fused_moe( + hidden, + mxfp8_weights.routed_w1, + mxfp8_weights.routed_w2, + all_weight, + all_ids, + block_size_M=64, + **kwargs, + ) + + +def test_fhmoe_rejects_unmarked_weight_layout(): + tensor = torch.empty(0) + weight = torch.empty(0, dtype=dtypes.fp8) + with pytest.raises(ValueError, match="is_shuffled=True"): + fused_moe( + tensor, + weight, + weight, + tensor, + tensor, + shared_w1=weight, + shared_w2=weight, + shared_w1_scale=tensor, + shared_w2_scale=tensor, + shared_expert_id=0, + ) + + +def _mock_gfx950_fhmoe_metadata(monkeypatch: pytest.MonkeyPatch): + import importlib + + fhmoe_module = importlib.import_module("aiter.fhmoe") + fused_moe_module = importlib.import_module("aiter.fused_moe") + monkeypatch.setattr(fhmoe_module, "get_gfx_runtime", lambda: "gfx950") monkeypatch.setattr(fused_moe_module, "get_cu_num", lambda: 256) monkeypatch.setattr(fused_moe_module, "get_gfx_runtime", lambda: "gfx950") - monkeypatch.delenv("AITER_BYPASS_TUNE_CONFIG", raising=False) + monkeypatch.setattr(dtypes, "fp8", torch.float8_e4m3fn) fused_moe_module.get_2stage_cfgs.cache_clear() fused_moe_module.cfg_2stages_by_file.clear() return fused_moe_module +def _mock_dsv4_i384_fhmoe_metadata(monkeypatch: pytest.MonkeyPatch): + fused_moe_module = _mock_gfx950_fhmoe_metadata(monkeypatch) + monkeypatch.delenv("AITER_BYPASS_TUNE_CONFIG", raising=False) + return fused_moe_module + + +def _supports_hy4_config(fhmoe_module, max_tokens: int, **overrides) -> bool: + request = { + "model_dim": 6144, + "inter_dim": 256, + "experts": 257, + "topk": 9, + "dtype": torch.bfloat16, + "hidden_pad": 0, + "intermediate_pad": 0, + "gate_mode": GateMode.INTERLEAVE, + "doweight_stage1": False, + } + request.update(overrides) + return fhmoe_module.supports_mxfp8_fhmoe_config(max_tokens, **request) + + +def test_gfx950_metadata_mock_normalizes_fp8_alias( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(dtypes, "fp8", torch.float8_e4m3fnuz) + _mock_gfx950_fhmoe_metadata(monkeypatch) + assert dtypes.fp8 == torch.float8_e4m3fn + + @pytest.mark.parametrize( ("overrides", "expected"), ( @@ -942,6 +1915,7 @@ def _mock_dsv4_i384_fhmoe_metadata(monkeypatch: pytest.MonkeyPatch): ({"intermediate_pad": 128}, False), ({"gate_mode": GateMode.SEPARATED}, False), ({"doweight_stage1": True}, False), + ({"routed_mxfp4": False}, False), ), ) def test_dsv4_i384_fhmoe_config_scope( @@ -960,6 +1934,7 @@ def test_dsv4_i384_fhmoe_config_scope( "intermediate_pad": 0, "gate_mode": GateMode.INTERLEAVE, "doweight_stage1": False, + "routed_mxfp4": True, } args.update(overrides) @@ -1001,7 +1976,8 @@ def test_dsv4_i384_fhmoe_capability_follows_csv( assert reader.fieldnames is not None fieldnames = reader.fieldnames - row_4096 = dict(rows[-1]) + dsv4_rows = [row for row in rows if row["inter_dim"] == "384"] + row_4096 = dict(dsv4_rows[-1]) row_4096["token"] = "4096" complete_path = tmp_path / "complete.csv" with complete_path.open("w", newline="") as config: @@ -1028,7 +2004,7 @@ def test_dsv4_i384_fhmoe_capability_follows_csv( with duplicate_path.open("w", newline="") as config: writer = csv.DictWriter(config, fieldnames=fieldnames) writer.writeheader() - writer.writerows([*rows, rows[-1]]) + writer.writerows([*rows, dsv4_rows[-1]]) monkeypatch.setattr( fhmoe, "_dsv4_i384_fhmoe_config_file", lambda: str(duplicate_path) ) @@ -1046,6 +2022,75 @@ def test_dsv4_i384_fhmoe_capability_follows_csv( ) assert not fhmoe.supports_dsv4_i384_fhmoe(2048) + mismatched_path = tmp_path / "mismatched_block.csv" + mismatched_rows = [dict(row) for row in rows] + mismatched_row = next( + row + for row in mismatched_rows + if row["inter_dim"] == "384" and row["token"] == "1" + ) + mismatched_row["block_m"] = "64" if mismatched_row["block_m"] != "64" else "32" + with mismatched_path.open("w", newline="") as config: + writer = csv.DictWriter(config, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(mismatched_rows) + monkeypatch.setattr( + fhmoe, + "_dsv4_i384_fhmoe_config_file", + lambda: str(mismatched_path), + ) + assert not fhmoe.supports_dsv4_i384_fhmoe(1) + + +@pytest.mark.parametrize( + "incompatible_kind", + ("wfp8", "f16", "invalid_sbm", "invalid_kwave"), +) +def test_dsv4_i384_fhmoe_capability_rejects_incompatible_kernels( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + incompatible_kind: str, +): + import csv + import importlib + + fhmoe = importlib.import_module("aiter.fhmoe") + _mock_dsv4_i384_fhmoe_metadata(monkeypatch) + source = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" + with source.open(newline="") as config: + reader = csv.DictReader(config) + rows = list(reader) + assert reader.fieldnames is not None + fieldnames = reader.fieldnames + + dsv4 = next( + row for row in rows if row["inter_dim"] == "384" and row["token"] == "1" + ) + hy4 = next(row for row in rows if row["inter_dim"] == "256" and row["token"] == "1") + if incompatible_kind == "wfp8": + dsv4["kernelName1"] = hy4["kernelName1"] + dsv4["kernelName2"] = hy4["kernelName2"] + elif incompatible_kind == "f16": + dsv4["kernelName2"] = dsv4["kernelName2"].replace("_bf16_", "_f16_") + elif incompatible_kind == "invalid_sbm": + dsv4["kernelName2"] = "flydsl_moe2_afp8_wfp4_bf16_t64x128x128_atomic_sbm32" + else: + dsv4["kernelName1"] = dsv4["kernelName1"].replace("_kw4_fp8", "_kw3_fp8") + + incompatible = tmp_path / f"{incompatible_kind}.csv" + with incompatible.open("w", newline="") as config: + writer = csv.DictWriter(config, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + monkeypatch.setattr( + fhmoe, + "_dsv4_i384_fhmoe_config_file", + lambda: str(incompatible), + ) + fhmoe._supports_dsv4_i384_fhmoe_config.cache_clear() + + assert not fhmoe.supports_dsv4_i384_fhmoe(1) + def test_dsv4_i384_fhmoe_capability_fails_closed( monkeypatch: pytest.MonkeyPatch, tmp_path: Path @@ -1083,11 +2128,26 @@ def test_dsv4_i384_fhmoe_capability_rejects_config_bypass( "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw4_fp8", "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic", ), + ( + 4, + "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w4_gui_kw2", + "flydsl_moe2_afp8_wfp4_bf16_t32x128x128_reduce_persist", + ), + ( + 8, + "flydsl_moe1_afp8_wfp4_bf16_t32x64x256_w3_gui", + "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_reduce_bnt2_persist", + ), ( 16, "flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w2_gui_fp8", "flydsl_moe2_afp8_wfp4_bf16_t32x128x128_atomic_bnt2", ), + ( + 512, + "flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8", + "flydsl_moe2_afp8_wfp4_bf16_t32x256x128_atomic_bnt2_persist", + ), ( 1536, "flydsl_moe1_afp8_wfp4_bf16_t64x128x256_w3_bnt0_gui", @@ -1100,44 +2160,341 @@ def test_dsv4_i384_fhmoe_capability_rejects_config_bypass( ), ), ) -def test_dsv4_i384_fhmoe_uses_dedicated_config( - monkeypatch: pytest.MonkeyPatch, - num_tokens: int, - expected_stage1: str, - expected_stage2: str, +def test_dsv4_i384_fhmoe_uses_dedicated_config( + monkeypatch: pytest.MonkeyPatch, + num_tokens: int, + expected_stage1: str, + expected_stage2: str, +): + fused_moe_module = _mock_gfx950_fhmoe_metadata(monkeypatch) + + config_path = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" + + metadata = fused_moe_module.get_2stage_cfgs( + fused_moe_module.get_padded_M(num_tokens), + 7168, + 384, + 385, + 7, + torch.bfloat16, + dtypes.fp8, + dtypes.fp4x2, + aiter.QuantType.per_1x32, + True, + aiter.ActivationType.Silu, + False, + 0, + 0, + True, + GateMode.INTERLEAVE, + config_file=str(config_path), + ) + + assert metadata.stage1.keywords["kernelName"] == expected_stage1 + assert metadata.stage2.keywords["kernelName"] == expected_stage2 + + +@pytest.mark.parametrize( + ("num_tokens", "expected_block_m"), + ( + (1, 32), + (512, 32), + (4096, 64), + (16384, 128), + (32768, 128), + (38836, 128), + ), +) +def test_hy4_mxfp8_fhmoe_uses_dedicated_config( + monkeypatch: pytest.MonkeyPatch, + num_tokens: int, + expected_block_m: int, +): + fused_moe_module = _mock_gfx950_fhmoe_metadata(monkeypatch) + config_path = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" + + metadata = fused_moe_module.get_2stage_cfgs( + fused_moe_module.get_padded_M(num_tokens), + 6144, + 256, + 257, + 9, + torch.bfloat16, + dtypes.fp8, + dtypes.fp8, + aiter.QuantType.per_1x32, + True, + aiter.ActivationType.Silu, + False, + 0, + 0, + True, + GateMode.INTERLEAVE, + config_file=str(config_path), + ) + + assert metadata.stage1.keywords["kernelName"].startswith("flydsl_moe1_afp8_wfp8_") + assert metadata.stage2.keywords["kernelName"].startswith("flydsl_moe2_afp8_wfp8_") + assert "_layout_" not in metadata.stage2.keywords["kernelName"] + assert metadata.fuse_quant == "fp8" + assert metadata.block_m == expected_block_m + + +def test_hy4_mxfp8_fhmoe_capability(monkeypatch: pytest.MonkeyPatch): + import importlib + + fhmoe = importlib.import_module("aiter.fhmoe") + _mock_gfx950_fhmoe_metadata(monkeypatch) + config_path = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" + monkeypatch.setattr(fhmoe, "_fhmoe_config_file", lambda: str(config_path)) + fhmoe._supports_mxfp8_fhmoe_config.cache_clear() + + assert fhmoe._config_has_mxfp8_fhmoe_geometry( + str(config_path), + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + ) + assert fhmoe._config_has_mxfp8_fhmoe_contract( + str(config_path), + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + dtype=torch.bfloat16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE, + doweight_stage1=False, + ) + assert not fhmoe._config_has_mxfp8_fhmoe_contract( + str(config_path), + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + dtype=torch.float16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE, + doweight_stage1=False, + ) + assert fhmoe._resolve_mxfp8_fhmoe_config_file( + 32, + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + dtype=torch.bfloat16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE, + doweight_stage1=False, + ) == str(config_path) + with pytest.raises(ValueError, match="partially matches"): + fhmoe._resolve_mxfp8_fhmoe_config_file( + 32, + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + dtype=torch.float16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE, + doweight_stage1=False, + ) + assert ( + fhmoe._resolve_mxfp8_fhmoe_config_file( + 32, + model_dim=256, + inter_dim=128, + experts=9, + topk=3, + dtype=torch.bfloat16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE, + doweight_stage1=False, + ) + is None + ) + assert _supports_hy4_config(fhmoe, 38836) + assert not _supports_hy4_config(fhmoe, 38837) + with pytest.raises(NotImplementedError, match="does not cover"): + fhmoe._resolve_mxfp8_fhmoe_config_file( + 38837, + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + dtype=torch.bfloat16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE, + doweight_stage1=False, + ) + assert not _supports_hy4_config(fhmoe, 50000) + monkeypatch.setenv("AITER_FLYDSL_STAGE2_FP8", "1") + assert _supports_hy4_config(fhmoe, 50000) + assert fhmoe._resolve_mxfp8_fhmoe_config_file( + 50000, + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + dtype=torch.bfloat16, + hidden_pad=0, + intermediate_pad=0, + gate_mode=GateMode.INTERLEAVE, + doweight_stage1=False, + ) == str(config_path) + monkeypatch.delenv("AITER_FLYDSL_STAGE2_FP8") + assert not _supports_hy4_config(fhmoe, 1 << 24) + assert not _supports_hy4_config(fhmoe, 0) + + +def test_mxfp8_fhmoe_capability_requires_live_gfx950( + monkeypatch: pytest.MonkeyPatch, +): + import importlib + + fhmoe = importlib.import_module("aiter.fhmoe") + monkeypatch.setattr(fhmoe, "get_gfx_runtime", lambda: "gfx942") + monkeypatch.setattr( + fhmoe, + "_fhmoe_config_file", + lambda: pytest.fail("off-arch capability must not probe the CSV"), + ) + assert not _supports_hy4_config(fhmoe, 32) + + +def test_reduce_output_limit_fits_uint32_byte_extent(): + from aiter.fhmoe_contract import ( + fhmoe_reduce_output_nbytes, + max_fhmoe_reduce_tokens, + ) + from aiter.fused_moe import _validate_flydsl_route_output_nbytes + + max_tokens = max_fhmoe_reduce_tokens(9, 6144, 2) + bytes_per_token = 9 * 6144 * 2 + assert max_tokens == 38836 + assert max_tokens * bytes_per_token <= (1 << 32) - 1 + assert (max_tokens + 1) * bytes_per_token > (1 << 32) - 1 + _validate_flydsl_route_output_nbytes( + fhmoe_reduce_output_nbytes(max_tokens, 9, 6144, 2) + ) + with pytest.raises(NotImplementedError, match="32-bit buffer-addressing"): + _validate_flydsl_route_output_nbytes( + fhmoe_reduce_output_nbytes(max_tokens + 1, 9, 6144, 2) + ) + + +def test_global_moe_large_m_uses_32768_floor_tier_until_131072(): + from aiter.fused_moe import get_padded_M + + assert get_padded_M(32768) == 32768 + assert get_padded_M(32769) == 32768 + assert get_padded_M(65536) == 32768 + assert get_padded_M(131071) == 32768 + assert get_padded_M(131072) == 131072 + assert get_padded_M(131073) == 131072 + + +def test_hy4_mxfp8_fhmoe_capability_fails_on_missing_bucket( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + import csv + import importlib + + fhmoe = importlib.import_module("aiter.fhmoe") + _mock_gfx950_fhmoe_metadata(monkeypatch) + source = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" + with source.open(newline="") as stream: + reader = csv.DictReader(stream) + assert reader.fieldnames is not None + rows = list(reader) + fields = reader.fieldnames + missing = tmp_path / "missing_hy4_m4096.csv" + with missing.open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + writer.writerows( + row + for row in rows + if not (row["model_dim"] == "6144" and row["token"] == "4096") + ) + monkeypatch.setattr(fhmoe, "_fhmoe_config_file", lambda: str(missing)) + fhmoe._supports_mxfp8_fhmoe_config.cache_clear() + + assert not _supports_hy4_config(fhmoe, 4096) + + +@pytest.mark.parametrize( + "incompatible_kind", + ( + "wfp4", + "f16", + "invalid_sbm", + "invalid_kwave", + ), +) +def test_hy4_mxfp8_fhmoe_capability_rejects_incompatible_kernels( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, incompatible_kind: str ): + import csv import importlib - fused_moe_module = importlib.import_module("aiter.fused_moe") - - config_path = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" - monkeypatch.setattr(fused_moe_module, "get_cu_num", lambda: 256) - monkeypatch.setattr(fused_moe_module, "get_gfx_runtime", lambda: "gfx950") - fused_moe_module.get_2stage_cfgs.cache_clear() - fused_moe_module.cfg_2stages_by_file.clear() - - metadata = fused_moe_module.get_2stage_cfgs( - fused_moe_module.get_padded_M(num_tokens), - 7168, - 384, - 385, - 7, - torch.bfloat16, - dtypes.fp8, - dtypes.fp4x2, - aiter.QuantType.per_1x32, - True, - aiter.ActivationType.Silu, - False, - 0, - 0, - True, - GateMode.INTERLEAVE, - config_file=str(config_path), + fhmoe = importlib.import_module("aiter.fhmoe") + _mock_gfx950_fhmoe_metadata(monkeypatch) + source = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" + with source.open(newline="") as stream: + reader = csv.DictReader(stream) + assert reader.fieldnames is not None + rows = list(reader) + fields = reader.fieldnames + dsv4 = next( + row for row in rows if row["model_dim"] == "7168" and row["token"] == "1" ) + incompatible = tmp_path / f"hy4_{incompatible_kind}_kernels.csv" + with incompatible.open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + for row in rows: + updated = dict(row) + if row["model_dim"] == "6144" and row["token"] == "1": + if incompatible_kind == "wfp4": + updated["kernelName1"] = dsv4["kernelName1"] + updated["kernelName2"] = dsv4["kernelName2"] + elif incompatible_kind == "f16": + updated["kernelName2"] = updated["kernelName2"].replace( + "_bf16_", "_f16_" + ) + elif incompatible_kind == "invalid_sbm": + updated["kernelName2"] = ( + "flydsl_moe2_afp8_wfp8_bf16_t64x128x128_atomic_sbm32" + ) + elif incompatible_kind == "invalid_kwave": + updated["kernelName1"] = updated["kernelName1"].replace( + "_kw2_fp8", "_kw5_fp8" + ) + writer.writerow(updated) + + monkeypatch.setattr(fhmoe, "_fhmoe_config_file", lambda: str(incompatible)) + fhmoe._supports_mxfp8_fhmoe_config.cache_clear() + + assert not _supports_hy4_config(fhmoe, 1) + + +def test_hy4_mxfp8_fhmoe_capability_rejects_config_bypass( + monkeypatch: pytest.MonkeyPatch, +): + import importlib - assert metadata.stage1.keywords["kernelName"] == expected_stage1 - assert metadata.stage2.keywords["kernelName"] == expected_stage2 + fhmoe = importlib.import_module("aiter.fhmoe") + monkeypatch.setenv("AITER_BYPASS_TUNE_CONFIG", "1") + assert not _supports_hy4_config(fhmoe, 1) @pytest.mark.parametrize("num_tokens", (768, 2049)) @@ -1147,9 +2504,8 @@ def test_dsv4_i384_fhmoe_config_requires_exact_bucket( num_tokens: int, ): import csv - import importlib - fused_moe_module = importlib.import_module("aiter.fused_moe") + fused_moe_module = _mock_gfx950_fhmoe_metadata(monkeypatch) source_path = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" with source_path.open(newline="") as config: reader = csv.DictReader(config) @@ -1164,11 +2520,7 @@ def test_dsv4_i384_fhmoe_config_requires_exact_bucket( writer.writeheader() writer.writerows(row for row in rows if row["token"] != missing_token) - monkeypatch.setattr(fused_moe_module, "get_cu_num", lambda: 256) - monkeypatch.setattr(fused_moe_module, "get_gfx_runtime", lambda: "gfx950") monkeypatch.setenv("AITER_ONLINE_TUNE", "1") - fused_moe_module.get_2stage_cfgs.cache_clear() - fused_moe_module.cfg_2stages_by_file.clear() with pytest.raises(NotImplementedError, match="requires an exact tuned config"): fused_moe_module.get_2stage_cfgs( @@ -1205,7 +2557,8 @@ def test_dsv4_i384_fhmoe_config_has_true_shapes(): with ordinary_path.open(newline="") as f: ordinary_rows = list(csv.DictReader(f)) - assert {int(row["token"]) for row in rows} == { + dsv4_rows = [row for row in rows if int(row["inter_dim"]) == 384] + assert {int(row["token"]) for row in dsv4_rows} == { 1, 2, 4, @@ -1219,13 +2572,12 @@ def test_dsv4_i384_fhmoe_config_has_true_shapes(): 1024, 2048, } - assert all(int(row["inter_dim"]) == 384 for row in rows) - assert all(int(row["shared_expert_id"]) == 384 for row in rows) - assert all(int(row["hidden_pad"]) == 0 for row in rows) - assert all(int(row["intermediate_pad"]) == 0 for row in rows) - assert all(row["gate_mode"] == "GateMode.INTERLEAVE" for row in rows) - assert all(row["kernelName1"].startswith("flydsl_") for row in rows) - assert all(row["kernelName2"].startswith("flydsl_") for row in rows) + assert all(int(row["shared_expert_id"]) == 384 for row in dsv4_rows) + assert all(int(row["hidden_pad"]) == 0 for row in dsv4_rows) + assert all(int(row["intermediate_pad"]) == 0 for row in dsv4_rows) + assert all(row["gate_mode"] == "GateMode.INTERLEAVE" for row in dsv4_rows) + assert all(row["kernelName1"].startswith("flydsl_") for row in dsv4_rows) + assert all(row["kernelName2"].startswith("flydsl_") for row in dsv4_rows) ordinary_m16 = next( row @@ -1248,7 +2600,9 @@ def test_fhmoe_aot_manifest_covers_native_i384(): ) config_path = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" ordinary_jobs = parse_csv(str(ordinary_path)) - dedicated_jobs = parse_csv(str(config_path)) + base_jobs = parse_csv(str(config_path)) + dedicated_jobs = [job for job in base_jobs if job["inter_dim"] == 384] + hy4_jobs = [job for job in base_jobs if job["inter_dim"] == 256] ordinary_fhmoe_jobs = [ job for job in ordinary_jobs if job.get("shared_expert_id", -1) >= 0 ] @@ -1257,6 +2611,7 @@ def test_fhmoe_aot_manifest_covers_native_i384(): assert len(dedicated_jobs) == 24 assert all(job["inter_dim"] == 384 for job in dedicated_jobs) assert all(job["shared_expert_id"] == 384 for job in dedicated_jobs) + assert all(job["b_dtype"] == "fp4" for job in dedicated_jobs) assert all(not job.get("enable_bias", False) for job in dedicated_jobs) assert {job["token_num"] for job in dedicated_jobs} == { 1, @@ -1272,14 +2627,44 @@ def test_fhmoe_aot_manifest_covers_native_i384(): 1024, 2048, } - for job in dedicated_jobs: + assert len(hy4_jobs) == 48 + assert all(job["shared_expert_id"] == 256 for job in hy4_jobs) + assert all(job["b_dtype"] == "fp8" for job in hy4_jobs) + hy4_stage1_jobs = [job for job in hy4_jobs if job["stage"] == 1] + assert len(hy4_stage1_jobs) == 32 + assert {job["clamp_shared"] for job in hy4_stage1_jobs} == {True, False} + assert all("clamp_shared" not in job for job in hy4_jobs if job["stage"] == 2) + assert {job["token_num"] for job in hy4_jobs} == { + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + 32768, + } + for job in base_jobs: params = get_flydsl_kernel_params(job["kernel_name"]) assert params is not None assert job.get("xcd_swizzle", 0) == params.get("xcd_swizzle", 0) + # HY4's FP8/FP8 rows are intentionally non-persistent. Restored DSV4 + # FP8/FP4 rows retain their legacy persist suffixes; the shared runtime + # normalizes FP8-A grid scheduling in resolve_flydsl_grid_y_persist_m. + if job["stage"] == 2 and job["a_dtype"] == "fp8" and job["b_dtype"] == "fp8": + assert not job.get("persist", False) if job["stage"] == 1: - assert 384 % job["tile_n"] == 0 + assert job["inter_dim"] % job["tile_n"] == 0 else: - assert 384 % job["tile_k"] == 0 + assert job["inter_dim"] % job["tile_k"] == 0 m2048_names = { job["kernel_name"] for job in dedicated_jobs if job["token_num"] == 2048 @@ -1290,16 +2675,200 @@ def test_fhmoe_aot_manifest_covers_native_i384(): } +def test_fhmoe_aot_import_is_safe_in_lightweight_mode(): + import subprocess + import sys + + root = Path(__file__).resolve().parents[1] + env = os.environ.copy() + env["AITER_AOT_IMPORT"] = "1" + code = """ +import aiter +assert not hasattr(aiter, "ActivationType") +from aiter.fhmoe_contract import max_fhmoe_reduce_tokens +from aiter.aot.flydsl.fhmoe import _FHMoEAOTBackend +assert max_fhmoe_reduce_tokens(9, 6144, 2) == 38836 +assert _FHMoEAOTBackend(shared_expert_id=256).shared_expert_id == 256 +""" + result = subprocess.run( + [sys.executable, "-c", code], + cwd=root, + env=env, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_fhmoe_aot_cache_is_retrievable_in_run_only_mode(tmp_path: Path): + import subprocess + import sys + + if get_gfx() != "gfx950": + pytest.skip("FHMoE AOT cache retrieval requires gfx950") + root = Path(__file__).resolve().parents[1] + cache = tmp_path / "cache" + env = os.environ.copy() + env["FLYDSL_RUNTIME_CACHE_DIR"] = str(cache) + compile_env = dict(env, AITER_AOT_IMPORT="1") + compile_code = """ +from aiter.aot.flydsl.moe import compile_one_config, parse_csv +from aiter.jit.core import AITER_CONFIGS +jobs = [ + job for job in parse_csv(AITER_CONFIGS.AITER_CONFIG_FHMOE_FILE) + if job["model_dim"] == 6144 and job["token_num"] == 1 +] +assert len(jobs) == 3 +for job in jobs: + assert compile_one_config(**job)["compile_time"] is not None +""" + compiled = subprocess.run( + [sys.executable, "-c", compile_code], + cwd=root, + env=compile_env, + capture_output=True, + text=True, + check=False, + ) + assert compiled.returncode == 0, compiled.stderr + + run_env = dict(env, FLYDSL_RUNTIME_RUN_ONLY="1") + retrieve_code = """ +from aiter.ops.flydsl.fhmoe import ( + compile_flydsl_fhmoe_stage1, + compile_flydsl_fhmoe_stage2, +) +compile_flydsl_fhmoe_stage1( + model_dim=6144, inter_dim=256, experts=257, topk=9, + tile_m=32, tile_n=64, tile_k=256, doweight_stage1=False, + a_dtype="fp8", b_dtype="fp8", out_dtype="fp8", act="silu", + persist_m=1, use_async_copy=True, k_batch=1, waves_per_eu=4, + b_nt=2, gate_mode="interleave", model_dim_pad=0, inter_dim_pad=0, + enable_bias=False, a_scale_one=False, xcd_swizzle=0, k_wave=2, + v2_output_layout=False, shared_expert_id=256, clamp_shared=False, +) +compile_flydsl_fhmoe_stage1( + model_dim=6144, inter_dim=256, experts=257, topk=9, + tile_m=32, tile_n=64, tile_k=256, doweight_stage1=False, + a_dtype="fp8", b_dtype="fp8", out_dtype="fp8", act="silu", + persist_m=1, use_async_copy=True, k_batch=1, waves_per_eu=4, + b_nt=2, gate_mode="interleave", model_dim_pad=0, inter_dim_pad=0, + enable_bias=False, a_scale_one=False, xcd_swizzle=0, k_wave=2, + v2_output_layout=False, shared_expert_id=256, clamp_shared=True, +) +compile_flydsl_fhmoe_stage2( + model_dim=6144, inter_dim=256, experts=257, topk=9, + tile_m=32, tile_n=128, tile_k=128, doweight_stage2=True, + a_dtype="fp8", b_dtype="fp8", out_dtype="bf16", accumulate=True, + enable_bias=False, model_dim_pad=0, inter_dim_pad=0, persist_m=1, + sort_block_m=0, waves_per_eu=None, use_async_copy=False, + cu_num_mul=1, b_nt=0, xcd_swizzle=0, shared_expert_id=256, + use_global_a=False, +) +""" + retrieved = subprocess.run( + [sys.executable, "-c", retrieve_code], + cwd=root, + env=run_env, + capture_output=True, + text=True, + check=False, + ) + assert retrieved.returncode == 0, retrieved.stderr + + +def test_fhmoe_aot_csv_preserves_padding(tmp_path: Path): + import csv + + from aiter.aot.flydsl.moe import parse_csv + + source = Path(__file__).resolve().parents[1] / "aiter/configs/tuned_fhmoe.csv" + with source.open(newline="") as stream: + reader = csv.DictReader(stream) + assert reader.fieldnames is not None + row = next( + row for row in reader if row["model_dim"] == "6144" and row["token"] == "1" + ) + fields = reader.fieldnames + row["hidden_pad"] = "64" + row["intermediate_pad"] = "128" + padded = tmp_path / "padded_fhmoe.csv" + with padded.open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + writer.writerow(row) + + jobs = parse_csv(str(padded)) + assert jobs + assert all(job["gfx"] == "gfx950" for job in jobs) + assert all(job["model_dim_pad"] == 64 for job in jobs) + assert all(job["inter_dim_pad"] == 128 for job in jobs) + + +@pytest.mark.parametrize("stage", (1, 2)) +def test_fhmoe_aot_forwards_padding_to_compiler( + monkeypatch: pytest.MonkeyPatch, stage: int +): + from aiter.aot.flydsl import moe as aot_moe + + forwarded = {} + + class Backend: + def build_stage1_args(self, *_args, **_kwargs): + return () + + def build_stage2_args(self, *_args, **_kwargs): + return () + + def compile_stage1(self, **kwargs): + forwarded.update(kwargs) + return object() + + def compile_stage2(self, **kwargs): + forwarded.update(kwargs) + return object() + + monkeypatch.setattr(aot_moe, "_run_compiled", lambda *_args: None) + aot_moe._precompile_to_cache( + stage=stage, + model_dim=256, + inter_dim=128, + experts=9, + topk=3, + tile_m=32, + tile_n=128, + tile_k=128, + a_dtype="fp8", + b_dtype="fp8", + out_dtype="bf16", + gate_mode="interleave", + mode="atomic", + cu_num=256, + token_num=1, + block_m=32, + model_dim_pad=64, + inter_dim_pad=128, + _aot_backend=Backend(), + ) + + assert forwarded["model_dim_pad"] == 64 + assert forwarded["inter_dim_pad"] == 128 + + def test_fhmoe_aot_precompile_keeps_native_i384(monkeypatch: pytest.MonkeyPatch): from aiter.aot.flydsl import fhmoe as aot_fhmoe from aiter.aot.flydsl import moe as aot_moe forwarded = {} + compile_arch = [] def precompile(**kwargs): forwarded.update(kwargs) + compile_arch.append(os.environ.get("FLYDSL_GPU_ARCH")) monkeypatch.setattr(aot_moe, "_precompile_to_cache", precompile) + monkeypatch.setenv("FLYDSL_GPU_ARCH", "gfx1250") aot_fhmoe.precompile_fhmoe_to_cache( experts=385, shared_expert_id=384, @@ -1312,6 +2881,348 @@ def precompile(**kwargs): assert forwarded["inter_dim"] == 384 assert forwarded["_aot_backend"].shared_expert_id == 384 + assert compile_arch == ["gfx950"] + assert os.environ["FLYDSL_GPU_ARCH"] == "gfx1250" + + +def test_fhmoe_aot_stage2_scale_size_uses_padded_inter_dim(): + from aiter.aot.flydsl.moe import _mx_w2_scale_numel + + assert _mx_w2_scale_numel(385, 7168, 384) == 385 * 7168 * (512 // 32) + assert _mx_w2_scale_numel(257, 6144, 256) == 257 * 6144 * (256 // 32) + + +def test_fhmoe_aot_precompile_accepts_hy4_mxfp8(monkeypatch: pytest.MonkeyPatch): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + from aiter.aot.flydsl import moe as aot_moe + + forwarded = {} + + def precompile(**kwargs): + forwarded.update(kwargs) + + monkeypatch.setattr(aot_moe, "_precompile_to_cache", precompile) + aot_fhmoe.precompile_fhmoe_to_cache( + experts=257, + shared_expert_id=256, + a_dtype="fp8", + b_dtype="fp8", + cu_num=256, + stage=1, + model_dim=6144, + inter_dim=256, + topk=9, + gate_mode="interleave", + ) + + assert forwarded["b_dtype"] == "fp8" + assert forwarded["inter_dim"] == 256 + assert forwarded["_aot_backend"].shared_expert_id == 256 + + +@pytest.mark.parametrize("cu_num", (-1, 0, 256)) +def test_fhmoe_aot_precompile_accepts_hy4_mxfp8_stage2( + monkeypatch: pytest.MonkeyPatch, cu_num: int +): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + from aiter.aot.flydsl import moe as aot_moe + + forwarded = {} + + def precompile(**kwargs): + forwarded.update(kwargs) + + monkeypatch.setattr(aot_moe, "_precompile_to_cache", precompile) + aot_fhmoe.precompile_fhmoe_to_cache( + experts=257, + shared_expert_id=256, + a_dtype="fp8", + b_dtype="fp8", + cu_num=cu_num, + stage=2, + model_dim=6144, + inter_dim=256, + topk=9, + ) + + assert forwarded["stage"] == 2 + assert forwarded["b_dtype"] == "fp8" + assert forwarded["cu_num"] == 256 + + +def test_fhmoe_aot_compile_job_forwards_explicit_target_gfx( + monkeypatch: pytest.MonkeyPatch, +): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + from aiter.aot.flydsl import moe as aot_moe + + forwarded = {} + + def precompile(**kwargs): + forwarded.update(kwargs) + + monkeypatch.setattr(aot_fhmoe, "precompile_fhmoe_to_cache", precompile) + result = aot_moe.compile_one_config( + kernel_name="flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic", + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + cu_num=256, + gfx="gfx950", + shared_expert_id=256, + stage=2, + a_dtype="fp8", + b_dtype="fp8", + ) + + assert result["compile_arch"] == "gfx950" + assert forwarded["gfx"] == "gfx950" + assert forwarded["cu_num"] == 256 + + +def test_fhmoe_aot_precompile_rejects_separated_mxfp8(): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + + with pytest.raises(ValueError, match="interleaved gate/up layout"): + aot_fhmoe.precompile_fhmoe_to_cache( + experts=257, + shared_expert_id=256, + a_dtype="fp8", + b_dtype="fp8", + cu_num=256, + stage=1, + model_dim=6144, + inter_dim=256, + topk=9, + ) + + +def test_fhmoe_aot_precompile_rejects_unknown_cu_count(): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + + with pytest.raises(ValueError, match="requires target gfx950"): + aot_fhmoe.precompile_fhmoe_to_cache( + experts=257, + shared_expert_id=256, + a_dtype="fp8", + b_dtype="fp8", + cu_num=999, + stage=2, + model_dim=6144, + inter_dim=256, + topk=9, + ) + + +def test_fhmoe_aot_precompile_rejects_non_gfx950_target(): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + + with pytest.raises(ValueError, match="gfx='gfx1250'.*cu_num=256"): + aot_fhmoe.precompile_fhmoe_to_cache( + experts=257, + shared_expert_id=256, + a_dtype="fp8", + b_dtype="fp8", + cu_num=256, + gfx="gfx1250", + stage=2, + model_dim=6144, + inter_dim=256, + topk=9, + ) + + +def test_fhmoe_aot_precompile_rejects_non_256_target_cu(): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + + with pytest.raises(ValueError, match="gfx='gfx942'.*cu_num=304"): + aot_fhmoe.precompile_fhmoe_to_cache( + experts=257, + shared_expert_id=256, + a_dtype="fp8", + b_dtype="fp8", + cu_num=304, + gfx="gfx942", + stage=2, + model_dim=6144, + inter_dim=256, + topk=9, + ) + + +def test_fhmoe_aot_precompile_rejects_invalid_stage(): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + + with pytest.raises(ValueError, match="stage must be 1 or 2"): + aot_fhmoe.precompile_fhmoe_to_cache( + experts=257, + shared_expert_id=256, + cu_num=256, + stage=0, + ) + + +def test_fhmoe_aot_clamp_default_and_explicit_unclamped_specialization( + monkeypatch: pytest.MonkeyPatch, +): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + from aiter.ops.flydsl import fhmoe as ops_fhmoe + + monkeypatch.setattr( + ops_fhmoe, + "compile_flydsl_fhmoe_stage1", + lambda **kwargs: kwargs, + ) + backend = aot_fhmoe._FHMoEAOTBackend(shared_expert_id=256) + kwargs = { + "model_dim": 6144, + "inter_dim": 256, + "experts": 257, + "topk": 9, + "b_dtype": "fp8", + "model_dim_pad": 0, + "inter_dim_pad": 0, + "gate_mode": "interleave", + "doweight_stage1": False, + } + + assert backend.compile_stage1(**kwargs)["clamp_shared"] is True + assert ( + backend.compile_stage1(**(kwargs | {"clamp_shared": False}))["clamp_shared"] + is False + ) + + +@pytest.mark.parametrize("clamp_shared", (True, False)) +def test_fhmoe_aot_precompile_forwards_clamp_shared( + monkeypatch: pytest.MonkeyPatch, + clamp_shared: bool, +): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + from aiter.aot.flydsl import moe as aot_moe + + captured = {} + monkeypatch.setattr( + aot_fhmoe._FHMoEAOTBackend, + "compile_stage1", + lambda _self, **kwargs: captured.update(kwargs) or object(), + ) + monkeypatch.setattr(aot_moe, "_run_compiled", lambda *_args, **_kwargs: None) + + aot_moe._precompile_to_cache( + stage=1, + model_dim=6144, + inter_dim=256, + experts=257, + topk=9, + tile_m=32, + tile_n=64, + tile_k=256, + a_dtype="fp8", + b_dtype="fp8", + out_dtype="fp8", + token_num=1, + gate_mode="interleave", + k_wave=2, + shared_expert_id=256, + clamp_shared=clamp_shared, + _aot_backend=aot_fhmoe._FHMoEAOTBackend(shared_expert_id=256), + ) + + assert captured["clamp_shared"] is clamp_shared + + +def test_fhmoe_aot_precompile_accepts_generic_mxfp8_clamp_modes( + monkeypatch: pytest.MonkeyPatch, +): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + from aiter.aot.flydsl import moe as aot_moe + + forwarded = {} + + def precompile(**kwargs): + forwarded.update(kwargs) + + monkeypatch.setattr(aot_moe, "_precompile_to_cache", precompile) + aot_fhmoe.precompile_fhmoe_to_cache( + experts=9, + shared_expert_id=8, + a_dtype="fp8", + b_dtype="fp8", + cu_num=256, + stage=1, + model_dim=256, + inter_dim=128, + topk=3, + gate_mode="interleave", + ) + + assert forwarded["b_dtype"] == "fp8" + assert forwarded["_aot_backend"].shared_expert_id == 8 + aot_fhmoe.precompile_fhmoe_to_cache( + experts=9, + shared_expert_id=8, + a_dtype="fp8", + b_dtype="fp8", + cu_num=256, + stage=1, + model_dim=256, + inter_dim=128, + topk=3, + gate_mode="interleave", + clamp_shared=False, + ) + assert forwarded["clamp_shared"] is False + + +def test_fhmoe_aot_keeps_mxfp4_shared_expert_clamped(): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + + with pytest.raises(ValueError, match="MXFP8 routed weights"): + aot_fhmoe.precompile_fhmoe_to_cache( + experts=385, + shared_expert_id=384, + a_dtype="fp8", + b_dtype="fp4", + cu_num=256, + stage=1, + model_dim=7168, + inter_dim=384, + topk=7, + gate_mode="interleave", + clamp_shared=False, + ) + + +def test_fhmoe_aot_keeps_dsv4_shape_mxfp8_generic( + monkeypatch: pytest.MonkeyPatch, +): + from aiter.aot.flydsl import fhmoe as aot_fhmoe + from aiter.aot.flydsl import moe as aot_moe + + forwarded = {} + + def precompile(**kwargs): + forwarded.update(kwargs) + + monkeypatch.setattr(aot_moe, "_precompile_to_cache", precompile) + aot_fhmoe.precompile_fhmoe_to_cache( + experts=385, + shared_expert_id=384, + a_dtype="fp8", + b_dtype="fp8", + cu_num=256, + stage=1, + model_dim=7168, + inter_dim=384, + topk=7, + gate_mode="interleave", + ) + + assert forwarded["b_dtype"] == "fp8" + assert forwarded["model_dim"] == 7168 + assert forwarded["_aot_backend"].shared_expert_id == 384 def test_fhmoe_aot_stage1_forwards_optional_swiglu_abi( @@ -1432,7 +3343,10 @@ def test_a4w4_routed_fp8_shared_heterogeneous_path( activation_quantization="fp4", ) shared_high = _torch_shared_reference( - hidden, a4_weights, "fp4", _swiglu_limit(profile) + hidden, + a4_weights, + "fp4", + _shared_swiglu_limit(_swiglu_limit(profile), clamp_shared=True), ) assert torch.isfinite(actual).all() error = _rel_l2(actual, routed_high + shared_high) diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index dd7e67bc804..7b0d5260e87 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -14,11 +14,16 @@ import aiter from aiter import dtypes from aiter.aot.flydsl.common import override_env, run_only_env +from aiter.fhmoe_contract import FHMOE_BUFFER_MAX_BYTES from aiter.fused_moe import ( + _flydsl_route_output_nbytes, + _is_flydsl_stage2, + _stage2_runtime_uses_route_reduce, fused_moe, fused_topk, get_2stage_cfgs, get_padded_M, + resolve_activation_dtype, torch_moe_stage1, torch_moe_stage2, ) @@ -1052,6 +1057,86 @@ def _moe_2stage_reference_workspace_gib(kwargs): ) +def _moe_2stage_route_workspace_skip_reason(kwargs): + """Return a reason when effective route-reduce exceeds 32-bit addressing.""" + token = kwargs["token"] + model_dim = kwargs["model_dim"] + topk = kwargs["topk"] + q_dtype_a = resolve_activation_dtype( + kwargs["qType"], + kwargs["WQDType"], + activation=kwargs["actType"], + gate_mode=kwargs["gateMode"], + M=token, + hidden_dtype=kwargs["dtype"], + gfx=get_gfx(), + ) + if q_dtype_a is None: + return None + is_shuffled = not ( + kwargs["WQDType"] == dtypes.fp4x2 + and kwargs["AQDType"] == dtypes.fp4x2 + and not kwargs.get("preshuffle", False) + ) + supports_expert_bias = ( + kwargs["qType"] == aiter.QuantType.per_1x32 + and kwargs["WQDType"] == dtypes.fp4x2 + and kwargs["AQDType"] in (dtypes.bf16, dtypes.fp16, dtypes.fp8) + and kwargs["actType"] != aiter.ActivationType.Situv2 + ) + try: + metadata = get_2stage_cfgs( + get_padded_M(token), + model_dim, + kwargs["inter_dim"], + kwargs["E"], + topk, + kwargs["dtype"], + q_dtype_a, + kwargs["WQDType"], + kwargs["qType"], + kwargs["use_g1u1"], + kwargs["actType"], + kwargs["doweight_stage1"], + kwargs.get("hidden_pad", 0), + kwargs.get("intermediate_pad", 0), + is_shuffled, + kwargs["gateMode"], + situ_beta=kwargs.get("beta", 1.0), + situ_linear_beta=kwargs.get("linear_beta", 1.0), + swiglu_limit=kwargs.get("swiglu_limit"), + opus_weights_shuffled=is_shuffled, + has_stage1_bias=supports_expert_bias, + has_stage2_bias=( + supports_expert_bias and not kwargs.get("disable_stage2_bias", False) + ), + ) + except Exception as exc: # noqa: BLE001 + aiter.logger.warning( + "moe_2stage: failed to resolve stage2 mode during preflight; " + "continuing to runtime validation: %s", + exc, + ) + return None + if not _is_flydsl_stage2(metadata.stage2): + return None + if not _stage2_runtime_uses_route_reduce(metadata.stage2, token, model_dim): + return None + route_bytes = _flydsl_route_output_nbytes( + metadata.stage2, + token, + topk, + model_dim, + torch.empty((), dtype=kwargs["dtype"]).element_size(), + ) + if route_bytes <= FHMOE_BUFFER_MAX_BYTES: + return None + return ( + f"FlyDSL stage2 reduce workspace {route_bytes / (1024**3):.1f} GiB " + "exceeds the current 32-bit buffer-addressing limit" + ) + + def _format_moe_2stage_case(kwargs): return ( f"token={kwargs['token']}, dim=({kwargs['model_dim']},{kwargs['inter_dim']}), " @@ -1062,6 +1147,10 @@ def _format_moe_2stage_case(kwargs): def _moe_2stage_skip_reason(kwargs): + route_workspace_reason = _moe_2stage_route_workspace_skip_reason(kwargs) + if route_workspace_reason is not None: + return route_workspace_reason + # Set AITER_MOE_2STAGE_MAX_FREE_MEMORY_FRACTION=0 to run manually. max_free_memory_fraction = float( os.environ.get("AITER_MOE_2STAGE_MAX_FREE_MEMORY_FRACTION", "0.8") diff --git a/op_tests/tuners/tune_fhmoe.py b/op_tests/tuners/tune_fhmoe.py new file mode 100644 index 00000000000..7d84e96a777 --- /dev/null +++ b/op_tests/tuners/tune_fhmoe.py @@ -0,0 +1,1291 @@ +# SPDX-License-Identifier: MIT +"""Offline FlyDSL tuner for fused heterogeneous MoE (FHMoE). + +The tuner measures the complete two-stage FHMoE call, including sorting, +stage-1 activation/quantization and stage-2 accumulation. It updates only rows +matching the requested contract and preserves every other row in the CSV. Each +selected full-M result must also match an independent FP32 reference on sampled +token rows before its kernels can be written. + +Candidate ranking remains a full-M comparison against the current kernel. A +full FP32 oracle for every token at large M would make tuning prohibitively +slow, so only the final winner receives the expensive independent check. The +tuner retains raw weights for experts reached by representative token rows, +computes the complete routed-plus-shared FP32 result for those rows, and compares +it with the corresponding rows from the winner's full-M output. + +Example (HY4 MXFP8 routed / FP8 shared on gfx950): + + python op_tests/tuners/tune_fhmoe.py \ + -i aiter/configs/tuned_fhmoe.csv \ + -o aiter/configs/tuned_fhmoe.csv \ + --tokens 1,2,4,8,16,32,64,128,256,512,1024,2048,4096,8192,16384,32768 +""" + +from __future__ import annotations + +import argparse +import csv +import fcntl +import hashlib +import math +import os +import stat +import tempfile +from collections import Counter +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.nn.functional as F +import triton + +import aiter +from aiter import dtypes +from aiter.fhmoe import _use_fhmoe_wrappers +from aiter.fhmoe_contract import max_fhmoe_reduce_tokens +from aiter.fused_moe import ( + _fused_moe_impl, + cfg_2stages_by_file, + get_2stage_cfgs, +) +from aiter.jit.utils.chip_info import get_cu_num, get_gfx_runtime +from aiter.ops.flydsl.moe_common import GateMode +from aiter.ops.flydsl.moe_kernels import ( + get_flydsl_kernel_params, + get_flydsl_stage1_kernels, + get_flydsl_stage2_kernels, + runtime_swiglu_limit, +) +from aiter.ops.shuffle import ( + shuffle_scale, + shuffle_scale_a16w4, + shuffle_weight_a16w4, +) +from aiter.utility import fp4_utils +from aiter.utility.mx_types import MxDtypeInt + +HY4_FHMOE_MAX_TOKENS = max_fhmoe_reduce_tokens(9, 6144, 2) + + +@dataclass(frozen=True) +class Contract: + gfx: str = "gfx950" + cu_num: int = 256 + model_dim: int = 6144 + inter_dim: int = 256 + routed_experts: int = 256 + routed_topk: int = 8 + shared_expert_id: int = 256 + act_type: str = "ActivationType.Silu" + dtype: str = "torch.bfloat16" + q_dtype_a: str = "torch.float8_e4m3fn" + q_dtype_w: str = "torch.float8_e4m3fn" + q_type: str = "QuantType.per_1x32" + use_g1u1: int = 1 + doweight_stage1: int = 0 + hidden_pad: int = 0 + intermediate_pad: int = 0 + gate_mode: str = "GateMode.INTERLEAVE" + ksplit: int = 0 + + @property + def experts(self) -> int: + return self.routed_experts + 1 + + @property + def topk(self) -> int: + return self.routed_topk + 1 + + +@dataclass +class Tensors: + hidden: torch.Tensor + topk_weights: torch.Tensor + topk_ids: torch.Tensor + w1: torch.Tensor + w2: torch.Tensor + w1_scale: torch.Tensor + w2_scale: torch.Tensor + shared_w1: torch.Tensor + shared_w2: torch.Tensor + shared_w1_scale: torch.Tensor + shared_w2_scale: torch.Tensor + output: torch.Tensor + # Minimal raw state retained for the final sampled FP32 acceptance check. + # Candidate timing still uses the full tensors above. + sample_rows: tuple[int, ...] + sample_expert_ids: tuple[int, ...] + sample_w1: torch.Tensor + sample_w2: torch.Tensor + sample_w1_scale: torch.Tensor + sample_w2_scale: torch.Tensor + sample_shared_w1: torch.Tensor + sample_shared_w2: torch.Tensor + sample_shared_w1_scale: torch.Tensor + sample_shared_w2_scale: torch.Tensor + + +CSV_FIELDS = ( + "gfx", + "cu_num", + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "shared_expert_id", + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "hidden_pad", + "intermediate_pad", + "gate_mode", + "block_m", + "ksplit", + "kernelName1", + "kernelName2", +) +PROFILE_FIELDS = ("token", "runtime_token", "stage", "kernel", "ms", "error") +_PAIR_BEAM_WIDTH = 4 + + +def _parse_tokens(value: str) -> list[int]: + tokens = sorted({int(token) for token in value.split(",") if token.strip()}) + if not tokens or any(token <= 0 for token in tokens): + raise argparse.ArgumentTypeError("tokens must be positive comma-separated ints") + return tokens + + +def _parse_nonnegative_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed < 0: + raise argparse.ArgumentTypeError("value must be finite and nonnegative") + return parsed + + +def _parse_positive_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError("value must be finite and positive") + return parsed + + +def _validate_benchmark_token(benchmark_token: int | None, tokens: list[int]) -> None: + if benchmark_token is None: + oversized = [token for token in tokens if token > HY4_FHMOE_MAX_TOKENS] + if oversized: + raise ValueError( + "HY4 tuning tokens must not exceed " + f"{HY4_FHMOE_MAX_TOKENS}, got {oversized}" + ) + return + from aiter.fused_moe import get_padded_M + + if benchmark_token <= 0: + raise ValueError("--benchmark-token must be positive") + if benchmark_token > HY4_FHMOE_MAX_TOKENS: + raise ValueError(f"--benchmark-token must not exceed {HY4_FHMOE_MAX_TOKENS}") + if len(tokens) != 1: + raise ValueError("--benchmark-token requires exactly one --tokens bucket") + if get_padded_M(benchmark_token) != tokens[0]: + raise ValueError( + f"benchmark M={benchmark_token} resolves to bucket " + f"{get_padded_M(benchmark_token)}, not {tokens[0]}" + ) + + +def _read_rows(path: Path) -> tuple[list[str], list[dict[str, str]]]: + with path.open(newline="") as stream: + reader = csv.DictReader(stream) + if reader.fieldnames is None: + raise ValueError(f"{path} has no CSV header") + rows = list(reader) + return list(reader.fieldnames), rows + + +def _write_rows( + path: Path, + fields: list[str], + rows: list[dict[str, str]], + *, + durable: bool = False, +) -> None: + if path.is_symlink(): + raise ValueError(f"refusing to replace symlink output: {path}") + path.parent.mkdir(parents=True, exist_ok=True) + previous_mode = path.stat().st_mode if path.exists() else 0o644 + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + "w", + newline="", + dir=path.parent, + prefix=f".{path.name}.", + delete=False, + ) as stream: + temporary = Path(stream.name) + writer = csv.DictWriter(stream, fieldnames=fields, lineterminator="\n") + writer.writeheader() + writer.writerows(rows) + os.fchmod(stream.fileno(), previous_mode & 0o777) + if durable: + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + if durable: + directory_fd = os.open(path.parent, os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + if temporary is not None and temporary.exists(): + temporary.unlink() + + +def _paths_alias(first: Path, second: Path) -> bool: + if first.exists() and second.exists(): + try: + return os.path.samefile(first, second) + except OSError: + pass + return first.resolve(strict=False) == second.resolve(strict=False) + + +def _validate_output_paths( + input_path: Path, output_path: Path, profile_path: Path | None +) -> None: + if output_path.is_symlink(): + raise ValueError(f"FHMoE tuner output must not be a symlink: {output_path}") + if profile_path is None: + return + if profile_path.is_symlink(): + raise ValueError(f"FHMoE tuner profile must not be a symlink: {profile_path}") + if _paths_alias(profile_path, input_path) or _paths_alias( + profile_path, output_path + ): + raise ValueError("FHMoE tuner profile must differ from input and output") + + +def _lock_directory() -> Path: + directory = Path("/tmp") / f"aiter-fhmoe-locks-{os.getuid()}" + directory.mkdir(mode=0o700, parents=True, exist_ok=True) + info = directory.lstat() + if not stat.S_ISDIR(info.st_mode) or info.st_uid != os.getuid(): + raise RuntimeError(f"unsafe FHMoE lock directory: {directory}") + os.chmod(directory, 0o700) + return directory + + +@contextmanager +def _output_lock(path: Path): + lock_id = hashlib.sha256(str(path.resolve()).encode()).hexdigest()[:16] + lock_path = _lock_directory() / f"{lock_id}.lock" + flags = os.O_CREAT | os.O_RDWR | os.O_CLOEXEC + flags |= getattr(os, "O_NOFOLLOW", 0) + descriptor = os.open(lock_path, flags, 0o600) + try: + info = os.fstat(descriptor) + if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid(): + raise RuntimeError(f"unsafe FHMoE lock file: {lock_path}") + fcntl.flock(descriptor, fcntl.LOCK_EX) + yield + finally: + fcntl.flock(descriptor, fcntl.LOCK_UN) + os.close(descriptor) + + +def _write_profile_rows( + profile_path: Path, + rows: list[dict[str, str | float]], +) -> None: + """Durably replace the measured token/runtime pair under a lock.""" + normalized = [ + {field: row.get(field, "") for field in PROFILE_FIELDS} for row in rows + ] + keys = {(str(row["token"]), str(row["runtime_token"])) for row in normalized} + with _output_lock(profile_path): + existing: list[dict[str, str | float]] = [] + if profile_path.exists(): + fields, existing_rows = _read_rows(profile_path) + legacy_fields = ["token", "stage", "kernel", "ms", "error"] + if fields == legacy_fields: + for row in existing_rows: + row["runtime_token"] = row["token"] + elif fields != list(PROFILE_FIELDS): + raise ValueError( + f"cannot merge FHMoE profile rows into {profile_path}: " + "CSV columns changed" + ) + existing = [ + row + for row in existing_rows + if (row["token"], row["runtime_token"]) not in keys + ] + _write_rows( + profile_path, + list(PROFILE_FIELDS), + [*existing, *normalized], + durable=True, + ) + + +def _profile_commit_failure_message( + output_path: Path, + profile_path: Path, +) -> str: + return ( + f"profile rows failed to commit to {profile_path}; the current token " + f"was not committed to {output_path}" + ) + + +def _validate_complete_row_keys( + expected_rows: list[dict[str, str]], + current_rows: list[dict[str, str]], + fields: list[str], + source: Path, +) -> None: + """Require the full input key multiset while allowing kernel updates.""" + key_fields = [ + field for field in fields if field not in ("kernelName1", "kernelName2") + ] + + def key(row: dict[str, str]) -> tuple[str, ...]: + return tuple(row[field] for field in key_fields) + + expected_keys = Counter(key(row) for row in expected_rows) + current_keys = Counter(key(row) for row in current_rows) + if current_keys != expected_keys: + missing = sum((expected_keys - current_keys).values()) + extra = sum((current_keys - expected_keys).values()) + raise RuntimeError( + f"existing output row/key set differs from input for {source}: " + f"missing={missing}, extra={extra}; use the output as --input only " + "if it contains the complete table" + ) + + +def _merge_tuned_output_rows( + input_path: Path, + output_path: Path, + fields: list[str], + rows: list[dict[str, str]], + selected: dict[int, int], + baseline_rows: dict[int, dict[str, str]], + contract: Contract, +) -> list[dict[str, str]]: + """Build a compare-and-swap merge while the caller holds the output lock.""" + source = output_path if output_path.exists() else input_path + latest_fields, latest_rows = _read_rows(source) + if latest_fields != fields: + raise ValueError(f"cannot merge FHMoE rows into {source}: CSV columns changed") + _validate_complete_row_keys(rows, latest_rows, fields, source) + for token, tuned_index in selected.items(): + matches = [ + index + for index, row in enumerate(latest_rows) + if _matches(row, contract, token) + ] + if len(matches) != 1: + raise ValueError( + f"expected one current HY4 FHMoE row for M={token}, " + f"found {len(matches)}" + ) + latest_index = matches[0] + if latest_rows[latest_index] != baseline_rows[token]: + raise RuntimeError( + f"FHMoE row M={token} changed during tuning; refusing stale update" + ) + updated = dict(latest_rows[latest_index]) + updated["kernelName1"] = rows[tuned_index]["kernelName1"] + updated["kernelName2"] = rows[tuned_index]["kernelName2"] + latest_rows[latest_index] = updated + return latest_rows + + +def _write_tuned_output( + input_path: Path, + output_path: Path, + fields: list[str], + rows: list[dict[str, str]], + selected: dict[int, int], + baseline_rows: dict[int, dict[str, str]], + contract: Contract, +) -> None: + """Compare-and-swap selected rows under a lock.""" + if output_path.is_symlink(): + raise ValueError(f"FHMoE tuner output must not be a symlink: {output_path}") + with _output_lock(output_path): + latest_rows = _merge_tuned_output_rows( + input_path, + output_path, + fields, + rows, + selected, + baseline_rows, + contract, + ) + _write_rows(output_path, fields, latest_rows, durable=True) + + +def _write_profile_and_tuned_output( + input_path: Path, + output_path: Path, + profile_path: Path, + fields: list[str], + rows: list[dict[str, str]], + profile_rows: list[dict[str, str | float]], + selected: dict[int, int], + baseline_rows: dict[int, dict[str, str]], + contract: Contract, +) -> None: + """Commit profile before output while serializing stale-update validation.""" + if output_path.is_symlink(): + raise ValueError(f"FHMoE tuner output must not be a symlink: {output_path}") + with _output_lock(output_path): + latest_rows = _merge_tuned_output_rows( + input_path, + output_path, + fields, + rows, + selected, + baseline_rows, + contract, + ) + try: + _write_profile_rows(profile_path, profile_rows) + except Exception as error: + raise RuntimeError( + _profile_commit_failure_message(output_path, profile_path) + ) from error + _write_rows(output_path, fields, latest_rows, durable=True) + + +def _validate_existing_output( + output_path: Path, + fields: list[str], + input_rows: list[dict[str, str]], + baseline_rows: dict[int, dict[str, str]], + contract: Contract, +) -> None: + if not output_path.exists(): + return + output_fields, output_rows = _read_rows(output_path) + if output_fields != fields: + raise ValueError(f"existing output has different CSV columns: {output_path}") + _validate_complete_row_keys(input_rows, output_rows, fields, output_path) + for token, baseline in baseline_rows.items(): + matches = [row for row in output_rows if _matches(row, contract, token)] + if len(matches) != 1 or matches[0] != baseline: + raise RuntimeError( + f"existing output row M={token} differs from input; use the " + "output as --input to resume, or remove it" + ) + + +def _matches(row: dict[str, str], contract: Contract, token: int) -> bool: + return ( + row["gfx"] == contract.gfx + and int(row["cu_num"]) == contract.cu_num + and int(row["token"]) == token + and int(row["model_dim"]) == contract.model_dim + and int(row["inter_dim"]) == contract.inter_dim + and int(row["expert"]) == contract.experts + and int(row["topk"]) == contract.topk + and int(row["shared_expert_id"]) == contract.shared_expert_id + and row["act_type"] == contract.act_type + and row["dtype"] == contract.dtype + and row["q_dtype_a"] == contract.q_dtype_a + and row["q_dtype_w"] == contract.q_dtype_w + and row["q_type"] == contract.q_type + and int(row["use_g1u1"]) == contract.use_g1u1 + and int(row["doweight_stage1"]) == contract.doweight_stage1 + and int(row["hidden_pad"]) == contract.hidden_pad + and int(row["intermediate_pad"]) == contract.intermediate_pad + and row["gate_mode"] == contract.gate_mode + and int(row["ksplit"]) == contract.ksplit + ) + + +def _select_rows( + rows: list[dict[str, str]], contract: Contract, tokens: list[int] +) -> dict[int, int]: + selected: dict[int, int] = {} + requested = set(tokens) + for index, row in enumerate(rows): + token = int(row["token"]) + if token not in requested or not _matches(row, contract, token): + continue + if token in selected: + raise ValueError( + f"multiple HY4 FHMoE rows match the complete contract for M={token}" + ) + selected[token] = index + return selected + + +def _effective_sort_block(name: str) -> int: + params = get_flydsl_kernel_params(name) + if params is None: + return -1 + return int(params.get("sort_block_m", 0) or params["tile_m"]) + + +def stage1_candidates(block_m: int, quick: bool = False) -> list[str]: + names: set[str] = set() + for name, params in get_flydsl_stage1_kernels("fp8", "fp8", "bf16").items(): + if ( + int(params["tile_m"]) != block_m + or int(params["tile_k"]) != 256 + or 256 % int(params["tile_n"]) != 0 + or int(params.get("k_batch", 1)) != 1 + or params.get("gate_mode") != "interleave" + ): + continue + candidate = f"{name}_fp8" + if get_flydsl_kernel_params(candidate) is None: + continue + if quick and ( + int(params["tile_n"]) not in (64, 128) + or int(params.get("k_wave", 1)) != 1 + or int(params.get("xcd_swizzle", 0)) != 0 + ): + continue + names.add(candidate) + return sorted(names) + + +def stage2_candidates(block_m: int, quick: bool = False) -> list[str]: + names: set[str] = set() + for name, params in get_flydsl_stage2_kernels("fp8", "fp8", "bf16").items(): + tile_m = int(params["tile_m"]) + if ( + block_m % tile_m != 0 + or 256 % int(params["tile_k"]) != 0 + or int(params["tile_n"]) != 128 + ): + continue + candidate = name if tile_m == block_m else f"{name}_sbm{block_m}" + parsed = get_flydsl_kernel_params(candidate) + if ( + parsed is None + or _effective_sort_block(candidate) != block_m + or parsed.get("persist") is True + or int(parsed.get("b_nt", 0)) != 0 + ): + continue + if quick and int(parsed.get("xcd_swizzle", 0)) != 0: + continue + names.add(candidate) + return sorted(names) + + +def _random_fp8( + shape: tuple[int, ...], generator: torch.Generator, device: torch.device +) -> torch.Tensor: + return ( + torch.randn(shape, dtype=torch.bfloat16, device=device, generator=generator) + .mul_(0.02) + .to(torch.float8_e4m3fn) + ) + + +def _varying_e8m0_scales( + shape: tuple[int, int, int], device: torch.device +) -> torch.Tensor: + """Generate deterministic 0.25/0.5/1/2 scales across E, N, and K groups.""" + experts, rows, groups = shape + expert_index = torch.arange(experts, device=device).remainder_(4).to(torch.uint8) + row_index = torch.arange(rows, device=device).remainder_(4).to(torch.uint8) + group_index = torch.arange(groups, device=device).remainder_(4).to(torch.uint8) + pattern = ( + expert_index[:, None, None] + + row_index[None, :, None] + + group_index[None, None, :] + ).remainder_(4) + return pattern.add_(0x7D) + + +def _sample_rows(token: int) -> tuple[int, ...]: + """Choose cheap representative rows without weakening the full-M timing run. + + The first 32 rows collectively route through all 256 HY4 experts. The + midpoint and final row additionally cover distant sorting positions. + """ + candidates = (*range(min(token, 32)), token // 2, token - 1) + return tuple(sorted({row for row in candidates if 0 <= row < token})) + + +def _sample_expert_ids(contract: Contract, rows: tuple[int, ...]) -> tuple[int, ...]: + """Return only the routed weights needed by the sampled token rows.""" + return tuple( + sorted( + { + (row * contract.routed_topk + slot) % contract.routed_experts + for row in rows + for slot in range(contract.routed_topk) + } + ) + ) + + +def _make_tensors( + contract: Contract, token: int, seed: int, device: torch.device +) -> Tensors: + generator = torch.Generator(device=device).manual_seed(seed) + e, h, i = contract.experts, contract.model_dim, contract.inter_dim + + raw_w1 = _random_fp8((e, 2 * i, h), generator, device) + raw_w2 = _random_fp8((e, h, i), generator, device) + raw_s1 = _varying_e8m0_scales((e, 2 * i, h // 32), device) + raw_s2 = _varying_e8m0_scales((e, h, i // 32), device) + sample_rows = _sample_rows(token) + sample_expert_ids = _sample_expert_ids(contract, sample_rows) + sample_indices = torch.tensor(sample_expert_ids, dtype=torch.long, device=device) + sample_w1 = raw_w1.index_select(0, sample_indices).clone() + sample_w2 = raw_w2.index_select(0, sample_indices).clone() + sample_w1_scale = raw_s1.index_select(0, sample_indices).clone() + sample_w2_scale = raw_s2.index_select(0, sample_indices).clone() + w1 = shuffle_weight_a16w4(raw_w1, 16, True).contiguous() + w2 = shuffle_weight_a16w4(raw_w2, 16, False).contiguous() + w1_scale = shuffle_scale_a16w4( + raw_s1.view(-1, raw_s1.shape[-1]).view(dtypes.fp8_e8m0), e, True + ).contiguous() + w2_scale = shuffle_scale( + raw_s2.view(-1, raw_s2.shape[-1]).view(dtypes.fp8_e8m0) + ).contiguous() + del raw_w1, raw_w2, raw_s1, raw_s2 + + shared_w1 = _random_fp8((1, 2 * i, h), generator, device) + shared_w2 = _random_fp8((1, h, i), generator, device) + shared_s1 = _varying_e8m0_scales((1, 2 * i, h // 32), device) + shared_s2 = _varying_e8m0_scales((1, h, i // 32), device) + sample_shared_w1 = shared_w1.clone() + sample_shared_w2 = shared_w2.clone() + sample_shared_w1_scale = shared_s1.clone() + sample_shared_w2_scale = shared_s2.clone() + shared_w1 = shuffle_weight_a16w4(shared_w1, 16, True).contiguous() + shared_w2 = shuffle_weight_a16w4(shared_w2, 16, False).contiguous() + shared_w1_scale = shuffle_scale_a16w4( + shared_s1.view(-1, shared_s1.shape[-1]).view(dtypes.fp8_e8m0), 1, True + ).contiguous() + shared_w2_scale = shuffle_scale( + shared_s2.view(-1, shared_s2.shape[-1]).view(dtypes.fp8_e8m0) + ).contiguous() + + hidden = torch.randn( + (token, h), dtype=torch.bfloat16, device=device, generator=generator + ) + row = torch.arange(token, device=device, dtype=torch.int32)[:, None] + slot = torch.arange(contract.routed_topk, device=device, dtype=torch.int32)[None, :] + routed_ids = (row * contract.routed_topk + slot) % contract.routed_experts + routed_weights = torch.arange( + 1, + contract.routed_topk + 1, + dtype=torch.float32, + device=device, + ).repeat(token, 1) + routed_weights.mul_(2.5 / routed_weights.sum(dim=1, keepdim=True)) + shared_ids = torch.full( + (token, 1), contract.shared_expert_id, dtype=torch.int32, device=device + ) + shared_weights = torch.ones((token, 1), dtype=torch.float32, device=device) + return Tensors( + hidden=hidden, + topk_weights=torch.cat((routed_weights, shared_weights), dim=1), + topk_ids=torch.cat((routed_ids, shared_ids), dim=1), + w1=w1, + w2=w2, + w1_scale=w1_scale, + w2_scale=w2_scale, + shared_w1=shared_w1, + shared_w2=shared_w2, + shared_w1_scale=shared_w1_scale, + shared_w2_scale=shared_w2_scale, + output=torch.empty((token, h), dtype=torch.bfloat16, device=device), + sample_rows=sample_rows, + sample_expert_ids=sample_expert_ids, + sample_w1=sample_w1, + sample_w2=sample_w2, + sample_w1_scale=sample_w1_scale, + sample_w2_scale=sample_w2_scale, + sample_shared_w1=sample_shared_w1, + sample_shared_w2=sample_shared_w2, + sample_shared_w1_scale=sample_shared_w1_scale, + sample_shared_w2_scale=sample_shared_w2_scale, + ) + + +def _candidate_row(base: dict[str, str], stage1: str, stage2: str) -> dict[str, str]: + row = dict(base) + row["kernelName1"] = stage1 + row["kernelName2"] = stage2 + return row + + +def _write_candidate(path: Path, row: dict[str, str]) -> None: + _write_rows(path, list(CSV_FIELDS), [row]) + get_2stage_cfgs.cache_clear() + cfg_2stages_by_file.clear() + + +def _runner( + tensors: Tensors, contract: Contract, config_file: Path, swiglu_limit: float +): + # This tuner targets HY4's model-level unclamped shared SiLU semantics. + clamp_shared = False + + def run(): + return _fused_moe_impl( + hidden_states=tensors.hidden, + w1=tensors.w1, + w2=tensors.w2, + topk_weight=tensors.topk_weights, + topk_ids=tensors.topk_ids, + activation=aiter.ActivationType.Silu.value, + quant_type=aiter.QuantType.per_1x32.value, + w1_scale=tensors.w1_scale, + w2_scale=tensors.w2_scale, + dtype=torch.bfloat16, + swiglu_limit=swiglu_limit, + gate_mode=GateMode.INTERLEAVE.value, + output=tensors.output, + _q_dtype_a=dtypes.fp8, + _metadata_transform=_use_fhmoe_wrappers, + _metadata_config_file=str(config_file), + _stage1_extra_args={ + "shared_w1": tensors.shared_w1, + "shared_w1_scale": tensors.shared_w1_scale, + "shared_expert_id": contract.shared_expert_id, + "swiglu_limit": swiglu_limit, + "clamp_shared": clamp_shared, + }, + _stage2_extra_args={ + "shared_w2": tensors.shared_w2, + "shared_w2_scale": tensors.shared_w2_scale, + "shared_expert_id": contract.shared_expert_id, + }, + ) + + return run + + +def _time_candidate( + run, + warmup_ms: float, + rep_ms: float, + stability_tolerance: float = 8e-3, +) -> tuple[float, torch.Tensor]: + result = run() + torch.accelerator.synchronize() + if not torch.isfinite(result).all(): + return float("inf"), result + reference = result.detach().clone() + timing = float(triton.testing.do_bench(run, warmup=warmup_ms, rep=rep_ms)) + for _ in range(3): + verification = run() + torch.accelerator.synchronize() + if not torch.isfinite(verification).all(): + return float("inf"), verification + if _relative_l2(verification, reference) > stability_tolerance: + return float("inf"), verification + return timing, reference + + +def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + denominator = torch.linalg.vector_norm(expected.float()).clamp_min(1e-12) + return float( + torch.linalg.vector_norm(actual.float() - expected.float()) / denominator + ) + + +def _fp8_group_quant_dequant( + x: torch.Tensor, *, fused_intermediate: bool = False +) -> torch.Tensor: + shape = x.shape + blocks = x.float().view(-1, 32) + amax = blocks.abs().amax(dim=1) + scale = ( + fp4_utils.f32_to_fused_moe_mxfp8_scale(amax) + if fused_intermediate + else fp4_utils.f32_to_mx_e8m0_scale(amax, dtype=MxDtypeInt.FP8_E4M3) + ) + scale_f32 = fp4_utils.e8m0_to_f32(scale).view(-1, 1) + quant = (blocks / scale_f32).to(torch.float8_e4m3fn) + return (quant.float() * scale_f32).view(shape) + + +def _dequant_fp8_weight(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + scale_f32 = fp4_utils.e8m0_to_f32(scale.view(dtypes.fp8_e8m0)) + return weight.float() * scale_f32.repeat_interleave(32, dim=-1) + + +def _routed_silu( + gate: torch.Tensor, up: torch.Tensor, swiglu_limit: float +) -> torch.Tensor: + activation_limit = runtime_swiglu_limit(swiglu_limit, "silu") + gate = gate.clamp(max=activation_limit) + up = up.clamp(min=-activation_limit, max=activation_limit) + return F.silu(gate) * up + + +@torch.no_grad() +def _sampled_fp32_reference( + tensors: Tensors, + contract: Contract, + swiglu_limit: float, +) -> torch.Tensor: + """Compute an independent routed-plus-shared FP32 oracle for sampled rows. + + This reproduces MXFP8 activation/weight dequantization and the model's BF16 + rounding points, but performs the GEMMs and SiLU in FP32. Work scales with + the number of sampled rows and their experts rather than full M. + """ + rows = torch.tensor( + tensors.sample_rows, dtype=torch.long, device=tensors.hidden.device + ) + hidden = _fp8_group_quant_dequant(tensors.hidden.index_select(0, rows).float()) + routed_ids = tensors.topk_ids.index_select(0, rows)[:, : contract.routed_topk] + routed_weights = tensors.topk_weights.index_select(0, rows)[ + :, : contract.routed_topk + ] + expanded_hidden = hidden[:, None, :].expand(-1, contract.routed_topk, -1) + slot_output = torch.zeros( + (*routed_ids.shape, contract.model_dim), + dtype=torch.float32, + device=hidden.device, + ) + expert_storage = { + expert_id: index for index, expert_id in enumerate(tensors.sample_expert_ids) + } + + for expert_id in torch.unique(routed_ids).tolist(): + mask = routed_ids == expert_id + storage = expert_storage[int(expert_id)] + w1 = _dequant_fp8_weight( + tensors.sample_w1[storage], tensors.sample_w1_scale[storage] + ) + gate, up = F.linear(expanded_hidden[mask], w1).chunk(2, dim=-1) + intermediate = _routed_silu(gate, up, swiglu_limit) + intermediate = _fp8_group_quant_dequant(intermediate, fused_intermediate=True) + w2 = _dequant_fp8_weight( + tensors.sample_w2[storage], tensors.sample_w2_scale[storage] + ) + slot_output[mask] = F.linear(intermediate, w2) + + routed_output = (slot_output * routed_weights[..., None]).sum(dim=1) + shared_w1 = _dequant_fp8_weight( + tensors.sample_shared_w1[0], tensors.sample_shared_w1_scale[0] + ) + shared_gate_up = F.linear(hidden, shared_w1).to(torch.bfloat16).float() + shared_gate, shared_up = shared_gate_up.chunk(2, dim=-1) + shared_intermediate = F.silu(shared_gate) * shared_up + shared_intermediate = _fp8_group_quant_dequant( + shared_intermediate.to(torch.bfloat16).float(), + fused_intermediate=True, + ) + shared_w2 = _dequant_fp8_weight( + tensors.sample_shared_w2[0], tensors.sample_shared_w2_scale[0] + ) + shared_output = F.linear(shared_intermediate, shared_w2) + shared_weight = tensors.topk_weights.index_select(0, rows)[:, contract.routed_topk] + return routed_output + shared_output * shared_weight[:, None] + + +def tune_row( + row: dict[str, str], + tensors: Tensors, + contract: Contract, + config_file: Path, + *, + quick: bool, + warmup_ms: float, + rep_ms: float, + tolerance: float, + swiglu_limit: float, + min_improvement_pct: float, + stability_tolerance: float = 8e-3, + baseline_only: bool = False, +) -> tuple[dict[str, str], list[dict[str, str | float]]]: + block_m = int(row["block_m"]) + current_s1, current_s2 = row["kernelName1"], row["kernelName2"] + profile: list[dict[str, str | float]] = [] + current_s1_params = get_flydsl_kernel_params(current_s1) + current_s2_params = get_flydsl_kernel_params(current_s2) + if current_s1_params is not None and ( + int(current_s1_params["tile_m"]) != block_m + or int(current_s1_params.get("k_batch", 1)) != 1 + or current_s1_params.get("gate_mode") != "interleave" + ): + raise ValueError( + f"invalid baseline stage1 kernel for block_m={block_m}: {current_s1}" + ) + if current_s2_params is not None and ( + _effective_sort_block(current_s2) != block_m + or current_s2_params.get("persist") is True + or int(current_s2_params.get("b_nt", 0)) != 0 + ): + raise ValueError( + f"invalid baseline stage2 kernel for block_m={block_m}: {current_s2}" + ) + + _write_candidate(config_file, row) + current_run = _runner(tensors, contract, config_file, swiglu_limit) + current_ms, reference = _time_candidate( + current_run, warmup_ms, rep_ms, stability_tolerance + ) + if not torch.isfinite(reference).all(): + raise RuntimeError( + f"M={row['token']}: baseline FHMoE kernels produced non-finite output: " + f"{current_s1}, {current_s2}" + ) + if not math.isfinite(current_ms): + raise RuntimeError( + f"M={row['token']}: baseline FHMoE kernels produced an unstable " + f"or non-finite timing: {current_s1}, {current_s2}" + ) + profile.append( + {"stage": "baseline", "kernel": f"{current_s1}|{current_s2}", "ms": current_ms} + ) + if baseline_only: + print( + f"M={row['token']}: baseline {current_ms:.4f} ms " + f"({current_s1}, {current_s2})", + flush=True, + ) + return row, profile + + def benchmark_candidate( + stage: str, + kernel: str, + candidate_row: dict[str, str], + ) -> tuple[float, float]: + try: + _write_candidate(config_file, candidate_row) + run = _runner(tensors, contract, config_file, swiglu_limit) + ms, output = _time_candidate(run, warmup_ms, rep_ms, stability_tolerance) + error = _relative_l2(output, reference) + except Exception as candidate_error: # noqa: BLE001 + ms = float("inf") + error = float("inf") + print( + f"M={row['token']}: rejected {stage} candidate {kernel!r}: " + f"{candidate_error}", + flush=True, + ) + profile.append({"stage": stage, "kernel": kernel, "ms": ms, "error": error}) + return ms, error + + best_s2, best_ms = current_s2, current_ms + stage2_scores: list[tuple[float, str]] = [] + for candidate in sorted(set(stage2_candidates(block_m, quick)) | {current_s2}): + candidate_row = _candidate_row(row, current_s1, candidate) + ms, error = benchmark_candidate( + "stage2", + candidate, + candidate_row, + ) + if error <= tolerance and math.isfinite(ms): + stage2_scores.append((ms, candidate)) + if error <= tolerance and ms < best_ms: + best_s2, best_ms = candidate, ms + + best_s1 = current_s1 + stage1_scores: list[tuple[float, str]] = [] + for candidate in sorted(set(stage1_candidates(block_m, quick)) | {current_s1}): + candidate_row = _candidate_row(row, candidate, best_s2) + ms, error = benchmark_candidate( + "stage1", + candidate, + candidate_row, + ) + if error <= tolerance and math.isfinite(ms): + stage1_scores.append((ms, candidate)) + if error <= tolerance and ms < best_ms: + best_s1, best_ms = candidate, ms + + measured_pairs = {(current_s1, candidate) for _, candidate in stage2_scores} | { + (candidate, best_s2) for _, candidate in stage1_scores + } + pair_stage1 = sorted(stage1_scores) + pair_stage2 = sorted(stage2_scores) + if quick: + pair_stage1 = pair_stage1[:_PAIR_BEAM_WIDTH] + pair_stage2 = pair_stage2[:_PAIR_BEAM_WIDTH] + for _, stage1 in pair_stage1: + for _, stage2 in pair_stage2: + if (stage1, stage2) in measured_pairs: + continue + candidate_row = _candidate_row(row, stage1, stage2) + kernel = f"{stage1}|{stage2}" + ms, error = benchmark_candidate( + "pair", + kernel, + candidate_row, + ) + if error <= tolerance and ms < best_ms: + best_s1, best_s2, best_ms = stage1, stage2, ms + + improvement_pct = ( + 100.0 * (current_ms - best_ms) / current_ms if current_ms > 0 else 0.0 + ) + if improvement_pct < min_improvement_pct: + best_s1, best_s2, best_ms = current_s1, current_s2, current_ms + tuned = _candidate_row(row, best_s1, best_s2) + print( + f"M={row['token']}: {current_ms:.4f} ms -> {best_ms:.4f} ms " + f"({current_s1}, {current_s2}) -> ({best_s1}, {best_s2})", + flush=True, + ) + return tuned, profile + + +def _validate_sampled_reference( + row: dict[str, str], + tensors: Tensors, + contract: Contract, + config_file: Path, + *, + swiglu_limit: float, + tolerance: float, +) -> float: + """Accept the full-M winner only when sampled rows match the FP32 oracle.""" + _write_candidate(config_file, row) + actual = _runner(tensors, contract, config_file, swiglu_limit)() + torch.accelerator.synchronize() + rows = torch.tensor( + tensors.sample_rows, dtype=torch.long, device=tensors.hidden.device + ) + actual_sample = actual.index_select(0, rows).float() + expected = _sampled_fp32_reference(tensors, contract, swiglu_limit) + if not torch.isfinite(actual_sample).all() or not torch.isfinite(expected).all(): + raise RuntimeError( + f"M={row['token']}: sampled FP32 validation produced non-finite output" + ) + error = _relative_l2(actual_sample, expected) + if error > tolerance: + raise RuntimeError( + f"M={row['token']}: sampled FP32 validation failed with " + f"relative L2 {error:.6f} > {tolerance:.6f}" + ) + print( + f"M={row['token']}: sampled FP32 validation passed for " + f"{len(tensors.sample_rows)} rows (relative L2={error:.6f})", + flush=True, + ) + return error + + +def _validate_tuning_environment() -> None: + for name in ("AITER_FLYDSL_FORCE_REDUCE", "AITER_FLYDSL_STAGE2_FP8"): + if os.environ.get(name, "0") not in ("", "0"): + raise RuntimeError( + f"FHMoE canonical tuning requires {name}=0; environment " + "overrides would benchmark a different implementation" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("-i", "--input", type=Path, required=True) + parser.add_argument("-o", "--output", type=Path, required=True) + parser.add_argument( + "--tokens", + type=_parse_tokens, + default=_parse_tokens( + "1,2,4,8,16,32,64,128,256,512,1024,2048,4096,8192,16384,32768" + ), + ) + parser.add_argument( + "--benchmark-token", + type=int, + help=( + f"actual M (1..{HY4_FHMOE_MAX_TOKENS}) used to tune exactly one " + "selected CSV bucket; " + "it must resolve to that bucket, for example --tokens 32768 " + f"--benchmark-token {HY4_FHMOE_MAX_TOKENS}" + ), + ) + parser.add_argument( + "--warmup-ms", "--warmup", type=_parse_positive_float, default=30.0 + ) + parser.add_argument( + "--rep-ms", "--iters", type=_parse_positive_float, default=120.0 + ) + parser.add_argument( + "--mp", + type=int, + default=1, + help="accepted for tuner-infrastructure compatibility; FHMoE tunes one GPU", + ) + parser.add_argument( + "--tolerance", + type=_parse_nonnegative_float, + default=8e-3, + help="maximum baseline-relative candidate-output L2 error", + ) + parser.add_argument( + "--stability-tolerance", + type=_parse_nonnegative_float, + default=8e-3, + help="maximum repeat-output relative L2 variation (not timing jitter)", + ) + parser.add_argument( + "--sampled-reference-tolerance", + type=_parse_nonnegative_float, + default=4e-2, + ) + parser.add_argument( + "--min-improvement-pct", type=_parse_nonnegative_float, default=1.0 + ) + parser.add_argument("--swiglu-limit", type=_parse_nonnegative_float, default=10.0) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument("--quick", action="store_true") + parser.add_argument( + "--baseline-only", + action="store_true", + help="Benchmark current CSV rows without searching or replacing kernels.", + ) + parser.add_argument( + "--profile", + type=Path, + help=( + "incremental CSV profile with columns " + "token,runtime_token,stage,kernel,ms,error" + ), + ) + args = parser.parse_args() + + _validate_output_paths(args.input, args.output, args.profile) + _validate_tuning_environment() + if args.mp != 1: + raise ValueError("FHMoE tuning currently supports --mp 1 only") + _validate_benchmark_token(args.benchmark_token, args.tokens) + if get_gfx_runtime() != "gfx950": + raise RuntimeError("FHMoE tuning currently requires gfx950") + contract = Contract() + cu_num = get_cu_num() + if cu_num != contract.cu_num: + raise RuntimeError(f"FHMoE tuning requires {contract.cu_num} CUs, got {cu_num}") + fields, rows = _read_rows(args.input) + if fields != list(CSV_FIELDS): + raise ValueError(f"unexpected FHMoE columns: {fields}") + + selected = _select_rows(rows, contract, args.tokens) + missing = sorted(set(args.tokens) - selected.keys()) + if missing: + raise ValueError(f"missing HY4 FHMoE rows for tokens: {missing}") + baseline_rows = {token: dict(rows[index]) for token, index in selected.items()} + if not args.baseline_only and not _paths_alias(args.input, args.output): + _validate_existing_output(args.output, fields, rows, baseline_rows, contract) + + device = torch.device("cuda") + with tempfile.TemporaryDirectory(prefix="fhmoe-tune-") as directory: + candidate_file = Path(directory) / "candidate.csv" + for token in args.tokens: + index = selected[token] + benchmark_token = ( + args.benchmark_token if args.benchmark_token is not None else token + ) + if benchmark_token != token: + print( + f"tuning CSV bucket M={token} with runtime M={benchmark_token}", + flush=True, + ) + tensors = _make_tensors( + contract, + benchmark_token, + args.seed + benchmark_token, + device, + ) + tuned, row_profile = tune_row( + rows[index], + tensors, + contract, + candidate_file, + quick=args.quick, + warmup_ms=args.warmup_ms, + rep_ms=args.rep_ms, + tolerance=args.tolerance, + swiglu_limit=args.swiglu_limit, + min_improvement_pct=args.min_improvement_pct, + stability_tolerance=args.stability_tolerance, + baseline_only=args.baseline_only, + ) + # Candidate ranking above compares complete full-M outputs for speed. + # Run the independent FP32 oracle only once, on the selected winner. + _validate_sampled_reference( + tuned, + tensors, + contract, + candidate_file, + swiglu_limit=args.swiglu_limit, + tolerance=args.sampled_reference_tolerance, + ) + rows[index] = tuned + for profile_row in row_profile: + profile_row["token"] = token + profile_row["runtime_token"] = benchmark_token + selected_row = {token: index} + selected_baseline = {token: baseline_rows[token]} + if not args.baseline_only and args.profile: + _write_profile_and_tuned_output( + args.input, + args.output, + args.profile, + fields, + rows, + row_profile, + selected_row, + selected_baseline, + contract, + ) + elif args.profile: + try: + _write_profile_rows(args.profile, row_profile) + except Exception as error: + raise RuntimeError( + _profile_commit_failure_message( + args.output, + args.profile, + ) + ) from error + elif not args.baseline_only: + _write_tuned_output( + args.input, + args.output, + fields, + rows, + selected_row, + selected_baseline, + contract, + ) + del tensors + torch.cuda.empty_cache() + + if args.baseline_only: + _write_tuned_output( + args.input, + args.output, + fields, + rows, + {}, + baseline_rows, + contract, + ) + print( + f"processed {len(selected)} HY4 FHMoE rows and wrote {args.output} " + f"on {get_gfx_runtime()} ({get_cu_num()} CUs)", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/tuning_tests/README.md b/op_tests/tuning_tests/README.md index 0029a54514f..3325f07cd6a 100644 --- a/op_tests/tuning_tests/README.md +++ b/op_tests/tuning_tests/README.md @@ -38,6 +38,7 @@ runtime budgets. | `batched_a8w8` | `csrc/ck_batched_gemm_a8w8/batched_gemm_a8w8_tune.py` | `a8w8_tuned_batched_gemm.csv` | ✓ | ✓ | | `batched_bf16` | `csrc/ck_batched_gemm_bf16/batched_gemm_bf16_tune.py` | `bf16_tuned_batched_gemm.csv` | ✓ | ✓ + shape_grouped | | `fmoe` | `csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py` | `tuned_fmoe.csv` + model_configs | ✓ | ✓ (bf16/fp8/int8/gelu) | +| `fhmoe` | `op_tests/tuners/tune_fhmoe.py` | `tuned_fhmoe.csv` | — | ✓ (HY4 MXFP8/FP8) | | `gradlib_bf16` | `gradlib/gradlib/gemm_tuner.py` | `bf16_tuned_gemm.csv` | ✓ | ✓ (hipBLASLt/ASM/FlyDSL) | | `gdn_k5_opt` | `csrc/gdn_k5/chunk_gdn_h_opt_tune.py` | `model_configs/*_chunk_gdn_h_opt_tuned.csv` | ✓ | ✓ (shape-only varlen smoke) | @@ -65,6 +66,7 @@ python3 -m unittest op_tests.tuning_tests.test_csv_validation \ op_tests.tuning_tests.test_mp_tuner_logic \ op_tests.tuning_tests.test_mixed_mxfp_tuning \ op_tests.tuning_tests.test_online_tune -v +python3 -m pytest op_tests/tuning_tests/test_fhmoe_tuner.py -v # Level 2: pipeline smoke (~10min) python3 -m unittest op_tests.tuning_tests.test_tune_pipeline -v @@ -73,12 +75,78 @@ python3 -m unittest op_tests.tuning_tests.test_tune_pipeline -v python3 -m unittest op_tests.tuning_tests.test_run_config -v # Everything -python3 -m unittest discover -s op_tests/tuning_tests -v +python3 -m pytest op_tests/tuning_tests -v +``` + +### Tuning HY4 heterogeneous MoE + +```bash +python op_tests/tuners/tune_fhmoe.py \ + -i aiter/configs/tuned_fhmoe.csv \ + -o aiter/configs/tuned_fhmoe.csv \ + --tokens 1,2,4,8,16,32,64,128,256,512,1024,2048,4096,8192,16384,32768 \ + --profile /tmp/hy4_fhmoe_profile.csv +``` + +Add `--quick` for the bounded pair-beam search used by tuning pipeline smoke +tests. Quick mode intentionally excludes nonzero XCD-swizzle candidates and +does not prove the production optimum. Without it, the tuner evaluates the full +valid stage-1 x stage-2 product; the default all-bucket sweep can take hours, so +exhaustive tuning should normally target selected `--tokens`. The tuner +optimizes the isolated FHMoE call; only commit changed rows after a matched +serving-level A/B test because full-model scheduling can favor a different +kernel. + +For a floor bucket that serves larger runtime shapes, keep the CSV key in +`--tokens` and provide the representative runtime M separately. For example: + +```bash +python op_tests/tuners/tune_fhmoe.py \ + -i aiter/configs/tuned_fhmoe.csv \ + -o aiter/configs/tuned_fhmoe.csv \ + --tokens 32768 \ + --benchmark-token 38836 +``` + +`--benchmark-token` requires exactly one `--tokens` bucket, must be positive +and no greater than 38836, and must resolve to the selected bucket. HY4 uses +the global power-of-two tiers below 32768 and the 32768 row through its safety +limit. Larger values are rejected until reduce output supports addressing +beyond 4 GiB. + +Candidate error checks are baseline-relative regression guards, not independent +oracles. Before updating each CSV row, the final selected full-M output is also +checked against an independent FP32 reference for representative token rows +covering every routed expert for M>=32 (and every expert reachable by smaller +buckets) plus the shared expert. Use +`--sampled-reference-tolerance` to adjust that final acceptance threshold. +`--tolerance` controls baseline-relative candidate accuracy, while +`--stability-tolerance` independently bounds repeat-output relative L2 +(default `8e-3`); it is not a timing-jitter threshold. + +Profile CSV columns are +`token,runtime_token,stage,kernel,ms,error`. Rows are committed incrementally +and replaced by `(token, runtime_token)`, so interrupted or resumed runs retain +completed measurements. Each token's profile rows are committed before its +tuned output row, so a profile-write failure leaves that token uncommitted and +safe to retry. Legacy five-column profiles are migrated with `runtime_token=token`. + +When input and output are different files, resume from the partial output, not +the original input, and select only unfinished buckets: + +```bash +python op_tests/tuners/tune_fhmoe.py \ + -i /tmp/partial_tuned_fhmoe.csv \ + -o /tmp/partial_tuned_fhmoe.csv \ + --tokens \ + --profile /tmp/hy4_fhmoe_profile.csv ``` ### Running individual tuner tests -Each tuner in `test_tune_pipeline.py` has two variants: `_mp1` (single GPU) and `_mp_default` (all GPUs). +Most tuners in `test_tune_pipeline.py` have `_mp1` (single GPU) and +`_mp_default` (all GPUs) variants. FHMoE is gfx950-only and defines `_mp1` +because its tuner intentionally runs on one GPU. ```bash # Run a specific tuner (both mp1 and mp_default) diff --git a/op_tests/tuning_tests/test_config_shape_collision.py b/op_tests/tuning_tests/test_config_shape_collision.py index 3bbd3777b1f..c7c52af3b17 100644 --- a/op_tests/tuning_tests/test_config_shape_collision.py +++ b/op_tests/tuning_tests/test_config_shape_collision.py @@ -393,10 +393,73 @@ def test_fhmoe(self): ) with open(merged, newline="") as f: rows = list(csv.DictReader(f)) + dsv4_rows = [ + row + for row in rows + if ( + row["model_dim"], + row["inter_dim"], + row["expert"], + row["topk"], + row["shared_expert_id"], + row["q_dtype_w"], + ) + == ( + "7168", + "384", + "385", + "7", + "384", + "torch.float4_e2m1fn_x2", + ) + ] + hy4_rows = [ + row + for row in rows + if ( + row["model_dim"], + row["inter_dim"], + row["expert"], + row["topk"], + row["shared_expert_id"], + row["q_dtype_w"], + ) + == ( + "6144", + "256", + "257", + "9", + "256", + "torch.float8_e4m3fn", + ) + ] self.assertEqual( - {int(row["token"]) for row in rows}, + {int(row["token"]) for row in dsv4_rows}, {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048}, ) + self.assertEqual( + {int(row["token"]) for row in hy4_rows}, + { + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + 32768, + }, + ) + self.assertEqual(len(dsv4_rows), 12) + self.assertEqual(len(hy4_rows), 16) def test_grouped_fmoe(self): self._check_family("AITER_CONFIG_GROUPED_FMOE", "tuned_grouped_fmoe") diff --git a/op_tests/tuning_tests/test_fhmoe_tuner.py b/op_tests/tuning_tests/test_fhmoe_tuner.py new file mode 100644 index 00000000000..afcb0df96c0 --- /dev/null +++ b/op_tests/tuning_tests/test_fhmoe_tuner.py @@ -0,0 +1,1202 @@ +# SPDX-License-Identifier: MIT + +import argparse +import multiprocessing +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from aiter.utility import fp4_utils +from op_tests.tuners import tune_fhmoe +from op_tests.tuners.tune_fhmoe import ( + PROFILE_FIELDS, + Contract, + _fp8_group_quant_dequant, + _matches, + _parse_nonnegative_float, + _parse_positive_float, + _parse_tokens, + _profile_commit_failure_message, + _read_rows, + _routed_silu, + _sample_expert_ids, + _sample_rows, + _select_rows, + _time_candidate, + _validate_benchmark_token, + _validate_existing_output, + _validate_output_paths, + _validate_sampled_reference, + _validate_tuning_environment, + _varying_e8m0_scales, + _write_profile_rows, + _write_rows, + _write_tuned_output, + stage1_candidates, + stage2_candidates, +) + + +def _concurrent_tuned_output_update( + barrier, + input_path, + output_path, + fields, + original, + token, + index, + kernel, +): + stale = [dict(row) for row in original] + stale[index]["kernelName1"] = kernel + barrier.wait(timeout=60) + _write_tuned_output( + input_path, + output_path, + fields, + stale, + {token: index}, + {token: dict(original[index])}, + Contract(), + ) + + +def test_parse_tokens(): + assert _parse_tokens("32,1,32,16") == [1, 16, 32] + with pytest.raises(argparse.ArgumentTypeError): + _parse_tokens("0") + + +@pytest.mark.parametrize("value", ("nan", "inf", "-1")) +def test_nonnegative_float_rejects_invalid_values(value: str): + with pytest.raises(argparse.ArgumentTypeError, match="finite and nonnegative"): + _parse_nonnegative_float(value) + + +@pytest.mark.parametrize("value", ("nan", "inf", "0", "-1")) +def test_positive_float_rejects_invalid_values(value: str): + with pytest.raises(argparse.ArgumentTypeError, match="finite and positive"): + _parse_positive_float(value) + + +def test_benchmark_token_respects_hy4_limit(): + _validate_benchmark_token(38836, [32768]) + with pytest.raises(ValueError, match="must not exceed 38836"): + _validate_benchmark_token(38837, [32768]) + _validate_benchmark_token(None, [32768]) + with pytest.raises(ValueError, match="tokens must not exceed 38836"): + _validate_benchmark_token(None, [32768, 131072]) + + +@pytest.mark.parametrize( + "name", ("AITER_FLYDSL_FORCE_REDUCE", "AITER_FLYDSL_STAGE2_FP8") +) +def test_canonical_tuning_rejects_runtime_overrides( + monkeypatch: pytest.MonkeyPatch, name: str +): + monkeypatch.setenv(name, "1") + with pytest.raises(RuntimeError, match=f"{name}=0"): + _validate_tuning_environment() + + +def test_large_m_samples_cover_boundary_experts(): + rows = _sample_rows(38836) + assert rows[:32] == tuple(range(32)) + assert rows[-2:] == (19418, 38835) + expert_ids = _sample_expert_ids(Contract(), rows) + assert expert_ids == tuple(range(256)) + + +def test_reference_scales_vary_across_all_weight_axes(): + scales = _varying_e8m0_scales((2, 3, 4), torch.device("cpu")) + assert set(scales.unique().tolist()) == {0x7D, 0x7E, 0x7F, 0x80} + assert scales[1, 0, 0] != scales[0, 0, 0] + assert scales[0, 1, 0] != scales[0, 0, 0] + assert scales[0, 0, 1] != scales[0, 0, 0] + + +def test_zero_swiglu_limit_keeps_routed_reference_unclamped(): + gate = torch.tensor([[2.0, -2.0]]) + up = torch.tensor([[3.0, -3.0]]) + expected = torch.nn.functional.silu(gate) * up + assert torch.equal(_routed_silu(gate, up, 0), expected) + + +def test_sampled_reference_uses_fused_moe_mxfp8_rounding(): + amax = torch.tensor([383.0, 384.0, 400.0, 448.0]) + scales = fp4_utils.f32_to_fused_moe_mxfp8_scale(amax).view(torch.uint8) + assert scales.tolist() == [0x7F, 0x80, 0x80, 0x80] + + +def test_input_and_fused_intermediate_use_distinct_mxfp8_rounding( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(tune_fhmoe.dtypes, "fp8", torch.float8_e4m3fnuz) + values = torch.zeros(32) + values[0] = 1.5 + values[1] = 2**-17 + + input_qdq = _fp8_group_quant_dequant(values) + fused_qdq = _fp8_group_quant_dequant(values, fused_intermediate=True) + + assert input_qdq[1] != 0 + assert fused_qdq[1] == 0 + + +def test_hy4_rows_and_candidate_kernels_are_complete(): + root = Path(__file__).resolve().parents[2] + fields, rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + contract = Contract() + expected_tokens = { + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + 32768, + } + actual_tokens = { + int(row["token"]) for row in rows if _matches(row, contract, int(row["token"])) + } + assert actual_tokens == expected_tokens + assert fields[-2:] == ["kernelName1", "kernelName2"] + + for block_m in (32, 64, 128): + quick_stage1 = stage1_candidates(block_m, quick=True) + assert quick_stage1 + assert all( + int(tune_fhmoe.get_flydsl_kernel_params(name).get("xcd_swizzle", 0)) == 0 + for name in quick_stage1 + ) + full_stage2 = stage2_candidates(block_m, quick=False) + assert full_stage2 + assert stage2_candidates(block_m, quick=True) + assert all( + tune_fhmoe.get_flydsl_kernel_params(name).get("persist") is not True + and int(tune_fhmoe.get_flydsl_kernel_params(name).get("b_nt", 0)) == 0 + for name in full_stage2 + ) + + +def test_matches_requires_complete_hy4_contract(): + root = Path(__file__).resolve().parents[2] + _, rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + contract = Contract() + row = next( + row for row in rows if row["model_dim"] == "6144" and row["token"] == "32" + ) + assert _matches(row, contract, 32) + + mismatches = { + "gfx": "gfx942", + "cu_num": "304", + "token": "64", + "model_dim": "4096", + "inter_dim": "128", + "expert": "256", + "topk": "8", + "shared_expert_id": "255", + "act_type": "ActivationType.Swiglu", + "dtype": "torch.float16", + "q_dtype_a": "torch.float8_e4m3fnuz", + "q_dtype_w": "torch.float4_e2m1fn_x2", + "q_type": "QuantType.per_Token", + "use_g1u1": "0", + "doweight_stage1": "1", + "hidden_pad": "1", + "intermediate_pad": "1", + "gate_mode": "GateMode.SEPARATED", + "ksplit": "1", + } + for field, value in mismatches.items(): + assert not _matches(row | {field: value}, contract, 32), field + + +def test_select_rows_rejects_duplicate_complete_contract(): + root = Path(__file__).resolve().parents[2] + _, rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + row = next( + row for row in rows if row["model_dim"] == "6144" and row["token"] == "32" + ) + with pytest.raises(ValueError, match="multiple HY4 FHMoE rows.*M=32"): + _select_rows([row, dict(row)], Contract(), [32]) + + +def test_csv_roundtrip_preserves_unrelated_rows(tmp_path: Path): + root = Path(__file__).resolve().parents[2] + fields, rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + output = tmp_path / "tuned_fhmoe.csv" + _write_rows(output, fields, rows) + roundtrip_fields, roundtrip_rows = _read_rows(output) + assert roundtrip_fields == fields + assert roundtrip_rows == rows + assert b"\r\n" not in output.read_bytes() + + +def test_tuned_output_merges_concurrent_row_updates(tmp_path: Path): + root = Path(__file__).resolve().parents[2] + fields, all_rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + contract = Contract() + original = [ + dict(row) + for row in all_rows + if _matches(row, contract, int(row["token"])) and row["token"] in ("1", "2") + ] + input_path = tmp_path / "input.csv" + output_path = tmp_path / "output.csv" + _write_rows(input_path, fields, original) + + context = multiprocessing.get_context("spawn") + barrier = context.Barrier(2) + common = (barrier, input_path, output_path, fields, original) + first = context.Process( + target=_concurrent_tuned_output_update, + args=(*common, 1, 0, "first_update"), + ) + second = context.Process( + target=_concurrent_tuned_output_update, + args=(*common, 2, 1, "second_update"), + ) + first.start() + second.start() + processes = (first, second) + try: + for process in processes: + process.join(120) + assert not any( + process.is_alive() for process in processes + ), "concurrent tuner update workers timed out: " + ", ".join( + process.name for process in processes if process.is_alive() + ) + assert first.exitcode == 0 + assert second.exitcode == 0 + finally: + for process in processes: + if process.is_alive(): + process.terminate() + for process in processes: + process.join(5) + + _, merged = _read_rows(output_path) + assert merged[0]["kernelName1"] == "first_update" + assert merged[1]["kernelName1"] == "second_update" + + +def test_tuner_main_commits_completed_tokens_and_resumes( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + import sys + + root = Path(__file__).resolve().parents[2] + fields, rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + input_path = tmp_path / "input.csv" + output_path = tmp_path / "output.csv" + _write_rows(input_path, fields, rows) + + fail_token_two = True + + def fake_tune_row(row, *_args, **_kwargs): + token = int(row["token"]) + if fail_token_two and token == 2: + raise RuntimeError("simulated interruption") + tuned = dict(row) + tuned["kernelName1"] = f"tuned_token_{token}" + return tuned, [] + + monkeypatch.setattr(tune_fhmoe, "get_gfx_runtime", lambda: "gfx950") + monkeypatch.setattr(tune_fhmoe, "get_cu_num", lambda: Contract().cu_num) + monkeypatch.setattr(tune_fhmoe, "_make_tensors", lambda *_args: object()) + monkeypatch.setattr(tune_fhmoe, "tune_row", fake_tune_row) + monkeypatch.setattr( + tune_fhmoe, + "_validate_sampled_reference", + lambda *_args, **_kwargs: 0.0, + ) + monkeypatch.setattr(tune_fhmoe.torch.cuda, "empty_cache", lambda: None) + monkeypatch.delenv("AITER_FLYDSL_FORCE_REDUCE", raising=False) + monkeypatch.delenv("AITER_FLYDSL_STAGE2_FP8", raising=False) + + monkeypatch.setattr( + sys, + "argv", + [ + "tune_fhmoe.py", + "-i", + str(input_path), + "-o", + str(output_path), + "--tokens", + "1,2", + "--quick", + ], + ) + with pytest.raises(RuntimeError, match="simulated interruption"): + tune_fhmoe.main() + + _, partial_rows = _read_rows(output_path) + token_one = next(row for row in partial_rows if _matches(row, Contract(), 1)) + token_two = next(row for row in partial_rows if _matches(row, Contract(), 2)) + assert token_one["kernelName1"] == "tuned_token_1" + assert token_two["kernelName1"] != "tuned_token_2" + + fail_token_two = False + monkeypatch.setattr( + sys, + "argv", + [ + "tune_fhmoe.py", + "-i", + str(output_path), + "-o", + str(output_path), + "--tokens", + "2", + "--quick", + ], + ) + tune_fhmoe.main() + + _, resumed_rows = _read_rows(output_path) + token_one = next(row for row in resumed_rows if _matches(row, Contract(), 1)) + token_two = next(row for row in resumed_rows if _matches(row, Contract(), 2)) + assert token_one["kernelName1"] == "tuned_token_1" + assert token_two["kernelName1"] == "tuned_token_2" + + +def test_tuner_main_does_not_commit_token_when_profile_write_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + import sys + + root = Path(__file__).resolve().parents[2] + fields, rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + input_path = tmp_path / "input.csv" + output_path = tmp_path / "output.csv" + profile_path = tmp_path / "profile.csv" + _write_rows(input_path, fields, rows) + + def fake_tune_row(row, *_args, **_kwargs): + tuned = dict(row) + tuned["kernelName1"] = "tuned_token_1" + profile = [ + {"stage": "baseline", "kernel": "candidate", "ms": 1.0, "error": 0.0} + ] + return tuned, profile + + def fail_profile_write(*_args): + raise OSError("profile write failed") + + monkeypatch.setattr(tune_fhmoe, "get_gfx_runtime", lambda: "gfx950") + monkeypatch.setattr(tune_fhmoe, "get_cu_num", lambda: Contract().cu_num) + monkeypatch.setattr(tune_fhmoe, "_make_tensors", lambda *_args: object()) + monkeypatch.setattr(tune_fhmoe, "tune_row", fake_tune_row) + monkeypatch.setattr( + tune_fhmoe, + "_validate_sampled_reference", + lambda *_args, **_kwargs: 0.0, + ) + monkeypatch.setattr( + tune_fhmoe, + "_write_profile_rows", + fail_profile_write, + ) + monkeypatch.setattr(tune_fhmoe.torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr( + sys, + "argv", + [ + "tune_fhmoe.py", + "-i", + str(input_path), + "-o", + str(output_path), + "--tokens", + "1", + "--profile", + str(profile_path), + "--quick", + ], + ) + + with pytest.raises(RuntimeError, match="current token was not committed"): + tune_fhmoe.main() + + assert not output_path.exists() + + +def test_tuner_main_uses_representative_runtime_token( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + import sys + + root = Path(__file__).resolve().parents[2] + fields, rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + input_path = tmp_path / "input.csv" + output_path = tmp_path / "output.csv" + profile_path = tmp_path / "profile.csv" + _write_rows(input_path, fields, rows) + observed_tokens = [] + + def make_tensors(_contract, token, *_args): + observed_tokens.append(token) + return object() + + monkeypatch.setattr(tune_fhmoe, "get_gfx_runtime", lambda: "gfx950") + monkeypatch.setattr(tune_fhmoe, "get_cu_num", lambda: Contract().cu_num) + monkeypatch.setattr(tune_fhmoe, "_make_tensors", make_tensors) + monkeypatch.setattr( + tune_fhmoe, + "tune_row", + lambda row, *_args, **_kwargs: ( + dict(row), + [{"stage": "baseline", "kernel": "stage1|stage2", "ms": 1.0}], + ), + ) + monkeypatch.setattr( + tune_fhmoe, + "_validate_sampled_reference", + lambda *_args, **_kwargs: 0.0, + ) + monkeypatch.setattr(tune_fhmoe.torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr( + sys, + "argv", + [ + "tune_fhmoe.py", + "-i", + str(input_path), + "-o", + str(output_path), + "--tokens", + "32768", + "--benchmark-token", + "38836", + "--profile", + str(profile_path), + "--quick", + ], + ) + + tune_fhmoe.main() + + assert observed_tokens == [38836] + _, output_rows = _read_rows(output_path) + assert any(_matches(row, Contract(), 32768) for row in output_rows) + profile_fields, profile_rows = _read_rows(profile_path) + assert profile_fields == [ + "token", + "runtime_token", + "stage", + "kernel", + "ms", + "error", + ] + assert profile_rows[0]["token"] == "32768" + assert profile_rows[0]["runtime_token"] == "38836" + + +def test_tuned_output_rejects_stale_same_token_update(tmp_path: Path): + root = Path(__file__).resolve().parents[2] + fields, all_rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + contract = Contract() + original = [ + dict(row) + for row in all_rows + if row["token"] == "32" and _matches(row, contract, 32) + ] + input_path = tmp_path / "input.csv" + output_path = tmp_path / "output.csv" + _write_rows(input_path, fields, original) + baseline = {32: dict(original[0])} + + first = [dict(row) for row in original] + first[0]["kernelName1"] = "first_winner" + _write_tuned_output( + input_path, output_path, fields, first, {32: 0}, baseline, contract + ) + + stale = [dict(row) for row in original] + stale[0]["kernelName1"] = "stale_winner" + with pytest.raises(RuntimeError, match="changed during tuning"): + _write_tuned_output( + input_path, output_path, fields, stale, {32: 0}, baseline, contract + ) + + _, preserved = _read_rows(output_path) + assert preserved[0]["kernelName1"] == "first_winner" + + +def test_existing_output_must_match_input_baseline(tmp_path: Path): + root = Path(__file__).resolve().parents[2] + fields, all_rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + contract = Contract() + baseline = next( + dict(row) + for row in all_rows + if row["token"] == "32" and _matches(row, contract, 32) + ) + output_path = tmp_path / "output.csv" + _write_rows(output_path, fields, [baseline]) + _validate_existing_output(output_path, fields, [baseline], {32: baseline}, contract) + + changed = dict(baseline) + changed["kernelName1"] = "newer_output" + _write_rows(output_path, fields, [changed]) + with pytest.raises(RuntimeError, match="differs from input"): + _validate_existing_output( + output_path, fields, [baseline], {32: baseline}, contract + ) + + +def test_existing_output_requires_complete_input_key_set(tmp_path: Path): + root = Path(__file__).resolve().parents[2] + fields, all_rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + contract = Contract() + selected = next( + dict(row) + for row in all_rows + if row["token"] == "32" and _matches(row, contract, 32) + ) + unrelated = next( + dict(row) + for row in all_rows + if row["token"] == "64" and _matches(row, contract, 64) + ) + input_rows = [selected, unrelated] + output_path = tmp_path / "output.csv" + + concurrently_updated = dict(unrelated) + concurrently_updated["kernelName1"] = "concurrent_unrelated_update" + _write_rows(output_path, fields, [selected, concurrently_updated]) + _validate_existing_output( + output_path, + fields, + input_rows, + {32: selected}, + contract, + ) + + _write_rows(output_path, fields, [selected]) + with pytest.raises(RuntimeError, match="row/key set differs.*missing=1"): + _validate_existing_output( + output_path, + fields, + input_rows, + {32: selected}, + contract, + ) + + tuned_rows = [dict(selected), dict(unrelated)] + tuned_rows[0]["kernelName1"] = "tuned_selected" + with pytest.raises(RuntimeError, match="row/key set differs.*missing=1"): + _write_tuned_output( + input_path=tmp_path / "input.csv", + output_path=output_path, + fields=fields, + rows=tuned_rows, + selected={32: 0}, + baseline_rows={32: selected}, + contract=contract, + ) + _, preserved = _read_rows(output_path) + assert preserved == [selected] + + +def test_profile_path_must_not_alias_input_or_output(tmp_path: Path): + input_path = tmp_path / "input.csv" + output_path = tmp_path / "output.csv" + input_path.write_text("x\n") + output_path.write_text("x\n") + + with pytest.raises(ValueError, match="must differ"): + _validate_output_paths(input_path, output_path, input_path) + with pytest.raises(ValueError, match="must differ"): + _validate_output_paths(input_path, output_path, output_path) + + hardlink = tmp_path / "profile-hardlink.csv" + hardlink.hardlink_to(output_path) + with pytest.raises(ValueError, match="must differ"): + _validate_output_paths(input_path, output_path, hardlink) + + +def test_profile_commit_failure_message_reports_uncommitted_token(tmp_path: Path): + output = tmp_path / "output.csv" + profile = tmp_path / "profile.csv" + + message = _profile_commit_failure_message( + output, + profile, + ) + assert f"profile rows failed to commit to {profile}" in message + assert f"current token was not committed to {output}" in message + + +def test_profile_rows_merge_and_replace_runtime_measurement(tmp_path: Path): + profile_path = tmp_path / "profile.csv" + first = { + "token": 32, + "runtime_token": 32, + "stage": "baseline", + "kernel": "first", + "ms": 1.0, + } + second = { + "token": 32768, + "runtime_token": 38836, + "stage": "baseline", + "kernel": "second", + "ms": 2.0, + } + replacement = dict(first, kernel="replacement", ms=0.5) + + _write_profile_rows(profile_path, [first]) + _write_profile_rows(profile_path, [second]) + _write_profile_rows(profile_path, [replacement]) + + fields, rows = _read_rows(profile_path) + assert fields == list(PROFILE_FIELDS) + assert {(row["token"], row["runtime_token"]) for row in rows} == { + ("32", "32"), + ("32768", "38836"), + } + token_32 = next(row for row in rows if row["token"] == "32") + assert token_32["kernel"] == "replacement" + assert token_32["ms"] == "0.5" + + +def test_profile_rows_migrate_legacy_schema(tmp_path: Path): + profile_path = tmp_path / "profile.csv" + _write_rows( + profile_path, + ["token", "stage", "kernel", "ms", "error"], + [ + { + "token": "32", + "stage": "baseline", + "kernel": "legacy", + "ms": "1.0", + "error": "", + } + ], + ) + + _write_profile_rows( + profile_path, + [ + { + "token": 32768, + "runtime_token": 38836, + "stage": "baseline", + "kernel": "new", + "ms": 2.0, + } + ], + ) + + fields, rows = _read_rows(profile_path) + assert fields == list(PROFILE_FIELDS) + legacy = next(row for row in rows if row["token"] == "32") + assert legacy["runtime_token"] == "32" + + +def test_output_symlink_is_rejected(tmp_path: Path): + input_path = tmp_path / "input.csv" + target = tmp_path / "target.csv" + output_link = tmp_path / "output.csv" + input_path.write_text("x\n") + target.write_text("x\n") + output_link.symlink_to(target) + + with pytest.raises(ValueError, match="must not be a symlink"): + _validate_output_paths(input_path, output_link, None) + + +def test_baseline_only_empty_selection_preserves_same_token_update(tmp_path: Path): + root = Path(__file__).resolve().parents[2] + fields, all_rows = _read_rows(root / "aiter/configs/tuned_fhmoe.csv") + contract = Contract() + original = [ + dict(row) + for row in all_rows + if row["token"] == "32" and _matches(row, contract, 32) + ] + path = tmp_path / "in_place.csv" + _write_rows(path, fields, original) + stale = [dict(row) for row in original] + + concurrent = [dict(row) for row in original] + concurrent[0]["kernelName1"] = "concurrent_update" + _write_rows(path, fields, concurrent) + + _write_tuned_output(path, path, fields, stale, {}, {}, contract) + + _, preserved = _read_rows(path) + assert preserved[0]["kernelName1"] == "concurrent_update" + + +def test_baseline_only_preserves_current_kernels( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "stage1", + "kernelName2": "stage2", + } + monkeypatch.setattr(tune_fhmoe, "_write_candidate", lambda *_args: None) + monkeypatch.setattr(tune_fhmoe, "_runner", lambda *_args: object()) + monkeypatch.setattr( + tune_fhmoe, + "_time_candidate", + lambda *_args: (1.25, torch.ones(1)), + ) + + tuned, profile = tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=True, + warmup_ms=1, + rep_ms=1, + tolerance=0, + swiglu_limit=10, + min_improvement_pct=1, + baseline_only=True, + ) + + assert tuned == row + assert profile == [{"stage": "baseline", "kernel": "stage1|stage2", "ms": 1.25}] + + +def test_tune_row_rejects_aliased_stage2_baseline(tmp_path: Path): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "flydsl_moe1_afp8_wfp8_bf16_t32x64x256_w3_gui_fp8", + "kernelName2": "flydsl_moe2_afp8_wfp8_bf16_t32x128x128_atomic_bnt2", + } + with pytest.raises(ValueError, match="invalid baseline stage2"): + tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=True, + warmup_ms=1, + rep_ms=1, + tolerance=0, + swiglu_limit=10, + min_improvement_pct=1, + baseline_only=True, + ) + + +@pytest.mark.parametrize("baseline_only", (False, True)) +def test_nonfinite_baseline_fails_before_tuning( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + baseline_only: bool, +): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "stage1", + "kernelName2": "stage2", + } + monkeypatch.setattr(tune_fhmoe, "_write_candidate", lambda *_args: None) + monkeypatch.setattr(tune_fhmoe, "_runner", lambda *_args: object()) + monkeypatch.setattr( + tune_fhmoe, + "_time_candidate", + lambda *_args: (float("inf"), torch.tensor([float("nan")])), + ) + + with pytest.raises(RuntimeError, match="baseline FHMoE kernels.*non-finite"): + tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=True, + warmup_ms=1, + rep_ms=1, + tolerance=0, + swiglu_limit=10, + min_improvement_pct=1, + baseline_only=baseline_only, + ) + + +@pytest.mark.parametrize("baseline_only", (False, True)) +def test_unstable_baseline_timing_fails_before_tuning( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + baseline_only: bool, +): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "stage1", + "kernelName2": "stage2", + } + monkeypatch.setattr(tune_fhmoe, "_write_candidate", lambda *_args: None) + monkeypatch.setattr(tune_fhmoe, "_runner", lambda *_args: object()) + monkeypatch.setattr( + tune_fhmoe, + "_time_candidate", + lambda *_args: (float("inf"), torch.ones(1)), + ) + + with pytest.raises(RuntimeError, match="unstable or non-finite timing"): + tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=True, + warmup_ms=1, + rep_ms=1, + tolerance=0, + swiglu_limit=10, + min_improvement_pct=1, + baseline_only=baseline_only, + ) + + +def test_tune_row_selects_only_fast_valid_candidates( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "stage1", + "kernelName2": "stage2", + } + active: dict[str, str] = {} + + def write_candidate(_path, candidate): + active.clear() + active.update(candidate) + + timings = { + ("stage1", "stage2"): (1.0, torch.ones(2)), + ("stage1", "stage2_fast"): (0.8, torch.ones(2)), + ("stage1", "stage2_wrong"): (0.5, torch.zeros(2)), + ("stage1", "stage2_nan"): ( + float("inf"), + torch.tensor([float("nan"), 1.0]), + ), + ("stage1_fast", "stage2"): (0.9, torch.ones(2)), + ("stage1_fast", "stage2_fast"): (0.6, torch.ones(2)), + } + monkeypatch.setattr(tune_fhmoe, "_write_candidate", write_candidate) + monkeypatch.setattr( + tune_fhmoe, + "_runner", + lambda *_args: (active["kernelName1"], active["kernelName2"]), + ) + + def time_candidate(run, *_args): + if run == ("stage1", "stage2_broken"): + raise RuntimeError("compile failed") + return timings[run] + + monkeypatch.setattr(tune_fhmoe, "_time_candidate", time_candidate) + monkeypatch.setattr( + tune_fhmoe, + "stage2_candidates", + lambda *_args: [ + "stage2_broken", + "stage2_fast", + "stage2_wrong", + "stage2_nan", + ], + ) + monkeypatch.setattr(tune_fhmoe, "stage1_candidates", lambda *_args: ["stage1_fast"]) + + tuned, profile = tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=True, + warmup_ms=1, + rep_ms=1, + tolerance=0.1, + swiglu_limit=10, + min_improvement_pct=1, + ) + + assert tuned["kernelName1"] == "stage1_fast" + assert tuned["kernelName2"] == "stage2_fast" + broken = next(row for row in profile if row["kernel"] == "stage2_broken") + assert broken["ms"] == float("inf") + assert broken["error"] == float("inf") + + +def test_tune_row_keeps_original_below_minimum_improvement( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "stage1", + "kernelName2": "stage2", + } + active: dict[str, str] = {} + + def write_candidate(_path, candidate): + active.clear() + active.update(candidate) + + monkeypatch.setattr(tune_fhmoe, "_write_candidate", write_candidate) + monkeypatch.setattr( + tune_fhmoe, + "_runner", + lambda *_args: (active["kernelName1"], active["kernelName2"]), + ) + monkeypatch.setattr( + tune_fhmoe, + "_time_candidate", + lambda run, *_: ( + (0.995 if run == ("stage1", "candidate") else 1.0), + torch.ones(1), + ), + ) + monkeypatch.setattr(tune_fhmoe, "stage2_candidates", lambda *_args: ["candidate"]) + monkeypatch.setattr(tune_fhmoe, "stage1_candidates", lambda *_args: []) + + tuned, _ = tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=True, + warmup_ms=1, + rep_ms=1, + tolerance=0, + swiglu_limit=10, + min_improvement_pct=1, + ) + + assert tuned == row + + +def test_pair_beam_finds_joint_kernel_improvement( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "stage1", + "kernelName2": "stage2", + } + active: dict[str, str] = {} + + def write_candidate(_path, candidate): + active.clear() + active.update(candidate) + + timings = { + ("stage1", "stage2"): 1.0, + ("stage1", "stage2_joint"): 1.1, + ("stage1_joint", "stage2"): 1.1, + ("stage1_joint", "stage2_joint"): 0.5, + } + monkeypatch.setattr(tune_fhmoe, "_write_candidate", write_candidate) + monkeypatch.setattr( + tune_fhmoe, + "_runner", + lambda *_args: (active["kernelName1"], active["kernelName2"]), + ) + monkeypatch.setattr( + tune_fhmoe, + "_time_candidate", + lambda run, *_args: (timings[run], torch.ones(1)), + ) + monkeypatch.setattr( + tune_fhmoe, "stage2_candidates", lambda *_args: ["stage2_joint"] + ) + monkeypatch.setattr( + tune_fhmoe, "stage1_candidates", lambda *_args: ["stage1_joint"] + ) + + tuned, _ = tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=True, + warmup_ms=1, + rep_ms=1, + tolerance=0, + swiglu_limit=10, + min_improvement_pct=1, + ) + + assert tuned["kernelName1"] == "stage1_joint" + assert tuned["kernelName2"] == "stage2_joint" + + +def test_full_search_checks_pairs_outside_quick_beam( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +): + row = { + "token": "32", + "block_m": "32", + "kernelName1": "stage1", + "kernelName2": "stage2", + } + active: dict[str, str] = {} + stage1 = [f"stage1_{index}" for index in range(5)] + stage2 = [f"stage2_{index}" for index in range(5)] + + def write_candidate(_path, candidate): + active.clear() + active.update(candidate) + + def timing(run): + s1, s2 = run + if (s1, s2) == (stage1[-1], stage2[-1]): + return 0.1 + if s1 == "stage1" and s2 == "stage2": + return 10.0 + if s1 == "stage1": + return 1.0 + stage2.index(s2) + if s2 == "stage2": + return 1.0 + stage1.index(s1) + return 9.0 + + monkeypatch.setattr(tune_fhmoe, "_write_candidate", write_candidate) + monkeypatch.setattr( + tune_fhmoe, + "_runner", + lambda *_args: (active["kernelName1"], active["kernelName2"]), + ) + monkeypatch.setattr( + tune_fhmoe, + "_time_candidate", + lambda run, *_args: (timing(run), torch.ones(1)), + ) + monkeypatch.setattr(tune_fhmoe, "stage1_candidates", lambda *_args: stage1) + monkeypatch.setattr(tune_fhmoe, "stage2_candidates", lambda *_args: stage2) + + tuned, _ = tune_fhmoe.tune_row( + row, + tensors=object(), + contract=Contract(), + config_file=tmp_path / "candidate.csv", + quick=False, + warmup_ms=1, + rep_ms=1, + tolerance=0, + swiglu_limit=10, + min_improvement_pct=1, + ) + + assert tuned["kernelName1"] == stage1[-1] + assert tuned["kernelName2"] == stage2[-1] + + +def test_time_candidate_rejects_nonfinite_repeat( + monkeypatch: pytest.MonkeyPatch, +): + outputs = iter((torch.ones(1), torch.tensor([float("nan")]))) + monkeypatch.setattr( + tune_fhmoe.triton.testing, + "do_bench", + lambda *_args, **_kwargs: 1.0, + ) + + timing, _ = _time_candidate(lambda: next(outputs), warmup_ms=1, rep_ms=1) + + assert timing == float("inf") + + +def test_time_candidate_uses_independent_stability_tolerance( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + tune_fhmoe.triton.testing, + "do_bench", + lambda *_args, **_kwargs: 1.0, + ) + + def measure(stability_tolerance: float) -> float: + outputs = iter( + ( + torch.ones(1), + torch.full((1,), 1.01), + torch.full((1,), 1.01), + torch.full((1,), 1.01), + ) + ) + timing, _ = _time_candidate( + lambda: next(outputs), + warmup_ms=1, + rep_ms=1, + stability_tolerance=stability_tolerance, + ) + return timing + + assert measure(2e-2) == 1.0 + assert measure(5e-3) == float("inf") + + +@pytest.mark.parametrize("expected_value,should_pass", ((1.0, True), (0.0, False))) +def test_sampled_reference_controls_final_acceptance( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + expected_value: float, + should_pass: bool, +): + tensors = SimpleNamespace( + sample_rows=(0, 1), + hidden=torch.zeros((2, 2)), + ) + monkeypatch.setattr(tune_fhmoe, "_write_candidate", lambda *_args: None) + monkeypatch.setattr( + tune_fhmoe, + "_runner", + lambda *_args: lambda: torch.ones((2, 2)), + ) + monkeypatch.setattr( + tune_fhmoe, + "_sampled_fp32_reference", + lambda *_args: torch.full((2, 2), expected_value), + ) + kwargs = { + "row": {"token": "32"}, + "tensors": tensors, + "contract": Contract(), + "config_file": tmp_path / "candidate.csv", + "swiglu_limit": 10, + "tolerance": 0.1, + } + + if should_pass: + assert _validate_sampled_reference(**kwargs) == 0 + else: + with pytest.raises(RuntimeError, match="sampled FP32 validation failed"): + _validate_sampled_reference(**kwargs) diff --git a/op_tests/tuning_tests/test_tune_pipeline.py b/op_tests/tuning_tests/test_tune_pipeline.py index 88f6a312960..47f99adff0b 100644 --- a/op_tests/tuning_tests/test_tune_pipeline.py +++ b/op_tests/tuning_tests/test_tune_pipeline.py @@ -33,6 +33,12 @@ def _gpu_available(): return False +def _get_gfx_runtime(): + from aiter.jit.utils.chip_info import get_gfx_runtime + + return get_gfx_runtime() + + def _get_platform_dtypes(): """Return (fp8_str, quant_type_str) based on GPU arch.""" try: @@ -286,6 +292,72 @@ def setUpClass(cls): "timeout": 1800, "timeout_mp1": 2400, }, + "fhmoe": { + "script": "op_tests/tuners/tune_fhmoe.py", + "header": [ + "gfx", + "cu_num", + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "shared_expert_id", + "act_type", + "dtype", + "q_dtype_a", + "q_dtype_w", + "q_type", + "use_g1u1", + "doweight_stage1", + "hidden_pad", + "intermediate_pad", + "gate_mode", + "block_m", + "ksplit", + "kernelName1", + "kernelName2", + ], + "shapes": [ + ( + "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", + ) + ], + "keys": [ + "gfx", + "cu_num", + "token", + "model_dim", + "inter_dim", + "expert", + "topk", + "shared_expert_id", + ], + "extra_args": ["--tokens", "32", "--quick"], + "timeout": 900, + "timeout_mp1": 900, + }, "gdn_k5_opt": { "script": "csrc/gdn_k5/chunk_gdn_h_opt_tune.py", "header": [ @@ -503,6 +575,11 @@ def _run_one(self, name, mp=1): self.assertIn( key, df.columns, f"{name} ({mp_label}): missing column {key}" ) + if name == "fhmoe": + self.assertIn("kernelName1", df.columns) + self.assertIn("kernelName2", df.columns) + self.assertTrue(df["kernelName1"].fillna("").str.strip().ne("").all()) + self.assertTrue(df["kernelName2"].fillna("").str.strip().ne("").all()) for _, row in df.iterrows(): us = float(row.get("us", -1)) self.assertNotEqual( @@ -584,6 +661,11 @@ def test_fmoe_mp1(self): def test_fmoe_mp_default(self): self._run_one("fmoe", mp=None) + def test_fhmoe_mp1(self): + if _get_gfx_runtime() != "gfx950": + self.skipTest("FHMoE tuning requires gfx950") + self._run_one("fhmoe", mp=1) + def test_a6w6_blockscale_mp1(self): self._run_one("a6w6_blockscale", mp=1)