[Fangzhou-Ai]

Summary

• Add a dedicated fused heterogeneous MoE (FHMoE) path for DeepSeek-V4 checkpoints whose 384 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 as logical expert 384, while keeping shared-expert fields out of the ordinary fused_moe_ custom-op schema and ordinary kernel compiler APIs.
• Share the mixed-MoE kernel builders between ordinary MoE and FHMoE through mixed_moe_gemm_2stage_common.py, leaving thin per-path facades.
• Support the current main launcher, XCD-swizzle, runtime, and AOT contracts, including the current stage-1 situ_beta, situ_linear_beta, k_batch_intra_block, and pass_swiglu_limit arguments.

Stack and review scope

This PR is the feature base and targets current main. #4314 contains the DeepSeek-V4 C1/C2/C32/C64 tuning rows and is stacked on it:

main aa6ee6102 -> #4269 2a4cb1339 -> #4314 57e848f53

This PR contains four commits across ten paths:

1. feat(flydsl): add fused heterogeneous MoE
2. refactor(flydsl): share MoE and FHMoE kernels
3. fix(flydsl): preserve ordinary MoE stage1 arguments
4. fix(flydsl): forward FHMoE AOT stage1 ABI flag

The feature, runtime wrappers, AOT adapter, GPU-kernel facade, and tests remain isolated in fhmoe.py / test_fhmoe.py files. Existing MoE modules contain only the narrow dispatch and shared-builder bridges needed by the new path. Tuning CSV changes remain in #4314.

Why heterogeneous fusion

The checkpoint stores routed experts as MXFP4 but the shared expert as FP8. Quantizing the shared expert to MXFP4 permits ordinary homogeneous fusion but adds avoidable weight error; launching a separate FP8 shared kernel preserves precision but gives up much of the fusion gain. FHMoE keeps one fused schedule while selecting native FP8 weight and scale resources only for the final logical shared expert.

Existing callers preserve their previous behavior when the shared-expert arguments are omitted. Compatibility coverage checks that omitted and explicit-default calls are bitwise identical and that the ordinary custom-op/compiler schemas have not acquired FHMoE-only fields.

Rebase and conflict resolution

The branch was rebuilt on main at aa6ee6102. The textual conflict was confined to mixed_moe_gemm_2stage.py: current main had moved FlyDSL buffer_ops and vector imports into the vendored kernel namespace, while this PR had moved the implementation into the shared builder. The resolution keeps the thin facade and carries the current-main imports into mixed_moe_gemm_2stage_common.py.

The rebase also forward-ports current launcher changes:

• ordinary stage-1 forwarding retains situ_beta, situ_linear_beta, and k_batch_intra_block;
• runtime and AOT jobs preserve parsed XCD-swizzle values;
• _FHMoEAOTBackend.build_stage1_args explicitly accepts and forwards pass_swiglu_limit to the shared stage-1 argument builder.

The last item was caught by the Python 3.12 prebuilt-wheel CI job after the first rebase. It is adapter-protocol drift in a newly added file, not a kernel/runtime failure. The fix keeps an explicit signature rather than absorbing future protocol changes through **kwargs, and a CPU-only regression passes False to prove that the value is forwarded rather than supplied by the downstream default.

Validation

Focused tests on the current head:

/home/amd/vllm/.venv/bin/python -m pytest \
  op_tests/test_fhmoe.py::test_fhmoe_aot_stage1_forwards_optional_swiglu_abi \
  op_tests/test_fhmoe.py::test_fhmoe_aot_jobs_preserve_xcd_swizzling -q

Result: 2 passed.

The exact DSV4 FHMoE stage-1 job that failed in wheel CI was also compiled through compile_one_config with COMPILE_ONLY=1, fake tensors, and target gfx950:

[flydsl] COMPILE_ONLY=1, compilation succeeded (arch=gfx950)
[OK] flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w4_gui_fp8
     model_dim=7168 inter_dim=512 E=385 topk=7

Additional checks:

