[fxmarty-amd]

Breaking: AITER FSE with MiniMax M3 and DeepSeek V4 now requires VLLM_ROCM_USE_AITER=1, like all other models. Updating the recipe https://recipes.vllm.ai/MiniMaxAI/MiniMax-M3 for that. AITER will be enabled by default with: #62092

Disclosure

AI assistance was used. The changes were reviewed and tested manually.

Overview

Allow compatible layers to use AITER fused shared experts (FSE) while other layers retain higher-precision shared experts.

Current vLLM rejects the entire model in case FSE can't be enabled for all layers, and various places assume all-or-nothing selection.

Motivation: certain layers are less sensitive to MXFP4 shared expert quantization than others, and fusion can be enabled for them.

Claims

• Enable per-layer FSE through the existing resolve_layer_fused_shared_expert API for models migrated to layer-scoped checkpoint loading.
• Redirect checkpoint weights only for layers that allocated fused shared-expert slots; preserve separate MXFP8 shared experts alongside MXFP4 routed experts.
• Explicitly reject mixed FSE with EPLB, expert parallelism, and the nixl_ep all2all backend, which assume uniform expert geometry across layers.
• Migrate the model loaders previously left to other PRs (MiniMax M3, DeepSeek V4, DeepSeek V2, DeepSeek V3.2, GLM5Next, AXK1, GLM4 Lite and the applicable MTP loaders) and route the EP/DP eligibility policy through explicit call-site flags (disable_under_ep) instead of model-type lookups.
• Report the resolved per-layer policy at startup: vllm serve logs the MoE layers that use FSE, those that keep a separate shared expert, and one line per distinct skip reason (deduplicated, not one per layer).

Validation

Unit tests

Run on MI350:

python -m pytest tests/model_executor/layers/test_fused_shared_expert.py

• main: 105 passed, 14 warnings in 62.09s (0:01:02)
• this PR: 166 passed, 3 skipped, 14 warnings in 76.12s (0:01:16)

Test changes in test_fused_shared_expert.py:

• test_models_fse_init now covers 13 architectures (including MiniMax M3, 2-layer DeepSeek V4 and Qwen4 Exp), both MXFP4-only and mixed MXFP4/MXFP8 shared experts, MTP loading, and the EPLB / expert-parallelism / nixl_ep mixed-FSE rejections.
• test_models_fse_init also runs off ROCm: where AITER is unavailable nothing fuses, so it asserts the unfused outcome (no appended shared-expert slot, weights kept in the separate shared expert, no skip-reason warning) instead of skipping.
• New focused tests cover 0-based ModuleList vs absolute ModuleDict index spaces, the expert-parallelism and nixl_ep rejection cases, and the DeepSeek V4 mapper fallback (mixed / all-excluded, .scale names).
• test_models_fse_init now uses single_rank_parallel_state instead of dist_init (~60% faster for that test), with shared fake-model, config, and MTP builders to avoid duplication.

End-to-end

Offline quantization

End-to-end mixed-precision checkpoint, using the Quark quantize_quark.py example script to quantize deepseek-ai/DeepSeek-V2-Lite to MXFP4 while keeping the shared experts of layers 5, 6 and 7 in BF16 (published as fxmarty-amd/DeepSeek-V2-Lite-mixed-shared-expert-mxfp4-testing):

python quantize_quark.py \
    --model_dir <DeepSeek-V2-Lite> \
    --quant_scheme mxfp4 \
    --exclude_layers "lm_head" "*self_attn*" "*mlp.gate" \
                     "*layers.5.mlp.shared_experts*" \
                     "*layers.6.mlp.shared_experts*" \
                     "*layers.7.mlp.shared_experts*" \
    --model_export hf_format --data_type bfloat16 \
    --num_calib_data 8 --seq_len 128 --skip_evaluation \
    --output_dir <mixed-mxfp4-checkpoint>

