[Test] Replace GEMM backend e2e matrices with layer-level unit tests (#33596)

python/sglang/srt/layers/quantization/fp8.py

@@ -484,7 +484,6 @@ class Fp8LinearMethod(LinearMethodBase):
         output_partition_sizes: List[int],
         skip_block_quant_check: bool = False,
     ):
-        tp_size = get_parallel().tp_size
         block_n, block_k = (
             quant_config.weight_block_size[0],
             quant_config.weight_block_size[1],
@@ -495,6 +494,7 @@ class Fp8LinearMethod(LinearMethodBase):
                 "Skipping block quantization checks for weight partition."
             )
         else:
+            tp_size = get_parallel().tp_size
             # Required by row parallel
             if tp_size > 1 and input_size // input_size_per_partition == tp_size:
                 if input_size_per_partition % block_k != 0:

test/registered/quant/test_fp8_blockwise_gemm.py

@@ -1,145 +0,0 @@
-import unittest
-from types import SimpleNamespace
-from urllib.parse import urlparse
-
-from sglang.srt.utils import get_device_sm, kill_process_tree
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.run_eval import run_eval
-from sglang.test.test_utils import (
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    popen_launch_server,
-    try_cached_model,
-)
-
-register_cuda_ci(est_time=430, stage="extra-b", runner_config="4-gpu-b200")
-
-MODEL_PATH = "Qwen/Qwen3-4B-Instruct-2507-FP8"
-MXFP8_MODEL_PATH = "zianglih/Qwen3-4B-Instruct-2507-MXFP8"
-
-
-class FP8BlockwiseGemmBase:
-    backend = None
-
-    @classmethod
-    def setUpClass(cls):
-        if cls.backend is None:
-            raise NotImplementedError("Subclass must set 'backend' attribute")
-        cls.model = try_cached_model(MODEL_PATH)
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        other_args = [
-            "--trust-remote-code",
-            "--fp8-gemm-backend",
-            cls.backend,
-        ]
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=other_args,
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        parsed_url = urlparse(self.base_url)
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            api="completion",
-            max_tokens=512,
-            num_examples=1319,
-            num_threads=200,
-            num_shots=8,
-        )
-        metrics = run_eval(args)
-        print(metrics)
-
-        self.assertGreaterEqual(metrics["score"], 0.8)
-
-
-class MXFP8GemmBase:
-    backend = None
-
-    @classmethod
-    def setUpClass(cls):
-        if cls.backend is None:
-            raise NotImplementedError("Subclass must set 'backend' attribute")
-        cls.model = try_cached_model(MXFP8_MODEL_PATH)
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        other_args = [
-            "--trust-remote-code",
-            "--fp8-gemm-backend",
-            cls.backend,
-            # TODO: pin tc_piecewise — default `breakable` prefill runs MXFP8 RMSNorm in bf16, hurting accuracy; unrelated to BCG-default change.
-            "--cuda-graph-backend-prefill",
-            "tc_piecewise",
-        ]
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=other_args,
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        parsed_url = urlparse(self.base_url)
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            api="completion",
-            max_tokens=512,
-            num_examples=1319,
-            num_threads=200,
-            num_shots=8,
-        )
-        metrics = run_eval(args)
-        print(metrics)
-
-        self.assertGreaterEqual(metrics["score"], 0.8)
-
-
-class TestFP8BlockwiseGemmTriton(FP8BlockwiseGemmBase, unittest.TestCase):
-    backend = "triton"
-
-
-class TestFP8BlockwiseGemmDeepGemm(FP8BlockwiseGemmBase, unittest.TestCase):
-    backend = "deep_gemm"
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestFP8BlockwiseGemmFlashinferTrtllm(FP8BlockwiseGemmBase, unittest.TestCase):
-    backend = "flashinfer_trtllm"
-
-
-@unittest.skipIf(get_device_sm() != 90, "Test requires CUDA SM 90")
-class TestFP8BlockwiseGemmFlashinferDeepGemm(FP8BlockwiseGemmBase, unittest.TestCase):
-    backend = "flashinfer_deepgemm"
-
-
-@unittest.skip("Currently PCG capture takes too long to complete, disable until fixed")
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestMXFP8GemmTriton(MXFP8GemmBase, unittest.TestCase):
-    backend = "triton"
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestMXFP8GemmFlashinferTrtllm(MXFP8GemmBase, unittest.TestCase):
-    backend = "flashinfer_trtllm"
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestMXFP8GemmFlashinferCutlass(MXFP8GemmBase, unittest.TestCase):
-    backend = "flashinfer_cutlass"
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/quant/test_fp8_gemm_sm120.py

@@ -1,86 +0,0 @@
-import unittest
-from types import SimpleNamespace
-from urllib.parse import urlparse
-
-from sglang.srt.utils import get_device_sm, kill_process_tree
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
-from sglang.test.test_utils import (
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    popen_launch_server,
-    try_cached_model,
-)
-
-register_cuda_ci(est_time=146, stage="extra-a", runner_config="1-gpu-small")
-
-PERTENSOR_MODEL_PATH = "nvidia/Llama-3.1-8B-Instruct-FP8"
-BLOCKWISE_MODEL_PATH = "Qwen/Qwen3-4B-Instruct-2507-FP8"
-
-
-class FP8GemmSM120Base:
-    model_path = None
-    backend = None
-    quantization = None
-
-    @classmethod
-    def setUpClass(cls):
-        if cls.backend is None:
-            raise NotImplementedError("Subclass must set 'backend' attribute")
-        cls.model = try_cached_model(cls.model_path)
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        other_args = [
-            "--trust-remote-code",
-            "--fp8-gemm-backend",
-            cls.backend,
-            "--cuda-graph-backend-prefill=disabled",
-        ]
-        if cls.quantization:
-            other_args += ["--quantization", cls.quantization]
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=other_args,
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        if hasattr(cls, "process"):
-            kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        parsed_url = urlparse(self.base_url)
-        args = SimpleNamespace(
-            num_shots=self.num_shots,
-            data_path=None,
-            num_questions=1319,
-            max_new_tokens=512,
-            parallel=200,
-            host=parsed_url.hostname,
-            port=parsed_url.port,
-        )
-        metrics = run_eval_few_shot_gsm8k(args)
-        print(f"{metrics=}")
-        self.assertGreaterEqual(metrics["accuracy"], self.accuracy_threshold)
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestFP8PerTensorGemmSM120Auto(FP8GemmSM120Base, unittest.TestCase):
-    model_path = PERTENSOR_MODEL_PATH
-    backend = "auto"
-    quantization = "modelopt_fp8"
-    num_shots = 5
-    accuracy_threshold = 0.73
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestFP8BlockwiseGemmSM120Auto(FP8GemmSM120Base, unittest.TestCase):
-    model_path = BLOCKWISE_MODEL_PATH
-    backend = "auto"
-    num_shots = 8
-    accuracy_threshold = 0.87
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/quant/test_nvfp4_gemm.py

@@ -1,86 +0,0 @@
-import unittest
-from types import SimpleNamespace
-from urllib.parse import urlparse
-
-from sglang.srt.utils import get_device_sm, kill_process_tree
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.run_eval import run_eval
-from sglang.test.test_utils import (
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    popen_launch_server,
-    try_cached_model,
-)
-
-register_cuda_ci(est_time=350, stage="base-c", runner_config="4-gpu-b200")
-
-MODEL_PATH = "nvidia/Llama-3.1-8B-Instruct-NVFP4"
-
-
-class FP4GemmBase:
-    backend = None
-
-    @classmethod
-    def setUpClass(cls):
-        if cls.backend is None:
-            raise NotImplementedError("Subclass must set 'backend' attribute")
-        cls.model = try_cached_model(MODEL_PATH)
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        other_args = [
-            "--trust-remote-code",
-            "--quantization",
-            "modelopt_fp4",
-            "--fp4-gemm-backend",
-            cls.backend,
-        ]
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=other_args,
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_gsm8k(self):
-        parsed_url = urlparse(self.base_url)
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            api="completion",
-            max_tokens=512,
-            num_examples=1319,
-            num_threads=200,
-        )
-        metrics = run_eval(args)
-        print(metrics)
-
-        # TODO: restore 0.64 once the BCG-prefill RMSNorm fp32 fix lands.
-        self.assertGreater(metrics["score"], 0.63)
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestFP4GemmFlashinferCutlass(FP4GemmBase, unittest.TestCase):
-    backend = "flashinfer_cutlass"
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestFP4GemmFlashinferCudnn(FP4GemmBase, unittest.TestCase):
-    backend = "flashinfer_cudnn"
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestFP4GemmFlashinferTrtllm(FP4GemmBase, unittest.TestCase):
-    backend = "flashinfer_trtllm"
-
-
-@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
-class TestFP4GemmFlashinferCutedsl(FP4GemmBase, unittest.TestCase):
-    backend = "flashinfer_cutedsl"
-
-
-if __name__ == "__main__":
-    unittest.main()

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

@@ -0,0 +1,245 @@
+"""Numerics for the FP8 dense-linear GEMM backends (--fp8-gemm-backend).
+
+Runs the quant-method layer path (create_weights ->
+process_weights_after_loading -> apply) against a dequantized-reference
+matmul, covering the per-backend weight preparation (e.g. UE8M0 scale requant
+for DeepGEMM, per-backend MXFP8 scale packing) and the GEMM dispatch.
+Three formats: FP8 blockwise (Fp8LinearMethod), MXFP8 (Fp8LinearMethod with
+use_mxfp8), and per-tensor FP8 (ModelOptFp8LinearMethod, auto dispatch).
+The backend set adapts to the device SM version, so the same file covers
+Hopper (SM90), B200-class (SM100/103), and consumer Blackwell (SM120).
+"""
+
+import unittest
+from unittest import mock
+
+import torch
+
+from sglang.srt.layers.quantization import fp8_utils
+from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
+from sglang.srt.layers.quantization.fp8_utils import Fp8GemmRunnerBackend
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptFp8Config,
+    ModelOptFp8LinearMethod,
+)
+from sglang.srt.utils import get_device_sm
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="4-gpu-b200")
+register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large")
+
+FP8_MAX = 448.0
+
+# (M, N, K), N and K multiples of the (128, 128) weight block.
+FP8_BLOCK_SHAPES = [
+    (64, 512, 512),
+    (5, 384, 896),
+    (128, 1024, 1024),
+]
+
+# (M, N, K); K must be a multiple of 256 (flashinfer trtllm mxfp8 requirement).
+MXFP8_SHAPES = [
+    (64, 512, 512),
+    (5, 384, 768),
+]
+
+# (M, N, K); per-tensor has no block-alignment constraints.
+PER_TENSOR_SHAPES = [
+    (64, 512, 512),
+    (5, 384, 896),
+]
+
+
+def _fp8_block_backends():
+    sm = get_device_sm()
+    if 100 <= sm < 110:
+        return ["triton", "deep_gemm", "flashinfer_trtllm", "flashinfer_cutlass"]
+    if sm >= 120:
+        # cutlass is the SM120-only explicit backend; the trtllm / deepgemm
+        # kernels do not support consumer Blackwell.
+        return ["triton", "cutlass"]
+    if sm == 90:
+        # flashinfer_deepgemm (swapAB) is SM90-only.
+        return ["triton", "deep_gemm", "flashinfer_deepgemm"]
+    return []
+
+
+def _mxfp8_backends():
+    # MXFP8 linear is validated on SM100/103 only.
+    if 100 <= get_device_sm() < 110:
+        return ["triton", "flashinfer_trtllm", "flashinfer_cutlass"]
+    return []
+
+
+def _quantize_fp8_blockwise(w: torch.Tensor, block: int = 128):
+    """Per (block, block) tile fp8 quantization; returns checkpoint-format
+    (w_fp8 [N, K], scale_inv fp32 [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-12)
+    scale = 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, 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."""
+    n, k = w.shape
+    groups = w.float().reshape(n, k // block, block)
+    amax = groups.abs().amax(dim=-1).clamp(min=1e-12)
+    exp = torch.ceil(torch.log2(amax / FP8_MAX)).clamp(min=-127, max=127)
+    scale = torch.pow(2.0, exp)
+    w_fp8 = (groups / scale[..., None]).to(torch.float8_e4m3fn)
+    w_dequant = (w_fp8.float() * scale[..., None]).reshape(n, k)
+    scale_e8m0 = (exp + 127).to(torch.uint8)
+    return w_fp8.reshape(n, k), scale_e8m0, w_dequant
+
+
+def _create_weights(method, n: int, k: int, device: str = "cuda"):
+    layer = torch.nn.Module()
+    kwargs = {}
+    if isinstance(method, Fp8LinearMethod):
+        # The shape check reads TP world size (needs distributed init); skip it here.
+        kwargs["skip_block_quant_check"] = True
+    method.create_weights(
+        layer,
+        input_size_per_partition=k,
+        output_partition_sizes=[n],
+        input_size=k,
+        output_size=n,
+        params_dtype=torch.bfloat16,
+        weight_loader=lambda *args, **kw: None,
+        **kwargs,
+    )
+    return layer.to(device)
+
+
+class _LinearBackendCheck(CustomTestCase):
+    def _check_backend(self, backend: str, allowed, shapes, build_layer):
+        if backend not in allowed:
+            self.skipTest(f"{backend} not in SM{get_device_sm()} backend set")
+        torch.manual_seed(7)
+        for m, n, k in shapes:
+            with self.subTest(backend=backend, shape=(m, n, k)):
+                with mock.patch.object(
+                    fp8_utils,
+                    "FP8_GEMM_RUNNER_BACKEND",
+                    Fp8GemmRunnerBackend(backend),
+                ):
+                    method, layer, w_dequant = build_layer(n, k)
+                    method.process_weights_after_loading(layer)
+
+                    x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
+                    out = method.apply(layer, x)
+
+                    ref = x.float() @ w_dequant.T
+                    self.assertEqual(out.shape, (m, n))
+                    cos = torch.nn.functional.cosine_similarity(
+                        out.float().flatten(), ref.flatten(), dim=0
+                    ).item()
+                    self.assertGreater(cos, 0.99)
+                    # atol covers single-element UE8M0 scale-rounding outliers
+                    # (deep_gemm); a wrong kernel/layout fails by orders more.
+                    torch.testing.assert_close(out.float(), ref, rtol=5e-2, atol=1e-1)
+
+
+@unittest.skipIf(get_device_sm() < 90, "FP8 GEMM backends require SM90+")
+class TestFp8BlockwiseLinearBackends(_LinearBackendCheck):
+    @staticmethod
+    def _build_layer(n: int, k: int):
+        quant_config = Fp8Config(
+            is_checkpoint_fp8_serialized=True,
+            activation_scheme="dynamic",
+            weight_block_size=[128, 128],
+        )
+        method = Fp8LinearMethod(quant_config)
+        layer = _create_weights(method, n, k)
+        w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10
+        w_fp8, scale_inv, w_dequant = _quantize_fp8_blockwise(w)
+        layer.weight.data.copy_(w_fp8)
+        layer.weight_scale_inv.data.copy_(scale_inv)
+        return method, layer, w_dequant
+
+    def _run(self, backend: str):
+        self._check_backend(
+            backend, _fp8_block_backends(), FP8_BLOCK_SHAPES, self._build_layer
+        )
+
+    def test_triton(self):
+        self._run("triton")
+
+    def test_deep_gemm(self):
+        self._run("deep_gemm")
+
+    def test_flashinfer_trtllm(self):
+        self._run("flashinfer_trtllm")
+
+    def test_flashinfer_cutlass(self):
+        self._run("flashinfer_cutlass")
+
+    def test_flashinfer_deepgemm(self):
+        self._run("flashinfer_deepgemm")
+
+    def test_cutlass(self):
+        self._run("cutlass")
+
+
+@unittest.skipIf(get_device_sm() < 90, "FP8 GEMM backends require SM90+")
+class TestMxfp8LinearBackends(_LinearBackendCheck):
+    @staticmethod
+    def _build_layer(n: int, k: int):
+        quant_config = Fp8Config(
+            is_checkpoint_fp8_serialized=True,
+            activation_scheme="dynamic",
+            use_mxfp8=True,
+        )
+        method = Fp8LinearMethod(quant_config)
+        layer = _create_weights(method, n, k)
+        w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10
+        w_fp8, scale_e8m0, w_dequant = _quantize_mxfp8(w)
+        layer.weight.data.copy_(w_fp8)
+        layer.weight_scale_inv.data.copy_(scale_e8m0)
+        return method, layer, w_dequant
+
+    def _run(self, backend: str):
+        self._check_backend(backend, _mxfp8_backends(), MXFP8_SHAPES, self._build_layer)
+
+    def test_triton(self):
+        self._run("triton")
+
+    def test_flashinfer_trtllm(self):
+        self._run("flashinfer_trtllm")
+
+    def test_flashinfer_cutlass(self):
+        self._run("flashinfer_cutlass")
+
+
+@unittest.skipIf(get_device_sm() < 90, "FP8 GEMM backends require SM90+")
+class TestModeloptFp8PerTensorLinear(_LinearBackendCheck):
+    """Per-tensor FP8 (ModelOptFp8LinearMethod, static scales) on the auto
+    dispatch path -- the checkpoint style of nvidia/*-FP8 models."""
+
+    @staticmethod
+    def _build_layer(n: int, k: int):
+        quant_config = ModelOptFp8Config(is_checkpoint_fp8_serialized=True)
+        method = ModelOptFp8LinearMethod(quant_config)
+        layer = _create_weights(method, n, k)
+        w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10
+        scale = (w.float().abs().max() / FP8_MAX).clamp(min=1e-12)
+        w_fp8 = (w.float() / scale).to(torch.float8_e4m3fn)
+        layer.weight.data.copy_(w_fp8)
+        layer.weight_scale.data.fill_(scale)
+        layer.input_scale.data.fill_(1.0 / FP8_MAX)
+        w_dequant = w_fp8.float() * scale
+        return method, layer, w_dequant
+
+    def test_auto(self):
+        self._check_backend("auto", ["auto"], PER_TENSOR_SHAPES, self._build_layer)
+
+
+if __name__ == "__main__":
+    unittest.main()

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

@@ -0,0 +1,175 @@
+"""Numerics for the NVFP4 dense-linear GEMM backends (--fp4-gemm-backend).
+
+Runs ModelOptFp4LinearMethod end to end (create_weights ->
+process_weights_after_loading -> apply) for each SM100 backend choice and
+checks the output against a dequantized-reference matmul. This covers both
+the per-backend weight preparation (padding / interleave / TRTLLM shuffle)
+and the GEMM kernel dispatch.
+"""
+
+import unittest
+from unittest import mock
+
+import torch
+from flashinfer import fp4_quantize
+
+from sglang.srt.layers.quantization import fp4_utils
+from sglang.srt.layers.quantization.fp4_utils import Fp4GemmRunnerBackend
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptFp4Config,
+    ModelOptFp4LinearMethod,
+)
+from sglang.srt.utils import get_device_sm
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="4-gpu-b200")
+
+kE2M1ToFloat = torch.tensor(
+    [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32
+)
+FLOAT8_E4M3_MAX = 448.0
+FLOAT4_E2M1_MAX = 6.0
+
+# (M, N, K). The second shape hits the padding paths: N=160 is not a multiple
+# of 128 (TRTLLM shuffle pad) and K=336 is neither a multiple of 32 (CUTLASS
+# K pad) nor K/16 a multiple of 4 (TRTLLM scale pad).
+SHAPES = [
+    (64, 256, 512),
+    (5, 160, 336),
+    (128, 1024, 1024),
+]
+
+BACKENDS = [
+    "flashinfer_cutedsl",
+    "flashinfer_cutlass",
+    "flashinfer_cudnn",
+    "flashinfer_trtllm",
+]
+
+
+def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
+    m_tiles = (m + 128 - 1) // 128
+    f = block_size * 4
+    k_tiles = (k + f - 1) // f
+    tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
+    tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
+    out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
+    # Crop the K-tile padding too: k // block_size scale columns, not k.
+    return out[0:m, 0 : k // block_size]
+
+
+def break_fp4_bytes(a, dtype):
+    assert a.dtype == torch.uint8
+    m, n = a.shape
+    a_flat = a.flatten()
+    high = (a_flat & 0xF0) >> 4
+    low = a_flat & 0x0F
+    combined = torch.stack((low, high), dim=1).flatten()
+    signs = (combined & 0x08).to(torch.bool)
+    abs_vals = (combined & 0x07).to(torch.long)
+    kE2M1 = kE2M1ToFloat.to(device=a.device)
+    values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
+    return values.reshape(m, n * 2).to(dtype=dtype)
+
+
+def dequantize_nvfp4_to_dtype(
+    tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16
+):
+    assert tensor_fp4.dtype == torch.uint8
+    m, packed_k = tensor_fp4.shape
+    k = packed_k * 2
+    tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32)
+    tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
+    tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
+    tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
+    tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
+    out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
+    return out.to(dtype=dtype)
+
+
+def _make_quantized_layer(n: int, k: int, device: str = "cuda"):
+    """Build a linear layer holding NVFP4 checkpoint-format weights; returns
+    (method, layer, w_dequant) with w_dequant the fp32 quant->dequant reference."""
+    quant_config = ModelOptFp4Config(
+        is_checkpoint_nvfp4_serialized=True,
+        group_size=16,
+        use_per_token_activation=False,
+    )
+    method = ModelOptFp4LinearMethod(quant_config)
+    layer = torch.nn.Module()
+    method.create_weights(
+        layer,
+        input_size_per_partition=k,
+        output_partition_sizes=[n],
+        input_size=k,
+        output_size=n,
+        params_dtype=torch.bfloat16,
+        weight_loader=lambda *args, **kwargs: None,
+    )
+    layer = layer.to(device)
+
+    w = torch.randn((n, k), device=device, dtype=torch.bfloat16) / 10
+    w_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w.abs().max().to(torch.float32)
+    w_q, w_sf_swizzled = fp4_quantize(w, w_gs)
+    w_sf_linear = convert_swizzled_to_linear(
+        w_sf_swizzled.view(torch.float8_e4m3fn), n, k, 16
+    )
+    w_dequant = dequantize_nvfp4_to_dtype(
+        w_q, w_sf_swizzled, w_gs, torch.float32, device
+    )
+
+    layer.weight.data.copy_(w_q)
+    layer.weight_scale.data.copy_(w_sf_linear)
+    layer.weight_scale_2.data.fill_(1.0 / w_gs)
+    # Calibrated activation amax stand-in (inputs are randn/10).
+    layer.input_scale.data.fill_(1.0 / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX))
+    return method, layer, w_dequant
+
+
+@unittest.skipIf(get_device_sm() < 100, "NVFP4 dense GEMM backends require SM100+")
+class TestNvFp4LinearBackends(CustomTestCase):
+    def _run_backend(self, backend: str):
+        torch.manual_seed(7)
+        for m, n, k in SHAPES:
+            with self.subTest(backend=backend, shape=(m, n, k)):
+                with mock.patch.object(
+                    fp4_utils,
+                    "FP4_GEMM_RUNNER_BACKEND",
+                    Fp4GemmRunnerBackend(backend),
+                ):
+                    method, layer, w_dequant = _make_quantized_layer(n, k)
+                    method.process_weights_after_loading(layer)
+
+                    x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
+                    out = method.apply(layer, x)
+
+                    x_gs = layer.input_scale_inv.data.float()
+                    x_q, x_sf = fp4_quantize(x, x_gs)
+                    x_dequant = dequantize_nvfp4_to_dtype(
+                        x_q, x_sf, x_gs, torch.float32, x.device
+                    )
+                    ref = x_dequant @ w_dequant.T
+
+                    self.assertEqual(out.shape, (m, n))
+                    cos = torch.nn.functional.cosine_similarity(
+                        out.float().flatten(), ref.flatten(), dim=0
+                    ).item()
+                    self.assertGreater(cos, 0.99)
+                    torch.testing.assert_close(out.float(), ref, rtol=5e-2, atol=5e-2)
+
+    def test_flashinfer_cutedsl(self):
+        self._run_backend("flashinfer_cutedsl")
+
+    def test_flashinfer_cutlass(self):
+        self._run_backend("flashinfer_cutlass")
+
+    def test_flashinfer_cudnn(self):
+        self._run_backend("flashinfer_cudnn")
+
+    def test_flashinfer_trtllm(self):
+        self._run_backend("flashinfer_trtllm")
+
+
+if __name__ == "__main__":
+    unittest.main()