[Quantization] ModelOpt MXFP8: small-M linear opts + env exclude-modules

tests/kernels/quantization/test_mxfp8_bf16_fallback.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness test for the opt-in small-M BF16 fallback in the FlashInfer
+CUTLASS MXFP8 linear kernel (VLLM_MXFP8_BF16_FALLBACK_SMALL_M).
+
+At M < 128 the kernel skips FlashInfer's mm_mxfp8 (which pads M up to a
+128-row tile) and instead does a plain BF16 matmul against a dequantized
+copy of the weight. This checks the fallback output matches the reference
+dequantize-then-matmul it is meant to reproduce.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.has_device_capability(100),
+    reason="FlashInfer CUTLASS MXFP8 kernel requires sm_100 (Blackwell)",
+)
+
+
+class _Layer:
+    """Minimal stand-in for the linear layer the kernel mutates."""
+
+
+@pytest.mark.parametrize("M", [1, 16, 17, 32, 64, 127])
+@pytest.mark.parametrize("N", [256])
+@pytest.mark.parametrize("K", [256])
+def test_bf16_fallback_matches_dequant_matmul(monkeypatch, M, N, K):
+    monkeypatch.setenv("VLLM_MXFP8_BF16_FALLBACK_SMALL_M", "1")
+
+    from vllm.model_executor.kernels.linear.mxfp8.flashinfer import (
+        FlashInferCutlassMxfp8LinearKernel,
+    )
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        dequant_mxfp8_to_bf16,
+        mxfp8_e4m3_quantize,
+    )
+
+    torch.manual_seed(0)
+    w = torch.randn(N, K, dtype=torch.bfloat16, device="cuda")
+    w_fp8, w_scale = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=False)
+
+    layer = _Layer()
+    layer.weight = w_fp8
+    layer.weight_scale = w_scale
+
+    kernel = FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearLayerConfig())
+    kernel.process_weights_after_loading(layer)
+    assert hasattr(layer, "weight_bf16"), "fallback weight was not cached"
+
+    x = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
+    out = kernel.apply_weights(layer, x)
+
+    ref = x @ dequant_mxfp8_to_bf16(w_fp8, w_scale).t()
+    assert out.shape == (M, N)
+    torch.testing.assert_close(out.float(), ref.float(), rtol=2e-2, atol=2e-2)

