minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-04 15:50:41 -07:00
committed by GitHub
parent b0fd31ba07
commit a0b3f1dde6
6 changed files with 421 additions and 318 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py
+245
View File
@@ -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
+175
View File
@@ -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()