diff --git a/docs/features/quantization/online.md b/docs/features/quantization/online.md index d3a4edd2710f..478bc160e839 100644 --- a/docs/features/quantization/online.md +++ b/docs/features/quantization/online.md @@ -86,12 +86,22 @@ vllm serve --quantization-config '{"moe":{"activation":"mxfp8"}}' vllm serve --quantization-config.moe.activation mxfp8 ``` -### Activation overrides on already-quantized checkpoints +### Activation overrides -For checkpoint-quantized models, `quantization_config` lets you pick an -activation format independently of the baked-in weights. The supported -overrides are checkpoint-specific; today this is wired up for MXFP4 MoE -checkpoints (gpt-oss) where you can opt into FP8 activations: +`quantization_config` lets you select an activation format independently of +the weight format. For online quantization, it can be used to override the default +activation quantization. + +For example, to enable MXFP4 MoE weights along MXFP8 activations: + +```bash +vllm serve Qwen/Qwen3.5-35B-A3B --quantization online \ + --quantization-config '{"moe":{"weight":"mxfp4","activation":"mxfp8"}}' +``` + +For pre-quantized models, the supported activation overrides are specific to each +weight quantization type; today this is for example available for MXFP4 +MoE checkpoints where you can opt into e.g. FP8 activations: ```bash vllm serve openai/gpt-oss-20b --quantization-config.moe.activation mxfp8 diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index fbd5cbc56392..59086b008e5a 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -24,6 +24,7 @@ from vllm.config.model import ModelConfig from vllm.config.quantization import ( QuantizationConfigArgs, + QuantSpec, resolve_quantization_config, ) from vllm.config.vllm import VllmConfig @@ -35,7 +36,13 @@ MarlinMxfp8LinearKernel, ) from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoEFactory +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, RoutedExperts +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, LinearBase, @@ -75,6 +82,7 @@ ) from vllm.model_executor.layers.quantization.online.mxfp8 import ( Mxfp8OnlineLinearMethod, + Mxfp8OnlineMoEMethod, ) from vllm.model_executor.layers.quantization.online.nvfp4 import ( Nvfp4OnlineMoEMethod, @@ -91,7 +99,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( amax_for_moe_weight_quant, amax_for_tp_weight_quant, + kFp8Dynamic128Sym, + kFp8DynamicTensorSym, + kFp8DynamicTokenSym, + kFp8Static128BlockSym, + kFp8StaticChannelSym, + kFp8StaticTensorSym, kMxfp8Dynamic, + kMxfp8Static, weight_amax, ) from vllm.model_executor.layers.vocab_parallel_embedding import ( @@ -124,6 +139,161 @@ def on_gfx942() -> bool: ) +@pytest.mark.parametrize( + ( + "method_cls", + "weight_key", + "activation_quant_key", + "has_activation_override", + "raises", + ), + [ + ( + Fp8PerTensorOnlineLinearMethod, + kFp8StaticTensorSym, + kFp8DynamicTensorSym, + True, + False, + ), + ( + Fp8PerBlockOnlineLinearMethod, + kFp8Static128BlockSym, + kFp8DynamicTensorSym, + True, + False, + ), + ( + Fp8PtpcOnlineLinearMethod, + kFp8StaticChannelSym, + kFp8DynamicTensorSym, + True, + False, + ), + ( + Fp8PerBlockOnlineLinearMethod, + kFp8Static128BlockSym, + kFp8Dynamic128Sym, + False, + False, + ), + ( + Fp8PtpcOnlineLinearMethod, + kFp8StaticChannelSym, + kFp8DynamicTokenSym, + False, + False, + ), + (Fp8PerTensorOnlineLinearMethod, kFp8StaticTensorSym, None, True, True), + (Fp8PerBlockOnlineLinearMethod, kFp8Static128BlockSym, None, True, True), + (Fp8PtpcOnlineLinearMethod, kFp8StaticChannelSym, None, True, True), + (Mxfp8OnlineLinearMethod, kMxfp8Static, None, True, True), + ], +) +def test_online_linear_methods_resolve_activation_quant_key( + default_vllm_config, + method_cls, + weight_key, + activation_quant_key, + has_activation_override, + raises, +) -> None: + class TestLinear(LinearBase): + def __init__(self): + torch.nn.Module.__init__(self) + + spec = QuantSpec(weight=weight_key) + if has_activation_override: + spec = QuantSpec(weight=weight_key, activation=activation_quant_key) + args = QuantizationConfigArgs(linear=spec) + default_vllm_config.model_config = SimpleNamespace( + quantization_config=args, + dtype=torch.bfloat16, + ) + + if raises: + with pytest.raises(NotImplementedError, match="activation=null"): + OnlineQuantizationConfig(args).get_quant_method(TestLinear(), "linear") + else: + method = OnlineQuantizationConfig(args).get_quant_method(TestLinear(), "linear") + assert method is not None + assert isinstance(method, method_cls) + assert method.activation_quant_key == activation_quant_key + + +@pytest.mark.parametrize( + ("method_cls", "weight_key", "activation_quant_key", "has_activation_override"), + [ + ( + Fp8PerTensorOnlineMoEMethod, + kFp8StaticTensorSym, + kFp8DynamicTensorSym, + True, + ), + ( + Fp8PerBlockOnlineMoEMethod, + kFp8Static128BlockSym, + kFp8Dynamic128Sym, + True, + ), + ( + Fp8PtpcOnlineMoEMethod, + kFp8StaticChannelSym, + kFp8DynamicTokenSym, + True, + ), + (Fp8PerTensorOnlineMoEMethod, kFp8StaticTensorSym, kFp8DynamicTensorSym, False), + (Fp8PerBlockOnlineMoEMethod, kFp8Static128BlockSym, kFp8Dynamic128Sym, False), + (Fp8PtpcOnlineMoEMethod, kFp8StaticChannelSym, kFp8DynamicTokenSym, False), + ], +) +def test_online_moe_methods_resolve_activation_quant_key( + default_vllm_config, + method_cls, + weight_key, + activation_quant_key, + has_activation_override, +) -> None: + class TestRoutedExperts(RoutedExperts): + def __init__(self): + torch.nn.Module.__init__(self) + self.moe_config = FusedMoEConfig( + num_experts=8, + experts_per_token=2, + hidden_dim=256, + intermediate_size=256, + num_local_experts=8, + num_logical_experts=8, + activation=MoEActivation.SILU, + device=current_platform.device_type, + routing_method=RoutingMethodType.Renormalize, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=torch.bfloat16, + ) + + spec = QuantSpec(weight=weight_key) + if has_activation_override: + spec = QuantSpec(weight=weight_key, activation=activation_quant_key) + args = QuantizationConfigArgs(moe=spec) + default_vllm_config.model_config = SimpleNamespace( + quantization_config=args, + dtype=torch.bfloat16, + ) + layer = TestRoutedExperts() + method = OnlineQuantizationConfig(args).get_quant_method(layer, "experts") + + assert method is not None + assert isinstance(method, method_cls) + assert method.activation_quant_key == activation_quant_key + + +def test_online_mxfp8_moe_rejects_explicit_null_activation() -> None: + with pytest.raises(NotImplementedError, match="activation=null"): + Mxfp8OnlineMoEMethod( + moe=object(), + activation_quant_key=None, + ) + + def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( monkeypatch, ) -> None: @@ -512,30 +682,33 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None: ("model_name,quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls"), [ # simple case - quantization='fp8_per_tensor' - ( + pytest.param( GRANITE_MODEL_NAME, "fp8_per_tensor", None, Fp8PerTensorOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, + id="fp8_per_tensor", ), # simple case - quantization='fp8_per_block' - ( + pytest.param( GRANITE_MODEL_NAME, "fp8_per_block", None, Fp8PerBlockOnlineLinearMethod, Fp8PerBlockOnlineMoEMethod, + id="fp8_per_block", ), - ( + pytest.param( GRANITE_MODEL_NAME, "fp8_per_channel", None, Fp8PtpcOnlineLinearMethod, Fp8PtpcOnlineMoEMethod, + id="fp8_per_channel", ), # quantization='online' with per-layer-kind overrides - ( + pytest.param( GRANITE_MODEL_NAME, "online", { @@ -544,9 +717,10 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None: }, Fp8PerBlockOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, + id="per_layer_kind_overrides", ), # quantization='online' with per-layer target patterns - ( + pytest.param( GRANITE_MODEL_NAME, "online", { @@ -557,9 +731,10 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None: }, Fp8PerBlockOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, + id="targets", ), # ignore with direct layer name - ( + pytest.param( GRANITE_MODEL_NAME, "fp8_per_tensor", # qkv_proj is fused from q_proj/k_proj/v_proj. The shard regex @@ -567,13 +742,15 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None: {"ignore": ["model.layers.1.self_attn.o_proj", "re:.*[qkv]_proj"]}, Fp8PerTensorOnlineLinearMethod, Fp8PerTensorOnlineMoEMethod, + id="ignore", ), - ( + pytest.param( GRANITE_MODEL_NAME, "mxfp4", None, Mxfp4OnlineLinearMethod, Mxfp4OnlineMoEMethod, + id="mxfp4", ), pytest.param( PARTIALLY_PREQUANTIZED_MODEL_NAME, @@ -583,16 +760,37 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None: CompressedTensorsMoEMethod, id="partially_prequantized_checkpoint", ), - ], - ids=[ - "fp8_per_tensor", - "fp8_per_block", - "fp8_per_channel", - "per_layer_kind_overrides", - "targets", - "ignore", - "mxfp4", - "partially_prequantized_checkpoint", + pytest.param( + GRANITE_MODEL_NAME, + "mxfp4", + {"moe": {"weight": "mxfp4", "activation": "mxfp8"}}, + Mxfp4OnlineLinearMethod, + Mxfp4OnlineMoEMethod, + id="mxfp4_activation_override", + ), + pytest.param( + GRANITE_MODEL_NAME, + "mxfp4", + {"moe": {"weight": "mxfp4", "activation": None}}, + Mxfp4OnlineLinearMethod, + Mxfp4OnlineMoEMethod, + id="mxfp4_null_activation_override", + ), + pytest.param( + GRANITE_MODEL_NAME, + "online", + { + "targets": { + "*experts*": { + "weight": "mxfp4", + "activation": "mxfp8", + }, + } + }, + UnquantizedLinearMethod, + Mxfp4OnlineMoEMethod, + id="targets_activation_override", + ), ], ) @pytest.mark.parametrize( @@ -619,7 +817,18 @@ def test_online_quantization( # TODO: Relax this condition once there is a native MXFP4_MXFP4 # linear/moe backend supported on cuda. - if quant_scheme == "mxfp4" and not (on_gfx950() or on_gfx942()): + target_specs = ( + online_quant_args.get("targets", {}) + if isinstance(online_quant_args, dict) + else {} + ) + uses_mxfp4_targets = any( + isinstance(spec, dict) and spec.get("weight") == "mxfp4" + for spec in target_specs.values() + ) + if (quant_scheme == "mxfp4" or uses_mxfp4_targets) and not ( + on_gfx950() or on_gfx942() + ): pytest.skip("mxfp4 online quantization is only tested on AMD gfx942, gfx950.") if current_platform.is_rocm(): @@ -689,6 +898,16 @@ def test_online_quantization( assert isinstance(o_proj.quant_method, expected_linear_cls) assert isinstance(moe._quant_method, expected_moe_cls) + moe_args = ( + online_quant_args.get("moe") if isinstance(online_quant_args, dict) else None + ) + if isinstance(moe_args, dict) and moe_args.get("activation") == "mxfp8": + assert moe._quant_method.activation_quant_key == quant_utils.kMxfp8Dynamic + elif isinstance(moe_args, dict) and "activation" in moe_args: + assert moe._quant_method.activation_quant_key is None + elif isinstance(target_specs.get("*experts*"), dict): + assert moe._quant_method.activation_quant_key == quant_utils.kMxfp8Dynamic + if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance( o_proj.quant_method.kernel, MarlinMxfp8LinearKernel ): @@ -701,6 +920,8 @@ def test_online_quantization( assert o_proj.weight.dtype == MXFP8_VALUE_DTYPE elif quant_scheme == "mxfp4": assert o_proj.weight.dtype == torch.uint8 + elif expected_linear_cls is UnquantizedLinearMethod: + assert o_proj.weight.dtype == torch.bfloat16 elif current_platform.is_cuda() or current_platform.is_xpu(): assert o_proj.weight.dtype == torch.float8_e4m3fn elif current_platform.is_rocm(): @@ -944,15 +1165,32 @@ def test_online_quantization_targets_reject_unsupported_layer() -> None: config.get_quant_method(lm_head, "lm_head") -def test_log_online_quantization(default_vllm_config, monkeypatch) -> None: - config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="fp8_per_tensor")) +@pytest.mark.parametrize( + ("target_quantization", "expected_target_summary"), + [ + ("mxfp4", "mxfp4"), + ( + {"weight": "mxfp4", "activation": "mxfp8"}, + '{"weight":"mxfp4","activation":"mxfp8"}', + ), + ], +) +def test_log_online_quantization( + default_vllm_config, + monkeypatch, + target_quantization, + expected_target_summary, +) -> None: + config = OnlineQuantizationConfig( + QuantizationConfigArgs(targets={"*experts*": target_quantization}) + ) config.quantized_layers = { "model.layers.0.mlp.down_proj": ("linear", "fp8_per_tensor", None), "model.layers.1.mlp.down_proj": ("linear", "fp8_per_tensor", None), - "model.layers.0.self_attn.qkv_proj": ( + "model.layers.0.block_sparse_moe.experts": ( "targets", - "mxfp4", - r"re:.*qkv_proj.*", + expected_target_summary, + "*experts*", ), } default_vllm_config.quant_config = config @@ -968,9 +1206,9 @@ def record_info(message: str, *args: object) -> None: log_online_quantization(default_vllm_config) assert logged_messages == [ - "Quantized 3 layers of types: mlp.down_proj: 2 (from linear: " - "fp8_per_tensor); self_attn.qkv_proj: 1 (from targets: " - "re:.*qkv_proj.*, mxfp4)" + "Quantized 3 layers of types: block_sparse_moe.experts: 1 (from targets: " + f"*experts*, {expected_target_summary}); mlp.down_proj: 2 (from linear: " + "fp8_per_tensor)" ] diff --git a/tests/quantization/test_quantization_config_args.py b/tests/quantization/test_quantization_config_args.py index 0bf203ac4516..8a85b83513cb 100644 --- a/tests/quantization/test_quantization_config_args.py +++ b/tests/quantization/test_quantization_config_args.py @@ -7,6 +7,7 @@ import pytest from vllm.config.quantization import ( + _ONLINE_SHORTHANDS, QUANT_KEY_NAMES, QuantizationConfigArgs, QuantSpec, @@ -23,15 +24,32 @@ kFp8Static128BlockSym, kFp8StaticTensorSym, kInt8StaticChannelSym, + kMxfp4Dynamic, + kMxfp4Static, kMxfp8Dynamic, + kMxfp8Static, ) # ---- QuantSpec ------------------------------------------------------------ +def test_args_resolve_ambiguous_strings_by_field(): + args = QuantizationConfigArgs(moe={"weight": "mxfp8", "activation": "mxfp8"}) + assert args.moe == QuantSpec(weight=kMxfp8Static, activation=kMxfp8Dynamic) + + assert QuantSpec(weight="mxfp8").weight is kMxfp8Static + args = QuantizationConfigArgs(moe=QuantSpec(weight="mxfp8")) + assert args.moe is not None + assert args.moe.weight is kMxfp8Static + + args = QuantizationConfigArgs(moe={"weight": "mxfp4", "activation": "mxfp4"}) + assert args.moe == QuantSpec(weight=kMxfp4Static, activation=kMxfp4Dynamic) + assert QuantSpec(weight="mxfp4").weight is kMxfp4Static + + def test_quant_spec_resolves_string_to_quant_key(): spec = QuantSpec(weight="mxfp8", activation="fp8_per_token") - assert spec.weight == kMxfp8Dynamic + assert spec.weight == kMxfp8Static assert spec.activation == kFp8DynamicTokenSym @@ -41,6 +59,13 @@ def test_quant_spec_accepts_quant_key_directly(): assert spec.activation is None +def test_quant_spec_equality_ignores_explicit_field_metadata(): + assert QuantSpec(weight=kFp8StaticTensorSym) == QuantSpec( + weight=kFp8StaticTensorSym, + activation=None, + ) + + def test_quant_spec_string_representation_uses_quantization_name(): assert str(QuantSpec(weight="mxfp4")) == "mxfp4" @@ -78,7 +103,16 @@ def test_args_string_shorthand_missing_slot_raises(): def test_args_accepts_dict_form(): args = QuantizationConfigArgs(moe={"activation": "mxfp8"}) - assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) + assert args.moe is not None + assert args.moe.weight is None + assert args.moe.activation is kMxfp8Dynamic + + +def test_online_shorthands_do_not_set_activation_defaults(): + for args in _ONLINE_SHORTHANDS.values(): + for spec in (args.linear, args.moe): + if spec is not None: + assert "activation" not in spec.fields_set def test_targets_reject_non_string_keys(): @@ -114,7 +148,9 @@ def test_resolve_quantization_config_only(): # checkpoint that needs an activation override), it's returned as-is. args = resolve_quantization_config(None, {"moe": {"activation": "mxfp8"}}) assert args.linear is None - assert args.moe == QuantSpec(weight=None, activation=kMxfp8Dynamic) + assert args.moe is not None + assert args.moe.weight is None + assert args.moe.activation is kMxfp8Dynamic def test_resolve_merges_explicit_over_shorthand(): @@ -124,8 +160,22 @@ def test_resolve_merges_explicit_over_shorthand(): "fp8_per_tensor", {"linear": "fp8_per_block"}, ) + assert args.linear == QuantSpec(weight=kFp8Static128BlockSym) assert args.moe == QuantSpec(weight=kFp8StaticTensorSym) + assert "activation" not in args.linear.fields_set + + +def test_resolve_preserves_explicit_null_activation(): + args = resolve_quantization_config( + "fp8_per_tensor", + {"linear": {"activation": None}}, + ) + assert args.linear == QuantSpec( + weight=kFp8StaticTensorSym, + activation=None, + ) + assert args.linear.fields_set == frozenset({"weight", "activation"}) def test_resolve_quantization_config_with_checkpoint_quantization(): @@ -137,10 +187,9 @@ def test_resolve_quantization_config_with_checkpoint_quantization(): def test_quant_key_names_round_trip(): - # Every advertised name should round-trip through QuantSpec without error - # and produce the same QuantKey it maps to. + # Every advertised name should round-trip through QuantSpec activation + # parsing without error and produce the same QuantKey it maps to. for name, expected in QUANT_KEY_NAMES.items(): - assert QuantSpec(weight=name).weight == expected, name assert QuantSpec(activation=name).activation == expected, name diff --git a/vllm/config/quantization.py b/vllm/config/quantization.py index 7d32a1730837..a3cbf7ea271f 100644 --- a/vllm/config/quantization.py +++ b/vllm/config/quantization.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Annotated, Any +from dataclasses import field +from typing import Annotated, Any, cast import regex as re from pydantic import ( @@ -26,6 +27,7 @@ kMxfp4Dynamic, kMxfp4Static, kMxfp8Dynamic, + kMxfp8Static, kNvfp4Static, ) @@ -42,6 +44,16 @@ "int8_per_channel_static": kInt8StaticChannelSym, } +# Ambiguous format names select the appropriate dynamic/static behavior +# for either activation (dynamic) or weight (static). +# Explicit ``*_static`` and ``*_dynamic`` names above retain +# their field-independent meanings. +# TODO: possibly deprecate the ``*_static`` and ``*_dynamic`` variants. +_WEIGHT_QUANT_KEY_NAMES: dict[str, QuantKey] = { + "mxfp8": kMxfp8Static, + "mxfp4": kMxfp4Static, +} + def _coerce_quant_key(v: Any) -> QuantKey | None: if v is None or isinstance(v, QuantKey): @@ -68,21 +80,55 @@ def _coerce_quant_key(v: Any) -> QuantKey | None: ), ] +_UNSET = cast(QuantKey | None, object()) + @config class QuantSpec: """Quantization spec for one layer kind (linear or MoE). - `None` on either side means the method class falls back to its own default - (typically inherited from the checkpoint, or unquantized for online). + An omitted activation key lets the quantization implementation choose its + default. An explicitly configured ``activation: null`` requests no + activation quantization; methods that do not support that request raise an + error. """ - weight: QuantKeyField = None + weight: QuantKeyField = _UNSET """Weight quantization key, or a name from QUANT_KEY_NAMES.""" - activation: QuantKeyField = None + activation: QuantKeyField = _UNSET """Activation quantization key, or a name from QUANT_KEY_NAMES.""" + _fields_set: frozenset[str] = field(init=False, repr=False, compare=False) + """Names explicitly provided when constructing this spec.""" + + @field_validator("weight", mode="before") + @classmethod + def _coerce_weight_quant_key(cls, v: Any) -> Any: + if isinstance(v, str) and v in _WEIGHT_QUANT_KEY_NAMES: + return _WEIGHT_QUANT_KEY_NAMES[v] + return v + + def __post_init__(self) -> None: + """ + We need a way to distinguish cases where `activation` is set by the + user or is a default `None`, as `_ONLINE_SHORTHANDS` do not hold the + default activation quant key, but may be overridden by users, including + to `null`. + """ + fields_set = set() + for field_name in ("weight", "activation"): + if getattr(self, field_name) is _UNSET: + setattr(self, field_name, None) + else: + fields_set.add(field_name) + self._fields_set = frozenset(fields_set) + + @property + def fields_set(self) -> frozenset[str]: + """Names explicitly provided when constructing this spec.""" + return self._fields_set + def __str__(self) -> str: def quant_key_str(quant_key: QuantKey | None) -> str: if quant_key is None: @@ -90,12 +136,16 @@ def quant_key_str(quant_key: QuantKey | None) -> str: return next( ( name - for name, known_quant_key in QUANT_KEY_NAMES.items() + for name, known_quant_key in ( + *QUANT_KEY_NAMES.items(), + *_WEIGHT_QUANT_KEY_NAMES.items(), + ) if known_quant_key == quant_key ), str(quant_key), ) + # TODO: QuantKey itself should define `__str__`, instead of having this logic. return quant_key_str(self.weight) @@ -117,11 +167,12 @@ class QuantizationConfigArgs: """Layers to skip quantization for. Online quantization also supports fnmatch-style patterns.""" - targets: dict[str, str] | None = None + targets: dict[str, str | QuantSpec] | None = None """Per-layer online quantization overrides, keyed by exact layer name or regex patterns with a `re:`, or fnmatch-style patterns for online - quantization, mapping to an online shorthand name (see - `_ONLINE_SHORTHANDS`). A layer that matches no pattern is left unquantized. + quantization. A target can map to an online shorthand name (see + `_ONLINE_SHORTHANDS`) or a quantization spec containing `weight` and + `activation` keys. A layer that matches no pattern is left unquantized. Mutually exclusive with `linear` and `moe`. """ @@ -130,6 +181,8 @@ class QuantizationConfigArgs: def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any: if not isinstance(v, str): return v + + # e.g. `--quantization-config.moe mxfp4` field_name = info.field_name assert field_name is not None if v in _ONLINE_SHORTHANDS: @@ -148,16 +201,25 @@ def _validate_targets(cls, v: Any) -> Any: return v if not isinstance(v, dict): raise TypeError(f"targets must be a dict, got {type(v).__name__}") - for pattern, shorthand in v.items(): + targets: dict[str, str | QuantSpec] = {} + for pattern, target in v.items(): if not isinstance(pattern, str): raise ValueError( f"targets keys must be strings, got {type(pattern).__name__}" ) - if not isinstance(shorthand, str) or shorthand not in _ONLINE_SHORTHANDS: + if isinstance(target, str): + if target not in _ONLINE_SHORTHANDS: + raise ValueError( + f"targets[{pattern}] = {target} is not a valid " + f"online shorthand name; expected one of " + f"{sorted(_ONLINE_SHORTHANDS)}" + ) + elif isinstance(target, dict): + target = QuantSpec(**target) + elif not isinstance(target, QuantSpec): raise ValueError( - f"targets[{pattern}] = {shorthand} is not a valid " - f"online shorthand name; expected one of " - f"{sorted(_ONLINE_SHORTHANDS)}" + f"targets[{pattern}] must be an online shorthand name or " + f"a quantization spec, got {type(target).__name__}" ) if pattern.startswith("re:"): try: @@ -166,7 +228,8 @@ def _validate_targets(cls, v: Any) -> Any: raise ValueError( f"targets key {pattern} is not a valid regex: {e}" ) from e - return v + targets[pattern] = target + return targets @model_validator(mode="after") def _validate_targets_exclusivity(self) -> "QuantizationConfigArgs": @@ -200,14 +263,16 @@ def _validate_targets_exclusivity(self) -> "QuantizationConfigArgs": moe=QuantSpec(weight=kFp8StaticChannelSym), ), "mxfp8": QuantizationConfigArgs( - linear=QuantSpec(weight=kMxfp8Dynamic), - moe=QuantSpec(weight=kMxfp8Dynamic), + linear=QuantSpec(weight=kMxfp8Static), + moe=QuantSpec(weight=kMxfp8Static), ), "mxfp4": QuantizationConfigArgs( linear=QuantSpec(weight=kMxfp4Static), moe=QuantSpec(weight=kMxfp4Static), ), # INT8 weight-only on MoE; linear stays unquantized (no `linear` field). + # TODO: this is broken since at least #41566, as Int8OnlineMoEMethod + # defaults to activation_quant_key=kInt8DynamicTokenSym. "int8_per_channel_weight_only": QuantizationConfigArgs( moe=QuantSpec(weight=kInt8StaticChannelSym), ), @@ -270,9 +335,31 @@ def resolve_quantization_config( if base is None: return quantization_config + def merge_spec( + base_spec: QuantSpec | None, + override_spec: QuantSpec | None, + ) -> QuantSpec | None: + if override_spec is None: + return base_spec + if base_spec is None: + return override_spec + + weight = ( + override_spec.weight + if "weight" in override_spec.fields_set + else base_spec.weight + ) + + if "activation" in override_spec.fields_set: + return QuantSpec(weight=weight, activation=override_spec.activation) + if "activation" in base_spec.fields_set: + return QuantSpec(weight=weight, activation=base_spec.activation) + + return QuantSpec(weight=weight) + return QuantizationConfigArgs( - linear=quantization_config.linear or base.linear, - moe=quantization_config.moe or base.moe, + linear=merge_spec(base.linear, quantization_config.linear), + moe=merge_spec(base.moe, quantization_config.moe), ignore=quantization_config.ignore or base.ignore, targets=quantization_config.targets or base.targets, ) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 1c03dcd468df..51b01a922b1d 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -222,7 +222,10 @@ from vllm.model_executor.kernels.linear.scaled_mm.zentorch import ( ZentorchInt8ScaledMMLinearKernel, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp8Dynamic, +) from vllm.platforms import PlatformEnum, current_platform logger = init_logger(__name__) @@ -846,9 +849,17 @@ def choose_mp_linear_kernel( ) -def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel: +def init_mxfp8_linear_kernel( + activation_quant_key: QuantKey | None = None, +) -> Mxfp8LinearKernel: """Select and instantiate the best MXFP8 linear kernel for the current platform.""" + if activation_quant_key not in (None, kMxfp8Dynamic): + raise NotImplementedError( + "MXFP8 linear oracle supports does not implement activation key override " + f"yet, got activation_quant_key={activation_quant_key}. " + "Please open an issue." + ) config = Mxfp8LinearLayerConfig() platform = current_platform._enum diff --git a/vllm/model_executor/layers/fused_moe/oracle/base.py b/vllm/model_executor/layers/fused_moe/oracle/base.py index 06be6becf83f..6bd5f2b755be 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/base.py +++ b/vllm/model_executor/layers/fused_moe/oracle/base.py @@ -27,7 +27,9 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config import get_current_vllm_config from vllm.config.kernel import MoEBackend +from vllm.config.quantization import QuantizationConfigArgs from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEQuantConfig, @@ -52,6 +54,15 @@ class MoEKernelOracle(ABC, Generic[BackendT]): (e.g. `make_quant_config` raises on the unquantized oracle). """ + @staticmethod + def get_user_moe_activation_override() -> tuple["QuantKey | None", bool]: + """Return the configured MoE activation key and whether it was set.""" + args = get_current_vllm_config().model_config.quantization_config + if not isinstance(args, QuantizationConfigArgs) or args.moe is None: + return None, False + + return args.moe.activation, "activation" in args.moe.fields_set + @abstractmethod def backend_enum_cls(self) -> type[BackendT]: """Return the concrete `Enum` class enumerating this oracle's diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 1fa2f63e5d21..15af199afb77 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -7,9 +7,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import envs -from vllm.config import get_current_vllm_config from vllm.config.kernel import MoEBackend -from vllm.config.quantization import QuantizationConfigArgs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, @@ -28,6 +26,7 @@ mxfp4_w4a16_moe_quant_config, ocp_mx_moe_quant_config, ) +from vllm.model_executor.layers.fused_moe.oracle.base import MoEKernelOracle from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( swap_w13_to_w31, ) @@ -396,21 +395,15 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None: return None # BF16 activation -def _user_moe_activation_override() -> QuantKey | None: - """User's MoE activation override from quantization_config, or None.""" - args = get_current_vllm_config().model_config.quantization_config - if not isinstance(args, QuantizationConfigArgs) or args.moe is None: - return None - return args.moe.activation - - def _resolve_activation_key( model_activation_key: QuantKey | None, ) -> QuantKey | None: """Combine the model-supplied activation key with the user override. Raises on conflict (both set and disagreeing).""" - user_override = _user_moe_activation_override() - if user_override is None: + user_override, has_user_override = ( + MoEKernelOracle.get_user_moe_activation_override() + ) + if not has_user_override: return model_activation_key if model_activation_key is None or model_activation_key == user_override: return user_override diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 2a03eacea8e1..9ddafbd71d48 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -9,6 +9,7 @@ backend_to_kernel_cls, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, kMxfp8Dynamic, kMxfp8Static, ) @@ -106,6 +107,7 @@ def _select_kernel_cls( def select_mxfp8_moe_backend( config: FusedMoEConfig, + activation_key: QuantKey | None = None, ) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: """Select the MXFP8 MoE backend and the best expert class. @@ -113,6 +115,12 @@ def select_mxfp8_moe_backend( A tuple of (fp8_backend, experts_cls). """ + if activation_key not in (None, kMxfp8Dynamic): + raise NotImplementedError( + "MXFP8 MoE only supports MXFP8 dynamic activations, got " + f"activation_key={activation_key}." + ) + runner_backend = config.moe_backend if runner_backend != "auto": backend = _BACKEND_NAME_MAP.get(runner_backend) diff --git a/vllm/model_executor/layers/quantization/online/base.py b/vllm/model_executor/layers/quantization/online/base.py index a16a33eaaa33..6462950d2161 100644 --- a/vllm/model_executor/layers/quantization/online/base.py +++ b/vllm/model_executor/layers/quantization/online/base.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from collections.abc import Mapping from enum import Enum from types import MappingProxyType @@ -10,12 +11,12 @@ from vllm.config.quantization import ( _ONLINE_SHORTHANDS, + QUANT_KEY_NAMES, QuantizationConfigArgs, QuantSpec, ) from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( - FusedMoEMethodBase, RoutedExperts, ) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( @@ -45,6 +46,7 @@ from vllm.model_executor.layers.quantization.online.int8 import ( Int8OnlineMoEMethod, ) +from vllm.model_executor.layers.quantization.online.moe_base import OnlineMoEMethodBase from vllm.model_executor.layers.quantization.online.mxfp4 import ( Mxfp4OnlineLinearMethod, Mxfp4OnlineMoEMethod, @@ -67,7 +69,7 @@ kFp8StaticTensorSym, kInt8StaticChannelSym, kMxfp4Static, - kMxfp8Dynamic, + kMxfp8Static, kNvfp4Static, ) @@ -82,14 +84,12 @@ class OnlineQuantizationSource(str, Enum): targets = "targets" -# Online dispatch tables, keyed by the QuantSpec.weight QuantKey. The -# corresponding method class handles the activation choice via its -# `supported_activation_quant` set. +# Online dispatch tables, keyed by the QuantSpec.weight QuantKey. _ONLINE_LINEAR_METHODS: dict[QuantKey, type] = { kFp8StaticTensorSym: Fp8PerTensorOnlineLinearMethod, kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod, kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod, - kMxfp8Dynamic: Mxfp8OnlineLinearMethod, + kMxfp8Static: Mxfp8OnlineLinearMethod, kMxfp4Static: Mxfp4OnlineLinearMethod, } @@ -97,7 +97,7 @@ class OnlineQuantizationSource(str, Enum): kFp8StaticTensorSym: Fp8PerTensorOnlineMoEMethod, kFp8Static128BlockSym: Fp8PerBlockOnlineMoEMethod, kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod, - kMxfp8Dynamic: Mxfp8OnlineMoEMethod, + kMxfp8Static: Mxfp8OnlineMoEMethod, kMxfp4Static: Mxfp4OnlineMoEMethod, kInt8StaticChannelSym: Int8OnlineMoEMethod, kNvfp4Static: Nvfp4OnlineMoEMethod, @@ -106,7 +106,7 @@ class OnlineQuantizationSource(str, Enum): def _find_matching_targets( prefix: str, - targets: Mapping[str, str], + targets: Mapping[str, str | QuantSpec], fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), ) -> list[str]: per_shard_matches = find_matching_patterns( @@ -127,8 +127,8 @@ def _find_matching_targets( ) matched_patterns = [next(iter(matches)) for matches in per_shard_matches] - quant_key_strs = {targets[pattern] for pattern in matched_patterns} - if len(quant_key_strs) > 1: + target_values = [targets[pattern] for pattern in matched_patterns] + if any(target != target_values[0] for target in target_values[1:]): raise ValueError( f"Found different quantization_config.targets values for the " f"shards of {prefix}: {matched_patterns}. vLLM requires all " @@ -237,19 +237,21 @@ def _get_method_cls( f"weight={spec.weight} is not supported; supported weight " f"keys: {sorted(str(k) for k in table)}" ) - # Online method classes pick their own activation format internally. - # Per-class activation overrides are not yet wired through; reject - # explicit overrides until the relevant method class opts in. - if spec.activation is not None: - raise ValueError( - f"activation override (activation={spec.activation}) is not " - f"yet supported for online {cls.__name__}" - ) return cls def resolve_quant_method_cls( self, layer: torch.nn.Module, prefix: str - ) -> tuple[OnlineQuantizationSource, str, str | None, QuantSpec, type] | None: + ) -> ( + tuple[ + OnlineQuantizationSource, + str, + str | None, + QuantSpec, + type, + QuantKey | None, + ] + | None + ): """Resolve quantization metadata and method class without instantiating it. Args: @@ -258,8 +260,8 @@ def resolve_quant_method_cls( Returns: A tuple of source, quantization key string, target pattern, spec, - and method class. Returns None when online quantization does not - apply to the layer. + method class, and activation quantization key. Returns None when + online quantization does not apply to the layer. """ quant_spec: QuantSpec | None if self.args.targets is not None: @@ -294,8 +296,28 @@ def resolve_quant_method_cls( quant_method_cls = self._get_method_cls(quant_spec, table, layer) 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 + + if isinstance(layer, RoutedExperts): + assert issubclass(quant_method_cls, OnlineMoEMethodBase) + else: + assert isinstance(layer, LinearBase) + assert issubclass(quant_method_cls, OnlineLinearBase) + + if "activation" not in quant_spec.fields_set: + activation_quant_key = quant_method_cls.default_activation_quant_key + else: + activation_quant_key = quant_spec.activation + + return ( + source, + quant_key_str, + target_pattern, + quant_spec, + quant_method_cls, + activation_quant_key, + ) def _resolve_targets_quant_method_metadata( self, prefix: str, layer: torch.nn.Module @@ -339,13 +361,36 @@ def _resolve_targets_quant_method_metadata( f"one target." ) target_pattern = matches[0] - quant_key_str = self.args.targets[target_pattern] - shorthand = _ONLINE_SHORTHANDS[quant_key_str] + target = self.args.targets[target_pattern] + if isinstance(target, str): + shorthand = _ONLINE_SHORTHANDS[target] + quant_key_str = target + quant_spec = ( + shorthand.linear if isinstance(layer, LinearBase) else shorthand.moe + ) + else: + quant_spec = target + target_config: dict[str, str | None] = {"weight": str(quant_spec)} + + # TODO: QuantKey itself should define `__str__`, + # instead of having this logic. + if "activation" in quant_spec.fields_set: + target_config["activation"] = ( + next( + ( + name + for name, known_quant_key in QUANT_KEY_NAMES.items() + if known_quant_key == quant_spec.activation + ), + str(quant_spec.activation), + ) + if quant_spec.activation is not None + else None + ) + quant_key_str = json.dumps(target_config, separators=(",", ":")) if isinstance(layer, LinearBase): - quant_spec = shorthand.linear table = _ONLINE_LINEAR_METHODS elif isinstance(layer, RoutedExperts): - quant_spec = shorthand.moe table = _ONLINE_MOE_METHODS else: raise ValueError( @@ -355,7 +400,7 @@ def _resolve_targets_quant_method_metadata( ) if quant_spec is None: raise ValueError( - f"targets pattern {target_pattern} = {quant_key_str} does " + f"targets pattern {target_pattern} = {target} does " f"not define a QuantSpec for {type(layer).__name__} layers " f"(matched at {prefix})." ) @@ -373,18 +418,28 @@ def get_quant_method( # `targets` takes precedence over `moe` and `linear` and is exclusive. resolved = self.resolve_quant_method_cls(layer, prefix) if resolved is not None: - source, quant_key_str, target_pattern, _, quant_method_cls = resolved + ( + source, + quant_key_str, + target_pattern, + _, + quant_method_cls, + activation_quant_key, + ) = resolved self.quantized_layers[prefix] = ( source.value, quant_key_str, target_pattern, ) if isinstance(layer, RoutedExperts): - assert issubclass(quant_method_cls, FusedMoEMethodBase) - return quant_method_cls(moe=layer.moe_config) + assert issubclass(quant_method_cls, OnlineMoEMethodBase) + return quant_method_cls( + moe=layer.moe_config, + activation_quant_key=activation_quant_key, + ) assert issubclass(quant_method_cls, OnlineLinearBase) - return quant_method_cls() + return quant_method_cls(activation_quant_key=activation_quant_key) if isinstance(layer, LinearBase): return UnquantizedLinearMethod() diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 0234702c8bab..43fd9ab3b862 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -12,11 +12,10 @@ 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 +from vllm.config.quantization import _UNSET from vllm.model_executor.kernels.linear import init_fp8_linear_kernel from vllm.model_executor.kernels.linear.scaled_mm import ( CutlassFP8ScaledMMLinearKernel, @@ -35,6 +34,7 @@ ) 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, @@ -113,14 +113,16 @@ def _is_tp_sharded(layer: Module, *, reduces_output_dim: bool = True) -> bool: class OnlineLinearBase(LinearMethodBase): - """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint + """Shared base for online linear methods. Loads fp16/bf16 checkpoint weights onto meta device and materializes them just-in-time.""" uses_meta_device: bool = True + default_activation_quant_key: QuantKey | None - def __init__(self): + def __init__(self, activation_quant_key: QuantKey | None): self.out_dtype = torch.get_default_dtype() self.input_dtype = get_current_vllm_config().model_config.dtype + self.activation_quant_key = activation_quant_key def create_weights( self, @@ -160,19 +162,31 @@ class Fp8PerTensorOnlineLinearMethod(OnlineLinearBase): """Online tensorwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" - def __init__(self): - super().__init__() + default_activation_quant_key = _UNSET + + def __init__( + self, + activation_quant_key: "QuantKey | None" = _UNSET, + ): + super().__init__(activation_quant_key) self.block_quant = False self.use_deep_gemm = False 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 + if self.activation_quant_key is _UNSET: + if cutlass_fp8_supported(): + self.activation_quant_key = kFp8DynamicTokenSym + else: + self.activation_quant_key = kFp8DynamicTensorSym + elif self.activation_quant_key is None: + # TODO: Remove once `init_fp8_linear_kernel` supports `None` + # activation quant key. + raise NotImplementedError( + "online FP8 per-tensor linear quantization does not support " + "activation=null" + ) def create_weights( self, @@ -194,6 +208,7 @@ def create_weights( **extra_weight_attrs, ) + 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, @@ -261,13 +276,22 @@ class Fp8PerBlockOnlineLinearMethod(OnlineLinearBase): """Online blockwise FP8 linear quantization. Loads fp16/bf16 weights and quantizes them per-block during loading.""" - def __init__(self): - super().__init__() + default_activation_quant_key = kFp8Dynamic128Sym + + def __init__( + self, + activation_quant_key: "QuantKey | None" = kFp8Dynamic128Sym, + ): + super().__init__(activation_quant_key) 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]), - ) + + # TODO: Remove once `init_fp8_linear_kernel` supports `None` + # activation quant key. + if self.activation_quant_key is None: + raise NotImplementedError( + "online FP8 per-block linear quantization does not support " + "activation=null" + ) self.weight_quant_key = create_fp8_quant_key( static=True, group_shape=GroupShape(*self.weight_block_size) ) @@ -293,6 +317,7 @@ def create_weights( ) layer.weight_block_size = self.weight_block_size + 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, @@ -346,7 +371,21 @@ class Fp8PtpcOnlineLinearMethod(OnlineLinearBase): """ weight_quant_key = kFp8StaticChannelSym - activation_quant_key = kFp8DynamicTokenSym + default_activation_quant_key = kFp8DynamicTokenSym + + def __init__( + self, + activation_quant_key: "QuantKey | None" = kFp8DynamicTokenSym, + ): + super().__init__(activation_quant_key) + + # TODO: Remove once `init_fp8_linear_kernel` supports `None` + # activation quant key. + if self.activation_quant_key is None: + raise NotImplementedError( + "online FP8 per-channel linear quantization does not support " + "activation=null" + ) def create_weights( self, @@ -368,6 +407,7 @@ def create_weights( **extra_weight_attrs, ) + 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, @@ -442,30 +482,18 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase): def __init__( self, - *, weight_block_size: list[int] | None, moe: FusedMoEConfig, - weight_key: "QuantKey | None" = None, - activation_key: "QuantKey | None" = None, + weight_key: "QuantKey", + activation_key: "QuantKey | None", allow_vllm_cutlass: bool = False, ): - super().__init__(moe) + super().__init__(moe, activation_quant_key=activation_key) self.weight_block_size = weight_block_size self.block_quant: bool = self.weight_block_size is not None self.weight_scale_name = ( "weight_scale_inv" if self.block_quant else "weight_scale" ) - - # Subclasses may pass explicit kernel keys (PTPC needs channelwise + - # per-token). - if weight_key is None or activation_key is None: - if self.block_quant: - weight_key = kFp8Static128BlockSym - activation_key = kFp8Dynamic128Sym - else: - weight_key = kFp8StaticTensorSym - activation_key = kFp8DynamicTensorSym - # Select Fp8 MoE backend self.fp8_backend, self.experts_cls = select_fp8_moe_backend( config=self.moe, @@ -553,14 +581,18 @@ class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase): """Online tensorwise FP8 MoE quantization. Loads fp16/bf16 weights and quantizes them per-tensor during loading.""" + default_activation_quant_key = kFp8DynamicTensorSym + def __init__( self, - *, moe: FusedMoEConfig, + activation_quant_key: "QuantKey | None" = kFp8DynamicTensorSym, ): super().__init__( weight_block_size=None, moe=moe, + weight_key=kFp8StaticTensorSym, + activation_key=activation_quant_key, ) def process_weights_after_loading(self, layer: Module) -> None: @@ -610,14 +642,18 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase): """Online blockwise FP8 MoE quantization. Loads fp16/bf16 weights and quantizes them per-block during loading.""" + default_activation_quant_key = kFp8Dynamic128Sym + def __init__( self, - *, moe: FusedMoEConfig, + activation_quant_key: "QuantKey | None" = kFp8Dynamic128Sym, ): super().__init__( weight_block_size=[128, 128], moe=moe, + weight_key=kFp8Static128BlockSym, + activation_key=activation_quant_key, ) def maybe_roundup_sizes( @@ -713,11 +749,12 @@ class Fp8PtpcOnlineMoEMethod(_Fp8OnlineMoEBase): per_act_token_quant: bool = True per_out_ch_quant: bool = True + default_activation_quant_key = kFp8DynamicTokenSym def __init__( self, - *, moe: FusedMoEConfig, + activation_quant_key: "QuantKey | None" = kFp8DynamicTokenSym, ): from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend @@ -725,7 +762,7 @@ def __init__( weight_block_size=None, moe=moe, weight_key=kFp8StaticChannelSym, - activation_key=kFp8DynamicTokenSym, + activation_key=activation_quant_key, allow_vllm_cutlass=True, ) # Reject backends whose make_fp8_moe_quant_config branch silently diff --git a/vllm/model_executor/layers/quantization/online/int8.py b/vllm/model_executor/layers/quantization/online/int8.py index 268345ca090d..342057c3cd55 100644 --- a/vllm/model_executor/layers/quantization/online/int8.py +++ b/vllm/model_executor/layers/quantization/online/int8.py @@ -10,6 +10,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig @@ -36,16 +37,18 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase): Loads fp16/bf16 weights and quantizes them per-row to int8 during loading. """ + default_activation_quant_key = kInt8DynamicTokenSym + def __init__( self, - *, moe: FusedMoEConfig, + activation_quant_key: "QuantKey | None" = kInt8DynamicTokenSym, ): - super().__init__(moe) + super().__init__(moe, activation_quant_key=activation_quant_key) self.int8_backend, self.experts_cls = select_int8_moe_backend( config=self.moe, weight_key=kInt8StaticChannelSym, - activation_key=kInt8DynamicTokenSym, + activation_key=activation_quant_key, ) def process_weights_after_loading(self, layer: Module) -> None: diff --git a/vllm/model_executor/layers/quantization/online/moe_base.py b/vllm/model_executor/layers/quantization/online/moe_base.py index 866c88b3d649..096bbc8ccea7 100644 --- a/vllm/model_executor/layers/quantization/online/moe_base.py +++ b/vllm/model_executor/layers/quantization/online/moe_base.py @@ -10,7 +10,9 @@ RoutedExperts, SharedExperts, ) +from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.model_executor.model_loader.reload.layerwise import ( initialize_online_processing, ) @@ -23,6 +25,15 @@ class OnlineMoEMethodBase(FusedMoEMethodBase): """ uses_meta_device: bool = True + default_activation_quant_key: QuantKey | None + + def __init__( + self, + moe: FusedMoEConfig, + activation_quant_key: QuantKey | None, + ) -> None: + super().__init__(moe) + self.activation_quant_key = activation_quant_key def create_weights( self, diff --git a/vllm/model_executor/layers/quantization/online/mxfp4.py b/vllm/model_executor/layers/quantization/online/mxfp4.py index 0115988e9265..400a3caa58fe 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp4.py +++ b/vllm/model_executor/layers/quantization/online/mxfp4.py @@ -15,6 +15,7 @@ FusedMoEQuantConfig, RoutedExperts, ) + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig @@ -80,10 +81,13 @@ class Mxfp4OnlineLinearMethod(OnlineLinearBase): FP4 with block-32 scales) during weight loading. """ - activation_quant_key = kMxfp4Dynamic + default_activation_quant_key = kMxfp4Dynamic - def __init__(self): - super().__init__() + def __init__( + self, + activation_quant_key: "QuantKey | None" = kMxfp4Dynamic, + ): + super().__init__(activation_quant_key) self.kernel = init_mxfp4_linear_kernel( activation_quant_key=self.activation_quant_key ) @@ -143,10 +147,14 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase): mxfp4_backend: Mxfp4MoeBackend experts_cls: "type[mk.FusedMoEExperts] | None" - activation_quant_key = kMxfp4Dynamic + default_activation_quant_key = kMxfp4Dynamic - def __init__(self, *, moe: FusedMoEConfig): - super().__init__(moe) + def __init__( + self, + moe: FusedMoEConfig, + activation_quant_key: "QuantKey | None" = kMxfp4Dynamic, + ): + super().__init__(moe=moe, activation_quant_key=activation_quant_key) self.weight_block_size: list[int] = [1, MXFP4_BLOCK_SIZE] self.weight_scale_name = "weight_scale" @@ -261,19 +269,14 @@ def _setup_kernel( def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> "FusedMoEQuantConfig | None": - # NOTE: unlike `QuarkOCP_MX_MoEMethod.get_fused_moe_quant_config`, this - # does not branch on `self.mxfp4_backend in TRITON_BACKENDS or - # AITER_MXFP4_FP8` to read `self.w13_precision_config`/ - # `self.w2_precision_config` instead of `layer.w13_weight_scale`/ - # `layer.w2_weight_scale`. For those backends, - # `convert_weight_to_mxfp4_moe_kernel_format` deletes the latter, so - # `getattr` below would raise `AttributeError`. This is currently - # unreachable because `activation_quant_key` is hardcoded to - # `kMxfp4Dynamic` - # TODO: When supporting online activation quant key override, - # fix this convert_weight_to_mxfp4_moe_kernel_format issue. - w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") - w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + if self.mxfp4_backend in TRITON_BACKENDS or self.mxfp4_backend in ( + Mxfp4MoeBackend.AITER_MXFP4_FP8, + ): + w1_scale = self.w13_precision_config + w2_scale = self.w2_precision_config + else: + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index 5165762d5b88..80291d829749 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -17,6 +17,7 @@ RoutedExperts, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend + from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( @@ -32,6 +33,7 @@ MXFP8_BLOCK_SIZE, mxfp8_e4m3_quantize, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -42,9 +44,23 @@ class Mxfp8OnlineLinearMethod(OnlineLinearBase): FP8 with block-32 scales) during weight loading. """ - def __init__(self): - super().__init__() - self.kernel = init_mxfp8_linear_kernel() + default_activation_quant_key = kMxfp8Dynamic + + def __init__( + self, + activation_quant_key: "QuantKey | None" = kMxfp8Dynamic, + ): + super().__init__(activation_quant_key) + # TODO: remove once `init_mxfp8_linear_kernel` truly supports None + # activation quant key (currently ignored) + if self.activation_quant_key is None: + raise NotImplementedError( + "online MXFP8 linear quantization does not support activation=null" + ) + + self.kernel = init_mxfp8_linear_kernel( + activation_quant_key=self.activation_quant_key + ) def create_weights( self, @@ -101,13 +117,27 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): fp8_backend: "Fp8MoeBackend" experts_cls: "type[mk.FusedMoEExperts] | None" + default_activation_quant_key = kMxfp8Dynamic - def __init__(self, *, moe: FusedMoEConfig): - super().__init__(moe) + def __init__( + self, + moe: FusedMoEConfig, + activation_quant_key: "QuantKey | None" = kMxfp8Dynamic, + ): + super().__init__(moe=moe, activation_quant_key=activation_quant_key) self.weight_block_size: list[int] = [1, MXFP8_BLOCK_SIZE] self.weight_scale_name = "weight_scale" + # TODO: remove once `select_mxfp8_moe_backend` truly supports None + # activation quant key (currently ignored) + if activation_quant_key is None: + raise NotImplementedError( + "online MXFP8 MoE quantization does not support activation=null" + ) - self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe) + self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend( + config=self.moe, + activation_key=activation_quant_key, + ) def create_weights( self, diff --git a/vllm/model_executor/layers/quantization/online/nvfp4.py b/vllm/model_executor/layers/quantization/online/nvfp4.py index 7abefce25732..21b66c34dca6 100644 --- a/vllm/model_executor/layers/quantization/online/nvfp4.py +++ b/vllm/model_executor/layers/quantization/online/nvfp4.py @@ -23,6 +23,7 @@ FLOAT4_E2M1_MAX, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, amax_for_moe_weight_quant, kNvfp4Dynamic, kNvfp4Static, @@ -87,20 +88,22 @@ class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase): (SM100) only. """ + default_activation_quant_key = kNvfp4Dynamic + def __init__( self, - *, moe: FusedMoEConfig, + activation_quant_key: "QuantKey | None" = kNvfp4Dynamic, ): if not current_platform.is_device_capability_family(100): raise ValueError( "nvfp4_per_token online quantization requires a Blackwell (SM100) GPU." ) - super().__init__(moe) + super().__init__(moe=moe, activation_quant_key=activation_quant_key) self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend( config=self.moe, weight_key=kNvfp4Static, - activation_key=kNvfp4Dynamic, + activation_key=activation_quant_key, ) def process_weights_after_loading(self, layer: Module) -> None: