[Quant] Serve 32-wide-K ue8m0 block-FP8 linears through the FlashInfer MXFP8 GEMMs (#40039)

test/registered/quant/test_fp8_utils.py

@@ -1,9 +1,10 @@
 import unittest
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import torch
 
+from sglang.srt.layers.quantization import fp8_utils
 from sglang.srt.layers.quantization.fp8 import (
     Fp8MoEMethod,
     _is_cuda,
@@ -11,8 +12,13 @@ from sglang.srt.layers.quantization.fp8 import (
     _is_hip,
 )
 from sglang.srt.layers.quantization.fp8_utils import (
+    Fp8GemmRunnerBackend,
+    Mxfp8DenseGemmBackend,
+    block_fp8_scale_to_mxfp8_e8m0,
+    can_serve_block_fp8_as_mxfp8,
     inverse_transform_scale_ue8m0,
     quant_weight_ue8m0,
+    resolve_block_fp8_mxfp8_backend,
     transform_scale_ue8m0,
 )
 from sglang.srt.runtime_context import get_platform
@@ -104,6 +110,68 @@ class TestInverseTransformScaleUe8m0(CustomTestCase):
             )
 
 
+class TestBlockFp8AsMxfp8(CustomTestCase):
+    def test_block_scale_to_e8m0_matches_reference(self):
+        # Sibling classes leave torch's default device on cuda; stay on cpu.
+        gen = torch.Generator().manual_seed(0)
+        n, k, block_n = 100, 256, 32  # 4 scale rows, the last one partial
+        exps = torch.randint(-20, 21, (4, k // 32), generator=gen, device="cpu")
+        got = block_fp8_scale_to_mxfp8_e8m0(
+            torch.exp2(exps.float()), (n, k), [block_n, 32]
+        )
+        ref = (exps + 127).to(torch.uint8).repeat_interleave(block_n, dim=0)[:n]
+        self.assertTrue(torch.equal(got, ref))
+        with self.assertRaises(ValueError):  # 128-wide K block is not MXFP8
+            block_fp8_scale_to_mxfp8_e8m0(
+                torch.ones(2, 8, device="cpu"), (64, 1024), [32, 128]
+            )
+        with self.assertRaises(ValueError):  # not a power of two
+            block_fp8_scale_to_mxfp8_e8m0(
+                torch.full((2, 8), 1.5, device="cpu"), (64, 256), [32, 32]
+            )
+
+    def test_serve_gate(self):
+        platform = MagicMock()
+        platform.is_blackwell = True
+        cutedsl = next(b for b in Mxfp8DenseGemmBackend if b.is_flashinfer_cutedsl())
+        with (
+            patch.object(fp8_utils, "_is_cuda", True),
+            patch.object(fp8_utils, "get_platform", return_value=platform),
+            patch.object(fp8_utils, "is_flashinfer_available", return_value=True),
+            patch.object(
+                fp8_utils, "resolve_mxfp8_dense_gemm_backend", return_value=cutedsl
+            ),
+        ):
+            for name, expected in (
+                ("flashinfer_cutedsl", True),
+                ("flashinfer_cutlass", True),
+                ("flashinfer_trtllm", False),
+                ("triton", False),
+                ("auto", False),
+            ):
+                with (
+                    self.subTest(backend=name),
+                    patch.object(
+                        fp8_utils, "FP8_GEMM_RUNNER_BACKEND", Fp8GemmRunnerBackend(name)
+                    ),
+                ):
+                    self.assertEqual(
+                        can_serve_block_fp8_as_mxfp8([32, 32], "ue8m0"), expected
+                    )
+                    self.assertEqual(
+                        resolve_block_fp8_mxfp8_backend().is_unsupported(), not expected
+                    )
+            with patch.object(
+                fp8_utils,
+                "FP8_GEMM_RUNNER_BACKEND",
+                Fp8GemmRunnerBackend.FLASHINFER_CUTEDSL,
+            ):
+                self.assertFalse(can_serve_block_fp8_as_mxfp8([128, 128], "ue8m0"))
+                self.assertFalse(can_serve_block_fp8_as_mxfp8([32, 32], None))
+                platform.is_blackwell = False
+                self.assertFalse(can_serve_block_fp8_as_mxfp8([32, 32], "ue8m0"))
+
+
 class TestApplyFp8LinearScaleDispatch(CustomTestCase):
     @classmethod
     def setUpClass(cls):
@@ -268,6 +336,7 @@ class TestApplyFp8LinearScaleDispatch(CustomTestCase):
         native_method = native_fp8.Fp8LinearMethod.__new__(native_fp8.Fp8LinearMethod)
         native_method.use_marlin = False
         native_method.use_mxfp8 = False
+        native_method.block_fp8_as_mxfp8 = False
         native_method.block_quant = False
         native_method.cutlass_fp8_supported = True
         native_method.use_per_token_if_dynamic = False

test/registered/unit/layers/quantization/test_flashinfer_trtllm_fp8_fallback.py

@@ -22,12 +22,20 @@ from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
+import functools
 import unittest
 from unittest.mock import MagicMock, patch
 
 import torch
 
+import sglang.srt.layers.quantization.fp8 as fp8
 import sglang.srt.layers.quantization.fp8_utils as fp8_utils
+from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
+from sglang.srt.layers.quantization.fp8_utils import (
+    Fp8GemmRunnerBackend,
+    dispatch_w8a8_block_fp8_linear,
+    triton_w8a8_block_fp8_linear,
+)
 from sglang.test.test_utils import CustomTestCase
 
 BLOCK_SIZE = [128, 128]
@@ -100,3 +108,45 @@ class TestFlashinferTrtllmFp8Fallback(CustomTestCase):
 
 if __name__ == "__main__":
     unittest.main(verbosity=3)
+
+
+class TestBlockSizeDispatch(CustomTestCase):
+    """Non-128-wide K blocks dispatch to Triton regardless of --fp8-gemm-backend;
+    128-wide blocks keep the backend choice."""
+
+    def test_128_wide_k_blocks_keep_the_backend_choice(self):
+        for name in ("triton", "deep_gemm"):
+            with (
+                self.subTest(backend=name),
+                patch.object(
+                    fp8_utils, "FP8_GEMM_RUNNER_BACKEND", Fp8GemmRunnerBackend(name)
+                ),
+            ):
+                default = dispatch_w8a8_block_fp8_linear()
+                self.assertIs(dispatch_w8a8_block_fp8_linear([128, 128]), default)
+                self.assertIs(dispatch_w8a8_block_fp8_linear([1, 128]), default)
+                fn = dispatch_w8a8_block_fp8_linear([32, 32], act_scale_ue8m0=True)
+                self.assertIsInstance(fn, functools.partial)
+                self.assertIs(fn.func, triton_w8a8_block_fp8_linear)
+                self.assertEqual(fn.keywords, {"act_scale_ue8m0": True})
+
+    def test_method_dispatches_on_the_effective_block_size(self):
+        # An MXFP8 checkpoint converted to block-fp8 at load time is a [128, 128]
+        # weight; dispatching on the pre-conversion [1, 32] would pick Triton.
+        with (
+            patch.object(
+                fp8_utils, "FP8_GEMM_RUNNER_BACKEND", Fp8GemmRunnerBackend.TRITON
+            ),
+            patch.object(fp8, "_mxfp8_to_block_fp8_required", True),
+        ):
+            method = Fp8LinearMethod(
+                Fp8Config(
+                    is_checkpoint_fp8_serialized=True,
+                    use_mxfp8=True,
+                    weight_block_size=[1, 32],
+                    scale_fmt="ue8m0",
+                )
+            )
+            self.assertTrue(method.convert_mxfp8_to_block)
+            self.assertIs(method.w8a8_block_fp8_linear, triton_w8a8_block_fp8_linear)
+            self.assertFalse(method.block_fp8_as_mxfp8)

test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py

@@ -1,8 +1,9 @@
 """Numerics for the FP8 dense-linear GEMM backends (--fp8-gemm-backend).
 
-Real layer path vs a dequantized-reference matmul, in three formats: FP8
-blockwise, MXFP8, and per-tensor FP8 (auto dispatch). Backend sets adapt to
-the device SM, so one file covers SM90 / SM100 / SM120.
+Real layer path vs a dequantized-reference matmul, in four formats: FP8
+blockwise, MXFP8, 32-wide-K ue8m0 block FP8 served as MXFP8, and per-tensor
+FP8 (auto dispatch). Backend sets adapt to the device SM, so one file covers
+SM90 / SM100 / SM120.
 """
 
 import unittest
@@ -14,6 +15,7 @@ from sglang.srt.layers.quantization import fp8_utils
 from sglang.srt.layers.quantization.fp8 import Fp8Config
 from sglang.srt.layers.quantization.fp8_utils import Fp8GemmRunnerBackend
 from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp8Config
+from sglang.srt.layers.quantization.mxfp8_input import Mxfp8SwizzledInput
 from sglang.srt.utils import get_device_sm
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.layer_ut_utils import (
@@ -43,6 +45,12 @@ MXFP8_SHAPES = [
     (5, 384, 768),
 ]
 
+# (M, N, K); N % 64 == 0 and K % 128 == 0 for the FlashInfer MXFP8 scale swizzle.
+BLOCK32_SHAPES = [
+    (64, 512, 512),
+    (5, 384, 768),
+]
+
 # (M, N, K); per-tensor has no block-alignment constraints.
 PER_TENSOR_SHAPES = [
     (64, 512, 512),
@@ -88,6 +96,27 @@ def _quantize_fp8_blockwise(w: torch.Tensor, block: int = 128):
     return w_fp8.reshape(n, k), scale, w_dequant
 
 
+def _block32_backends():
+    # The block-fp8-as-MXFP8 route takes the FlashInfer CUTLASS / CuTe-DSL MXFP8
+    # kernels only, on SM100/103.
+    if get_device_sm() in (100, 103):
+        return ["flashinfer_cutlass", "flashinfer_cutedsl"]
+    return []
+
+
+def _quantize_fp8_block32_ue8m0(w: torch.Tensor, block: int = 32):
+    """Per (block, block) tile fp8 quantization with power-of-two scales; returns
+    checkpoint-format (w_fp8 [N, K], scale e8m0 [N/block, K/block]) and the
+    dequant reference."""
+    n, k = w.shape
+    tiles = w.float().reshape(n // block, block, k // block, block)
+    amax = tiles.abs().amax(dim=(1, 3)).clamp(min=1e-30)
+    scale = torch.exp2(torch.ceil(torch.log2(amax / FP8_MAX)))
+    w_fp8 = (tiles / scale[:, None, :, None]).to(torch.float8_e4m3fn)
+    w_dequant = (w_fp8.float() * scale[:, None, :, None]).reshape(n, k)
+    return w_fp8.reshape(n, k), scale.to(torch.float8_e8m0fnu), w_dequant
+
+
 def _quantize_mxfp8(w: torch.Tensor, block: int = 32):
     """Per (1, block) group e8m0 quantization; returns checkpoint-format
     (w_fp8 [N, K], scale uint8 [N, K/block]) and the dequant reference."""
@@ -238,6 +267,78 @@ class TestMxfp8LinearBackends(_LinearBackendCheck):
             is_backend_supported.assert_called_once_with("cute-dsl", 107)
 
 
+class TestBlockFp8AsMxfp8Linear(_LinearBackendCheck):
+    """A 32-wide-K ue8m0 block-fp8 weight served through the MXFP8 GEMMs."""
+
+    @staticmethod
+    def _build_layer(n: int, k: int, keep_plain_weight_layout: bool = False):
+        quant_config = Fp8Config(
+            is_checkpoint_fp8_serialized=True,
+            activation_scheme="dynamic",
+            weight_block_size=[32, 32],
+            scale_fmt="ue8m0",
+        )
+        layer = _make_linear(quant_config, n, k)
+        if keep_plain_weight_layout:
+            layer.keep_plain_weight_layout = True
+        w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10
+        w_fp8, scale_e8m0, w_dequant = _quantize_fp8_block32_ue8m0(w)
+        load_linear_weights(layer, weight=w_fp8, weight_scale_inv=scale_e8m0)
+        return layer, w_dequant
+
+    def _run(self, backend: str):
+        self._check_backend(
+            backend, _block32_backends(), BLOCK32_SHAPES, self._build_layer
+        )
+
+    def test_flashinfer_cutlass(self):
+        self._run("flashinfer_cutlass")
+
+    def test_flashinfer_cutedsl(self):
+        self._run("flashinfer_cutedsl")
+
+    def test_mxfp8_view_and_swizzled_input(self):
+        if "flashinfer_cutedsl" not in _block32_backends():
+            self.skipTest(f"cutedsl not in SM{get_device_sm()} backend set")
+        from sglang.kernels.ops.attention.dsv4.wo_a_bf16 import (
+            _quantize_partial,
+            _wo_a_reduce,
+        )
+
+        torch.manual_seed(7)
+        with mock.patch.object(
+            fp8_utils,
+            "FP8_GEMM_RUNNER_BACKEND",
+            Fp8GemmRunnerBackend.FLASHINFER_CUTEDSL,
+        ):
+            n, k = 512, 2048
+            layer, _ = self._build_layer(n, k)
+            layer.quant_method.process_weights_after_loading(layer)
+            self.assertTrue(layer.quant_method.block_fp8_as_mxfp8)
+            self.assertTrue(layer.block_fp8_mxfp8_ready)
+            # Block scales stay in place for the Triton fallback and raw readers.
+            self.assertEqual(tuple(layer.weight_scale_inv.shape), (n // 32, k // 32))
+            self.assertIsNotNone(layer.weight_scale_inv_swizzled)
+
+            # A prequantized 128x4-swizzled MXFP8 activation must give the same
+            # output as the bf16 input the layer quantizes itself.
+            rows = 6
+            partial = torch.randn(8, rows, 2, k // 2, device="cuda")
+            bf16 = torch.empty(rows, k, dtype=torch.bfloat16, device="cuda")
+            _wo_a_reduce[(rows * 8,)](partial, bf16, rows * k, num_warps=4)
+            q, s = _quantize_partial(partial)
+            swizzled = layer.quant_method.apply(layer, Mxfp8SwizzledInput(q, s))
+            plain = layer.quant_method.apply(layer, bf16)
+            torch.testing.assert_close(swizzled, plain, rtol=0, atol=0)
+
+            # A layer that keeps the plain weight layout has no MXFP8 view.
+            plain_layer, _ = self._build_layer(n, k, keep_plain_weight_layout=True)
+            plain_layer.quant_method.process_weights_after_loading(plain_layer)
+            self.assertFalse(plain_layer.block_fp8_mxfp8_ready)
+            with self.assertRaises(ValueError):
+                plain_layer.quant_method.apply(plain_layer, Mxfp8SwizzledInput(q, s))
+
+
 @unittest.skipIf(get_device_sm() < 90, "FP8 GEMM backends require SM90+")
 class TestModeloptFp8PerTensorLinear(_LinearBackendCheck):
     """Per-tensor FP8 (ModelOptFp8LinearMethod, static scales) on the auto