feat(quantization): add native SM110 MXFP6

aphrodite/config/model.py

@@ -219,7 +219,8 @@ class ModelConfig:
     determine the data type of the weights."""
     quantization_config: dict[str, Any] | QuantizationConfigArgs | None = None
     """User-facing quantization configuration. Carries per-layer-kind specs
-    (linear, moe) and ignore patterns; see :class:`QuantizationConfigArgs`.
+    (linear, moe), ignore patterns, and ordered precision overrides; see
+    :class:`QuantizationConfigArgs`.
     Auto-populated from the matching online shorthand when `quantization` is
     one of the values in `ONLINE_QUANT_SHORTHAND_NAMES`."""
     allow_deprecated_quantization: bool = False

aphrodite/config/quantization.py

@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: disable-error-code=call-arg
 
-from typing import Annotated, Any
+from typing import Annotated, Any, Literal
 
 from pydantic import Field, GetPydanticSchema, ValidationInfo, field_validator
 from pydantic_core import core_schema
@@ -17,6 +18,10 @@
     kFp8StaticTensorSym,
     kInt8StaticChannelSym,
     kMxfp4Dynamic,
+    kMxfp6E2m3Dynamic,
+    kMxfp6E2m3Static,
+    kMxfp6E3m2Dynamic,
+    kMxfp6E3m2Static,
     kMxfp8Dynamic,
     kNvfp4Static,
 )
@@ -31,6 +36,10 @@
     "fp8_per_block_dynamic": kFp8Dynamic128Sym,
     "mxfp8": kMxfp8Dynamic,
     "mxfp4": kMxfp4Dynamic,
+    "mxfp6_e2m3": kMxfp6E2m3Static,
+    "mxfp6_e2m3_dynamic": kMxfp6E2m3Dynamic,
+    "mxfp6_e3m2": kMxfp6E3m2Static,
+    "mxfp6_e3m2_dynamic": kMxfp6E3m2Dynamic,
     "int8_per_channel_static": kInt8StaticChannelSym,
 }
 
@@ -69,6 +78,32 @@ class QuantSpec:
     """Activation quantization key, or a name from QUANT_KEY_NAMES."""
 
 
+def _coerce_override_weight(v: Any) -> Any:
+    if v in (None, "bf16"):
+        return v
+    return _coerce_quant_key(v)
+
+
+OverrideWeightField = Annotated[
+    QuantKey | Literal["bf16"] | None,
+    GetPydanticSchema(lambda _src, _handler: core_schema.no_info_plain_validator_function(_coerce_override_weight)),
+]
+
+
+@config
+class QuantOverride:
+    """Ordered module-level override for online quantization."""
+
+    pattern: str = ""
+    """Exact module prefix or ``re:`` regular expression."""
+
+    weight: OverrideWeightField = None
+    """Replacement weight format; ``bf16`` leaves the module unquantized."""
+
+    activation: QuantKeyField = None
+    """Replacement activation format; omitted fields inherit the base spec."""
+
+
 @config
 class QuantizationConfigArgs:
     """User-facing quantization configuration.
@@ -86,6 +121,9 @@ class QuantizationConfigArgs:
     ignore: list[str] = Field(default_factory=list)
     """Layers to skip quantization for."""
 
+    overrides: list[QuantOverride] = Field(default_factory=list)
+    """Ordered module precision overrides. Later matching rules win."""
+
     @field_validator("linear", "moe", mode="before")
     @classmethod
     def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any:
@@ -122,6 +160,16 @@ def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any:
         linear=QuantSpec(weight=kMxfp8Dynamic),
         moe=QuantSpec(weight=kMxfp8Dynamic),
     ),