(--exclude_layers replaces the template defaults, hence the repeated deepseek_v2 entries; DeepSeek-V2-Lite's legacy remote code needs a small transformers-5 shim for is_torch_fx_available and _tied_weights_keys to run the example script.)

VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 vllm serve fxmarty-amd/DeepSeek-V2-Lite-mixed-shared-expert-mxfp4-testing on MI350 then reports the resolved per-layer policy:
WARNING ... VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be enabled - skipping shared-expert fusion: Quark excludes shared experts at <shared_expert>.
INFO ... AITER fused shared experts: enabled on MoE layer(s) [1, 2, 3, 4, 8, ..., 26]; disabled on MoE layer(s) [5, 6, 7].

Online shared experts quantization

amd/Qwen3.5-35B-A3B-MXFP4 keeps the shared experts in BF16 while the routed experts are MXFP4. Targeting a few layers' shared experts to mxfp4 online should fuse only those layers. This is the exact server command that produced the results below (mixed FSE, layers 10 to 20):

HIP_VISIBLE_DEVICES=2,3 \
VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 \
vllm serve amd/Qwen3.5-35B-A3B-MXFP4 \
  --served-model-name amd/Qwen3.5-35B-A3B-MXFP4 \
  --port 8101 \
  --gpu-memory-utilization 0.5 \
  --tensor-parallel-size 2 \
  --quantization-config '{"targets":{"re:.*layers\\.(1[0-9]|20)\\.mlp\\.shared_expert\\..*":"mxfp4"}}'

Captured on MI350 with HIP_VISIBLE_DEVICES=2,3 and TP=2 (both TP workers report the same policy):

(Worker_TP1 pid=4846) INFO 10-21 19:31:58 [base_loader.py:151] Quantizing 22 layers of types: mlp.shared_expert.down_proj: 11 (from targets: re:.*layers\.(1[0-9]|20)\.mlp\.shared_expert\..*, mxfp4); mlp.shared_expert.gate_up_proj: 11 (from targets: re:.*layers\.(1[0-9]|20)\.mlp\.shared_expert\..*, mxfp4)
(Worker_TP1 pid=4846) INFO 10-21 19:31:58 [utils.py:563] AITER fused shared experts: enabled on MoE layer(s) [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]; disabled on MoE layer(s) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39].
(Worker_TP0 pid=4834) WARNING 10-21 19:31:58 [utils.py:156] VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be enabled - skipping shared-expert fusion: Quark excludes shared experts at <shared_expert>.
(Worker_TP0 pid=4834) INFO 10-21 19:31:58 [utils.py:563] AITER fused shared experts: enabled on MoE layer(s) [10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]; disabled on MoE layer(s) [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39].

This confirms that the online mxfp4 targets make only the targeted layers FSE-compatible; every other layer keeps its BF16 shared expert, and the single skip reason is reported once instead of per layer.

Accuracy: all versus some shared experts quantized

