From a0b3f1dde6ae347a2632e4af60171f62fbf66503 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 4 Aug 2026 15:50:41 -0700 Subject: [PATCH] [Test] Replace GEMM backend e2e matrices with layer-level unit tests (#33596) --- python/sglang/srt/layers/quantization/fp8.py | 2 +- .../quant/test_fp8_blockwise_gemm.py | 145 ----------- test/registered/quant/test_fp8_gemm_sm120.py | 86 ------ test/registered/quant/test_nvfp4_gemm.py | 86 ------ .../test_fp8_blockwise_linear_backends.py | 245 ++++++++++++++++++ .../test_nvfp4_linear_backends.py | 175 +++++++++++++ 6 files changed, 421 insertions(+), 318 deletions(-) delete mode 100644 test/registered/quant/test_fp8_blockwise_gemm.py delete mode 100644 test/registered/quant/test_fp8_gemm_sm120.py delete mode 100644 test/registered/quant/test_nvfp4_gemm.py create mode 100644 test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py create mode 100644 test/registered/unit/layers/quantization/test_nvfp4_linear_backends.py diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 3e142d29d..86859e526 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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: diff --git a/test/registered/quant/test_fp8_blockwise_gemm.py b/test/registered/quant/test_fp8_blockwise_gemm.py deleted file mode 100644 index 664cc61dc..000000000 --- a/test/registered/quant/test_fp8_blockwise_gemm.py +++ /dev/null @@ -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() diff --git a/test/registered/quant/test_fp8_gemm_sm120.py b/test/registered/quant/test_fp8_gemm_sm120.py deleted file mode 100644 index e4400955c..000000000 --- a/test/registered/quant/test_fp8_gemm_sm120.py +++ /dev/null @@ -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() diff --git a/test/registered/quant/test_nvfp4_gemm.py b/test/registered/quant/test_nvfp4_gemm.py deleted file mode 100644 index 4e8e4b357..000000000 --- a/test/registered/quant/test_nvfp4_gemm.py +++ /dev/null @@ -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() diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py new file mode 100644 index 000000000..04f675bd1 --- /dev/null +++ b/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() diff --git a/test/registered/unit/layers/quantization/test_nvfp4_linear_backends.py b/test/registered/unit/layers/quantization/test_nvfp4_linear_backends.py new file mode 100644 index 000000000..1ca187407 --- /dev/null +++ b/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()