[Feature][FlyDSL] Add fused heterogeneous MXFP4/FP8 shared-expert MoE (FHMoE)

[Fangzhou-Ai]

Summary

• Add a dedicated fused heterogeneous MoE (FHMoE) path for checkpoints whose routed experts use MXFP4/E8M0 while the shared expert remains native FP8/E8M0.
• Fuse the shared expert into the same FlyDSL stage-1 and stage-2 schedules without adding shared-expert fields to the ordinary fused_moe_ custom-op schema or ordinary kernel compiler APIs.
• Keep FHMoE dispatch, runtime wrappers, GPU kernels, AOT handling, and tests in dedicated fhmoe.py / test_fhmoe.py files, with only narrow compatibility bridges in existing MoE modules.

Stack and review scope

This PR directly targets #4314's afz/dsv4-bm16-tuning branch:

#4315 XCD fix -> #4314 BM16 tuning -> #4269 FHMoE -> future FHMoE tuning

GitHub therefore displays exactly one FHMoE commit (cc36e67f9) and eight paths:

• dedicated: aiter/fhmoe.py, aiter/ops/flydsl/fhmoe.py, aiter/ops/flydsl/kernels/fhmoe.py, aiter/aot/flydsl/fhmoe.py, and op_tests/test_fhmoe.py
• compatibility bridges: aiter/fused_moe.py, aiter/ops/flydsl/moe_kernels.py, and aiter/aot/flydsl/moe.py

It contains no model-config CSV change, ordinary mixed-MoE GPU-kernel change, XCD helper/test change, or ordinary MoE boundary test.

XCD policy and follow-up tuning

The initial FHMoE enablement intentionally uses XCD0 for every batch size:

• dedicated runtime and compiler APIs expose no XCD parameter;
• runtime dispatch passes XCD0 through its private ordinary-MoE bridge;
• the dedicated GPU kernel contains no XCD workgroup-remap helper;
• all 32 cloned AOT jobs normalize their names and settings to XCD0.

This gives the feature one deterministic, validated scheduling contract. A follow-up FHMoE tuning PR will target this PR's upstream afz/dsv4-fmoe-bm16-post4 branch, benchmark the corrected XCD mapping from #4315 across relevant batch sizes, and select XCD settings only where they provide validated correctness and performance.

Why heterogeneous fusion

The checkpoint stores the 384 routed experts as MXFP4 but the shared expert as FP8. Converting the shared expert to MXFP4 recovers fusion performance but adds substantial weight error; scheduling a separate FP8 shared kernel preserves precision but loses much of the fusion gain. FHMoE keeps one fused schedule while selecting FP8 weight and scale resources only for the final logical shared expert.

Existing callers keep their prior behavior when the shared-expert arguments are omitted. Compatibility coverage verifies that omitted and explicit-default calls are bitwise identical.

Validation

All GPU validation below ran on gfx950 from final head cc36e67f9.

Exact DeepSeek-V4 TP8 dimensions across M=4,8,16,32,64:

HIP_VISIBLE_DEVICES=7 \
  AITER_HETERO_MOE_DSV4=1 \
  AITER_HETERO_MOE_FULL_SWEEP=1 \
  /home/amd/vllm/.venv/bin/python -m pytest op_tests/test_fhmoe.py -q

Result: 13 passed. The suite checks FP32 precision oracles, separately scheduled routed/shared controls, isolated routed and shared contributions, forced-reduce determinism, atomic-output bounds, graph replay, legacy compatibility, invalid contracts, the dedicated custom-op schema, and XCD0 runtime/AOT behavior.

The full small-profile sweep also passed:

HIP_VISIBLE_DEVICES=7 \
  AITER_HETERO_MOE_FULL_SWEEP=1 \
  /home/amd/vllm/.venv/bin/python -m pytest op_tests/test_fhmoe.py -q

Result: 13 passed.

Focused API/AOT/compiler-bridge coverage passed (3 passed), including a regression that invokes both shared generic dispatchers through the private XCD0 adapter while keeping the public FHMoE signatures XCD-free. Fresh-process imports succeeded in both fused_moe -> fhmoe and fhmoe -> fused_moe order, registering torch.ops.aiter.fhmoe_ without the old custom-op name.

Static checks passed on all eight FHMoE paths:

uvx ruff check <eight FHMoE paths>
/home/amd/vllm/.venv/bin/python -m py_compile <eight FHMoE paths>
git diff --check 1d5dc2691...HEAD

Model evaluation and serving measurements

Prior end-to-end validation of the same heterogeneous fused behavior completed all 1,319 GSM8K samples with 20-shot prompts, concurrency 64, max_tokens=2048, and temperature 0:

Strict match	Flexible extract
1270/1319 (96.285%)	1269/1319 (96.209%)

All requests completed without a server error or GPU-fault match. This model evaluation predates the isolation and naming refactor; the final head was revalidated with the exact-shape precision suite above.

Saved DeepSeek-V4-Pro TP8 8K-input/1K-output deployment measurements showed:

Concurrency	No-FHMoE output tok/s	FHMoE output tok/s	Gain
4	143.680	166.273	+15.72%
8	263.593	310.652	+17.85%
16	443.155	506.383	+14.27%
32	677.115	776.618	+14.70%
64	920.519	1031.343	+12.04%

All measured requests completed. These are deployment-oriented cross-version measurements, not an isolated same-base PR A/B; the final structural refactor was not reswept. The isolated BM16 tuning results are documented in #4314.

Duplicate-work check

Open-PR searches found no equivalent heterogeneous MXFP4-routed/FP8-shared fused implementation. The original #4269 is updated in place rather than opening a duplicate feature PR.

AI assistance

OpenAI Codex was used to audit related work, implement and organize the patch, run validation and benchmarks, organize the stacked history, and draft this description. The human submitter reviewed every changed line and is responsible for understanding and defending the change end-to-end.

[Fangzhou-Ai]

E2E GSM8K 8*MI355X DSV4 Pro TP=8 EP=1 test result, conc=64 n-shot=20