iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Quantization] Enable shared expert fusion compatibility with online shared_expert quantization (showcase: along Quark MXFP4 routed experts) (#55686)

Browse Source 
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
Signed-off-by: fxmarty-amd <felmarty@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>
This commit is contained in:
fxmarty-amd
2026-10-08 12:04:46 -04:00
committed by GitHub
co-authored by OpenAI Codex Misha Goin
parent e2414ec1db
commit 9f3d70e081
24 changed files with 984 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/layers/test_fused_shared_expert.py
+78 -11
View File
@@ -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"
)
tests/model_executor/model_loader/test_reload.py
+1
View File
@@ -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,
tests/models/kimi_k3/test_gfx942_int4.py
+3 -7
View File
@@ -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")
tests/quantization/test_auto_gptq.py
+1
View File
@@ -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,
tests/quantization/test_config_utils.py
+2 -3
View File
@@ -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(
tests/quantization/test_online.py
+3 -6
View File
@@ -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:
tests/quantization/test_online_shared_expert_fusion.py
+637
View File
@@ -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."
]
tests/quantization/test_quantization_config_args.py
+1 -5
View File
@@ -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")
tests/quantization/utils.py
+23 -6
View File
@@ -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)
vllm/model_executor/layers/fused_moe/config.py
+1
View File
@@ -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
vllm/model_executor/layers/fused_moe/layer.py
+3
View File
@@ -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
vllm/model_executor/layers/fused_moe/routed_experts.py
+56
View File
@@ -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
vllm/model_executor/layers/fused_moe/utils.py
+34
View File
@@ -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 "
vllm/model_executor/layers/quantization/base_config.py
+1 -1
View File
@@ -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:
vllm/model_executor/layers/quantization/online/base.py
+27 -17
View File
@@ -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] = (
vllm/model_executor/layers/quantization/online/fp8.py
+12 -11
View File
@@ -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,
vllm/model_executor/layers/quantization/online/mxfp8.py
+3
View File
@@ -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,
vllm/model_executor/layers/quantization/utils/config_utils.py
+92 -6
View File
@@ -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
vllm/model_executor/models/AXK1.py
+1
View File
@@ -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:
vllm/model_executor/models/deepseek_v2.py
+1
View File
@@ -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,
)
vllm/model_executor/models/glm4_moe.py
+1
View File
@@ -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:
vllm/model_executor/models/qwen3_next.py
+1
View File
@@ -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,
vllm/models/glm5next/common/model.py
+1
View File
@@ -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,
)
vllm/models/minimax_m3/amd/model.py
+1
View File
@@ -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",
)