387 lines
14 KiB
Python
387 lines
14 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
from sglang.srt.layers.quantization.fp8 import (
|
|
Fp8MoEMethod,
|
|
_is_cuda,
|
|
_is_gfx95_supported,
|
|
_is_hip,
|
|
)
|
|
from sglang.srt.layers.quantization.fp8_utils import (
|
|
inverse_transform_scale_ue8m0,
|
|
quant_weight_ue8m0,
|
|
transform_scale_ue8m0,
|
|
)
|
|
from sglang.srt.runtime_context import get_platform
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large")
|
|
|
|
|
|
class TestMxfp8MoeScaleLayout(CustomTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not (
|
|
(_is_cuda and get_platform().is_sm100) or (_is_hip and _is_gfx95_supported)
|
|
):
|
|
raise unittest.SkipTest(
|
|
"MXFP8 MoE quantization requires SM100 or ROCm gfx95"
|
|
)
|
|
|
|
def test_cutlass_serialized_scales_remain_expert_first(self):
|
|
class CutlassBackend:
|
|
def is_cutlass(self):
|
|
return True
|
|
|
|
def is_flashinfer_trtllm(self):
|
|
return False
|
|
|
|
def is_flashinfer_trtllm_routed(self):
|
|
return False
|
|
|
|
def is_deep_gemm(self):
|
|
return False
|
|
|
|
layer = SimpleNamespace(
|
|
w13_weight=torch.nn.Parameter(
|
|
torch.zeros((2, 64, 32), dtype=torch.float8_e4m3fn, device="cuda")
|
|
),
|
|
w2_weight=torch.nn.Parameter(
|
|
torch.zeros((2, 32, 32), dtype=torch.float8_e4m3fn, device="cuda")
|
|
),
|
|
w13_weight_scale_inv=torch.nn.Parameter(
|
|
torch.zeros((2, 64, 1), dtype=torch.uint8, device="cuda"),
|
|
requires_grad=False,
|
|
),
|
|
w2_weight_scale_inv=torch.nn.Parameter(
|
|
torch.zeros((2, 32, 1), dtype=torch.uint8, device="cuda"),
|
|
requires_grad=False,
|
|
),
|
|
)
|
|
method = object.__new__(Fp8MoEMethod)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.quantization.fp8.get_moe_runner_backend",
|
|
return_value=CutlassBackend(),
|
|
):
|
|
method._process_mxfp8_moe_weights(layer, quantize=False)
|
|
|
|
self.assertEqual(tuple(layer.w13_weight_scale_inv.shape), (2, 64, 1))
|
|
self.assertEqual(tuple(layer.w2_weight_scale_inv.shape), (2, 32, 1))
|
|
|
|
|
|
class TestInverseTransformScaleUe8m0(CustomTestCase):
|
|
def test_round_trip(self):
|
|
for _ in range(100):
|
|
weight_bf16 = torch.randn(
|
|
# DeepSeek V3 kv_b_proj
|
|
(32768, 512),
|
|
dtype=torch.bfloat16,
|
|
device="cuda",
|
|
)
|
|
|
|
weight_block_size = [128, 128]
|
|
|
|
qweight, sf_fp32_original = quant_weight_ue8m0(
|
|
weight_bf16, weight_block_size=weight_block_size
|
|
)
|
|
mn = qweight.shape[-2]
|
|
|
|
sf_packed_original = transform_scale_ue8m0(sf_fp32_original, mn=mn)
|
|
sf_fp32_recreated = inverse_transform_scale_ue8m0(sf_packed_original, mn=mn)
|
|
|
|
sf_packed_recreated = transform_scale_ue8m0(sf_fp32_recreated, mn=mn)
|
|
|
|
assert torch.all(sf_packed_original == sf_packed_recreated), (
|
|
f"{sf_packed_original=} {sf_packed_recreated}"
|
|
)
|
|
assert torch.all(sf_fp32_original == sf_fp32_recreated), (
|
|
f"{sf_fp32_original=} {sf_fp32_recreated}"
|
|
)
|
|
|
|
|
|
class TestApplyFp8LinearScaleDispatch(CustomTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA is not available")
|
|
torch.set_default_device("cuda")
|
|
|
|
@staticmethod
|
|
def _make_inputs(dtype=torch.bfloat16):
|
|
M, K, N = 8, 16, 32
|
|
input = torch.randn(M, K, dtype=dtype)
|
|
qinput = input.to(torch.float8_e4m3fn)
|
|
weight = torch.randn(N, K).to(torch.float8_e4m3fn).t()
|
|
input_scale = torch.tensor([0.05], dtype=torch.float32)
|
|
weight_scale = torch.linspace(0.01, 0.03, N, dtype=torch.float32)
|
|
return input, qinput, weight, input_scale, weight_scale
|
|
|
|
def test_native_scalar_a_static_prequant_and_dynamic_scale_shapes(self):
|
|
import sglang.srt.layers.quantization.fp8_utils as fp8_utils
|
|
|
|
exec_config = SimpleNamespace(
|
|
graph=SimpleNamespace(
|
|
cuda_graph_config=SimpleNamespace(
|
|
prefill=SimpleNamespace(tc_compiler="none")
|
|
)
|
|
)
|
|
)
|
|
for capability in (
|
|
"is_sm90",
|
|
"is_sm100",
|
|
"is_sm120",
|
|
):
|
|
with self.subTest(capability=capability):
|
|
input, qinput, weight, input_scale, weight_scale = self._make_inputs()
|
|
seen_scales = []
|
|
|
|
def fake_fp8_scaled_mm(
|
|
mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None
|
|
):
|
|
seen_scales.append(scales_a)
|
|
return torch.empty(
|
|
(mat_a.shape[0], mat_b.shape[1]),
|
|
dtype=out_dtype,
|
|
device=mat_a.device,
|
|
)
|
|
|
|
capabilities = {
|
|
"is_sm90": False,
|
|
"is_sm100": False,
|
|
"is_sm120": False,
|
|
}
|
|
capabilities[capability] = True
|
|
with (
|
|
patch.object(
|
|
fp8_utils,
|
|
"get_platform",
|
|
return_value=SimpleNamespace(**capabilities),
|
|
),
|
|
patch.object(
|
|
fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm
|
|
),
|
|
patch.object(fp8_utils, "get_exec", return_value=exec_config),
|
|
):
|
|
fp8_utils.apply_fp8_linear(
|
|
input,
|
|
weight,
|
|
weight_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=True,
|
|
)
|
|
fp8_utils.apply_fp8_linear(
|
|
input,
|
|
weight,
|
|
weight_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=True,
|
|
use_per_token_if_dynamic=True,
|
|
compressed_tensor_quant=True,
|
|
)
|
|
fp8_utils.apply_fp8_linear(
|
|
qinput,
|
|
weight,
|
|
weight_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=True,
|
|
pre_quant_output_dtype=input.dtype,
|
|
)
|
|
fp8_utils.apply_fp8_linear(
|
|
input,
|
|
weight,
|
|
weight_scale,
|
|
input_scale=None,
|
|
cutlass_fp8_supported=True,
|
|
use_per_token_if_dynamic=True,
|
|
compressed_tensor_quant=True,
|
|
)
|
|
|
|
self.assertEqual(seen_scales[0].numel(), 1)
|
|
self.assertEqual(seen_scales[1].numel(), 1)
|
|
self.assertIs(seen_scales[2], input_scale)
|
|
self.assertEqual(tuple(seen_scales[3].shape), (input.shape[0], 1))
|
|
|
|
def test_without_native_scalar_a_static_scale_is_repeated(self):
|
|
import sglang.srt.layers.quantization.fp8_utils as fp8_utils
|
|
|
|
input, qinput, weight, input_scale, weight_scale = self._make_inputs()
|
|
seen_scales = []
|
|
|
|
def fake_fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
|
|
seen_scales.append(scales_a)
|
|
return torch.empty(
|
|
(mat_a.shape[0], mat_b.shape[1]), dtype=out_dtype, device=mat_a.device
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
fp8_utils,
|
|
"get_platform",
|
|
return_value=SimpleNamespace(
|
|
is_sm90=False,
|
|
is_sm100=False,
|
|
is_sm120=False,
|
|
),
|
|
),
|
|
patch.object(fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm),
|
|
):
|
|
fp8_utils.apply_fp8_linear(
|
|
input,
|
|
weight,
|
|
weight_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=True,
|
|
)
|
|
fp8_utils.apply_fp8_linear(
|
|
qinput,
|
|
weight,
|
|
weight_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=True,
|
|
pre_quant_output_dtype=input.dtype,
|
|
)
|
|
|
|
self.assertEqual(tuple(seen_scales[0].shape), (input.shape[0], 1))
|
|
self.assertEqual(tuple(seen_scales[1].shape), (input.shape[0], 1))
|
|
|
|
def test_linear_methods_forward_fused_scalar_tuple(self):
|
|
import sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 as compressed_fp8
|
|
import sglang.srt.layers.quantization.fp8 as native_fp8
|
|
|
|
input, qinput, weight, input_scale, weight_scale = self._make_inputs(
|
|
torch.float16
|
|
)
|
|
|
|
class Layer:
|
|
pass
|
|
|
|
layer = Layer()
|
|
layer.weight = weight
|
|
layer.weight_scale = weight_scale
|
|
layer.input_scale = input_scale
|
|
|
|
native_method = native_fp8.Fp8LinearMethod.__new__(native_fp8.Fp8LinearMethod)
|
|
native_method.use_marlin = False
|
|
native_method.use_mxfp8 = False
|
|
native_method.block_quant = False
|
|
native_method.cutlass_fp8_supported = True
|
|
native_method.use_per_token_if_dynamic = False
|
|
|
|
compressed_method = compressed_fp8.CompressedTensorsW8A8Fp8.__new__(
|
|
compressed_fp8.CompressedTensorsW8A8Fp8
|
|
)
|
|
compressed_method.weight_block_size = None
|
|
|
|
fused_input = (qinput, input_scale, input.dtype)
|
|
with patch.object(native_fp8, "apply_fp8_linear") as native_apply:
|
|
native_apply.return_value = torch.empty(
|
|
(qinput.shape[0], weight.shape[1]), dtype=input.dtype
|
|
)
|
|
native_method.apply(layer, fused_input)
|
|
self.assertIs(native_apply.call_args.kwargs["input_scale"], input_scale)
|
|
self.assertEqual(
|
|
native_apply.call_args.kwargs["pre_quant_output_dtype"], input.dtype
|
|
)
|
|
|
|
with patch.object(compressed_fp8, "apply_fp8_linear") as compressed_apply:
|
|
compressed_apply.return_value = torch.empty(
|
|
(qinput.shape[0], weight.shape[1]), dtype=input.dtype
|
|
)
|
|
compressed_method.apply_weights(layer, fused_input)
|
|
self.assertIs(compressed_apply.call_args.kwargs["input_scale"], input_scale)
|
|
self.assertEqual(
|
|
compressed_apply.call_args.kwargs["pre_quant_output_dtype"],
|
|
input.dtype,
|
|
)
|
|
|
|
|
|
class TestApplyFp8LinearPrequantOutputDtype(CustomTestCase):
|
|
"""apply_fp8_linear with a pre-quantized fp8 activation must emit the
|
|
caller-supplied ``pre_quant_output_dtype`` (the model's activation dtype),
|
|
not the fp8 input dtype. Regression test for FP16 models where hardcoding
|
|
bf16 caused a query/key dtype mismatch in attention."""
|
|
|
|
DTYPES = [torch.float16, torch.bfloat16]
|
|
FP8_DTYPE = torch.float8_e4m3fn
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA is not available")
|
|
torch.set_default_device("cuda")
|
|
|
|
def _run(self, dtype):
|
|
from sglang.srt.layers.quantization.fp8_utils import (
|
|
apply_fp8_linear,
|
|
cutlass_fp8_supported,
|
|
)
|
|
|
|
torch.manual_seed(0)
|
|
M, K, N = 33, 512, 256
|
|
cf = cutlass_fp8_supported()
|
|
fp8_info = torch.finfo(self.FP8_DTYPE)
|
|
|
|
normed = torch.randn(M, K, dtype=dtype)
|
|
input_scale = torch.tensor([0.05], dtype=torch.float32)
|
|
# Per-channel fp8 weight in column-major (K, N) layout.
|
|
w = torch.randn(N, K, dtype=dtype) * 0.05
|
|
w_scale = (w.abs().amax(dim=1) / fp8_info.max).float()
|
|
weight = (
|
|
(w.float() / w_scale[:, None])
|
|
.clamp(fp8_info.min, fp8_info.max)
|
|
.to(self.FP8_DTYPE)
|
|
.t()
|
|
)
|
|
|
|
# Reference: non-pre-quantized input -> output dtype == input dtype.
|
|
ref = apply_fp8_linear(
|
|
input=normed,
|
|
weight=weight,
|
|
weight_scale=w_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=cf,
|
|
)
|
|
self.assertEqual(ref.dtype, dtype)
|
|
|
|
qinput = (
|
|
(normed.float() * input_scale.reciprocal())
|
|
.clamp(fp8_info.min, fp8_info.max)
|
|
.to(self.FP8_DTYPE)
|
|
)
|
|
|
|
# Pre-quantized input with the dtype propagated -> output matches dtype.
|
|
out = apply_fp8_linear(
|
|
input=qinput,
|
|
weight=weight,
|
|
weight_scale=w_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=cf,
|
|
pre_quant_output_dtype=dtype,
|
|
)
|
|
self.assertEqual(out.dtype, dtype)
|
|
self.assertTrue(torch.allclose(out.float(), ref.float(), atol=2e-2, rtol=2e-2))
|
|
|
|
# Without the dtype hint, the pre-quantized path falls back to bf16.
|
|
out_default = apply_fp8_linear(
|
|
input=qinput,
|
|
weight=weight,
|
|
weight_scale=w_scale,
|
|
input_scale=input_scale,
|
|
cutlass_fp8_supported=cf,
|
|
)
|
|
self.assertEqual(out_default.dtype, torch.bfloat16)
|
|
|
|
def test_prequant_output_dtype(self):
|
|
for dtype in self.DTYPES:
|
|
with self.subTest(dtype=dtype):
|
|
self._run(dtype)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|