[ROCm][Quantization][MOE] Enable fused shared experts for block-quantized FP8

[xuebwang-amd]

Purpose

To enable fp8 models such as DeepSeek-R1-0528 ("n_shared_experts": 1, "quant_method": "fp8")

On top of PR #51695.

PR #51695 (88b2bff2c63) centralised the "may this model fuse its shared experts into the routed grouped GEMM?" decision into is_shared_expert_quant_fse_compatible. That helper implements None, DeepseekV4FP8Config and QuarkConfig, and closes with a TODO plus a generic fallback for everything else.

Fp8Config therefore lands in the fallback. On ROCm with VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1, a block-quantized FP8 checkpoint (DeepSeek-V3/R1 and friends) now reports

shared-expert FSE quantization compatibility is not implemented for Fp8Config

is logged as VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be enabled: ... (vllm/model_executor/layers/fused_moe/utils.py:85-90, duplicated at vllm/models/deepseek_v4/amd/model.py:239-245 and
vllm/models/minimax_m3/amd/model.py:386-392) and falls back to running the shared expert as separate linears. This is not a silent failure -- it warns -- but before #51695 deepseek_v2.py gated fused shared experts on the env var alone, with no quantization check, so this configuration did take the fused path.

This PR fills that TODO for Fp8Config, restoring the fused path for the case that is actually safe, and rejecting the cases that are not.

What changes

Three gates, then compatible.

Gate 1 -- store_dtype == "mxfp4" (config_utils.py:152-157)

This gate is defensive. Folding an FP8 shared expert into MXFP4 expert slots is exactly the storage-format mismatch this helper exists to catch.

Gate 2 -- weight_block_size is None (config_utils.py:161-166)

Fused shared experts are enabled only for block-quantized FP8. Two distinct populations are rejected here:

• Serialized per-tensor checkpoints. Their scales are 0-D (AutoFP8) or size-1 (compressed-tensors) -- On a 0-D tensor that is an IndexError; on a size-1 tensor with n_shared_experts > 1 the following divisibility assert (deepseek_v2.py:1689-1692) fails instead. Reporting these compatible would turn a model that loads today into a load-time crash.
• Online (non-serialized) FP8. No scale tensors exist in the checkpoint, so nothing crashes -- this half is rejected because it is untested, not because broken.

Gate 3 -- asymmetric ignored_layers (config_utils.py:169-185)

ignored_layers + ignored_layers_match_mode is the only per-layer discriminator on Fp8Config; the whole class body (fp8.py:92-236) has no layer_quant_config / targets / overrides, and every other field (activation_scheme, weight_block_size, is_checkpoint_fp8_serialized, use_deep_gemm) is model-global and so cannot differ between routed and shared experts by construction. Once ignore-status agrees, both sides are driven by the same config object and reach the same FP8 scheme.

Unit Test Plan & Result

Plan

7 new test functions (9 parametrized cases) in tests/model_executor/layers/test_fused_shared_expert.py, modelled on the existing Quark compatibility tests in the same file:

test	asserts
test_block_fp8_shared_expert_fse_is_compatible	block FP8 -> (True, None)
test_per_tensor_fp8_shared_expert_fse_is_incompatible	gate 2
test_mxfp4_store_dtype_fp8_shared_expert_fse_is_incompatible	gate 1
test_fp8_shared_expert_fse_rejects_asymmetric_ignored_layers	gate 3, parametrized over routed-only / gate_up-only / down-only
test_fp8_shared_expert_fse_allows_symmetric_ignored_layers	symmetric exclusion -> compatible
test_fp8_shared_expert_fse_expands_packed_projections	packed_modules_mapping fused-shard expansion
test_fp8_shared_expert_fse_propagates_partial_shard_exclusion	limitation 4 above

Results

Passed

End-to-end Test Plan & Result (negative -> positive)

Plan

Two arms:

• negative = unmodified upstream/main
• positive = the same tree plus this PR changes

DeepSeek-R1-0528 is used to run the end-to-end validation:

• Checkpoint: DeepSeek-R1-0528, unmodified, all 61 layers. Its quantization_config is {"quant_method": "fp8", "activation_scheme": "dynamic", "fmt": "e4m3", "weight_block_size": [128, 128]} with n_shared_experts: 1, n_routed_experts: 256 and first_k_dense_replace: 3, so layers 3..60 -- 58 layers -- are MoE layers with a shared expert. It resolves to plain Fp8Config: not QuarkConfig, not CompressedTensorsConfig, and not DeepseekV4FP8Config. That is exactly the config class that had no branch, and it is the class every stock DeepSeek-V3/R1 FP8 release lands on.
• Hardware/env: 8x MI325X (gfx942, 256 GiB), TP8, torch 2.11.
• Environments: VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1.

Results

GSM8K

	negative (upstream/main)	positive (this PR)
"cannot be enabled" warnings	464 (58 MoE layers x 8 ranks)	0
fused shared experts	declined everywhere	active everywhere
GSM8K accuracy	0.9568 (1262/1319)	0.9538 (1258/1319)
invalid responses	0.000	0.000
output tokens/s	2352.7	2503.6 (+6.4%)
questions/s	24.23	26.04 (+7.5%)

Throughput (non-eager)

5 timed repeats per batch size, median reported:

batch	negative (upstream/main)	positive (this PR)	delta
1	90.9 tok/s	103.2 tok/s	+13.5%
16	985.7 tok/s	1153.8 tok/s	+17.1%
64	2799.6 tok/s	3222.8 tok/s	+15.1%

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[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.

[fxmarty-amd]

What model is it for?

[xuebwang-amd]

What model is it for?

DeepSeek-R1-0528 as one example model, see PR description.

[fxmarty-amd]

DeepSeek-R1-0528 as one example model, see PR description.

Makes sense, looking at https://huggingface.co/deepseek-ai/DeepSeek-R1-0528 it indeed uses the same dtype for routed experts and shared expert.

[fxmarty-amd]

LGTM great!

Just one comment

[fxmarty-amd]

store_dtype comes from #45200

specifically https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash

I think the message here is misleading see https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro-FP4-DFlash/blob/main/config.json#L130 - I would just use if quant_config.store_dtype is not None => not supported

[xuebwang-amd]

Good catch, agreed.
On MiMo-V2.5-Pro-FP4-DFlash, n_shared_experts: null means it never reaches this helper, so the gate stays defensive.
Updated, thanks.