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

Revert "Add flashinfer rmsnorm + quant fusion support SM90, SM100, SM120" (#33455)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-08-03 19:03:56 -07:00
committed by GitHub
parent cdff33d738
commit eb31a53338
14 changed files with 65 additions and 1032 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/layers/test_layernorm_fusion.py
-142
View File
@@ -1,142 +0,0 @@
import itertools
import unittest
import torch
from sglang.srt.layers.layernorm import RMSNorm
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large")
class TestRMSNormFp8QuantFusion(CustomTestCase):
DTYPES = [torch.bfloat16, torch.half]
NUM_TOKENS = [7, 83, 512]
HIDDEN_SIZES = [512, 4096]
ADD_RESIDUAL = [False, True]
SEED = 0
FP8_DTYPE = torch.float8_e4m3fn
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA is not available")
from sglang.srt.layers.layernorm import _flashinfer_rmsnorm_quant_available
if not _flashinfer_rmsnorm_quant_available:
raise unittest.SkipTest("flashinfer rmsnorm_quant is not available")
torch.set_default_device("cuda")
def _run_fusion_test(self, num_tokens, hidden_size, add_residual, dtype):
torch.manual_seed(self.SEED)
layer = RMSNorm(hidden_size).to(dtype=dtype)
layer.weight.data.normal_(mean=1.0, std=0.1)
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
residual = torch.randn_like(x) if add_residual else None
# Per-tensor reciprocal scale (as carried by a static FP8 linear).
scale = torch.tensor([0.05], dtype=torch.float32)
with torch.inference_mode():
ref = layer.forward_native(
x.clone(), residual.clone() if add_residual else None
)
normed_ref = ref[0] if add_residual else ref
residual_ref = ref[1] if add_residual else None
result = layer.forward_with_per_tensor_quant_fusion(
x.clone(), scale, residual.clone() if add_residual else None
)
if add_residual:
(q, s, out_dtype), r = result
else:
q, s, out_dtype = result
r = None
# Output contract.
self.assertEqual(q.dtype, self.FP8_DTYPE)
self.assertIs(s, scale)
self.assertEqual(out_dtype, dtype)
self.assertEqual(tuple(q.shape), (num_tokens, hidden_size))
if add_residual:
self.assertEqual(r.dtype, dtype)
self.assertTrue(
torch.allclose(r.float(), residual_ref.float(), atol=1e-2, rtol=1e-2)
)
# Numerical: dequantized (q * scale) matches the reference normed output
# within FP8 e4m3 precision.
deq = q.float() * scale
ref_flat = normed_ref.float().flatten()
cos = torch.nn.functional.cosine_similarity(deq.flatten(), ref_flat, dim=0)
self.assertGreater(cos.item(), 0.99)
rel_err = (
deq.flatten() - ref_flat
).abs().mean() / ref_flat.abs().mean().clamp_min(1e-6)
self.assertLess(rel_err.item(), 0.1)
def test_rms_norm_fp8_quant_fusion(self):
for params in itertools.product(
self.NUM_TOKENS,
self.HIDDEN_SIZES,
self.ADD_RESIDUAL,
self.DTYPES,
):
with self.subTest(
num_tokens=params[0],
hidden_size=params[1],
add_residual=params[2],
dtype=params[3],
):
self._run_fusion_test(*params)
def test_forward_cuda_quant_linear_dispatch(self):
"""forward_cuda routes to the fused path only when applicable."""
import sglang.srt.layers.layernorm as ln_mod
torch.manual_seed(self.SEED)
hidden_size, num_tokens = 512, 32
x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16)
residual = torch.randn_like(x)
scale = torch.tensor([0.05], dtype=torch.float32)
orig_static_scale = ln_mod._fp8_static_input_scale
ln_mod._fp8_static_input_scale = lambda linear: scale
try:
plain = RMSNorm(hidden_size).to(dtype=torch.bfloat16)
plain.weight.data.normal_(mean=1.0, std=0.1)
with torch.inference_mode():
# Plain norm -> fused (fp8, scale, dtype) + bf16 residual.
(q, s, out_dtype), r = plain(
x.clone(), residual.clone(), quant_linear=object()
)
# variance_size_override is incompatible -> must not fuse.
var_layer = RMSNorm(hidden_size, var_hidden_size=hidden_size // 2).to(
dtype=torch.bfloat16
)
var_out = var_layer(x.clone(), residual.clone(), quant_linear=object())
# cast_x_before_out_mul (HF semantics) is incompatible -> must not fuse.
cast_layer = RMSNorm(hidden_size, cast_x_before_out_mul=True).to(
dtype=torch.bfloat16
)
cast_out = cast_layer(
x.clone(), residual.clone(), quant_linear=object()
)
finally:
ln_mod._fp8_static_input_scale = orig_static_scale
self.assertEqual(q.dtype, self.FP8_DTYPE)
self.assertIs(s, scale)
self.assertEqual(out_dtype, torch.bfloat16)
self.assertEqual(r.dtype, torch.bfloat16)
self.assertEqual(var_out[0].dtype, torch.bfloat16)
self.assertEqual(cast_out[0].dtype, torch.bfloat16)
if __name__ == "__main__":
unittest.main()
test/registered/quant/test_fp8_utils.py
+1 -267
View File
@@ -1,6 +1,4 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import torch
@@ -12,7 +10,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large")
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large")
class TestInverseTransformScaleUe8m0(CustomTestCase):
@@ -45,269 +43,5 @@ class TestInverseTransformScaleUe8m0(CustomTestCase):
), 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_supported",
"_is_sm100_supported",
"_is_sm120_supported",
):
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_supported": False,
"_is_sm100_supported": False,
"_is_sm120_supported": False,
}
capabilities[capability] = True
with patch.multiple(fp8_utils, **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.multiple(
fp8_utils,
_is_sm90_supported=False,
_is_sm100_supported=False,
_is_sm120_supported=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()