+    "mxfp6": QuantizationConfigArgs(
+        linear=QuantSpec(weight=kMxfp6E2m3Static, activation=kMxfp8Dynamic),
+        moe=QuantSpec(weight=kMxfp6E2m3Static, activation=kMxfp8Dynamic),
+        overrides=[
+            QuantOverride(
+                pattern=r"re:(^|.*\.)(gate|router|shared_expert_gate|lm_head)$",
+                weight="bf16",
+            )
+        ],
+    ),
     # INT8 weight-only on MoE; linear stays unquantized (no `linear` field).
     "int8_per_channel_weight_only": QuantizationConfigArgs(
         moe=QuantSpec(weight=kInt8StaticChannelSym),
@@ -177,4 +225,5 @@ def resolve_quantization_config(
         linear=quantization_config.linear or base.linear,
         moe=quantization_config.moe or base.moe,
         ignore=quantization_config.ignore or base.ignore,
+        overrides=[*base.overrides, *quantization_config.overrides],
     )

aphrodite/engine/arg_utils.py

@@ -513,7 +513,7 @@ class EngineArgs:
     quantization: QuantizationMethods | str | None = ModelConfig.quantization
     quantization_config: "dict[str, Any] | QuantizationConfigArgs | None" = None
     """User-facing quantization configuration. Carries per-layer-kind
-    QuantSpecs (linear, moe) and ignore patterns; see
+    QuantSpecs (linear, moe), ignore patterns, and ordered precision overrides; see
     :class:`QuantizationConfigArgs`. Auto-populated from the matching online
     shorthand when `quantization` is one of the values in
     `ONLINE_QUANT_SHORTHAND_NAMES`."""

aphrodite/model_executor/kernels/linear/mxfp6/SOURCE.md

@@ -0,0 +1,18 @@
+<!-- SPDX-License-Identifier: Apache-2.0 -->
+<!-- SPDX-FileCopyrightText: Copyright contributors to the vLLM project -->
+
+# Vendored CuTe DSL source
+
+`cutedsl_kernel.py` is derived from NVIDIA CUTLASS commit
+`f94ec46f4f63f96003d6cfdf2014731e7672c281`:
+
+`examples/python/CuTeDSL/cute/blackwell/kernel/blockscaled_gemm/dense_blockscaled_gemm_persistent.py`
+
+Sonar integrates the kernel with its packed-weight ABI and Torch custom
+operator in `cutedsl.py`. The upstream command-line reference helpers remain
+in the vendored file to make future CUTLASS updates easier.
+
+`cutedsl_grouped_kernel.py` derives from the grouped block-scaled GEMM example
+at the same CUTLASS commit. It ports mixed-width operand and FP6 TMA-unpack
+support from the dense implementation. NVIDIA's grouped example only accepted
+same-width operands at that revision.

aphrodite/model_executor/kernels/linear/mxfp6/__init__.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .base import Mxfp6LinearKernel, Mxfp6LinearLayerConfig
+from .cutedsl import CutedslMxfp6LinearKernel
+from .cutedsl_grouped import cutedsl_grouped_mxfp6_gemm
+
+__all__ = [
+    "CutedslMxfp6LinearKernel",
+    "cutedsl_grouped_mxfp6_gemm",
+    "Mxfp6LinearKernel",
+    "Mxfp6LinearLayerConfig",
+]

aphrodite/model_executor/kernels/linear/mxfp6/base.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+
+
+@dataclass(frozen=True)
+class Mxfp6LinearLayerConfig:
+    weight_format: Literal["e2m3", "e3m2"] = "e2m3"
+    activation_format: Literal["mxfp8", "mxfp6_e2m3", "mxfp6_e3m2"] = "mxfp8"
+
+
+class Mxfp6LinearKernel(ABC):
+    def __init__(self, config: Mxfp6LinearLayerConfig) -> None:
+        supported, reason = self.is_supported()
+        if not supported:
+            raise ValueError(reason)
+        self.config = config
+
+    @classmethod
+    @abstractmethod
+    def is_supported(cls) -> tuple[bool, str | None]: ...
+
+    @classmethod
+    @abstractmethod
+    def can_implement_shape(cls, n: int, k: int) -> tuple[bool, str | None]: ...
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ...
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor: ...

aphrodite/model_executor/kernels/linear/mxfp6/cutedsl.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Native Blackwell MXFP8 x MXFP6 linear kernel."""
+
+from functools import lru_cache
+from typing import Literal
+
+import torch
+from torch.nn.parameter import Parameter
+
+from aphrodite.model_executor.layers.quantization.utils.mxfp8_utils import swizzle_mxfp8_scale
+from aphrodite.platforms import current_platform
+from aphrodite.utils.import_utils import has_cutedsl
+
+from .base import Mxfp6LinearKernel
+
+
+@lru_cache(maxsize=16)
+def _compile_gemm(activation_format: str, weight_format: str, output_dtype: torch.dtype):
+    import cuda.bindings.driver as cuda
+    import cutlass
+    from cutlass import utils
+
+    from .cutedsl_kernel import (
+        Sm100BlockScaledPersistentDenseGemmKernel,
+        scaled_mm,
+    )
+
+    weight_dtype = cutlass.Float6E2M3FN if weight_format == "e2m3" else cutlass.Float6E3M2FN
+    activation_dtype = {
+        "mxfp8": cutlass.Float8E4M3FN,
+        "mxfp6_e2m3": cutlass.Float6E2M3FN,
+        "mxfp6_e3m2": cutlass.Float6E3M2FN,
+    }[activation_format]
+    out_dtype = cutlass.BFloat16 if output_dtype == torch.bfloat16 else cutlass.Float16
+    cluster = (1, 1)
+    gemm = Sm100BlockScaledPersistentDenseGemmKernel(32, (128, 128), cluster)
+    max_clusters = utils.HardwareInfo().get_max_active_clusters(1)
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    return scaled_mm(
+        gemm,
+        activation_dtype,
+        weight_dtype,
+        out_dtype,
+        cutlass.Float8E8M0FNU,
+        "k",
+        "k",
+        "n",
+        max_clusters,
+        stream,
+    )
+
+
+@torch.library.custom_op("aphrodite::cutedsl_mxfp6_gemm", mutates_args={"out"})
+def _cutedsl_mxfp6_gemm(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    out: torch.Tensor,
+    activation_format: str,
+    weight_format: str,
+) -> None:
+    """Launch the CuTe DSL kernel behind an opaque Torch operator boundary."""
+    import cuda.bindings.driver as cuda
+    import cutlass
+    import cutlass.cute as cute
+    from cutlass.cute.runtime import make_ptr
+
+    from aphrodite.model_executor.layers.quantization.utils.mxfp6_online_utils import quantize_mxfp6_cuda
+    from aphrodite.model_executor.layers.quantization.utils.mxfp8_utils import mxfp8_e4m3_quantize
+
+    m, k = x.shape
+    n = out.shape[1]
+    if activation_format == "mxfp8":
+        x_q, x_scale = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True)
+        activation_dtype = cutlass.Float8E4M3FN
+    else:
+        activation_encoding: Literal["e2m3", "e3m2"] = "e2m3" if activation_format == "mxfp6_e2m3" else "e3m2"
+        x_q, x_scale = quantize_mxfp6_cuda(x, activation_encoding)
+        x_scale = swizzle_mxfp8_scale(x_scale, M=m, K=k)
+        activation_dtype = cutlass.Float6E2M3FN if activation_encoding == "e2m3" else cutlass.Float6E3M2FN
+    weight_dtype = cutlass.Float6E2M3FN if weight_format == "e2m3" else cutlass.Float6E3M2FN
+    out_dtype = cutlass.BFloat16 if out.dtype == torch.bfloat16 else cutlass.Float16
+    stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    compiled = _compile_gemm(activation_format, weight_format, out.dtype)
+    compiled(
+        make_ptr(
+            activation_dtype,
+            x_q.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
+        make_ptr(
+            weight_dtype,
+            weight.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
+        make_ptr(
+            cutlass.Float8E8M0FNU,
+            x_scale.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=32,
+        ),
+        make_ptr(
+            cutlass.Float8E8M0FNU,
+            weight_scale.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=32,
+        ),
+        make_ptr(
+            out_dtype,
+            out.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
+        (m, n, k, 1),
+        stream,
+    )
+
+
+class CutedslMxfp6LinearKernel(Mxfp6LinearKernel):
+    """Thor-native tcgen05 MXFP6 GEMM through NVIDIA CUTLASS DSL."""
+
+    @classmethod
+    def is_supported(cls) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "MXFP6 requires CUDA"
+        capability = current_platform.get_device_capability()
+        if capability is None or capability.to_int() != 110:
+            return False, "the initial native MXFP6 kernel requires SM110"
+        if not has_cutedsl():
+            return False, "MXFP6 requires nvidia-cutlass-dsl"
+        return True, None
+
+    @classmethod
+    def can_implement_shape(cls, n: int, k: int) -> tuple[bool, str | None]:
+        if n < 128 or n % 128:
+            return False, "output width must be a multiple of 128"
+        if k < 128 or k % 128:
+            return False, "input width must be a multiple of 128"
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        n = layer.mxfp6_logical_n
+        k = layer.mxfp6_logical_k
+        supported, reason = self.can_implement_shape(n, k)
+        if not supported:
+            raise ValueError(reason)
+        scales = swizzle_mxfp8_scale(layer.weight_scale.data, M=n, K=k)
+        layer.weight = Parameter(layer.weight.data.contiguous(), requires_grad=False)
+        layer.weight_scale = Parameter(scales.contiguous(), requires_grad=False)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        n = layer.mxfp6_logical_n
+        k = layer.mxfp6_logical_k
+        input_shape = x.shape
+        x_2d = x.reshape(-1, k)
+        m = x_2d.shape[0]
+        out = torch.empty((m, n), dtype=x.dtype, device=x.device)
+
+        _cutedsl_mxfp6_gemm(
+            x_2d,
+            layer.weight,
+            layer.weight_scale,
+            out,
+            self.config.activation_format,
+            self.config.weight_format,
+        )
+        if bias is not None:
+            out.add_(bias)
+        return out.view(*input_shape[:-1], n)