diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index d8d5c4fb03..78aa563925 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/tests/model_executor/layers/test_fused_shared_expert.py @@ -15,10 +15,12 @@ from torch import nn import vllm.config as vllm_config_module from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.fused_moe import utils as fused_moe_utils from vllm.model_executor.layers.fused_moe.layer import determine_expert_counts -from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +from vllm.model_executor.layers.quantization.experts_int8 import ExpertsInt8Config from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig from vllm.model_executor.layers.quantization.utils.config_utils import ( get_quark_ocp_mx_group_size, @@ -141,10 +143,6 @@ def get_deepseek_v4_quark_config(exclude: list[str]) -> dict[str, Any]: return quantization_config -def _stub_quant_config() -> QuantizationConfig: - return cast(QuantizationConfig, object()) - - class _StubGlm5NextAttention(nn.Module): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__() @@ -340,7 +338,7 @@ def test_resolve_layer_fused_shared_expert_skips_compatibility_when_disabled( ) assert not fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) @@ -355,7 +353,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter( assert ( fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) is False ) @@ -364,7 +362,7 @@ def test_resolve_layer_fused_shared_expert_normalizes_unavailable_aiter( def test_resolve_layer_fused_shared_expert_passes_module_prefixes( monkeypatch: pytest.MonkeyPatch, ) -> None: - quant_config = _stub_quant_config() + quant_config = ExpertsInt8Config() monkeypatch.setattr( fused_moe_utils.rocm_aiter_ops, "is_fusion_moe_shared_experts_enabled", @@ -413,7 +411,7 @@ def test_resolve_layer_fused_shared_expert_rejects_incompatible_quantization( ) assert not fused_moe_utils.resolve_layer_fused_shared_expert( - _stub_quant_config(), "model.layers.0.mlp" + ExpertsInt8Config(), "model.layers.0.mlp" ) assert "shared experts are excluded" in caplog.text @@ -423,6 +421,7 @@ def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix( ) -> None: class DeepseekV4Config: expert_dtype = "fp4" + online_quantization_config = None def _is_quark_mxfp4_ocp(self, hf_config: object) -> bool: return True @@ -1070,6 +1069,73 @@ def test_quark_shared_expert_fse_compatibility( ) +@pytest.mark.parametrize("exclude_shared_expert", [False, True]) +def test_online_targets_off_shared_expert_defer_to_checkpoint( + exclude_shared_expert: bool, +) -> None: + """Unrelated online targets preserve the checkpoint's FSE decision.""" + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig( + { + **_QUARK_FSE_CONFIG, + "exclude": ( + [f"{shared_prefix}.down_proj"] if exclude_shared_expert else [] + ), + } + ) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*linear_attn.out_proj": "mxfp4"}) + ) + + compatible, reason = is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) + + assert compatible is not exclude_shared_expert + assert reason == ( + f"Quark excludes shared experts at {shared_prefix}" + if exclude_shared_expert + else None + ) + + +def test_online_targeting_only_one_shared_projection_rejects_fse() -> None: + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []}) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={f"{shared_prefix}.down_proj": "mxfp4"}) + ) + + assert is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) == ( + False, + "online quantization targets only part of the shared expert at " + f"{shared_prefix}", + ) + + +def test_online_fp8_shared_expert_rejects_quark_mxfp4_fse() -> None: + shared_prefix = "model.layers.0.mlp.shared_expert" + quant_config = QuarkConfig({**_QUARK_FSE_CONFIG, "exclude": []}) + quant_config.online_quantization_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={f"{shared_prefix}*": "fp8_per_tensor"}) + ) + + assert is_shared_expert_quant_fse_compatible( + quant_config, + "model.layers.0.mlp.experts", + shared_prefix, + ) == ( + False, + "online shared-expert quantization keys do not match the routed expert keys", + ) + + def test_quark_shared_expert_fse_exclude_is_scoped_to_the_layer() -> None: """Excluding one layer's shared experts must not disable FSE elsewhere. @@ -1335,14 +1401,15 @@ def test_quark_packed_layer_config_must_match_global_config() -> None: def test_non_quark_shared_expert_fse_is_incompatible() -> None: compatible, reason = is_shared_expert_quant_fse_compatible( - _stub_quant_config(), + ExpertsInt8Config(), "model.layers.0.mlp.experts", "model.layers.0.mlp.shared_experts", ) assert not compatible assert reason == ( - "shared-expert FSE quantization compatibility is not implemented for object" + "shared-expert FSE quantization compatibility is not implemented for " + "ExpertsInt8Config" ) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index aaad16778f..66c5b49050 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -528,6 +528,7 @@ def test_padded_moe_reload_releases_each_layer( layer.quant_method = method layer.expert_map_manager = SimpleNamespace(map_global_to_local=lambda i: i) layer._loaded_expert_biases = set() + layer._fused_shared_expert_quantizer = None method.create_weights( layer, experts, diff --git a/tests/models/kimi_k3/test_gfx942_int4.py b/tests/models/kimi_k3/test_gfx942_int4.py index 20bef41eb8..e38a8ecff9 100644 --- a/tests/models/kimi_k3/test_gfx942_int4.py +++ b/tests/models/kimi_k3/test_gfx942_int4.py @@ -4,7 +4,7 @@ import sys from collections.abc import Callable from types import ModuleType, SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest @@ -91,9 +91,7 @@ def test_int4_per_group_32_is_not_an_online_moe_method(): QuantizationConfigArgs(moe={"weight": "int4_per_group_32"}) ) assert ( - config.resolve_quant_method_cls( - Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" - ) + config.resolve_quant_method_cls(RoutedExperts, "model.layers.0.mlp.experts") is None ) @@ -103,6 +101,4 @@ def test_unknown_online_moe_weight_still_raises(): QuantizationConfigArgs(moe={"weight": "fp8_per_token"}) ) with pytest.raises(ValueError, match="online quantization"): - config.resolve_quant_method_cls( - Mock(spec=RoutedExperts), "model.layers.0.mlp.experts" - ) + config.resolve_quant_method_cls(RoutedExperts, "model.layers.0.mlp.experts") diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index 7e54599c80..4f60a974ca 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -162,6 +162,7 @@ def test_routed_experts_loads_per_expert_biases(): class Loader: quant_config = None quant_method = object() + _fused_shared_expert_quantizer = None moe_config = SimpleNamespace( is_act_and_mul=True, tp_rank=0, diff --git a/tests/quantization/test_config_utils.py b/tests/quantization/test_config_utils.py index 6a0c27748f..a0cb12a9f8 100644 --- a/tests/quantization/test_config_utils.py +++ b/tests/quantization/test_config_utils.py @@ -2,11 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests quantization configuration matching utilities.""" -from unittest.mock import Mock - import pytest from vllm.config.quantization import QuantizationConfigArgs +from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( should_ignore_layer, ) @@ -58,7 +57,7 @@ def test_online_ignore_supports_fnmatch_patterns(): ) with pytest.raises(ValueError, match="matches both quantization_config.ignore"): - config.resolve_quant_method_cls(Mock(), layer_name) + config.resolve_quant_method_cls(LinearBase, layer_name) @pytest.mark.parametrize( diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index c6b9894009..ce13aa9a65 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -195,7 +195,6 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( config, ) - layer = Mock(spec=LinearBase) for prefix in ( "vision_model.encoder.layers.0.attn.qkv", "vision_model.encoder.layers.1.mlp.fc2", @@ -203,10 +202,10 @@ def test_qianfan_online_fp8_keeps_vision_layers_unquantized( "mlp1.1", "mlp1.3", ): - assert config.resolve_quant_method_cls(layer, prefix) is None + assert config.resolve_quant_method_cls(LinearBase, prefix) is None resolved = config.resolve_quant_method_cls( - layer, "language_model.model.layers.0.self_attn.qkv_proj" + LinearBase, "language_model.model.layers.0.self_attn.qkv_proj" ) assert resolved is not None and resolved[-1] is method_cls assert other_config.ignored_layers == [] @@ -1320,9 +1319,7 @@ def test_online_quantization_targets_ignore_collision() -> None: ) ) with pytest.raises(ValueError, match="matches both quantization_config.ignore"): - config.resolve_quant_method_cls( - Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" - ) + config.resolve_quant_method_cls(LinearBase, "model.layers.0.self_attn.o_proj") def test_online_quantization_targets_reject_unsupported_layer() -> None: diff --git a/tests/quantization/test_online_shared_expert_fusion.py b/tests/quantization/test_online_shared_expert_fusion.py new file mode 100644 index 0000000000..2dc3d5da76 --- /dev/null +++ b/tests/quantization/test_online_shared_expert_fusion.py @@ -0,0 +1,637 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests ROCm AITER fused shared experts with online quantization.""" + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from tests.quantization.utils import load_model_without_vllm_runner +from vllm._aiter_ops import rocm_aiter_ops +from vllm.config.model import ModelConfig +from vllm.config.quantization import ( + QuantizationConfigArgs, + resolve_quantization_config, +) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe.experts import rocm_aiter_moe +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig +from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig +from vllm.model_executor.layers.quantization.quark.quark_moe import ( + QuarkOCP_MX_MoEMethod, +) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + mxfp4_quantize, +) +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + record_metadata_for_reloading, +) +from vllm.model_executor.model_loader.utils import get_model_architecture +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), + reason="ROCm AITER shared-expert tests require ROCm.", +) + +_QUARK_MXFP4_CONFIG: dict[str, Any] = { + "global_quant_config": { + "input_tensors": { + "dtype": "fp4", + "is_dynamic": True, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "output_tensors": None, + "weight": { + "dtype": "fp4", + "is_dynamic": False, + "qscheme": "per_group", + "ch_axis": -1, + "group_size": 32, + "block_size": None, + "symmetric": None, + "round_method": "half_even", + "scale_type": "float", + "scale_format": "e8m0", + "scale_calculation_mode": "even", + "mx_element_dtype": None, + "observer_cls": "PerBlockMXObserver", + "is_scale_quant": False, + "enable_buffer_reuse": False, + "max_input_numel": 4194304, + }, + "bias": None, + "target_device": None, + }, + "algo_config": None, + "softmax_quant_spec": None, + "quant_method": "quark", + "layer_type_quant_config": {}, + "layer_quant_config": {}, + "exclude": ["re:^(?!.*\\.mlp\\.experts(?:\\.|$)).*$"], + "kv_cache_quant_config": {}, + "kv_cache_post_rope": False, + "quant_mode": "eager_mode", + "version": "0.12+9d3d471cdf1", + "export": { + "kv_cache_group": [], + "min_kv_scale": 0.0, + "pack_method": "reorder", + "weight_format": "real_quantized", + "weight_merge_groups": None, + }, +} + +_QUARK_MXFP4_LAYER_CONFIG: dict[str, Any] = _QUARK_MXFP4_CONFIG["global_quant_config"] +_QUARK_MXFP4_CONFIG = { + **_QUARK_MXFP4_CONFIG, + "global_quant_config": { + "input_tensors": {}, + "output_tensors": None, + "weight": {}, + "bias": None, + "target_device": None, + }, + "layer_quant_config": {"*mlp.experts*": _QUARK_MXFP4_LAYER_CONFIG}, +} + + +class _StubAttention(torch.nn.Module): + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__() + self.o_proj = SimpleNamespace(reduce_results=True) + + +@pytest.fixture(autouse=True) +def reset_aiter_shared_expert_topk_metadata(monkeypatch: pytest.MonkeyPatch): + """Isolate AITER's process-global shared-expert routing buffer.""" + monkeypatch.setattr(rocm_aiter_moe, "aiter_topK_meta_data", None) + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() + yield + rocm_aiter_moe.init_aiter_topK_meta_data.cache_clear() + + +def _write_minimal_moe_config(model_path: Path, architecture: str) -> None: + """Write a compact MoE config with one shared expert.""" + config = { + "architectures": [architecture], + "hidden_size": 256, + "intermediate_size": 512, + "moe_intermediate_size": 256, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 64, + "num_hidden_layers": 1, + "vocab_size": 256, + "max_position_embeddings": 64, + "q_lora_rank": 64, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "n_routed_experts": 4, + "num_experts": 4, + "num_experts_per_tok": 1, + "n_group": 1, + "topk_group": 1, + "n_shared_experts": 1, + "first_k_dense_replace": 0, + "moe_layer_freq": 1, + "rms_norm_eps": 1e-6, + "hidden_act": "silu", + "tie_word_embeddings": False, + "rope_theta": 10000, + "quantization_config": _QUARK_MXFP4_CONFIG, + } + if architecture == "AXK1ForCausalLM": + config["model_type"] = "axk1" + elif architecture == "MiniMaxM3SparseForCausalLM": + text_config = { + **config, + "model_type": "minimax_m3_text", + "dense_intermediate_size": 512, + "hidden_act": "swigluoai", + "num_local_experts": 4, + "n_shared_experts": 1, + "moe_layer_freq": [1], + "sparse_attention_config": None, + "quantization_config": { + **_QUARK_MXFP4_CONFIG, + "layer_quant_config": { + "*block_sparse_moe.experts*": _QUARK_MXFP4_LAYER_CONFIG + }, + "exclude": [r"re:^(?!.*\.block_sparse_moe\.experts(?:\.|$)).*$"], + }, + } + config = { + "architectures": [architecture], + "model_type": "minimax_m3_vl", + "text_config": text_config, + } + elif architecture == "Glm4MoeForCausalLM": + config["model_type"] = "glm4_moe" + elif architecture == "Glm4MoeLiteForCausalLM": + config["model_type"] = "glm4_moe_lite" + elif architecture in ("Qwen3NextForCausalLM", "Qwen3_5MoeForCausalLM"): + config.update( + model_type=( + "qwen3_next" + if architecture == "Qwen3NextForCausalLM" + else "qwen3_5_moe_text" + ), + layer_types=[ + "linear_attention" + if architecture == "Qwen3NextForCausalLM" + else "full_attention" + ], + head_dim=64, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + shared_expert_intermediate_size=256, + ) + elif architecture == "Qwen4ExpForCausalLM": + config.update( + model_type="qwen4_exp_text", + layer_types=["full_attention"], + head_dim=64, + linear_key_head_dim=64, + linear_value_head_dim=64, + linear_num_key_heads=4, + linear_num_value_heads=4, + shared_expert_intermediate_size=256, + eos_token_id=1, + hc_count=2, + hc_lowrank=4, + ple_layer_ids=[], + ) + elif architecture == "Glm5NextForCausalLM": + config.update( + model_type="glm5_next", + layer_types=["linear_attention"], + mlp_layer_types=["sparse"], + qk_rope_head_dim=0, + index_topk=1, + index_kpool=1, + pad_token_id=None, + scoring_func="sigmoid", + topk_method="noaux_tc", + mhc=False, + ) + else: + config["model_type"] = ( + "deepseek_v3" + if "V3" in architecture + else "deepseek_v32" + if architecture == "DeepseekV32ForCausalLM" + else "deepseek_v2" + ) + if architecture in ("DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"): + config.update( + index_topk=1, + index_kpool=1, + index_n_heads=1, + index_head_dim=32, + index_kv_lora_rank=32, + ) + (model_path / "config.json").write_text(json.dumps(config)) + + +@pytest.mark.parametrize( + "load_path", ["load_weights", "fused_gate_up", "weight_loader"] +) +def test_online_shared_expert_loads_bf16_weights_into_mxfp4_slot( + default_vllm_config, + dist_init, + load_path: str, +) -> None: + """A BF16 shared expert is quantized while routed MXFP4 weights are loaded. + + `weight_loader` covers DeepSeek-style model loaders, which bypass + `RoutedExperts.load_weights`. + """ + default_vllm_config.model_config = ModelConfig() + hidden_size = intermediate_size = 64 + num_routed_experts = 2 + device = current_platform.device_type + online_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*shared_expert*": "mxfp4"}) + ) + quant_config = QuarkConfig(_QUARK_MXFP4_CONFIG) + quant_config.online_quantization_config = online_config + + with torch.device(device): + runner = FusedMoEFactory( + num_experts=num_routed_experts, + top_k=1, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + n_shared_experts=1, + fuse_shared_experts=True, + shared_expert_prefix="model.layers.0.mlp.shared_expert", + prefix="model.layers.0.mlp.experts", + quant_config=quant_config, + ) + layer = runner.routed_experts + assert isinstance(layer.quant_method, QuarkOCP_MX_MoEMethod) + + routed_gate, _ = mxfp4_quantize( + torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + ) + routed_up, _ = mxfp4_quantize( + torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + ) + routed_down, _ = mxfp4_quantize( + torch.randn(hidden_size, intermediate_size, dtype=torch.bfloat16) + ) + shared_gate = torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + shared_up = torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16) + shared_down = torch.randn(hidden_size, intermediate_size, dtype=torch.bfloat16) + + weights = [ + ("0.gate_proj.weight", routed_gate), + ("0.up_proj.weight", routed_up), + ("0.down_proj.weight", routed_down), + ] + if load_path == "load_weights": + weights += [ + ("2.gate_proj.weight", shared_gate), + ("2.up_proj.weight", shared_up), + ("2.down_proj.weight", shared_down), + ] + elif load_path == "fused_gate_up": + weights += [ + ("2.gate_up_proj.weight", torch.cat([shared_gate, shared_up])), + ("2.down_proj.weight", shared_down), + ] + list(layer.load_weights(weights)) + + if load_path == "weight_loader": + for stem, shard_id, weight in ( + ("w13", "w1", shared_gate), + ("w13", "w3", shared_up), + ("w2", "w2", shared_down), + ): + param = getattr(layer, f"{stem}_weight") + param.weight_loader( + param, weight, f"{layer.layer_name}.{stem}_weight", shard_id, 2 + ) + + expected_gate, expected_gate_scale = mxfp4_quantize(shared_gate) + expected_up, expected_up_scale = mxfp4_quantize(shared_up) + expected_down, expected_down_scale = mxfp4_quantize(shared_down) + assert torch.equal(layer.w13_weight[2, :intermediate_size], expected_gate) + assert torch.equal(layer.w13_weight[2, intermediate_size:], expected_up) + assert torch.equal(layer.w2_weight[2], expected_down) + assert torch.equal( + layer.w13_weight_scale[2, :intermediate_size], expected_gate_scale + ) + assert torch.equal(layer.w13_weight_scale[2, intermediate_size:], expected_up_scale) + assert torch.equal(layer.w2_weight_scale[2], expected_down_scale) + + +def test_online_shared_expert_reload_compatibility( + default_vllm_config, + dist_init, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reload a BF16 shared weight into its MXFP4 slot via meta staging.""" + default_vllm_config.model_config = ModelConfig() + online_config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*shared_expert*": "mxfp4"}) + ) + quant_config = QuarkConfig(_QUARK_MXFP4_CONFIG) + quant_config.online_quantization_config = online_config + + with torch.device(current_platform.device_type): + layer = FusedMoEFactory( + num_experts=2, + top_k=1, + hidden_size=64, + intermediate_size=64, + n_shared_experts=1, + fuse_shared_experts=True, + shared_expert_prefix="model.layers.0.mlp.shared_expert", + prefix="model.layers.0.mlp.experts", + quant_config=quant_config, + ).routed_experts + monkeypatch.setattr( + layer.quant_method, "process_weights_after_loading", lambda _: None + ) + + routed_weights = [] + for expert_id in range(2): + for projection in ("gate_proj", "up_proj", "down_proj"): + weight, scale = mxfp4_quantize( + torch.randn(64, 64, dtype=torch.bfloat16) + ) + prefix = f"{expert_id}.{projection}" + routed_weights.extend( + [(f"{prefix}.weight", weight), (f"{prefix}.weight_scale", scale)] + ) + + shared_gate = torch.randn(64, 64, dtype=torch.bfloat16) + shared_up = torch.randn(64, 64, dtype=torch.bfloat16) + shared_down = torch.randn(64, 64, dtype=torch.bfloat16) + + def checkpoint_weights(gate: torch.Tensor): + return [ + *routed_weights, + ("2.gate_proj.weight", gate), + ("2.up_proj.weight", shared_up), + ("2.down_proj.weight", shared_down), + ] + + list(layer.load_weights(checkpoint_weights(shared_gate))) + old_weight = layer.w13_weight[2, :64].clone() + reloaded_gate = shared_gate + 1 + expected_weight, expected_scale = mxfp4_quantize(reloaded_gate) + assert not torch.equal(old_weight, expected_weight) + + record_metadata_for_reloading(layer) + initialize_layerwise_reload(layer) + + assert layer.w13_weight.is_meta + list(layer.load_weights(checkpoint_weights(reloaded_gate))) + finalize_layerwise_reload(layer, default_vllm_config.model_config) + assert not layer.w13_weight.is_meta + assert torch.equal(layer.w13_weight[2, :64], expected_weight) + assert torch.equal(layer.w13_weight_scale[2, :64], expected_scale) + + +# TODO: Add DeepseekV4ForConditionalGeneration once it uses +# resolve_layer_fused_shared_expert. +@pytest.mark.parametrize( + "architecture", + [ + "AXK1ForCausalLM", + "DeepseekForCausalLM", + "DeepseekV2ForCausalLM", + "DeepseekV32ForCausalLM", + "DeepseekV3ForCausalLM", + "Glm4MoeForCausalLM", + "Glm4MoeLiteForCausalLM", + "Glm5NextForCausalLM", + "GlmMoeDsaForCausalLM", + "MiniMaxM3SparseForCausalLM", + "Qwen3NextForCausalLM", + "Qwen3_5MoeForCausalLM", + "Qwen4ExpForCausalLM", + ], +) +def test_online_quantization( + architecture: str, + tmp_path: Path, + monkeypatch, + dist_init, + workspace_init, +) -> None: + """Online MXFP4 fuses each architecture's shared expert into its MoE.""" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") + rocm_aiter_ops.refresh_env_variables() + if architecture == "Glm5NextForCausalLM": + from vllm.models.glm5next.common import model as glm5_next_model + + for attention in ("Glm5NextLinearAttention", "Glm5NextMLAAttention"): + monkeypatch.setattr(glm5_next_model, attention, _StubAttention) + monkeypatch.setattr( + glm5_next_model, "_fused_shared_experts_tuned", lambda _: True + ) + elif architecture == "Glm4MoeLiteForCausalLM": + from vllm.model_executor.models import glm4_moe_lite + + monkeypatch.setattr(glm4_moe_lite, "Glm4MoeLiteAttention", _StubAttention) + monkeypatch.setattr(glm4_moe_lite, "Glm4MoeLiteMLAAttention", _StubAttention) + elif architecture == "DeepseekV32ForCausalLM": + from vllm.models.deepseek_v32.amd import model as deepseek_v32_model + + monkeypatch.setattr( + deepseek_v32_model, "DeepseekV32MLAAttention", _StubAttention + ) + elif architecture == "Qwen3_5MoeForCausalLM": + from vllm.model_executor.models import qwen3_5 + + monkeypatch.setattr(qwen3_5, "Qwen3NextAttention", _StubAttention) + elif architecture == "Qwen4ExpForCausalLM": + from vllm.models.qwen4_exp.amd import model as qwen4_exp_model + + monkeypatch.setattr(qwen4_exp_model, "Qwen3NextAttention", _StubAttention) + _write_minimal_moe_config(tmp_path, architecture) + expected_shared_expert_name = ( + "shared_expert" if architecture.startswith("Qwen") else "shared_experts" + ) + target_pattern = f"*{expected_shared_expert_name}*" + + logged_messages: list[str] = [] + logged_warnings: list[str] = [] + + def record_info(message: str, *args: object, **_kwargs: object) -> None: + if "Quantizing " in message and "of types" in message: + logged_messages.append(message % args) + + def record_warning(message: str, *args: object, **_kwargs: object) -> None: + logged_warnings.append(message % args) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.base_loader.logger.info", record_info + ) + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning + ) + + model, vllm_config = load_model_without_vllm_runner( + str(tmp_path), + dtype="bfloat16", + quantization="quark", + model_config_kwargs={ + "quantization_config": resolve_quantization_config( + "quark", {"targets": {target_pattern: "mxfp4"}} + ) + }, + model_loader_cls=DummyModelLoader, + ) + expected_model_cls, resolved_architecture = get_model_architecture( + vllm_config.model_config + ) + assert resolved_architecture == architecture + assert isinstance(model, expected_model_cls) + + fused_moes = [ + module + for module in model.modules() + if getattr(module, "is_fused_shared_expert_enabled", False) + and hasattr(module, "experts") + ] + assert len(fused_moes) == 1, f"{architecture} did not enable FSE" + moe = fused_moes[0] + shared_expert_name = ( + "shared_expert" if hasattr(moe, "shared_expert") else "shared_experts" + ) + assert shared_expert_name == expected_shared_expert_name + assert getattr(moe, shared_expert_name) is None + + routed_experts = moe.experts.routed_experts + assert isinstance(routed_experts.quant_method, QuarkOCP_MX_MoEMethod) + shared_expert_prefix = routed_experts.moe_config.shared_expert_prefix + assert shared_expert_prefix is not None + assert shared_expert_prefix.endswith(shared_expert_name) + assert routed_experts._fused_shared_expert_quantizer is not None + expert_map_manager = routed_experts.expert_map_manager + assert expert_map_manager.num_fused_shared_experts == 1 + assert expert_map_manager.map_global_to_local( + routed_experts.moe_config.num_experts + ) == (routed_experts.moe_config.num_experts) + expected_weight_dtype = { + Mxfp4MoeBackend.AITER_MXFP4_BF16: torch.float4_e2m1fn_x2, + Mxfp4MoeBackend.AITER_MXFP4_MXFP4: torch.float4_e2m1fn_x2, + Mxfp4MoeBackend.EMULATION: torch.uint8, + }[routed_experts.quant_method.mxfp4_backend] + assert routed_experts.w13_weight.dtype == expected_weight_dtype + assert routed_experts.w13_weight_scale.dtype == torch.uint8 + assert routed_experts.w2_weight.dtype == expected_weight_dtype + assert routed_experts.w2_weight_scale.dtype == torch.uint8 + assert vllm_config.quant_config is not None + assert vllm_config.quant_config.online_quantization_config is not None + assert not any( + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " + "cannot be enabled" in warning + for warning in logged_warnings + ) + if architecture == "MiniMaxM3SparseForCausalLM": + # TODO: Use resolve_layer_fused_shared_expert in MiniMax M3 and remove + # this branch once its virtual fused projections are registered. + assert logged_messages == ["Quantizing 0 layers of types: "] + else: + assert logged_messages == [ + "Quantizing 2 layers of types: " + f"mlp.{shared_expert_name}.down_proj: 1 " + f"(from targets: {target_pattern}, mxfp4); " + f"mlp.{shared_expert_name}.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] + + +def test_online_gate_up_proj_target_disables_fse( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + dist_init, + workspace_init, +) -> None: + """A partial shared-expert target quantizes while disabling FSE.""" + architecture = "DeepseekForCausalLM" + target_pattern = "*shared_experts.gate_up_proj*" + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "1") + rocm_aiter_ops.refresh_env_variables() + _write_minimal_moe_config(tmp_path, architecture) + + logged_messages: list[str] = [] + logged_warnings: list[str] = [] + + def record_info(message: str, *args: object, **_kwargs: object) -> None: + if "Quantizing " in message and "of types" in message: + logged_messages.append(message % args) + + def record_warning(message: str, *args: object, **_kwargs: object) -> None: + logged_warnings.append(message % args) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.base_loader.logger.info", record_info + ) + monkeypatch.setattr( + "vllm.model_executor.layers.fused_moe.utils.logger.warning", record_warning + ) + + model, _ = load_model_without_vllm_runner( + str(tmp_path), + dtype="bfloat16", + quantization="quark", + model_config_kwargs={ + "quantization_config": resolve_quantization_config( + "quark", {"targets": {target_pattern: "mxfp4"}} + ) + }, + model_loader_cls=DummyModelLoader, + ) + + moe = next( + module + for module in model.modules() + if hasattr(module, "is_fused_shared_expert_enabled") + and hasattr(module, "shared_experts") + ) + assert not moe.is_fused_shared_expert_enabled + assert moe.shared_experts is not None + assert logged_messages == [ + "Quantizing 1 layers of types: " + "mlp.shared_experts.gate_up_proj: 1 " + f"(from targets: {target_pattern}, mxfp4)" + ] + assert logged_warnings == [ + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but cannot be " + "enabled - skipping for this layer: online quantization targets only " + "part of the shared expert at model.layers.0.mlp.shared_experts." + ] diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 7e2e21a3ac..f24566f5fa 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for QuantizationConfigArgs parsing.""" -from unittest.mock import Mock - import pytest from tests.quantization.utils import quant_config_args, quant_spec @@ -244,6 +242,4 @@ def test_targets_reject_moe_only_shorthand_for_linear_layer(): ) with pytest.raises(ValueError, match="does not define a QuantSpec"): - config.resolve_quant_method_cls( - Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" - ) + config.resolve_quant_method_cls(LinearBase, "model.layers.0.self_attn.o_proj") diff --git a/tests/quantization/utils.py b/tests/quantization/utils.py index ed9776ad1c..d8446952f0 100644 --- a/tests/quantization/utils.py +++ b/tests/quantization/utils.py @@ -15,7 +15,11 @@ from vllm.config import ( VllmConfig, set_current_vllm_config, ) -from vllm.config.quantization import QuantizationConfigArgs, QuantSpec +from vllm.config.quantization import ( + QuantizationConfigArgs, + QuantSpec, + resolve_quantization_config, +) from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.platforms import current_platform @@ -40,10 +44,17 @@ def _limit_num_hidden_layers( original_load_weights = model.load_weights def should_load_weight(name: str) -> bool: - for prefix in ("model.layers.", "layers."): - if name.startswith(prefix): - layer_idx = int(name.removeprefix(prefix).split(".", 1)[0]) - return layer_idx < num_hidden_layers + layers_prefix = "layers." + if name.startswith(layers_prefix): + layer_name = name.removeprefix(layers_prefix) + elif ".layers." in name: + layer_name = name.split(".layers.", maxsplit=1)[1] + else: + return True + + layer_idx_str = layer_name.split(".", maxsplit=1)[0] + if layer_idx_str.isdigit(): + return int(layer_idx_str) < num_hidden_layers return True def load_weights(weights): @@ -71,13 +82,19 @@ def load_model_without_vllm_runner( quantization=quantization, **(model_config_kwargs or {}), ) + model_config.quantization_config = resolve_quantization_config( + model_config.quantization, model_config.quantization_config + ) vllm_config_args = dict(vllm_config_kwargs or {}) vllm_config_args.setdefault( "compilation_config", CompilationConfig(mode=CompilationMode.NONE) ) vllm_config = VllmConfig(model_config=model_config, **vllm_config_args) hf_overrides = (model_config_kwargs or {}).get("hf_overrides") or {} - num_hidden_layers = hf_overrides.get("num_hidden_layers") + text_config_overrides = hf_overrides.get("text_config") or {} + num_hidden_layers = text_config_overrides.get( + "num_hidden_layers", hf_overrides.get("num_hidden_layers") + ) with set_current_vllm_config(vllm_config): model_loader = model_loader_cls(vllm_config.load_config) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 158bc90769..f20af8a6f8 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1293,6 +1293,7 @@ class FusedMoEConfig: has_bias: bool = False is_lora_enabled: bool = False has_hash_routing: bool = False + shared_expert_prefix: str | None = None # When True, the MoE skips its final cross-rank all-reduce (and the separate # shared-expert reduce), returning the partial per-rank sum. The caller is diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 294c538c7b..6b7c64bdc4 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -122,6 +122,7 @@ def FusedMoEFactory( is_fused_checkpoint_transposed: bool = False, n_shared_experts: int | None = None, fuse_shared_experts: bool = False, + shared_expert_prefix: str | None = None, router_logits_dtype: torch.dtype | None = None, gate: torch.nn.Module | None = None, shared_experts: torch.nn.Module | None = None, @@ -197,6 +198,7 @@ def FusedMoEFactory( n_shared_experts: Number of shared experts to fuse into the routed grouped GEMM (ROCm; requires aiter FSE or the router-append path) fuse_shared_experts: Whether to enable shared-expert fusion. + shared_expert_prefix: Checkpoint prefix for the fused shared expert. router_logits_dtype: Data type for router logits buffers gate: Pre-configured gate module shared_experts: Pre-configured shared experts module @@ -360,6 +362,7 @@ def FusedMoEFactory( elastic_ep_max_dp_size=vllm_config.parallel_config.elastic_ep_max_dp_size, has_bias=has_bias, is_lora_enabled=vllm_config.lora_config is not None, + shared_expert_prefix=shared_expert_prefix, activation=moe_activation, device=vllm_config.device_config.device, routing_method=router.routing_method_type, # Not ideal diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index adc8a27010..ae5829a311 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) +from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, resolve_quant_method, @@ -153,6 +154,7 @@ class RoutedExperts(PluggableLayer): self.apply_router_weight_on_input = apply_router_weight_on_input # End random parameters self._loaded_expert_biases: set[str] = set() + self._fused_shared_expert_quantizer = self._get_fused_shared_expert_quantizer() self.quant_method = self._get_quant_method( self.layer_name, @@ -630,6 +632,31 @@ class RoutedExperts(PluggableLayer): # _to_scalar's reshape(()) would reject the size-2 weight_shape. param_data[expert_id] = loaded_weight + def _get_fused_shared_expert_quantizer( + self, + ) -> Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] | None: + """Return the quantizer for an online-quantized fused shared expert.""" + online_config = getattr(self.quant_config, "online_quantization_config", None) + prefix = self.moe_config.shared_expert_prefix + if ( + online_config is None + or prefix is None + or self.expert_map_manager.num_fused_shared_experts == 0 + ): + return None + # FSE compatibility checks require every shared-expert projection to + # use the routed experts' weight key. + resolved = online_config.resolve_quant_method_cls( + LinearBase, f"{prefix}.down_proj" + ) + if resolved is None: + return None + from vllm.model_executor.layers.quantization.online.base import ( + ONLINE_SHARED_EXPERT_QUANTIZERS, + ) + + return ONLINE_SHARED_EXPERT_QUANTIZERS[resolved[3].weight] + @overload def weight_loader( self, @@ -687,6 +714,35 @@ class RoutedExperts(PluggableLayer): return False if return_success else None # Hereafter, `expert_id` is local physical id + # A full-precision fused shared expert is quantized into the checkpoint + # layout, then loaded like a pre-quantized expert weight and scale. + if ( + self._fused_shared_expert_quantizer is not None + and self.moe_config.num_logical_experts + <= global_expert_id + < self.moe_config.num_logical_experts + + self.expert_map_manager.num_fused_shared_experts + and weight_name.endswith("_weight") + and loaded_weight.is_floating_point() + and loaded_weight.dtype != param.dtype + ): + weight, weight_scale = self._fused_shared_expert_quantizer( + loaded_weight.to(self.moe_config.device) + ) + stem = "w2" if shard_id == "w2" else "w13" + self.weight_loader( + getattr(self, f"{stem}_weight_scale"), + weight_scale, + f"{weight_name}_scale", + shard_id, + global_expert_id, + return_success=True, + ) + loaded = self.weight_loader( + param, weight, weight_name, shard_id, global_expert_id, True + ) + return loaded if return_success else None + # is_transposed: if the dim to shard the weight # should be flipped. Required by GPTQ/AWQ (K-first format). # should be whatever dimension intermediate_size_per_partition is diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index a98e181cae..c724a7c7cb 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -84,6 +84,40 @@ def resolve_layer_fused_shared_expert( else (True, None) ) is_fused_shared_expert_enabled = fse_requested and fse_compatible + + # In online shared_expert quantization is used, register it to + # `online_quantization_config.quantized_layers` ahead of weight loading. + if is_fused_shared_expert_enabled and quant_config is not None: + online_quantization_config = quant_config.online_quantization_config + if online_quantization_config is not None: + from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, + ) + + for projection_name in ("gate_up_proj", "down_proj"): + projection_prefix = f"{prefix}.{shared_expert_name}.{projection_name}" + quant_method_metadata = ( + online_quantization_config.resolve_quant_method_cls( + LinearBase, projection_prefix + ) + ) + + if quant_method_metadata is None: + continue + + source, quant_key_str, target_pattern, _, quant_method_cls = ( + quant_method_metadata + ) + + if quant_method_cls in (None, UnquantizedLinearMethod): + continue + + online_quantization_config.quantized_layers[projection_prefix] = ( + source.value, + quant_key_str, + target_pattern, + ) if fse_requested and not is_fused_shared_expert_enabled: logger.warning( "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled but " diff --git a/vllm/model_executor/layers/quantization/base_config.py b/vllm/model_executor/layers/quantization/base_config.py index 6c1f934739..3ec2e728e1 100644 --- a/vllm/model_executor/layers/quantization/base_config.py +++ b/vllm/model_executor/layers/quantization/base_config.py @@ -320,7 +320,7 @@ def resolve_quant_method( base_quant_method, (UnquantizedLinearMethod, UnquantizedFusedMoEMethod) ) online_target = quant_config.online_quantization_config.resolve_quant_method_cls( - layer, prefix + type(layer), prefix ) if checkpoint_is_quantized: diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index 3f35b56a31..511714837f 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Mapping +from collections.abc import Callable, Mapping from enum import Enum from types import MappingProxyType from typing import Any @@ -60,6 +60,7 @@ from vllm.model_executor.layers.quantization.utils.config_utils import ( find_matching_patterns, get_layer_name_after_index, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Static128BlockSym, @@ -104,6 +105,15 @@ _ONLINE_MOE_METHODS: dict[QuantKey, type] = { kNvfp4Static: Nvfp4OnlineMoEMethod, } +# Quantizers for full-precision fused shared-expert weights, keyed by the routed +# experts' weight QuantKey. Each shard is quantized to the checkpoint layout on +# arrival, so only block-local schemes are supported. +ONLINE_SHARED_EXPERT_QUANTIZERS: dict[ + QuantKey, Callable[[torch.Tensor], tuple[torch.Tensor, torch.Tensor]] +] = { + kMxfp4Static: mxfp4_quantize, +} + def _find_matching_targets( prefix: str, @@ -217,14 +227,14 @@ class OnlineQuantizationConfig(QuantizationConfig): self, spec: QuantSpec | None, table: dict[QuantKey, type], - layer: torch.nn.Module, + layer_cls: type[torch.nn.Module], ) -> type | None: """Resolve the online method class for a layer's quantization spec. Args: spec: Quantization specification to resolve. table: Mapping from weight quantization keys to method classes. - layer: Layer that will use the resolved method. + layer_cls: Type of the layer that will use the resolved method. Returns: The matching method class, or None when ``spec`` has no weight @@ -239,7 +249,7 @@ class OnlineQuantizationConfig(QuantizationConfig): cls = table.get(spec.weight) if cls is None: raise ValueError( - f"online quantization for {type(layer).__name__} with " + f"online quantization for {layer_cls.__name__} with " f"weight={spec.weight} is not supported; supported weight " f"keys: {sorted(str(k) for k in table)}" ) @@ -254,12 +264,12 @@ class OnlineQuantizationConfig(QuantizationConfig): return cls def resolve_quant_method_cls( - self, layer: torch.nn.Module, prefix: str + self, layer_cls: type[torch.nn.Module], prefix: str ) -> tuple[OnlineQuantizationSource, str, str | None, QuantSpec, type] | None: """Resolve quantization metadata and method class without instantiating it. Args: - layer: Layer for which to resolve online quantization. + layer_cls: Type of layer for which to resolve online quantization. prefix: Fully qualified layer name. Returns: @@ -271,17 +281,17 @@ class OnlineQuantizationConfig(QuantizationConfig): quant_spec: QuantSpec | None if self.args.targets is not None: resolved_pattern = self._resolve_targets_quant_method_metadata( - prefix, layer + prefix, layer_cls ) if resolved_pattern is None: return None source, quant_key_str, target_pattern, quant_spec, table = resolved_pattern else: - if isinstance(layer, LinearBase): + if issubclass(layer_cls, LinearBase): source = OnlineQuantizationSource.linear quant_spec = self.args.linear table = _ONLINE_LINEAR_METHODS - elif isinstance(layer, RoutedExperts): + elif issubclass(layer_cls, RoutedExperts): source = OnlineQuantizationSource.moe quant_spec = self.args.moe table = _ONLINE_MOE_METHODS @@ -298,14 +308,14 @@ class OnlineQuantizationConfig(QuantizationConfig): quant_key_str = str(quant_spec) target_pattern = None - quant_method_cls = self._get_method_cls(quant_spec, table, layer) + quant_method_cls = self._get_method_cls(quant_spec, table, layer_cls) if quant_method_cls is None: return None assert quant_spec is not None return source, quant_key_str, target_pattern, quant_spec, quant_method_cls def _resolve_targets_quant_method_metadata( - self, prefix: str, layer: torch.nn.Module + self, prefix: str, layer_cls: type[torch.nn.Module] ) -> ( tuple[OnlineQuantizationSource, str, str, QuantSpec, dict[QuantKey, type]] | None @@ -314,7 +324,7 @@ class OnlineQuantizationConfig(QuantizationConfig): Args: prefix: Fully qualified layer name. - layer: Layer matched against configured target patterns. + layer_cls: Type of layer matched against configured target patterns. Returns: A tuple of source, quantization key string, target pattern, spec, @@ -349,22 +359,22 @@ class OnlineQuantizationConfig(QuantizationConfig): target_pattern = matches[0] quant_key_str = self.args.targets[target_pattern] shorthand = _ONLINE_SHORTHANDS[quant_key_str] - if isinstance(layer, LinearBase): + if issubclass(layer_cls, LinearBase): quant_spec = shorthand.linear table = _ONLINE_LINEAR_METHODS - elif isinstance(layer, RoutedExperts): + elif issubclass(layer_cls, RoutedExperts): quant_spec = shorthand.moe table = _ONLINE_MOE_METHODS else: raise ValueError( f"Layer {prefix} was matched by quantization_config.targets " f"({target_pattern}), but online quantization is not supported for " - f"{type(layer).__name__}." + f"{layer_cls.__name__}." ) if quant_spec is None: raise ValueError( f"targets pattern {target_pattern} = {quant_key_str} does " - f"not define a QuantSpec for {type(layer).__name__} layers " + f"not define a QuantSpec for {layer_cls.__name__} layers " f"(matched at {prefix})." ) return ( @@ -379,7 +389,7 @@ class OnlineQuantizationConfig(QuantizationConfig): self, layer: torch.nn.Module, prefix: str ) -> "QuantizeMethodBase | None": # `targets` takes precedence over `moe` and `linear` and is exclusive. - resolved = self.resolve_quant_method_cls(layer, prefix) + resolved = self.resolve_quant_method_cls(type(layer), prefix) if resolved is not None: source, quant_key_str, target_pattern, _, quant_method_cls = resolved self.quantized_layers[prefix] = ( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 23a36c8fbc..f4662ad5a6 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -12,8 +12,6 @@ if TYPE_CHECKING: FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend - from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey - import vllm.envs as envs from vllm import _custom_ops as ops from vllm.config import get_current_vllm_config @@ -40,6 +38,7 @@ from vllm.model_executor.layers.quantization.online.moe_base import ( ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + QuantKey, amax_for_moe_weight_quant, amax_for_tp_weight_quant, create_fp8_quant_key, @@ -122,6 +121,7 @@ class OnlineLinearBase(LinearMethodBase): weights onto meta device and materializes them just-in-time.""" uses_meta_device: bool = True + activation_quant_key: QuantKey | None def __init__(self): self.out_dtype = torch.get_default_dtype() @@ -203,6 +203,10 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" + activation_quant_key = ( + kFp8DynamicTokenSym if cutlass_fp8_supported() else kFp8DynamicTensorSym + ) + def __init__(self): super().__init__() @@ -211,11 +215,6 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): self.use_marlin = False self.marlin_input_dtype = None self.weight_quant_key = kFp8StaticTensorSym - # Use per-token quantization for better perf if dynamic and cutlass - if cutlass_fp8_supported(): - self.activation_quant_key = kFp8DynamicTokenSym - else: - self.activation_quant_key = kFp8DynamicTensorSym def create_weights( self, @@ -237,6 +236,8 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): **extra_weight_attrs, ) + # TODO: init_fp8_linear_kernel should support activation_quant_key=None + assert self.activation_quant_key is not None self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, @@ -305,13 +306,11 @@ class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): """Online blockwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-block during loading.""" + activation_quant_key = kFp8Dynamic128Sym + def __init__(self): super().__init__() self.weight_block_size = [128, 128] - self.activation_quant_key = create_fp8_quant_key( - static=False, - group_shape=GroupShape(1, self.weight_block_size[0]), - ) self.weight_quant_key = create_fp8_quant_key( static=True, group_shape=GroupShape(*self.weight_block_size) ) @@ -337,6 +336,8 @@ class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): ) layer.weight_block_size = self.weight_block_size + # TODO: init_fp8_linear_kernel should support activation_quant_key=None + assert self.activation_quant_key is not None self.fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 9cddab85c1..ecc196e9f1 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -9,6 +9,7 @@ import torch from torch.nn import Module from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey if TYPE_CHECKING: import vllm.model_executor.layers.fused_moe.modular_kernel as mk @@ -42,6 +43,8 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase): FP8 with block-32 scales) during weight loading. """ + activation_quant_key: QuantKey | None = None + def create_weights( self, layer: torch.nn.Module, diff --git a/vllm/model_executor/layers/quantization/utils/config_utils.py b/vllm/model_executor/layers/quantization/utils/config_utils.py index eb50ed1a23..6e294adb3a 100644 --- a/vllm/model_executor/layers/quantization/utils/config_utils.py +++ b/vllm/model_executor/layers/quantization/utils/config_utils.py @@ -8,9 +8,13 @@ from typing import TYPE_CHECKING import regex as re +from vllm.logger import init_logger + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +logger = init_logger(__name__) + def get_quark_ocp_mx_group_size( quant_config: "QuantizationConfig | None", @@ -70,18 +74,100 @@ def is_shared_expert_quant_fse_compatible( A compatibility flag and, when incompatible, the reason. """ + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + from vllm.model_executor.layers.quantization.online.fp8 import OnlineLinearBase + from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, + ) + from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config + if projection_names is None: projection_names = ["gate_up_proj", "down_proj"] if quant_config is None: return True, None - from vllm.model_executor.layers.quantization.fp8 import Fp8Config - from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig - from vllm.model_executor.layers.quantization.utils.quant_utils import ( - is_layer_skipped, - ) - from vllm.models.deepseek_v4.quant_config import DeepseekV4FP8Config + online_quant_config = quant_config.online_quantization_config + shared_expert_targets = None + if online_quant_config is not None: + from vllm.model_executor.layers.linear import LinearBase + + online_quant_config.packed_modules_mapping = quant_config.packed_modules_mapping + shared_expert_targets = [ + online_quant_config.resolve_quant_method_cls( + LinearBase, f"{shared_expert_prefix}.{projection_name}" + ) + for projection_name in projection_names + ] + if all(target is None for target in shared_expert_targets): + shared_expert_targets = None + + if shared_expert_targets is not None: + from vllm.model_executor.layers.fused_moe import RoutedExperts + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.layers.quantization.online.base import ( + ONLINE_SHARED_EXPERT_QUANTIZERS, + ) + + # NOTE: online shared experts quantization check is only implemented for quark + # quant method at the moment. This can be extended here for other quant methods. + # Requirement is: + # - Retrieve routed weight quant key + # - Retrieve routed activation quant key + # and ensure the online quantization set on shared expert match them. + if isinstance(quant_config, QuarkConfig): + routed_weight_key, routed_activation_key, _ = ( + quant_config.get_quant_method_target(expert_prefix, RoutedExperts) + ) + else: + reason = ( + "online quantization FSE compatibility check is not implemented for " + f"{type(quant_config).__name__}" + ) + logger.warning_once(reason) + return False, reason + + if routed_weight_key not in ONLINE_SHARED_EXPERT_QUANTIZERS: + return ( + False, + "online shared-expert quantization is not supported for routed " + f"expert weight key {routed_weight_key}", + ) + + for shared_expert_target in shared_expert_targets: + if shared_expert_target is None: + return ( + False, + "online quantization targets only part of the shared expert at " + f"{shared_expert_prefix}", + ) + + _, _, _, shared_quant_spec, shared_method_cls = shared_expert_target + if shared_method_cls in (None, UnquantizedLinearMethod): + return ( + False, + "online quantization excludes FSE weights at " + f"{shared_expert_prefix}", + ) + assert issubclass(shared_method_cls, OnlineLinearBase) + + # NOTE: We can not rely on shared_quant_spec.activation_quant_key + # as _ONLINE_SHORTHANDS: dict[str, QuantSpec] in quantization.py + # does not define the default activation quant key. + # `online_quant_config.resolve_quant_method_cls` should return the + # activation quant key as well. Only some online methods define it + # on the class, so compare the weight key first. + if ( + shared_quant_spec.weight != routed_weight_key + or shared_method_cls.activation_quant_key != routed_activation_key + ): + return ( + False, + "online shared-expert quantization keys do not match the " + "routed expert keys", + ) + return True, None if isinstance(quant_config, DeepseekV4FP8Config): from vllm.config import get_current_vllm_config diff --git a/vllm/model_executor/models/AXK1.py b/vllm/model_executor/models/AXK1.py index da38035db9..d8ceb5083c 100644 --- a/vllm/model_executor/models/AXK1.py +++ b/vllm/model_executor/models/AXK1.py @@ -199,6 +199,7 @@ class AXK1MoE(nn.Module): if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 70d87529f6..a1a919fae6 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -396,6 +396,7 @@ class DeepseekV2MoE(nn.Module): if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", router_logits_dtype=self.gate.out_dtype, ) diff --git a/vllm/model_executor/models/glm4_moe.py b/vllm/model_executor/models/glm4_moe.py index 5e5974e209..4e541a4941 100644 --- a/vllm/model_executor/models/glm4_moe.py +++ b/vllm/model_executor/models/glm4_moe.py @@ -204,6 +204,7 @@ class Glm4MoE(nn.Module): config.n_shared_experts if self.is_fused_shared_expert_enabled else None ), fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index 19fc8c1562..9580894e22 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/vllm/model_executor/models/qwen3_next.py @@ -232,6 +232,7 @@ class Qwen3NextSparseMoeBlock(nn.Module): is_sequence_parallel=self.is_sequence_parallel, n_shared_experts=1 if self.shared_expert is None else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_expert", shared_expert_gate=self.shared_expert_gate if self.shared_expert is None else None, diff --git a/vllm/models/glm5next/common/model.py b/vllm/models/glm5next/common/model.py index 778738c137..26becc9619 100644 --- a/vllm/models/glm5next/common/model.py +++ b/vllm/models/glm5next/common/model.py @@ -307,6 +307,7 @@ class Glm5NextMoE(nn.Module): if self.is_fused_shared_expert_enabled else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", router_logits_dtype=self.gate.out_dtype, swiglu_limit=swiglu_limit, ) diff --git a/vllm/models/minimax_m3/amd/model.py b/vllm/models/minimax_m3/amd/model.py index 9105f97af8..3bf8928603 100644 --- a/vllm/models/minimax_m3/amd/model.py +++ b/vllm/models/minimax_m3/amd/model.py @@ -465,6 +465,7 @@ class MiniMaxM3MoE(nn.Module): self.n_shared_experts if self.is_fused_shared_expert_enabled else None ), fuse_shared_experts=self.is_fused_shared_expert_enabled, + shared_expert_prefix=f"{prefix}.shared_experts", quant_config=quant_config, prefix=f"{prefix}.experts", )