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

Signed-off-by: Felix Marty <[email protected]>

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;

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()

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)

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",

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:

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):

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)

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:

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:

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:

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:

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",
+]

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

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)

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:

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)

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)