gsm8k_platinum (1209 samples, 5 repeats) on MI350, and evaluated with a patched lm-evaluation-harness (https://github.com/fxmarty-amd/lm-evaluation-harness/tree/lm-eval-fixes at 30563cd).

lm_eval --model local-completions \
  --model_args "model=amd/Qwen3.5-35B-A3B-MXFP4,base_url=http://0.0.0.0:8101/v1/completions,num_concurrent=1000,timeout=99999,tokenized_requests=False,think_end_token='</think>'" \
  --tasks gsm8k_platinum \
  --gen_kwargs "max_gen_toks=8000" \
  --log_samples \
  --output_path <output-dir> \
  --apply_chat_template \
  --repeats 5 \
  --report_repeat_stats

Online mxfp4 targets on shared experts	FSE MoE layers	exact_match (strict, mean over 5 repeats)	exact_match (flexible, mean over 5 repeats)	mean gen tokens
none (baseline)	—	0.9843	0.9788	3543
.*layers\.(1[0-9]|20)\.mlp\.shared_expert\..*	[10..20]	0.9866	0.9798	3568
.*\.mlp\.shared_expert\..*	[0..39]	0.9835	0.9763	3873

The all-layers case logs no skip reason at all, as expected:

INFO ... AITER fused shared experts: enabled on MoE layer(s) [0, 1, ..., 39]; disabled on MoE layer(s) none.

Quantizing all shared experts is detrimental to generation length (~9% longer completions: median 3327 → 3767, p90 5820 → 6112), while the mixed per-layer policy preserves it (median 3327 → 3385, p90 5820 → 5789).

Details

Left to other PRs

• Move FHMOE's separate model-wide eligibility policy to this per-layer API:

  vllm/vllm/models/deepseek_v4/amd/model.py

  Lines 212 to 240 in 12b005e

  	def _heterogeneous_shared_expert_enabled(vllm_config: VllmConfig) -> bool:
  	config = vllm_config.model_config.hf_config
  	quant_config = vllm_config.quant_config
  	parallel_config = vllm_config.parallel_config
  	reasons: list[str] = []
  	if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
  	return False
  	if not on_gfx950():
  	reasons.append("the device is not gfx950")
  	if parallel_config.enable_expert_parallel:
  	reasons.append("expert parallelism is enabled")
  	if parallel_config.enable_eplb:
  	reasons.append("EPLB is enabled")
  	if parallel_config.tensor_parallel_size != 8:
  	reasons.append("tensor parallelism is not 8")
  	if parallel_config.data_parallel_size != 1:
  	reasons.append("data parallelism is enabled")
  	if parallel_config.prefill_context_parallel_size != 1:
  	reasons.append("prefill context parallelism is enabled")
  	if vllm_config.kernel_config.moe_backend != "aiter":
  	reasons.append("the MoE backend is not AITER")
  	if vllm_config.model_config.dtype != torch.bfloat16:
  	reasons.append("the model dtype is not BF16")
  	if not rocm_aiter_ops.fused_moe_supports_heterogeneous_shared_expert(1):
  	reasons.append(
  	"AITER lacks CSV-backed DSV4 I384 heterogeneous shared-expert support"
  	)

  and

  vllm/vllm/models/deepseek_v4/amd/model.py

  Lines 1077 to 1092 in 12b005e

  	self.fuse_heterogeneous_shared_expert = _heterogeneous_shared_expert_enabled(
  	vllm_config
  	)
  	self.use_mega_moe = vllm_config.kernel_config.moe_backend == "aiter_mega_moe"
  	self.start_layer, self.end_layer, self.layers = make_layers(
  	config.num_hidden_layers,
  	lambda prefix: DeepseekV4DecoderLayer(
  	vllm_config,
  	prefix=prefix,
  	topk_indices_buffer=self.topk_indices_buffer,
  	aux_stream_list=aux_stream_list,
  	fuse_heterogeneous_shared_expert=(
  	self.fuse_heterogeneous_shared_expert
  	),
  	),
  	prefix=f"{prefix}.layers",

• Make the DeepSeek V4 MTP loader and the DeepSeek V4.1 (AMD) loader consume fused_shared_expert_layers. They build MoE blocks through the per-layer resolver but still load as if unfused; construction now raises NotImplementedError if any MTP/V4.1 layer resolves fused, until this follow-up lands.
• Remove the fused-shared-expert logic from the NVIDIA model definitions, which cannot enable FSE.

Limitations

• Mixed FSE with EPLB or expert parallelism is unsupported.
• Routed and shared experts must have compatible quantization within each fused layer; heterogeneous precision inside one fused GEMM is unsupported.
• Fused layers must share routing-buffer geometry; mixed FSE is also rejected with nixl_ep, and Mori remains unsupported with FSE requested.
• Forward execution, graph replay, model evaluations, and performance are not validated. The NVIDIA model paths were updated too, although FSE is not supported there.

Compatibility changes

• Compared with main, migrated models accept mixed fusion eligibility with EPLB and expert parallelism off, retaining incompatible shared experts separately at their configured precision.
• Internal breaking changes: maybe_fuse_shared_experts() requires a fused_shared_expert_layers collection instead of the optional enabled boolean/environment fallback; migrated model-wide FSE booleans become layer sets; FusedMoEConfig.aiter_fmoe_shared_expert_enabled becomes is_fused_shared_expert_requested.
• is_model_fused_shared_expert_compatible() is removed; get_fused_shared_expert_layers() collects the fused layer indices and rejects mixed FSE with EPLB, expert parallelism, or nixl_ep.
• resolve_layer_fused_shared_expert() takes an explicit disable_under_ep flag (set by MiniMax M3 and DeepSeek V4) instead of branching on model type. DeepSeek V4 performs its own data_parallel_size > 1 fail-fast in DeepseekV4MoE.__init__. MiniMax M3 now enables FSE through the resolver (AITER-based) instead of the raw VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS flag.
• The DeepSeek V4 CPU model keeps its is_fused_shared_expert_enabled = False scaffolding; only the removed is_model_fused_shared_expert_compatible() call is replaced with that constant, so the broader dead-FSE cleanup can be split into a separate PR.

Related work

• This addresses per-layer policy, not the ModelOpt MXFP8 compatibility support in [ROCm][Quantization][MOE] Enable fused shared experts for MXFP8 checkpoints #57916; the duplicate-work search found no open PR implementing this same policy migration.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ Triggered Buildkite CI #97265 for commit 1f7e92677a8e.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fxmarty-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[fxmarty-amd]

/ci cancel

[github-actions]

✅ Requested cancellation of 1 CI build for fse-per-layer: #97265.

[fxmarty-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97276 for commit 2218bbadb19d.

[fxmarty-amd]

Note: updating vllm/triton_utils/dispatcher.py in 2b0d81d to fix the CI bug: buildkite/ci/pr/amd-mi300-model-executor where:

   if hasattr(value, "__dict__"):
       wrapper = lookup(vars(value).get("kernel"))

was failing with AttributeError: '_lib_fn_wrapper' object has no attribute 'get' in the CI, due to https://github.com/numba/llvmlite/blob/v0.47.0/llvmlite/binding/ffi.py#L153-L165

[fxmarty-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97316 for commit 5c5f4f61698c.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fxmarty-amd.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[fxmarty-amd]

Merged main (up to ab0e5ef) into fse-per-layer.

Conflicts:

• vllm/models/minimax_m3/amd/model.py: kept kMxfp8Dynamic from [ROCm][MiniMax-M3.1] Fuse MXFP8 activation quant into the Gemma input_layernorm for qkv_proj #63637. Dropped the is_shared_expert_quant_fse_compatible import, since the model-wide check is replaced by resolve_layer_fused_shared_expert.
• vllm/triton_utils/dispatcher.py: kept the vars()-is-a-dict guard from this branch (2b0d81d) on top of [Bugfix] Forward launch grid to Triton kernel overrides #63999's rewrite. [Bugfix] Forward launch grid to Triton kernel overrides #63999 still does hasattr(value, "__dict__") + vars(value).get("kernel"), which raises AttributeError: '_lib_fn_wrapper' object has no attribute 'get' once llvmlite is imported. I checked this directly against main's _rebind_impls.

Breakage git didn't flag, fixed in 01c144f: this branch removed from vllm import envs from the MiniMax M3 AMD model, while #63637 now uses envs.VLLM_BATCH_INVARIANT there. That raised NameError in all 18 test_online_shared_expert_fusion.py tests.

Other main commits: the other four that touch PR files (#63155, #63999, #63996, #62318) don't interact with FSE. No new code in main calls the APIs the PR removed or renamed.

Tests rerun (MI350, after rebuilding the C++ extensions)

python -m pytest tests/model_executor/layers/test_fused_shared_expert.py
# 168 passed, 3 skipped, 14 warnings in 114.20s (0:01:54)

python -m pytest tests/quantization/test_online_shared_expert_fusion.py
# 18 passed, 14 warnings in 42.01s

python -m pytest tests/model_executor/test_triton_dispatcher.py
# 12 passed, 1 skipped, 25 warnings in 4.46s

python -m pytest tests/models/test_deepseek_v4_aiter_mega_moe.py
# 16 passed, 14 warnings in 10.03s

python -m pytest tests/models/test_deepseek_v4_vl_rocm.py
# 14 passed, 14 warnings in 21.21s

python -m pytest tests/kernels/moe/test_fp8_tp_loading.py
# 24 passed, 6 skipped, 14 warnings in 6.05s

python -m pytest tests/model_executor/model_loader/test_reload.py
# partial: 77 passed, 0 failed (up to 77%); the DeepSeek-V3 cases did not complete
# because the HF download of inference-optimization/DeepSeek-V3-debug-multiply stalled

[fxmarty-amd]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97980 for commit fcd5cec526c9.

ℹ️ This PR is 3 commits behind upstream main. Merged upstream main into this branch as fcd5cec526c9 before starting CI. Pull the branch before pushing more changes.