[ROCm][Quantization][MOE] Enable fused shared experts for block-quantized FP8 (#53097)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Aug 29, 2026 at 00:23 UTC
06569a8696076eeae9558928b00f035ded8f8b60
2 files changed
+177
tests/model_executor/layers/test_fused_shared_expert.py
+133
@@ -15,6 +15,7 @@ import vllm.config as vllm_config_module
15
from vllm.config import VllmConfig, set_current_vllm_config
16
from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils
17
from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts
18
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config
19
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
20
from vllm.model_executor.layers.quantization.utils.config_utils import (
21
is_shared_expert_quant_fse_compatible,
@@ -710,3 +711,135 @@ def test_non_quark_shared_expert_fse_is_incompatible() -> None:
711
assert reason == (
712
"shared-expert FSE quantization compatibility is not implemented for object"
713
)
714
+
715
+
716
+def _fp8_config(**kwargs: Any) -> Fp8Config:
717
+ return Fp8Config(
718
+ is_checkpoint_fp8_serialized=True,
719
+ activation_scheme="dynamic",
720
+ weight_block_size=[128, 128],
721
+ **kwargs,
722
+ )
723
+
724
+
725
+def test_block_fp8_shared_expert_fse_is_compatible() -> None:
726
+ compatible, reason = is_shared_expert_quant_fse_compatible(
727
+ _fp8_config(),
728
+ "model.layers.0.mlp.experts",
729
+ "model.layers.0.mlp.shared_experts",
730
+ )
731
+
732
+ assert compatible
733
+ assert reason is None
734
+
735
+
736
+def test_per_tensor_fp8_shared_expert_fse_is_incompatible() -> None:
737
+ """Per-tensor scales are 0-D, so the shared-expert chunker cannot slice them."""
738
+ compatible, reason = is_shared_expert_quant_fse_compatible(
739
+ Fp8Config(is_checkpoint_fp8_serialized=True, activation_scheme="dynamic"),
740
+ "model.layers.0.mlp.experts",
741
+ "model.layers.0.mlp.shared_experts",
742
+ )
743
+
744
+ assert not compatible
745
+ assert reason == (
746
+ "FP8 shared-expert FSE is only implemented for block-quantized checkpoints"
747
+ )
748
+
749
+
750
+@pytest.mark.parametrize("store_dtype", ["mxfp4", "nvfp4"])
751
+def test_store_dtype_fp8_shared_expert_fse_is_incompatible(store_dtype: str) -> None:
752
+ """Any routed-expert storage override rules out fusing the shared expert."""
753
+ compatible, reason = is_shared_expert_quant_fse_compatible(
754
+ _fp8_config(store_dtype=store_dtype),
755
+ "model.layers.0.mlp.experts",
756
+ "model.layers.0.mlp.shared_experts",
757
+ )
758
+
759
+ assert not compatible
760
+ assert reason == (
761
+ f"FP8 stores routed experts as {store_dtype}, which is not supported "
762
+ "for fused shared experts at model.layers.0.mlp.shared_experts"
763
+ )
764
+
765
+
766
+@pytest.mark.parametrize(
767
+ "ignored_layers",
768
+ [
769
+ ["model.layers.0.mlp.experts"],
770
+ ["model.layers.0.mlp.shared_experts.gate_up_proj"],
771
+ ["model.layers.0.mlp.shared_experts.down_proj"],
772
+ ],
773
+)
774
+def test_fp8_shared_expert_fse_rejects_asymmetric_ignored_layers(
775
+ ignored_layers: list[str],
776
+) -> None:
777
+ compatible, reason = is_shared_expert_quant_fse_compatible(
778
+ _fp8_config(ignored_layers=ignored_layers),
779
+ "model.layers.0.mlp.experts",
780
+ "model.layers.0.mlp.shared_experts",
781
+ )
782
+
783
+ assert not compatible
784
+ assert reason == (
785
+ "FP8 ignores routed and shared experts inconsistently at "
786
+ "model.layers.0.mlp.shared_experts"
787
+ )
788
+
789
+
790
+def test_fp8_shared_expert_fse_allows_symmetric_ignored_layers() -> None:
791
+ """Both sides unquantized is still a consistent scheme."""
792
+ compatible, reason = is_shared_expert_quant_fse_compatible(
793
+ _fp8_config(
794
+ ignored_layers=[
795
+ "model.layers.0.mlp.experts",
796
+ "model.layers.0.mlp.shared_experts.gate_up_proj",
797
+ "model.layers.0.mlp.shared_experts.down_proj",
798
+ ]
799
+ ),
800
+ "model.layers.0.mlp.experts",
801
+ "model.layers.0.mlp.shared_experts",
802
+ )
803
+
804
+ assert compatible
805
+ assert reason is None
806
+
807
+
808
+def test_fp8_shared_expert_fse_expands_packed_projections() -> None:
809
+ """Both shards of a fused projection ignored is symmetric within that
810
+ projection, so the mismatch reported is against the routed experts."""
811
+ quant_config = _fp8_config(
812
+ ignored_layers=[
813
+ "model.layers.0.mlp.shared_experts.gate_proj",
814
+ "model.layers.0.mlp.shared_experts.up_proj",
815
+ ]
816
+ )
817
+ quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
818
+
819
+ compatible, reason = is_shared_expert_quant_fse_compatible(
820
+ quant_config,
821
+ "model.layers.0.mlp.experts",
822
+ "model.layers.0.mlp.shared_experts",
823
+ )
824
+
825
+ assert not compatible
826
+ assert reason == (
827
+ "FP8 ignores routed and shared experts inconsistently at "
828
+ "model.layers.0.mlp.shared_experts"
829
+ )
830
+
831
+
832
+def test_fp8_shared_expert_fse_propagates_partial_shard_exclusion() -> None:
833
+ """Half a fused projection excluded is rejected by `is_layer_skipped`
834
+ itself, exactly as it is in `Fp8Config.get_quant_method`."""
835
+ quant_config = _fp8_config(
836
+ ignored_layers=["model.layers.0.mlp.shared_experts.gate_proj"]
837
+ )
838
+ quant_config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
839
+
840
+ with pytest.raises(ValueError, match="some but not all shards"):
841
+ is_shared_expert_quant_fse_compatible(
842
+ quant_config,
843
+ "model.layers.0.mlp.experts",
844
+ "model.layers.0.mlp.shared_experts",
845
+ )
vllm/model_executor/layers/quantization/utils/config_utils.py
+44
@@ -32,7 +32,11 @@ def is_shared_expert_quant_fse_compatible(
32
if quant_config is None:
33
return True, None
34
35
+ from vllm.model_executor.layers.quantization.fp8 import Fp8Config
36
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
37
+ from vllm.model_executor.layers.quantization.utils.quant_utils import (
38
+ is_layer_skipped,
39
+ )
40
from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config
41
42
if isinstance(quant_config, DeepseekV4FP8Config):
@@ -144,6 +148,46 @@ def is_shared_expert_quant_fse_compatible(
148
f"shared experts at {shared_expert_prefix}",
149
)
150
151
+ if isinstance(quant_config, Fp8Config):
152
+ if quant_config.store_dtype is not None:
153
+ return (
154
+ False,
155
+ f"FP8 stores routed experts as {quant_config.store_dtype}, which "
156
+ f"is not supported for fused shared experts at "
157
+ f"{shared_expert_prefix}",
158
+ )
159
+
160
+ # Serialized per-tensor checkpoints store 0-D or size-1 scales, which
161
+ # the shared-expert weight chunker cannot slice into the appended expert
162
+ # slots; online FP8 is simply untested. Both lack a weight block size.
163
+ if quant_config.weight_block_size is None:
164
+ return (
165
+ False,
166
+ "FP8 shared-expert FSE is only implemented for block-quantized "
167
+ "checkpoints",
168
+ )
169
+
170
+ def is_ignored(layer_name: str) -> bool:
171
+ return is_layer_skipped(
172
+ prefix=layer_name,
173
+ ignored_layers=quant_config.ignored_layers,
174
+ fused_mapping=quant_config.packed_modules_mapping,
175
+ match_mode=quant_config.ignored_layers_match_mode,
176
+ )
177
+
178
+ expert_ignored = is_ignored(expert_prefix)
179
+ if any(
180
+ is_ignored(f"{shared_expert_prefix}.{projection_name}") != expert_ignored
181
+ for projection_name in projection_names
182
+ ):
183
+ return (
184
+ False,
185
+ "FP8 ignores routed and shared experts inconsistently at "
186
+ f"{shared_expert_prefix}",
187
+ )
188
+
189
+ return True, None
190
+
191
# TODO: Extend FSE support detection to other quantization methods. Typically,
192
# one would check that the experts and shared_experts use the same
193
# quantization config. This may be refactored as part of QuantizationConfig later.