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

[ROCm][Quark][7/N] Use MXFP4 linear kernel abstraction for emulation backend (#48949)

Browse Source 
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
This commit is contained in:
fxmarty-amd
2026-07-31 10:07:15 -07:00
committed by GitHub
parent d87d2ca747
commit e67a2e0a56
17 changed files with 824 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/quantization/test_mxfp4_kernel_selection.py
+118 -22
View File
@@ -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()
tests/kernels/quantization/test_mxfp6_kernel_selection.py
+150
View File
@@ -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)