tests/kernels/quantization/test_mxfp8_smallm_gemm.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness test for the Triton small-M MXFP8 GEMM (mxfp8_smallm_gemm).
+
+The kernel computes, per 32-wide block,
+    out[m, n] = sum_k in_fp8[m, k] * w_fp8[n, k] * 2^(in_sf + w_sf - 254)
+which is exactly a dequantize-then-matmul. We check it against that
+reference over the small-M (M < 128) range it is used for.
+"""
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Triton MXFP8 small-M GEMM requires a CUDA-like device",
+)
+
+
+@pytest.mark.parametrize("M", [1, 16, 17, 32, 64, 127])
+@pytest.mark.parametrize("N", [128, 256])
+@pytest.mark.parametrize("K", [128, 256])
+def test_mxfp8_smallm_gemm_matches_reference(M, N, K):
+    from vllm.model_executor.kernels.linear.mxfp8.triton_smallm import (
+        mxfp8_smallm_gemm,
+    )
+    from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+        dequant_mxfp8_to_bf16,
+        mxfp8_e4m3_quantize,
+    )
+
+    device = current_platform.device_type
+    torch.manual_seed(0)
+    x = torch.randn(M, K, dtype=torch.bfloat16, device=device)
+    w = torch.randn(N, K, dtype=torch.bfloat16, device=device)
+
+    x_fp8, x_scale = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False)
+    w_fp8, w_scale = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=False)
+
+    out = mxfp8_smallm_gemm(x_fp8, x_scale, w_fp8, w_scale)
+    assert out.shape == (M, N)
+    assert out.dtype == torch.bfloat16
+
+    x_deq = dequant_mxfp8_to_bf16(x_fp8, x_scale)
+    w_deq = dequant_mxfp8_to_bf16(w_fp8, w_scale)
+    ref = x_deq.float() @ w_deq.float().t()
+
+    torch.testing.assert_close(out.float(), ref, rtol=3e-2, atol=3e-2)

vllm/envs.py

@@ -117,6 +117,9 @@
     VLLM_DISABLE_PYNCCL: bool = False
     VLLM_USE_OINK_OPS: bool = False
     VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True
+    VLLM_MXFP8_BF16_FALLBACK_SMALL_M: bool = False
+    VLLM_MXFP8_TRITON_SMALLM: bool = False
+    VLLM_MODELOPT_EXTRA_EXCLUDE_MODULES: list[str] = []
     VLLM_ROCM_USE_AITER: bool = False
     VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
     VLLM_ROCM_USE_AITER_LINEAR: bool = True
@@ -1135,6 +1138,27 @@ def _resolve_rust_frontend_path() -> str | None:
         os.getenv("VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD", "True").lower()
         in ("true", "1")
     ),
+    # Cache a BF16-dequantized copy of MXFP8 linear weights so small-batch
+    # (M < 128) GEMMs can skip FlashInfer mm_mxfp8's 128-row tile padding.
+    # Opt-in: roughly doubles MXFP8 linear weight memory.
+    "VLLM_MXFP8_BF16_FALLBACK_SMALL_M": lambda: (
+        os.getenv("VLLM_MXFP8_BF16_FALLBACK_SMALL_M", "False").lower()
+        in ("true", "1")
+    ),
+    # Use a Triton MXFP8 GEMM for small-batch (M < 128) linear layers instead
+    # of FlashInfer mm_mxfp8, which pads M up to a 128-row tile. Takes
+    # precedence over VLLM_MXFP8_BF16_FALLBACK_SMALL_M when both are set.
+    "VLLM_MXFP8_TRITON_SMALLM": lambda: (
+        os.getenv("VLLM_MXFP8_TRITON_SMALLM", "False").lower() in ("true", "1")
+    ),
+    # Extra ModelOpt exclude_modules patterns (comma-separated fnmatch
+    # wildcards) appended to those from the checkpoint, so operators can keep
+    # specific modules in BF16 without re-exporting the checkpoint.
+    "VLLM_MODELOPT_EXTRA_EXCLUDE_MODULES": lambda: [
+        p.strip()
+        for p in os.environ.get("VLLM_MODELOPT_EXTRA_EXCLUDE_MODULES", "").split(",")
+        if p.strip()
+    ],
     "VLLM_ROCM_USE_AITER": lambda: (
         os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1")
     ),

vllm/model_executor/kernels/linear/mxfp8/flashinfer.py

@@ -4,8 +4,10 @@
 import torch
 from torch.nn.parameter import Parameter
 
+import vllm.envs as envs
 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
     MXFP8_BLOCK_SIZE,
+    dequant_mxfp8_to_bf16,
     mxfp8_e4m3_quantize,
     swizzle_mxfp8_scale,
 )
@@ -14,6 +16,7 @@
 from vllm.utils.flashinfer import has_flashinfer_cutedsl
 
 from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+from .triton_smallm import mxfp8_smallm_gemm
 
 
 class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
@@ -44,6 +47,22 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             weight_scale_swizzled.contiguous(), requires_grad=False
         )
 
+        # Optional small-M BF16 fallback (see apply_weights). mm_mxfp8 pads M
+        # up to a 128-row tile, so at small M most GEMM rows are wasted. Cache
+        # a BF16-dequantized weight to enable a plain matmul there instead.
+        if envs.VLLM_MXFP8_BF16_FALLBACK_SMALL_M:
+            weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale_2d)
+            layer.weight_bf16 = Parameter(
+                weight_bf16.contiguous(), requires_grad=False
+            )
+
+        # Optional small-M Triton GEMM (see apply_weights). The kernel indexes
+        # the scale row-major, so cache the un-swizzled [N, K/32] scale.
+        if envs.VLLM_MXFP8_TRITON_SMALLM:
+            layer.weight_scale_raw = Parameter(
+                weight_scale_2d.contiguous(), requires_grad=False
+            )
+
     def apply_weights(
         self,
         layer: torch.nn.Module,
@@ -57,6 +76,34 @@ def apply_weights(
 
         input_shape = x.shape
         input_2d = x.view(-1, K)
+        M = input_2d.shape[0]
+
+        # Small-M Triton MXFP8 GEMM: handle M < 128 natively instead of padding
+        # to mm_mxfp8's 128-row tile. Only active when weight_scale_raw was
+        # cached at load time (VLLM_MXFP8_TRITON_SMALLM). Takes precedence over
+        # the BF16 fallback below when both are enabled.
+        if M < 128 and hasattr(layer, "weight_scale_raw"):
+            input_mxfp8, input_scale = mxfp8_e4m3_quantize(
+                input_2d, is_sf_swizzled_layout=False
+            )
+            output = mxfp8_smallm_gemm(
+                input_mxfp8, input_scale, weight, layer.weight_scale_raw
+            )
+            if bias is not None:
+                output = output + bias
+            return output.view(*input_shape[:-1], N).to(out_dtype)
+
+        # Small-M BF16 fallback: avoid mm_mxfp8's 128-row tile padding when
+        # M < 128. Only active when weight_bf16 was cached at load time
+        # (VLLM_MXFP8_BF16_FALLBACK_SMALL_M); otherwise behaviour is unchanged.
+        if M < 128 and hasattr(layer, "weight_bf16"):
+            output = torch.matmul(
+                input_2d.to(torch.bfloat16), layer.weight_bf16.t()
+            )
+            if bias is not None:
+                output = output + bias
+            return output.view(*input_shape[:-1], N).to(out_dtype)
+
         min_dim = 128
 
         assert min_dim <= K, (