From e67a2e0a56b7ba535bae8bfd98c1c4e4fa742e63 Mon Sep 17 00:00:00 2001 From: fxmarty-amd Date: Fri, 31 Jul 2026 19:07:15 +0200 Subject: [PATCH] [ROCm][Quark][7/N] Use MXFP4 linear kernel abstraction for `emulation` backend (#48949) Signed-off-by: Felix Marty --- csrc/core/scalar_type.hpp | 3 + .../test_mxfp4_kernel_selection.py | 140 +++++++++++++--- .../test_mxfp6_kernel_selection.py | 150 ++++++++++++++++++ .../model_executor/kernels/linear/__init__.py | 100 +++++++++++- .../kernels/linear/mxfp4/aiter.py | 31 +++- .../kernels/linear/mxfp4/base.py | 8 +- .../kernels/linear/mxfp4/emulation.py | 106 +++++++++++++ .../kernels/linear/mxfp4/flashinfer.py | 5 + .../kernels/linear/mxfp4/humming.py | 14 +- .../kernels/linear/mxfp4/marlin.py | 15 +- .../kernels/linear/mxfp4/xpu.py | 7 +- .../kernels/linear/mxfp6/__init__.py | 12 ++ .../kernels/linear/mxfp6/base.py | 76 +++++++++ .../kernels/linear/mxfp6/emulation.py | 93 +++++++++++ .../schemes/compressed_tensors_w4a4_mxfp4.py | 7 +- .../quark/schemes/quark_ocp_mx.py | 113 +++++++------ .../layers/quantization/utils/quant_utils.py | 34 +++- 17 files changed, 824 insertions(+), 90 deletions(-) create mode 100644 tests/kernels/quantization/test_mxfp6_kernel_selection.py create mode 100644 vllm/model_executor/kernels/linear/mxfp4/emulation.py create mode 100644 vllm/model_executor/kernels/linear/mxfp6/__init__.py create mode 100644 vllm/model_executor/kernels/linear/mxfp6/base.py create mode 100644 vllm/model_executor/kernels/linear/mxfp6/emulation.py diff --git a/csrc/core/scalar_type.hpp b/csrc/core/scalar_type.hpp index b6f39ed795..e1e3490c2d 100644 --- a/csrc/core/scalar_type.hpp +++ b/csrc/core/scalar_type.hpp @@ -327,6 +327,8 @@ static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); +static inline constexpr auto kFE2M3f = + ScalarType::float_(2, 3, true, ScalarType::NAN_NONE); static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); static inline constexpr auto kFE8M0fnu = @@ -346,6 +348,7 @@ static inline constexpr auto kUint8b128 = kU8B128; static inline constexpr auto kFloat4_e2m1f = kFE2M1f; static inline constexpr auto kFloat6_e3m2f = kFE3M2f; +static inline constexpr auto kFloat6_e2m3f = kFE2M3f; static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn; static inline constexpr auto kFloat8_e5m2 = kFE5M2; static inline constexpr auto kFloat16_e8m7 = kFE8M7; diff --git a/tests/kernels/quantization/test_mxfp4_kernel_selection.py b/tests/kernels/quantization/test_mxfp4_kernel_selection.py index 20c171a87c..c3f8bc094d 100644 --- a/tests/kernels/quantization/test_mxfp4_kernel_selection.py +++ b/tests/kernels/quantization/test_mxfp4_kernel_selection.py @@ -12,15 +12,40 @@ import torch from vllm.model_executor.kernels.linear import ( AiterMxfp4LinearKernel, + EmulationMxfp4LinearKernel, + FlashInferMxFp4LinearKernel, + HummingMxFp4LinearKernel, + MarlinMxFp4LinearKernel, MxFp4LinearKernel, MxFp4LinearLayerConfig, + XPUMxFp4LinearKernel, init_mxfp4_linear_kernel, register_linear_kernel, ) +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + quant_dequant_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, + kMxfp6E2M3Dynamic, + kMxfp6E3M2Dynamic, +) from vllm.platforms import PlatformEnum pytestmark = pytest.mark.cpu_test +# Kernels that quantize activations themselves (true W4A4): they require an +# explicit MXFP4-dynamic activation key. +_TRUE_W4A4_KERNELS = [ + FlashInferMxFp4LinearKernel, + XPUMxFp4LinearKernel, + AiterMxfp4LinearKernel, +] + +# Weight-only (A16) kernels: they never quantize activations. They still accept +# MXFP4 activation keys as an intentional compatibility fallback. +_WEIGHT_ONLY_KERNELS = [MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel] + def test_can_implement_is_abstract(): """Test that can_implement()/is_supported() are properly defined.""" @@ -28,6 +53,97 @@ def test_can_implement_is_abstract(): assert hasattr(MxFp4LinearKernel, "is_supported") +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) +def test_true_w4a4_kernels_accept_dynamic_mxfp4_activation(kernel_cls): + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) + can_implement, reason = kernel_cls.can_implement(config) + assert can_implement, reason + + +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) +def test_true_w4a4_kernels_reject_unset_activation(kernel_cls): + """None means weight-only/unquantized activations, not dynamic MXFP4.""" + config = MxFp4LinearLayerConfig() + can_implement, reason = kernel_cls.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize("kernel_cls", _TRUE_W4A4_KERNELS) +def test_true_w4a4_kernels_reject_explicit_non_mxfp4_activation(kernel_cls): + """FlashInfer/XPU/Aiter quantize activations to MXFP4 internally, so an + explicit request for a different activation format must be rejected.""" + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) + can_implement, reason = kernel_cls.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS) +@pytest.mark.parametrize("activation_quant_key", [None, kMxfp4Dynamic]) +def test_weight_only_kernels_accept_unquantized_or_mxfp4_activation( + kernel_cls, activation_quant_key +): + """Marlin/Humming never quantize activations, so an unset activation key, + or one that already describes MXFP4-shaped data, is tolerated. When an + activation key is explicitly set, a warning must be logged noting that it + is ignored, since these kernels are weight-only (A16).""" + config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) + with patch(f"{kernel_cls.__module__}.logger.warning_once") as warning_once: + can_implement, reason = kernel_cls.can_implement(config) + assert can_implement, reason + + if activation_quant_key is None: + warning_once.assert_not_called() + else: + warning_once.assert_called_once() + message = warning_once.call_args.args[0] + assert "the requested activation quantization" in message + assert "is ignored" in message + + +@pytest.mark.parametrize("kernel_cls", _WEIGHT_ONLY_KERNELS) +def test_weight_only_kernels_reject_non_mxfp4_activation(kernel_cls): + config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp6E3M2Dynamic) + can_implement, reason = kernel_cls.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize( + "activation_quant_key", + [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic], +) +def test_emulation_kernel_accepts_any_config(activation_quant_key): + """EmulationMxfp4LinearKernel is the universal fallback: it must accept + every supported activation format.""" + config = MxFp4LinearLayerConfig(activation_quant_key=activation_quant_key) + with patch( + "vllm.model_executor.kernels.linear._get_linear_backend", + return_value="emulation", + ): + can_implement, reason = EmulationMxfp4LinearKernel.can_implement(config) + assert can_implement, reason + + +def test_emulation_kernel_derives_quant_dequant_func_from_config(): + """quant_dequant_func must be derived purely from the config's activation + QuantKey, not set externally.""" + with patch( + "vllm.model_executor.kernels.linear._get_linear_backend", + return_value="emulation", + ): + weight_only_config = MxFp4LinearLayerConfig() + kernel = EmulationMxfp4LinearKernel(weight_only_config) + x = torch.randn(4) + # identity for weight-only + assert torch.equal(kernel.quant_dequant_func(x), x) + + w4a4_config = MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) + kernel = EmulationMxfp4LinearKernel(w4a4_config) + assert kernel.quant_dequant_func is quant_dequant_mxfp4 + + def test_aiter_kernel_is_supported_requires_native_mx_support(): """AiterMxfp4LinearKernel must not be selected on platforms without native MX compute, even if AITER itself is importable.""" @@ -70,10 +186,10 @@ def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock) platform_mock._enum = PlatformEnum.OOT register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4") - kernel = init_mxfp4_linear_kernel() + kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) assert isinstance(kernel, OOTMxFp4LinearKernel) - assert kernel.config == MxFp4LinearLayerConfig() + assert kernel.config == MxFp4LinearLayerConfig(activation_quant_key=kMxfp4Dynamic) class UnsupportedMxFp4LinearKernel(MxFp4LinearKernel): @@ -108,23 +224,3 @@ def test_init_mxfp4_linear_kernel_raises_when_no_kernel_matches(platform_mock): with pytest.raises(ValueError, match="Failed to find a kernel"): init_mxfp4_linear_kernel() - - -@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.is_aiter_found_and_supported") -@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform") -@patch("vllm.model_executor.kernels.linear.current_platform") -def test_init_mxfp4_linear_kernel_raises_on_rocm_without_aiter( - linear_platform_mock, aiter_platform_mock, is_aiter_found_and_supported_mock -): - """On ROCm, the only registered MXFP4 linear kernel is AITER-based. - If AITER is not found/supported, no kernel should be selected.""" - linear_platform_mock._enum = PlatformEnum.ROCM - aiter_platform_mock.supports_mx.return_value = True - is_aiter_found_and_supported_mock.return_value = False - - with pytest.raises( - ValueError, - match="(?s)Failed to find a kernel.*" - "AITER not found or not supported on the current platform", - ): - init_mxfp4_linear_kernel() diff --git a/tests/kernels/quantization/test_mxfp6_kernel_selection.py b/tests/kernels/quantization/test_mxfp6_kernel_selection.py new file mode 100644 index 0000000000..2c7b565bea --- /dev/null +++ b/tests/kernels/quantization/test_mxfp6_kernel_selection.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MXFP6 linear kernel selection logic (CPU-only) + +Run `pytest tests/kernels/quantization/test_mxfp6_kernel_selection.py`. +""" + +from unittest.mock import patch + +import pytest +import torch + +from vllm.model_executor.kernels.linear import ( + EmulationMxfp6LinearKernel, + MxFp6LinearKernel, + MxFp6LinearLayerConfig, + init_mxfp6_linear_kernel, + register_linear_kernel, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, + kMxfp4Static, + kMxfp6E2M3Dynamic, + kMxfp6E2M3Static, + kMxfp6E3M2Dynamic, + kMxfp6E3M2Static, +) +from vllm.platforms import PlatformEnum + +pytestmark = pytest.mark.cpu_test + +# The only implementation available at the moment is software emulation. +_WEIGHT_QUANT_KEYS = [kMxfp6E3M2Static, kMxfp6E2M3Static] + + +def test_can_implement_is_abstract(): + """Test that can_implement()/is_supported() are properly defined.""" + assert hasattr(MxFp6LinearKernel, "can_implement") + assert hasattr(MxFp6LinearKernel, "is_supported") + + +def test_emulation_kernel_rejects_non_mxfp6_weights(): + """EmulationMxfp6LinearKernel must not implement a non-MXFP6 weight + format.""" + config = MxFp6LinearLayerConfig(weight_quant_key=kMxfp4Static) + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) + assert not can_implement + assert reason + + +@pytest.mark.parametrize("weight_quant_key", _WEIGHT_QUANT_KEYS) +@pytest.mark.parametrize( + "activation_quant_key", + [None, kMxfp4Dynamic, kMxfp6E3M2Dynamic, kMxfp6E2M3Dynamic], +) +def test_emulation_kernel_accepts_any_supported_config( + weight_quant_key, activation_quant_key +): + """EmulationMxfp6LinearKernel is the only backend today: it must accept + every supported weight/activation format combination.""" + config = MxFp6LinearLayerConfig( + weight_quant_key=weight_quant_key, activation_quant_key=activation_quant_key + ) + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) + assert can_implement, reason + + +@pytest.mark.parametrize("weight_quant_key", _WEIGHT_QUANT_KEYS) +def test_emulation_kernel_rejects_non_mxfp4_or_mxfp6_activation(weight_quant_key): + config = MxFp6LinearLayerConfig( + weight_quant_key=weight_quant_key, activation_quant_key=kMxfp4Static + ) + can_implement, reason = EmulationMxfp6LinearKernel.can_implement(config) + assert not can_implement + assert reason + + +class OOTMxFp6LinearKernel(MxFp6LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp6_linear_kernel_dispatches_to_registered_kernel(platform_mock): + """init_mxfp6_linear_kernel should select a registered kernel that + reports itself as supported/able to implement the given config, and + construct it with that exact config.""" + platform_mock._enum = PlatformEnum.OOT + register_linear_kernel(OOTMxFp6LinearKernel, PlatformEnum.OOT, "mxfp6") + + kernel = init_mxfp6_linear_kernel( + weight_quant_key=kMxfp6E3M2Static, activation_quant_key=kMxfp6E3M2Dynamic + ) + + assert isinstance(kernel, OOTMxFp6LinearKernel) + assert kernel.config == MxFp6LinearLayerConfig( + weight_quant_key=kMxfp6E3M2Static, activation_quant_key=kMxfp6E3M2Dynamic + ) + + +class UnsupportedMxFp6LinearKernel(MxFp6LinearKernel): + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return False, "never supported" + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + pass + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + pass + + +@patch("vllm.model_executor.kernels.linear.current_platform") +def test_init_mxfp6_linear_kernel_raises_when_no_kernel_matches(platform_mock): + platform_mock._enum = PlatformEnum.UNSPECIFIED + register_linear_kernel( + UnsupportedMxFp6LinearKernel, PlatformEnum.UNSPECIFIED, "mxfp6" + ) + + with pytest.raises(ValueError, match="Failed to find a kernel"): + init_mxfp6_linear_kernel(weight_quant_key=kMxfp6E3M2Static) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 2892f1aa71..aaf46875e8 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -77,6 +77,9 @@ from vllm.model_executor.kernels.linear.mxfp4 import ( from vllm.model_executor.kernels.linear.mxfp4.aiter import ( AiterMxfp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp4.emulation import ( + EmulationMxfp4LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp4.flashinfer import ( FlashInferMxFp4LinearKernel, ) @@ -89,6 +92,13 @@ from vllm.model_executor.kernels.linear.mxfp4.marlin import ( from vllm.model_executor.kernels.linear.mxfp4.xpu import ( XPUMxFp4LinearKernel, ) +from vllm.model_executor.kernels.linear.mxfp6 import ( + MxFp6LinearKernel, + MxFp6LinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mxfp6.emulation import ( + EmulationMxfp6LinearKernel, +) from vllm.model_executor.kernels.linear.mxfp8 import ( Mxfp8LinearKernel, Mxfp8LinearLayerConfig, @@ -294,6 +304,8 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = { "emulation": { EmulationMxfp8LinearKernel, EmulationNvFp4LinearKernel, + EmulationMxfp6LinearKernel, + EmulationMxfp4LinearKernel, }, "xpu": { XPUW8A8FP8LinearKernel, @@ -467,14 +479,25 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { ], } +_POSSIBLE_MXFP6_KERNELS: dict[PlatformEnum, list[type[MxFp6LinearKernel]]] = { + PlatformEnum.CUDA: [ + EmulationMxfp6LinearKernel, + ], + PlatformEnum.ROCM: [ + EmulationMxfp6LinearKernel, + ], +} + _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = { PlatformEnum.CUDA: [ FlashInferMxFp4LinearKernel, MarlinMxFp4LinearKernel, HummingMxFp4LinearKernel, + EmulationMxfp4LinearKernel, ], PlatformEnum.ROCM: [ AiterMxfp4LinearKernel, + EmulationMxfp4LinearKernel, ], PlatformEnum.XPU: [ XPUMxFp4LinearKernel, @@ -812,9 +835,15 @@ def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel: ) -def init_mxfp4_linear_kernel() -> MxFp4LinearKernel: +def init_mxfp4_linear_kernel( + activation_quant_key: QuantKey | None = None, +) -> MxFp4LinearKernel: """Select and instantiate the best MXFP4 linear kernel for the current platform.""" + config = MxFp4LinearLayerConfig( + activation_quant_key=activation_quant_key, + ) + linear_backend = _get_linear_backend() platform = current_platform._enum @@ -843,8 +872,13 @@ def init_mxfp4_linear_kernel() -> MxFp4LinearKernel: failure_reasons.append(f"{kernel_cls.__name__}: {reason}") continue + can_implement, reason = kernel_cls.can_implement(config) + if not can_implement: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + logger.info_once("Using %s for MXFP4 GEMM", kernel_cls.__name__) - return kernel_cls(MxFp4LinearLayerConfig()) + return kernel_cls(config) raise ValueError( "Failed to find a kernel that can implement the " @@ -852,6 +886,59 @@ def init_mxfp4_linear_kernel() -> MxFp4LinearKernel: ) +def init_mxfp6_linear_kernel( + weight_quant_key: QuantKey, + activation_quant_key: QuantKey | None = None, +) -> MxFp6LinearKernel: + """Select and instantiate the best MXFP6 linear kernel for the + current platform.""" + config = MxFp6LinearLayerConfig( + weight_quant_key=weight_quant_key, + activation_quant_key=activation_quant_key, + ) + + linear_backend = _get_linear_backend() + + platform = current_platform._enum + possible = list(_POSSIBLE_MXFP6_KERNELS.get(platform, [])) + + # Apply --linear-backend filtering when set. + if linear_backend != "auto": + filtered = _filter_kernels_by_backend(linear_backend, possible) + if not filtered: + raise ValueError( + f"--linear-backend={linear_backend} was requested but no " + f"'{linear_backend}' kernel exists for MXFP6 layers." + ) + possible = filtered + + failure_reasons = [] + for kernel_cls in possible: + if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS: + failure_reasons.append( + f" {kernel_cls.__name__} disabled by environment variable" + ) + continue + + is_supported, reason = kernel_cls.is_supported() + if not is_supported: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + can_implement, reason = kernel_cls.can_implement(config) + if not can_implement: + failure_reasons.append(f"{kernel_cls.__name__}: {reason}") + continue + + logger.info_once("Using %s for MXFP6 GEMM", kernel_cls.__name__) + return kernel_cls(config) + + raise ValueError( + "Failed to find a kernel that can implement the " + "MXFP6 linear layer. Reasons: \n" + "\n".join(failure_reasons) + ) + + def init_wfp8_a16_linear_kernel( weight_quant_key: QuantKey, activation_quant_key: QuantKey, @@ -1037,6 +1124,10 @@ def register_linear_kernel( if platform not in _POSSIBLE_MXFP4_KERNELS: _POSSIBLE_MXFP4_KERNELS[platform] = [] _POSSIBLE_MXFP4_KERNELS[platform].append(kernel_class) + elif kernel_type == "mxfp6": + if platform not in _POSSIBLE_MXFP6_KERNELS: + _POSSIBLE_MXFP6_KERNELS[platform] = [] + _POSSIBLE_MXFP6_KERNELS[platform].append(kernel_class) else: raise ValueError(f"Unrecognized kernel type: {kernel_type}") @@ -1091,7 +1182,12 @@ __all__ = [ "init_mxfp4_linear_kernel", "MxFp4LinearKernel", "MxFp4LinearLayerConfig", + "MxFp6LinearKernel", + "MxFp6LinearLayerConfig", + "init_mxfp6_linear_kernel", + "EmulationMxfp6LinearKernel", "AiterMxfp4LinearKernel", + "EmulationMxfp4LinearKernel", "FlashInferMxFp4LinearKernel", "MarlinMxFp4LinearKernel", "FlashInferCutedslMxfp8LinearKernel", diff --git a/vllm/model_executor/kernels/linear/mxfp4/aiter.py b/vllm/model_executor/kernels/linear/mxfp4/aiter.py index 54f8098b9a..f1336dd008 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/aiter.py +++ b/vllm/model_executor/kernels/linear/mxfp4/aiter.py @@ -4,11 +4,18 @@ import torch from torch.nn.parameter import Parameter +import vllm.envs as envs from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) from vllm.platforms import current_platform from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig +logger = init_logger(__name__) + # NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP # which can force the engine core to spawn instead of fork. # is_aiter_found_and_supported() checks platform + arch + library availability via @@ -131,12 +138,34 @@ class AiterMxfp4LinearKernel(MxFp4LinearKernel): ) -> tuple[bool, str | None]: if not current_platform.supports_mx(): return False, "current platform does not support native MXFP4 computation" + + from vllm._aiter_ops import is_aiter_found_and_supported + from vllm.model_executor.kernels.linear import _get_linear_backend + + linear_backend = _get_linear_backend() + + if ( + current_platform.is_rocm() + and current_platform.supports_mx() + and "AiterMxfp4LinearKernel" not in envs.VLLM_DISABLED_KERNELS + and linear_backend == "auto" + and not is_aiter_found_and_supported() + ): + logger.warning_once( + "This platform supports native MXFP4 W4A4 MOE " + "computation via AITER MOE backend, but AITER is not " + "found or not supported. Consider installing AITER: " + "https://github.com/ROCm/aiter." + ) + if is_aiter_found_and_supported(): return True, None return False, "AITER not found or not supported on the current platform" @classmethod - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "only supports MXFP4 dynamic activation" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp4/base.py b/vllm/model_executor/kernels/linear/mxfp4/base.py index 868faa4731..9c3088df52 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/base.py +++ b/vllm/model_executor/kernels/linear/mxfp4/base.py @@ -6,6 +6,8 @@ from dataclasses import dataclass import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + @dataclass class MxFp4LinearLayerConfig: @@ -13,9 +15,13 @@ class MxFp4LinearLayerConfig: All MXFP4 layers share the same structure: packed uint8 weights (2 FP4 values per byte) and per-block weight scales (group size 32). + + Attributes: + activation_quant_key: Identifies the activation quantization format, + or `None` when activations must not be quantized. """ - pass + activation_quant_key: QuantKey | None = None class MxFp4LinearKernel(ABC): diff --git a/vllm/model_executor/kernels/linear/mxfp4/emulation.py b/vllm/model_executor/kernels/linear/mxfp4/emulation.py new file mode 100644 index 0000000000..66436521b8 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp4/emulation.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +from functools import partial + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + dequant_mxfp4, + quant_dequant_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( + quant_dequant_mxfp6, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp6E2M3Dynamic, + kMxfp6E3M2Dynamic, +) +from vllm.platforms import current_platform + +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig + +logger = init_logger(__name__) + +_ACTIVATION_QUANT_DEQUANT_FUNCS: dict[ + QuantKey, Callable[[torch.Tensor], torch.Tensor] +] = { + kMxfp4Dynamic: quant_dequant_mxfp4, + kMxfp6E3M2Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e3m2"), + kMxfp6E2M3Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e2m3"), +} + + +class EmulationMxfp4LinearKernel(MxFp4LinearKernel): + """Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear).""" + + def __init__(self, config: MxFp4LinearLayerConfig) -> None: + super().__init__(config) + if config.activation_quant_key is None: + # no input Q/DQ for weight-only + self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = ( + lambda x: x + ) + else: + self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[ + config.activation_quant_key + ] + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key not in ( + None, + kMxfp4Dynamic, + kMxfp6E3M2Dynamic, + kMxfp6E2M3Dynamic, + ): + return False, "only supports MXFP4 or MXFP6 or unquantized activations" + + if ( + current_platform.is_rocm() + and current_platform.supports_mx() + and config.activation_quant_key != kMxfp4Dynamic + ): + logger.warning_once( + "The current platform supports native MXFP4/MXFP6 computation, " + f"but kernels for activation_quant_key={config.activation_quant_key} " + f"are not yet integrated in vLLM. Using EmulationMxfp4LinearKernel, " + "with simulated weight dequantization and activation " + "QDQ (quantize and dequantize), with the linear " + "layers computed in high precision." + ) + + if not current_platform.supports_mx(): + logger.warning_once( + "The current platform does not support native MXFP4 " + "computation. Using EmulationMxfp4LinearKernel, with simulated weight " + "dequantization and activation QDQ (quantize and dequantize), with " + "the linear layers computed in high precision." + ) + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + dq_w = dequant_mxfp4(layer.weight, layer.weight_scale, x.dtype) + qdq_x = self.quant_dequant_func(x) + return F.linear(qdq_x, dq_w, bias) diff --git a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py index 532a7d55e9..4491e85522 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp4/flashinfer.py @@ -7,6 +7,9 @@ from torch.nn.parameter import Parameter from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( swizzle_mxfp4_scales, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutedsl @@ -28,6 +31,8 @@ class FlashInferMxFp4LinearKernel(MxFp4LinearKernel): @classmethod def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "only supports MXFP4 dynamic activation" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py index d93f5d4815..f38b8ead63 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/humming.py +++ b/vllm/model_executor/kernels/linear/mxfp4/humming.py @@ -3,14 +3,18 @@ import torch +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.humming_utils import ( convert_linear_layer_to_humming_standard, prepare_humming_layer, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic from vllm.platforms import current_platform from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig +logger = init_logger(__name__) + class HummingMxFp4LinearKernel(MxFp4LinearKernel): """Humming GEMM Kernel for MXFP4.""" @@ -28,7 +32,15 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel): return True, None @classmethod - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key not in (None, kMxfp4Dynamic): + return False, "only supports MXFP4 dynamic or unquantized activations" + if config.activation_quant_key is not None: + logger.warning_once( + "HummingMxFp4LinearKernel is a weight-only (A16) kernel; " + "the requested activation quantization (%s) is ignored.", + config.activation_quant_key, + ) return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp4/marlin.py b/vllm/model_executor/kernels/linear/mxfp4/marlin.py index 3844075207..30ff0d58d4 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp4/marlin.py @@ -3,8 +3,13 @@ import torch +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp4Dynamic + from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig +logger = init_logger(__name__) + class MarlinMxFp4LinearKernel(MxFp4LinearKernel): @classmethod @@ -20,7 +25,15 @@ class MarlinMxFp4LinearKernel(MxFp4LinearKernel): return False, "Marlin FP4 not available" @classmethod - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key not in (None, kMxfp4Dynamic): + return False, "only supports MXFP4 dynamic or unquantized activations" + if config.activation_quant_key is not None: + logger.warning_once( + "MarlinMxFp4LinearKernel is a weight-only (A16) kernel; " + "the requested activation quantization (%s) is ignored.", + config.activation_quant_key, + ) return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp4/xpu.py b/vllm/model_executor/kernels/linear/mxfp4/xpu.py index 8d33939d2e..c499d0a17a 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/xpu.py +++ b/vllm/model_executor/kernels/linear/mxfp4/xpu.py @@ -6,6 +6,9 @@ import torch from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( xpu_mxfp4_quantize as quant_mxfp4, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform @@ -24,7 +27,9 @@ class XPUMxFp4LinearKernel(MxFp4LinearKernel): return True, None @classmethod - def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: + if config.activation_quant_key != kMxfp4Dynamic: + return False, "only supports MXFP4 dynamic activation" return True, None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/mxfp6/__init__.py b/vllm/model_executor/kernels/linear/mxfp6/__init__.py new file mode 100644 index 0000000000..bbdd6f0b52 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.model_executor.kernels.linear.mxfp6.base import ( + MxFp6LinearKernel, + MxFp6LinearLayerConfig, +) + +__all__ = [ + "MxFp6LinearKernel", + "MxFp6LinearLayerConfig", +] diff --git a/vllm/model_executor/kernels/linear/mxfp6/base.py b/vllm/model_executor/kernels/linear/mxfp6/base.py new file mode 100644 index 0000000000..b3ba208634 --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/base.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import torch + +from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey + + +@dataclass +class MxFp6LinearLayerConfig: + """Configuration for an MXFP6 linear layer. + + All MXFP6 layers share the same structure: packed uint8 weights (2 FP4 values per + byte) and per-block weight scales (group size 32). + + Attributes: + weight_quant_key: Identifies the weight quantization format. Can be + kMxfp6E2M3Static or kMxfp6E3M2Static. + activation_quant_key: Identifies the activation quantization format, + or `None` when activations must not be quantized. + """ + + weight_quant_key: QuantKey + activation_quant_key: QuantKey | None = None + + +class MxFp6LinearKernel(ABC): + """Base class for MXFP6 quantized linear kernels. + + Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc). + The kernel selection mechanism iterates over registered subclasses in + priority order,calling ``is_supported`` and ``can_implement`` to find the best + match for the current hardware. + """ + + def __init__(self, config: MxFp6LinearLayerConfig) -> None: + assert self.can_implement(config)[0] + assert self.is_supported()[0] + self.config = config + + @classmethod + @abstractmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + """Return whether this kernel can run on the current platform.""" + raise NotImplementedError + + @classmethod + @abstractmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + """Return whether this kernel can handle *config*.""" + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + """Transform weights into the format required by this kernel. + + Called once after checkpoint weights have been loaded onto the + device. Implementations should repack / swizzle / pad weights + and scales in-place on *layer*. + """ + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the quantized GEMM.""" + raise NotImplementedError diff --git a/vllm/model_executor/kernels/linear/mxfp6/emulation.py b/vllm/model_executor/kernels/linear/mxfp6/emulation.py new file mode 100644 index 0000000000..0ef68deb0d --- /dev/null +++ b/vllm/model_executor/kernels/linear/mxfp6/emulation.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Callable +from functools import partial + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter + +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + quant_dequant_mxfp4, +) +from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( + dequant_mxfp6, + quant_dequant_mxfp6, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp6E2M3Dynamic, + kMxfp6E2M3Static, + kMxfp6E3M2Dynamic, + kMxfp6E3M2Static, +) + +from .base import MxFp6LinearKernel, MxFp6LinearLayerConfig + +_WEIGHT_DEQUANT_FUNCS: dict[QuantKey, Callable[..., torch.Tensor]] = { + kMxfp6E3M2Static: partial(dequant_mxfp6, quant_dtype="fp6_e3m2"), + kMxfp6E2M3Static: partial(dequant_mxfp6, quant_dtype="fp6_e2m3"), +} + +_ACTIVATION_QUANT_DEQUANT_FUNCS: dict[ + QuantKey, Callable[[torch.Tensor], torch.Tensor] +] = { + kMxfp4Dynamic: quant_dequant_mxfp4, + kMxfp6E3M2Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e3m2"), + kMxfp6E2M3Dynamic: partial(quant_dequant_mxfp6, quant_dtype="fp6_e2m3"), +} + + +class EmulationMxfp6LinearKernel(MxFp6LinearKernel): + """Software emulation fallback for OCP MXFP4/MXFP6 (dequant + F.linear).""" + + def __init__(self, config: MxFp6LinearLayerConfig) -> None: + super().__init__(config) + self.dequant_func = _WEIGHT_DEQUANT_FUNCS[config.weight_quant_key] + if config.activation_quant_key is None: + # no input Q/DQ for weight-only + self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = ( + lambda x: x + ) + else: + self.quant_dequant_func = _ACTIVATION_QUANT_DEQUANT_FUNCS[ + config.activation_quant_key + ] + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + return True, None + + @classmethod + def can_implement(cls, config: MxFp6LinearLayerConfig) -> tuple[bool, str | None]: + if config.weight_quant_key not in ( + kMxfp6E2M3Static, + kMxfp6E3M2Static, + ): + return False, "only supports MXFP6 weights" + if config.activation_quant_key not in ( + None, + kMxfp4Dynamic, + kMxfp6E3M2Dynamic, + kMxfp6E2M3Dynamic, + ): + return False, "only supports MXFP4 or MXFP6 or unquantized activations" + + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype) + qdq_x = self.quant_dequant_func(x) + return F.linear(qdq_x, dq_w, bias) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py index 7b994e10e4..2eb8ba717b 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py @@ -9,6 +9,9 @@ from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel from vllm.model_executor.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsScheme, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp4Dynamic, +) from vllm.model_executor.parameter import ( GroupQuantScaleParameter, ModelWeightParameter, @@ -35,7 +38,9 @@ class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme): def __init__(self): self.group_size = 32 - self.kernel = init_mxfp4_linear_kernel() + self.kernel = init_mxfp4_linear_kernel( + activation_quant_key=kMxfp4Dynamic, + ) @classmethod def get_min_capability(cls) -> int: diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index 2d42babeb3..8b6edd6cae 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -3,25 +3,28 @@ from collections.abc import Callable from fractions import Fraction -from functools import partial from typing import Any import torch -import torch.nn.functional as F from vllm.logger import init_logger -from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel -from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( - dequant_mxfp4, - quant_dequant_mxfp4, -) -from vllm.model_executor.layers.quantization.utils.mxfp6_utils import ( - dequant_mxfp6, - quant_dequant_mxfp6, +from vllm.model_executor.kernels.linear import ( + MxFp4LinearKernel, + MxFp6LinearKernel, + init_mxfp4_linear_kernel, + init_mxfp6_linear_kernel, ) from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, - OCP_MX_Scheme, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp4Dynamic, + kMxfp4Static, + kMxfp6E2M3Dynamic, + kMxfp6E2M3Static, + kMxfp6E3M2Dynamic, + kMxfp6E3M2Static, ) from vllm.model_executor.parameter import ( GroupQuantScaleParameter, @@ -35,8 +38,22 @@ from .quark_scheme import QuarkScheme logger = init_logger(__name__) +_WEIGHT_QUANT_KEY_MAP: dict[str, QuantKey] = { + "mxfp4": kMxfp4Static, + "mxfp6_e3m2": kMxfp6E3M2Static, + "mxfp6_e2m3": kMxfp6E2M3Static, +} + +_ACTIVATION_QUANT_KEY_MAP: dict[str, QuantKey] = { + "mxfp4": kMxfp4Dynamic, + "mxfp6_e3m2": kMxfp6E3M2Dynamic, + "mxfp6_e2m3": kMxfp6E2M3Dynamic, +} + class QuarkOCP_MX(QuarkScheme): + ocp_mx_linear: MxFp6LinearKernel | MxFp4LinearKernel + def __init__( self, weight_quant_spec: dict[str, Any], @@ -49,35 +66,26 @@ class QuarkOCP_MX(QuarkScheme): self.weight_dtype = weight_quant_spec["dtype"].replace("fp", "mxfp") self.input_dtype: str | None = None if input_quant_spec is not None: - input_quant = input_quant_spec["dtype"] - if input_quant == "fp8_e4m3": - self.input_dtype = "fp8" - else: - self.input_dtype = input_quant.replace("fp", "mxfp") + self.input_dtype = input_quant_spec["dtype"].replace("fp", "mxfp") - self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype( - self.input_dtype, self.weight_dtype + if self.input_dtype not in [None, *_ACTIVATION_QUANT_KEY_MAP]: + raise ValueError( + f"Unsupported input_dtype={self.input_dtype} for QuarkOCP_MX. " + f"Supported activation dtypes are {_ACTIVATION_QUANT_KEY_MAP.keys()}, " + "or None for weight-only quantization." + ) + + self.weight_quant_key = _WEIGHT_QUANT_KEY_MAP[self.weight_dtype] + self.activation_quant_key = ( + _ACTIVATION_QUANT_KEY_MAP[self.input_dtype] + if self.input_dtype is not None + else None ) if self.weight_dtype == "mxfp4": self.packed_factor: int | Fraction = 2 - self.dequant_func = dequant_mxfp4 else: self.packed_factor = Fraction(numerator=8, denominator=6) - self.dequant_func = partial( - dequant_mxfp6, quant_dtype=self.weight_dtype.replace("mx", "") - ) - - if self.input_dtype is None: - self.quant_dequant_func: Callable[[torch.Tensor], torch.Tensor] = ( - lambda x: x - ) # no input Q/DQ for weight-only - elif self.input_dtype == "mxfp4": - self.quant_dequant_func = quant_dequant_mxfp4 - else: - self.quant_dequant_func = partial( - quant_dequant_mxfp6, quant_dtype=self.input_dtype.replace("mx", "") - ) if input_quant_spec is None: self.static_input_scales = False @@ -90,16 +98,6 @@ class QuarkOCP_MX(QuarkScheme): "implemented. Please open an issue." ) - # TODO: integrate (or test) mixed-precision kernel. - self.emulate = not current_platform.supports_mx() or ( - self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4" - ) - - # TODO: Move emulation code path as a kernel, and always - # use init_mxfp4_linear_kernel. - if not self.emulate: - self.ocp_mx_linear = init_mxfp4_linear_kernel() - if not current_platform.supports_mx(): logger.warning_once( "The current platform does not support native MXFP4/MXFP6 " @@ -151,17 +149,10 @@ class QuarkOCP_MX(QuarkScheme): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False) - if self.emulate: - if self.dynamic_mxfp4_quant: - self.process_dynamic_mxfp4_weights_after_loading(layer) - else: - layer.weight_scale = torch.nn.Parameter( - layer.weight_scale.data, requires_grad=False - ) - else: - if self.dynamic_mxfp4_quant: - self.process_dynamic_mxfp4_weights_after_loading(layer) - self.ocp_mx_linear.process_weights_after_loading(layer) + if self.dynamic_mxfp4_quant: + self.process_dynamic_mxfp4_weights_after_loading(layer) + + self.ocp_mx_linear.process_weights_after_loading(layer) def create_weights( self, @@ -218,14 +209,20 @@ class QuarkOCP_MX(QuarkScheme): ) layer.register_parameter("weight_scale", weight_scale) + if self.weight_quant_key == kMxfp4Static: + self.ocp_mx_linear = init_mxfp4_linear_kernel( + activation_quant_key=self.activation_quant_key, + ) + elif self.weight_quant_key in [kMxfp6E2M3Static, kMxfp6E3M2Static]: + self.ocp_mx_linear = init_mxfp6_linear_kernel( + weight_quant_key=self.weight_quant_key, + activation_quant_key=self.activation_quant_key, + ) + def apply_weights( self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - if self.emulate: - dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype) - qdq_x = self.quant_dequant_func(x) - return F.linear(qdq_x, dq_w, bias) return self.ocp_mx_linear.apply_weights(layer, x, bias) diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 453e08e612..6678aa15e2 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -90,6 +90,7 @@ class ScaleDesc: GroupShape.PER_CHANNEL: "per_channel", } group_shape = d.get(self.group_shape, str(self.group_shape)) + return ( f"{fx.graph.dtype_abbrs[self.dtype]}," f"{'static' if self.static else 'dynamic'},{group_shape}" @@ -106,15 +107,24 @@ class QuantKey: symmetric: symmetric if True, asymmetric if False """ - dtype: torch.dtype + # TODO: QuantKey.dtype is assumed to be `torch.dtype` in matcher_utils.py, + # but #37990 introduced e.g. `kInt4Static` that uses a `ScalarType` dtype, + # same for kMxfp6 that does not have a native torch representation. + # Logical dtype and storage (torch) dtype should be separated (see #48949). + dtype: torch.dtype | ScalarType scale: ScaleDesc scale2: ScaleDesc | None = None symmetric: bool = True def __str__(self): scale2_str = f"scale2({self.scale2})," if self.scale2 else "" + dtype_description = ( + fx.graph.dtype_abbrs[self.dtype] + if isinstance(self.dtype, torch.dtype) + else self.dtype + ) return ( - f"QuantKey({fx.graph.dtype_abbrs[self.dtype]}," + f"QuantKey({dtype_description}," f"scale({self.scale}),{scale2_str}" f"{'a' if not self.symmetric else ''}symmetric)" ) @@ -172,6 +182,26 @@ kMxfp8Dynamic = QuantKey(FP8_DTYPE, scale=kMxfp8DynamicGroupScale, symmetric=Tru kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True) +kMxfp6E3M2StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) +kMxfp6E3M2Static = QuantKey( + scalar_types.float6_e3m2f, scale=kMxfp6E3M2StaticGroupScale, symmetric=True +) + +kMxfp6E3M2DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) +kMxfp6E3M2Dynamic = QuantKey( + scalar_types.float6_e3m2f, scale=kMxfp6E3M2DynamicGroupScale, symmetric=True +) + +kMxfp6E2M3StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) +kMxfp6E2M3Static = QuantKey( + scalar_types.float6_e2m3f, scale=kMxfp6E2M3StaticGroupScale, symmetric=True +) + +kMxfp6E2M3DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) +kMxfp6E2M3Dynamic = QuantKey( + scalar_types.float6_e2m3f, scale=kMxfp6E2M3DynamicGroupScale, symmetric=True +) + # TODO: convert this to use SCALAR_TYPE. This is not right. kInt4StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1)) kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True)