[Test] Replace GEMM backend e2e matrices with layer-level unit tests (#33596)
This commit is contained in:
Liangsheng Yin
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user