uvx ruff==0.16.0 check <ten changed paths>
/home/amd/vllm/.venv/bin/python -m py_compile <ten changed paths>
git diff --check aa6ee6102...2a4cb1339

Result: all passed. GitHub's Ruff and Black checks also pass on the current head.

Earlier full GPU correctness validation covered both the small profile and exact DeepSeek-V4 TP8 dimensions at M=4,8,16,32,64, including the FP32 oracle, isolated routed/shared contributions, forced-reduce determinism, atomic bounds, graph replay, schema compatibility, invalid contracts, and runtime/AOT bridges: 13 passed in each profile.

The editable benchmark install was refreshed from the restacked tip without prebuilding kernels:

PREBUILD_KERNELS=0 AITER_USE_SYSTEM_TRITON=1 \
  uv pip install --python /home/amd/vllm/.venv/bin/python \
  --reinstall --no-cache --no-build-isolation --no-deps \
  --editable /home/amd/aiter-pr4269

Installed version: 0.1.20.dev34+g57e848f53.d20260730.

Model evaluation and serving performance

The current rebased runtime stack was tested with the repository's unchanged scripts on 8×gfx950, TP=8, expert parallelism disabled, and concurrency 4 for the performance point. The subsequent pass_swiglu_limit change is AOT-only and does not alter this runtime path.

The server log confirms the intended route: native-FP8 shared-expert fusion, internal E=385/top-k=7, no EP, and the tuned C1/C2 configurations on every rank.

GSM8K used all 1,319 samples, 20-shot prompts, concurrency 64, max_tokens=2048, and temperature 0:

Filter	Result
strict match	1269/1319 (96.209%)
flexible extract	1268/1319 (96.133%)

All 1,319 requests returned HTTP 200 with no retries or server errors. Against the immediately preceding same-script baseline, both metrics moved by 4/1319 answers (-0.303 percentage points); 13 extracted answers changed, with 7 correct-to-wrong and 3 wrong-to-correct.

InferenceX 8K-input/1K-output at concurrency 4 completed 40/40 requests:

Metric	Result
output throughput	206.858 tok/s
total throughput	1859.526 tok/s
mean TTFT	477.185 ms
mean TPOT	18.360 ms

Relative to the immediately preceding TP8/C4 run, output throughput improved 0.41%, TTFT improved 2.75%, and TPOT improved 0.38%.

Commands:

AITER_ROOT=/home/amd/aiter-pr4269 TP=8 PORT=8888 GPU_MEM_UTIL=0.8 \
  bash dsv4_rocm_bench/start_server.sh

PORT=8888 CONC=64 bash dsv4_rocm_bench/run_gsm8k.sh
PORT=8888 CONCS=4 bash dsv4_rocm_bench/run_inferencex_matrix.sh

Duplicate-work check

Current open-PR searches for heterogeneous MoE and MXFP4 FP8 shared expert found this feature PR, its tuning follow-up #4314, and #3739. #3739 is a config-only DeepSeek-V3 E=33/top-k9 tuning change where the fused shared expert uses the same MXFP4 representation as the routed experts; it does not implement this PR's heterogeneous MXFP4-routed/native-FP8-shared kernel or ABI. This updates the existing #4269 rather than opening a duplicate.

AI assistance

OpenAI Codex was used to audit related work, implement and organize the patch, resolve and verify the rebase, run tests and model evaluations, restack the branches, and draft this description. The human submitter reviewed the changes and is responsible for understanding and defending them end-to-end.

[Fangzhou-Ai]

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

[yadaish]

@Fangzhou-Ai aiter/ops/flydsl/kernels/fhmoe.py is duplicated from original moe implement which makes the moe really hard to maintain.

[Fangzhou-Ai]

@Fangzhou-Ai aiter/ops/flydsl/kernels/fhmoe.py is duplicated from original moe implement which makes the moe really hard to maintain.

Yes, it's a good point, I was thinking to separate the common code into a new file then both moe and fhmoe can import them to reduce duplication. Do you think it's a good idea here?

[Fangzhou-Ai]

@yadaish I have revised this PR with a new file contains the common code to minimize the future frictions between these 2 kernels.