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
benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
-185
View File
@@ -1,185 +0,0 @@
"""Microbenchmark: fused RMSNorm + static per-tensor FP8 quant, comparing the
flashinfer default kernels against the CuTe-DSL kernels and the unfused
baseline (RMSNorm followed by a separate static FP8 quant).
Providers:
unfused RMSNorm.forward_cuda + static_quant_fp8
fused flashinfer rmsnorm_quant / fused_add_rmsnorm_quant (default)
fused_cute flashinfer rmsnorm_quant_cute / fused_add_rmsnorm_quant_cute
All fused providers produce an ``(fp8, scale)`` activation (and updated residual
when a residual is supplied), matching what a downstream FP8 static-per-tensor
linear consumes. Covers the no-residual and residual (fused-add) cases across a
few hidden sizes so you can pick the fastest kernel per shape.
Run:
python benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
"""
import itertools
import numpy as np
import torch
import triton
from flashinfer.norm import fused_add_rmsnorm_quant, rmsnorm_quant
from flashinfer.testing import bench_gpu_time
from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
from sglang.srt.layers.layernorm import RMSNorm, _flashinfer_rmsnorm_quant_available
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for this benchmark")
if not _flashinfer_rmsnorm_quant_available:
raise RuntimeError(
"flashinfer rmsnorm_quant / fused_add_rmsnorm_quant is not available; "
"install flashinfer to benchmark the fused path"
)
try:
from flashinfer.norm import fused_add_rmsnorm_quant_cute, rmsnorm_quant_cute
_CUTE_AVAILABLE = True
except ImportError:
_CUTE_AVAILABLE = False
DEVICE = "cuda"
DTYPE = torch.bfloat16
FP8_DTYPE = torch.float8_e4m3fn
HIDDEN_SIZES = [4096, 8192]
# Per-tensor reciprocal scale (q = normed / scale); 0.05 keeps normed/scale well
# within the e4m3 range for unit-scale activations.
SCALE_VALUE = 0.05
def make_layer(hidden_size):
layer = RMSNorm(hidden_size).to(device=DEVICE, dtype=DTYPE)
layer.weight.data.normal_(mean=1.0, std=0.1)
return layer
def make_inputs(num_tokens, hidden_size, add_residual):
x = torch.randn(num_tokens, hidden_size, device=DEVICE, dtype=DTYPE)
residual = torch.randn_like(x) if add_residual else None
scale = torch.tensor([SCALE_VALUE], device=DEVICE, dtype=torch.float32)
return x, residual, scale
def run_unfused(layer, x, residual, scale):
out = layer(x, residual)
if residual is not None:
normed, residual_out = out
q, q_scale = static_quant_fp8(normed, scale)
return (q, q_scale), residual_out
q, q_scale = static_quant_fp8(out, scale)
return q, q_scale
def _run_fused(kernel, add_kernel, layer, x, residual, scale):
out = torch.empty_like(x, dtype=FP8_DTYPE)
if residual is not None:
# In-place: residual += x, then out = quant(rmsnorm(residual) * w).
add_kernel(out, x, residual, layer.weight.data, scale, layer.variance_epsilon)
return (out, scale), residual
kernel(out, x, layer.weight.data, scale, layer.variance_epsilon)
return out, scale
def run_fused_default(layer, x, residual, scale):
return _run_fused(rmsnorm_quant, fused_add_rmsnorm_quant, layer, x, residual, scale)
def run_fused_cute(layer, x, residual, scale):
return _run_fused(
rmsnorm_quant_cute, fused_add_rmsnorm_quant_cute, layer, x, residual, scale
)
RUNNERS = {
"unfused": run_unfused,
"fused": run_fused_default,
"fused_cute": run_fused_cute,
}
# (provider key, plot label, style)
_PROVIDERS = [
("unfused", "rmsnorm + static_quant_fp8 (unfused)", ("blue", "-")),
("fused", "rmsnorm_quant (fused, default)", ("green", "-")),
]
if _CUTE_AVAILABLE:
_PROVIDERS.append(
("fused_cute", "rmsnorm_quant_cute (fused, cute-dsl)", ("red", "-"))
)
def _bench_ms(fn, args, quantiles=(0.5, 0.2, 0.8)):
# Pass the GPU tensors as input_args so flashinfer's cold_l2_cache flush can
# find them; a zero-arg callable trips its "no GPU tensors found" warning and
# silently disables cold-L2 timing.
times = bench_gpu_time(
fn=fn,
input_args=args,
use_cuda_graph=True,
dry_run_time_ms=25,
repeat_time_ms=100,
)
return tuple(float(np.percentile(times, q * 100)) for q in quantiles)
def _check_correctness():
"""One-shot sanity check that every fused provider agrees with the unfused
baseline within FP8 precision."""
fused_providers = [p for p in RUNNERS if p != "unfused"]
for hidden_size, add_residual in itertools.product(HIDDEN_SIZES, [False, True]):
layer = make_layer(hidden_size)
x, residual, scale = make_inputs(64, hidden_size, add_residual)
with torch.inference_mode():
ref = run_unfused(
layer, x.clone(), residual.clone() if add_residual else None, scale
)
(uq, _), _ = ref if add_residual else (ref, None)
ref_deq = uq.float() * scale
for provider in fused_providers:
if provider == "fused_cute" and not _CUTE_AVAILABLE:
continue
with torch.inference_mode():
out = RUNNERS[provider](
layer, x.clone(), residual.clone() if add_residual else None, scale
)
(q, _), _ = out if add_residual else (out, None)
cos = torch.nn.functional.cosine_similarity(
(q.float() * scale).flatten(), ref_deq.flatten(), dim=0
).item()
assert (
cos > 0.99
), f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
print("correctness check passed (all fused providers vs unfused within FP8)")
configs = [
triton.testing.Benchmark(
x_names=["num_tokens"],
x_vals=[512, 1024, 2048, 4096, 8192, 16384],
x_log=False,
line_arg="provider",
line_vals=[p[0] for p in _PROVIDERS],
line_names=[p[1] for p in _PROVIDERS],
styles=[p[2] for p in _PROVIDERS],
ylabel="latency (ms)",
plot_name=f"rmsnorm_fp8_quant_h{hidden_size}_residual{add_residual}",
args={"hidden_size": hidden_size, "add_residual": add_residual},
)
for hidden_size, add_residual in itertools.product(HIDDEN_SIZES, [False, True])
]
@triton.testing.perf_report(configs)
def benchmark(num_tokens, hidden_size, add_residual, provider):
layer = make_layer(hidden_size)
x, residual, scale = make_inputs(num_tokens, hidden_size, add_residual)
return _bench_ms(RUNNERS[provider], (layer, x, residual, scale))
if __name__ == "__main__":
torch.manual_seed(0)
_check_correctness()
benchmark.run(print_data=True, show_plots=False)
python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py
+8 -22
View File
@@ -8,7 +8,9 @@ import torch
import triton
from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm
from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import per_tensor_quant_fp8
from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import (
per_tensor_quant_fp8,
)
from sglang.utils import is_in_ci
# Optional vLLM import
@@ -104,7 +106,7 @@ def sglang_scaled_fp8_quant(
if IS_CI:
batch_sizes = [1] # Single batch size for CI
else:
batch_sizes = [1, 2, 8, 16, 64, 128, 256, 512, 1024, 2048]
batch_sizes = [1, 16, 64, 128, 256, 512, 1024, 2048]
# Filter line_vals based on vLLM availability
if VLLM_AVAILABLE:
@@ -113,39 +115,24 @@ if VLLM_AVAILABLE:
"vllm-fp8-bf16",
"sglang-fp8-fp16",
"sglang-fp8-bf16",
"sglang-scalar-a-fp8-fp16",
"sglang-scalar-a-fp8-bf16",
]
line_names = [
"vllm-fp8-fp16",
"vllm-fp8-bf16",
"sglang-fp8-fp16",
"sglang-fp8-bf16",
"sglang-scalar-a-fp8-fp16",
"sglang-scalar-a-fp8-bf16",
]
styles = [
("green", "-"),
("green", "--"),
("blue", "-"),
("blue", "--"),
("red", "-"),
("red", "--"),
]
styles = [("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")]
else:
line_vals = [
"sglang-fp8-fp16",
"sglang-fp8-bf16",
"sglang-scalar-a-fp8-fp16",
"sglang-scalar-a-fp8-bf16",
]
line_names = [
"sglang-fp8-fp16",
"sglang-fp8-bf16",
"sglang-scalar-a-fp8-fp16",
"sglang-scalar-a-fp8-bf16",
]
styles = [("blue", "-"), ("blue", "--"), ("red", "-"), ("red", "--")]
styles = [("blue", "-"), ("blue", "--")]
@triton.testing.perf_report(
@@ -187,9 +174,8 @@ def benchmark(batch_size, provider, N, K):
lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype),
quantiles=quantiles,
)
elif "sglang" in provider:
a_scale = scale_a_scalar if "scalar-a" in provider else scale_a
a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, a_scale)
elif "sglang-fp8" in provider:
a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, scale_a)
b_fp8, scale_b_fp8 = sglang_scaled_fp8_quant(b, scale_b)
b_fp8 = b_fp8.t()
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu
+23 -70
View File
@@ -448,22 +448,19 @@ template <
typename MainloopScheduleType,
typename EpilogueScheduleType,
typename TileSchedulerType = void,
bool WithBias = false,
bool ScalarA = false>
bool WithBias = false>
struct DeviceGemmFp8RowwiseSm100 {
static_assert(std::is_same_v<ElementType, cutlass::float_e4m3_t>, "ElementType must be FP8(e4m3)");
using TileShape = CTAShape;
using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
using ElementComputeEpilogue = float;
using VectorScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
0,
TileShape,
ElementComputeEpilogue,
ElementComputeEpilogue,
cute::Stride<cute::Int<1>, cute::Int<0>, cute::Int<0>>>;
using ScalarScaleA = cutlass::epilogue::fusion::Sm90ScalarBroadcast<float>;
using ScaleA = std::conditional_t<ScalarA, ScalarScaleA, VectorScaleA>;
using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast<
0,
@@ -554,11 +551,7 @@ struct DeviceGemmFp8RowwiseSm100 {
auto* data_ptr = static_cast<T*>(tensor.data_ptr());
static_assert(
std::is_same_v<Descriptor, ScaleA> || std::is_same_v<Descriptor, ScaleB> || std::is_same_v<Descriptor, Bias>);
if constexpr (std::is_same_v<Descriptor, ScalarScaleA>) {
return Arguments{{}, {data_ptr}, {}};
} else {
return Arguments{data_ptr};
}
return Arguments{data_ptr};
}
public:
@@ -664,7 +657,7 @@ void launch_sm100_fp8_scaled_mm(
TORCH_CHECK(status == cutlass::Status::kSuccess)
}
template <typename OutType, bool ScalarA>
template <typename OutType>
void sm100_fp8_dispatch_bias(
torch::Tensor& out,
const torch::Tensor& a,
@@ -702,8 +695,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
true,
ScalarA>;
true>;
using BiasGemm256 = DeviceGemmFp8RowwiseSm100<
ElementInput,
ElementOutput,
@@ -713,8 +705,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
true,
ScalarA>;
true>;
using BiasGemm64 = DeviceGemmFp8RowwiseSm100<
ElementInput,
ElementOutput,
@@ -724,8 +715,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
true,
ScalarA>;
true>;
using BiasGemm16 = DeviceGemmFp8RowwiseSm100<
ElementInput,
ElementOutput,
@@ -735,8 +725,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
true,
ScalarA>;
true>;
// Gemm type without bias
using GemmDefault = DeviceGemmFp8RowwiseSm100<
@@ -748,8 +737,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
false,
ScalarA>;
false>;
using Gemm256 = DeviceGemmFp8RowwiseSm100<
ElementInput,
ElementOutput,
@@ -759,8 +747,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
false,
ScalarA>;
false>;
using Gemm64 = DeviceGemmFp8RowwiseSm100<
ElementInput,
ElementOutput,
@@ -770,8 +757,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
false,
ScalarA>;
false>;
using Gemm16 = DeviceGemmFp8RowwiseSm100<
ElementInput,
ElementOutput,
@@ -781,8 +767,7 @@ void sm100_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
false,
ScalarA>;
false>;
// next power of 2 (minimum 16)
uint32_t const m = a.size(0);
@@ -826,10 +811,7 @@ void sm100_fp8_dispatch_shape(
const torch::Tensor& scales_a,
const torch::Tensor& scales_b,
const c10::optional<torch::Tensor>& bias) {
if (scales_a.numel() == 1) {
return sm100_fp8_dispatch_bias<OutType, true>(out, a, b, scales_a, scales_b, bias);
}
return sm100_fp8_dispatch_bias<OutType, false>(out, a, b, scales_a, scales_b, bias);
return sm100_fp8_dispatch_bias<OutType>(out, a, b, scales_a, scales_b, bias);
}
template <
@@ -841,22 +823,19 @@ template <
typename MainloopScheduleType,
typename EpilogueScheduleType,
typename TileSchedulerType = void,
bool WithBias = false,
bool ScalarA = false>
bool WithBias = false>
struct DeviceGemmFp8RowwiseSm120 {
static_assert(std::is_same_v<ElementType, cutlass::float_e4m3_t>, "ElementType must be FP8(e4m3)");
using TileShape = CTAShape;
using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
using ElementComputeEpilogue = float;
using VectorScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
0,
TileShape,
ElementComputeEpilogue,
ElementComputeEpilogue,
cute::Stride<cute::Int<1>, cute::Int<0>, cute::Int<0>>>;
using ScalarScaleA = cutlass::epilogue::fusion::Sm90ScalarBroadcast<float>;
using ScaleA = std::conditional_t<ScalarA, ScalarScaleA, VectorScaleA>;
using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast<
0,
@@ -947,11 +926,7 @@ struct DeviceGemmFp8RowwiseSm120 {
auto* data_ptr = static_cast<T*>(tensor.data_ptr());
static_assert(
std::is_same_v<Descriptor, ScaleA> || std::is_same_v<Descriptor, ScaleB> || std::is_same_v<Descriptor, Bias>);
if constexpr (std::is_same_v<Descriptor, ScalarScaleA>) {
return Arguments{{}, {data_ptr}, {}};
} else {
return Arguments{data_ptr};
}
return Arguments{data_ptr};
}
public:
@@ -1057,7 +1032,7 @@ void launch_sm120_fp8_scaled_mm(
TORCH_CHECK(status == cutlass::Status::kSuccess)
}
template <typename OutType, bool ScalarA>
template <typename OutType>
void sm120_fp8_dispatch_bias(
torch::Tensor& out,
const torch::Tensor& a,
@@ -1085,8 +1060,7 @@ void sm120_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
true,
ScalarA>;
true>;
using GemmDefault = DeviceGemmFp8RowwiseSm120<
ElementInput,
@@ -1097,8 +1071,7 @@ void sm120_fp8_dispatch_bias(
MainloopScheduleType,
EpilogueScheduleType,
TileSchedulerType,
false,
ScalarA>;
false>;
if (bias) {
return launch_sm120_fp8_scaled_mm<BiasGemmDefault, true>(out, a, b, scales_a, scales_b, bias);
@@ -1115,10 +1088,7 @@ void sm120_fp8_dispatch_shape(
const torch::Tensor& scales_a,
const torch::Tensor& scales_b,
const c10::optional<torch::Tensor>& bias) {
if (scales_a.numel() == 1) {
return sm120_fp8_dispatch_bias<OutType, true>(out, a, b, scales_a, scales_b, bias);
}
return sm120_fp8_dispatch_bias<OutType, false>(out, a, b, scales_a, scales_b, bias);
return sm120_fp8_dispatch_bias<OutType>(out, a, b, scales_a, scales_b, bias);
}
#endif
@@ -1145,26 +1115,7 @@ torch::Tensor fp8_scaled_mm(
TORCH_CHECK(mat_b.scalar_type() == torch::kFloat8_e4m3fn, "mat_b must be Float8_e4m3fn");
TORCH_CHECK(out_dtype == torch::kHalf || out_dtype == torch::kBFloat16, "out_dtype must be Half or BFloat16");
auto sm_version = getSMVersion();
TORCH_CHECK(
scales_a.numel() == 1 || scales_a.numel() == mat_a.size(0),
"scales_a must contain either one scalar scale or one scale per row; got ",
scales_a.numel(),
" elements for M=",
mat_a.size(0));
bool scalar_a_scale_supported = false;
#if defined CUDA_VERSION && CUDA_VERSION >= 12000
scalar_a_scale_supported = sm_version == 90;
#endif
#if defined CUDA_VERSION && CUDA_VERSION >= 12080
scalar_a_scale_supported = scalar_a_scale_supported || sm_version >= 100;
#endif
TORCH_CHECK(
scales_a.numel() != 1 || mat_a.size(0) == 1 || scalar_a_scale_supported,
"scalar scales_a with M > 1 is unsupported on SM",
sm_version,
" for this build; got M=",
mat_a.size(0));
TORCH_CHECK(scales_a.numel() == mat_a.size(0), "size of scales_a is not matched");
TORCH_CHECK(scales_b.numel() == mat_b.size(1), "size of scales_b is not matched");
TORCH_CHECK(scales_a.is_contiguous(), "scales_a must be contiguous");
TORCH_CHECK(scales_b.is_contiguous(), "scales_b msut be contiguous");
@@ -1180,6 +1131,8 @@ torch::Tensor fp8_scaled_mm(
torch::Tensor out = torch::empty({mat_a.size(0), mat_b.size(1)}, mat_a.options().dtype(out_dtype));
TORCH_CHECK((out.size(1) * out.element_size()) % 16 == 0, "out must be multiple of 16 bytes for memory alignment");
auto sm_version = getSMVersion();
#if defined CUDA_VERSION && CUDA_VERSION >= 12080
if (sm_version >= 120) {
if (out_dtype == torch::kBFloat16) {
python/sglang/kernels/aot/tests/test_fp8_gemm.py
-87
View File
@@ -5,22 +5,6 @@ import torch
from sgl_kernel import fp8_scaled_mm
def _cuda_version_at_least(major, minor):
if torch.version.cuda is None:
return False
version = tuple(int(component) for component in torch.version.cuda.split(".")[:2])
return version >= (major, minor)
def _native_scalar_a_supported():
if not torch.cuda.is_available():
return False
capability = torch.cuda.get_device_capability()
if capability == (9, 0):
return _cuda_version_at_least(12, 0)
return capability[0] in (10, 12) and _cuda_version_at_least(12, 8)
def torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias):
o = torch.matmul(a.to(torch.float32), b.to(torch.float32))
o = o.to(torch.float32)
@@ -54,40 +38,6 @@ def _test_accuracy_once(M, N, K, with_bias, out_dtype, device):
print(f"M={M}, N={N}, K={K}, with_bias={with_bias}, out_dtype={out_dtype}: OK")
def _test_scalar_a_accuracy_once(M, N, K, with_bias, out_dtype, device):
fp8_info = torch.finfo(torch.float8_e4m3fn)
a_fp8 = (
torch.randn(M, K, dtype=torch.float32, device=device)
.clamp(min=fp8_info.min, max=fp8_info.max)
.to(torch.float8_e4m3fn)
)
b_fp8 = (
torch.randn(N, K, dtype=torch.float32, device=device)
.clamp(min=fp8_info.min, max=fp8_info.max)
.to(torch.float8_e4m3fn)
.t()
)
scale_a = torch.tensor([0.03125], device=device, dtype=torch.float32)
scale_a_repeated = scale_a.repeat(M)
# Resemble merged projections whose component matrices were quantized with
# different tensorwise scales before concatenation.
scale_b = torch.empty(N, device=device, dtype=torch.float32)
first_boundary = N // 3
second_boundary = 2 * N // 3
scale_b[:first_boundary] = 0.015625
scale_b[first_boundary:second_boundary] = 0.03125
scale_b[second_boundary:] = 0.0625
bias = torch.randn(N, device=device, dtype=out_dtype) if with_bias else None
expected = torch_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias)
actual = fp8_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias)
repeated = fp8_scaled_mm(a_fp8, b_fp8, scale_a_repeated, scale_b, out_dtype, bias)
torch.testing.assert_close(expected, actual, rtol=0.02, atol=1)
torch.testing.assert_close(repeated, actual, rtol=0, atol=0)
@pytest.mark.parametrize("M", [1, 128, 512, 1024, 4096])
@pytest.mark.parametrize("N", [16, 128, 512, 1024, 4096])
@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
@@ -142,43 +92,6 @@ def test_accuracy_sm90_swap_ab(shape_mn, K, with_bias, out_dtype):
_test_accuracy_once(M, N, K, with_bias, out_dtype, "cuda")
@pytest.mark.skipif(
not _native_scalar_a_supported(),
reason="native scalar A scales require a compatible SM90, SM100, or SM120 build",
)
@pytest.mark.parametrize("M", [1, 2, 8, 16, 64, 189])
@pytest.mark.parametrize("with_bias", [True, False])
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
def test_scalar_a_channelwise_b(M, with_bias, out_dtype):
_test_scalar_a_accuracy_once(M, 6144, 4096, with_bias, out_dtype, "cuda")
def test_rejects_invalid_a_scale_count():
M, N, K = 8, 128, 512
a = torch.randn(M, K, device="cuda").to(torch.float8_e4m3fn)
b = torch.randn(N, K, device="cuda").to(torch.float8_e4m3fn).t()
scale_a = torch.ones(2, device="cuda", dtype=torch.float32)
scale_b = torch.ones(N, device="cuda", dtype=torch.float32)
with pytest.raises(RuntimeError, match="scales_a must contain either"):
fp8_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16, None)
@pytest.mark.skipif(
not torch.cuda.is_available() or torch.cuda.get_device_capability() != (8, 9),
reason="SM89-specific scalar A validation",
)
def test_rejects_scalar_a_with_multiple_rows_on_sm89():
M, N, K = 8, 128, 512
a = torch.randn(M, K, device="cuda").to(torch.float8_e4m3fn)
b = torch.randn(N, K, device="cuda").to(torch.float8_e4m3fn).t()
scale_a = torch.ones(1, device="cuda", dtype=torch.float32)
scale_b = torch.ones(N, device="cuda", dtype=torch.float32)
with pytest.raises(RuntimeError, match="scalar scales_a with M > 1 is unsupported"):
fp8_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16, None)
PRODUCTION_LIKE_FP8_GEMM_CASES = [
(189, 4608, 8192, False, torch.bfloat16),
(3330, 256, 8192, False, torch.bfloat16),
python/sglang/srt/layers/layernorm.py
-145
View File
@@ -55,7 +55,6 @@ _is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
_is_xpu = is_xpu()
_flashinfer_layernorm_available = False
_flashinfer_rmsnorm_quant_available = False
if _is_cuda or _is_xpu or _is_musa:
if _is_flashinfer_available:
@@ -84,19 +83,8 @@ if _is_cuda or _is_xpu or _is_musa:
_flashinfer_layernorm_available = True
except (ImportError, AttributeError):
_flashinfer_layernorm_available = False
try:
from flashinfer.norm import (
fused_add_rmsnorm_quant as _flashinfer_fused_add_rmsnorm_quant,
)
from flashinfer.norm import rmsnorm_quant as _flashinfer_rmsnorm_quant
_flashinfer_rmsnorm_quant_available = True
except (ImportError, AttributeError):
_flashinfer_rmsnorm_quant_available = False
else:
_flashinfer_layernorm_available = False
_flashinfer_rmsnorm_quant_available = False
from sgl_kernel import (
fused_add_rmsnorm,
@@ -169,7 +157,6 @@ if _is_cuda:
logger = logging.getLogger(__name__)
if _is_npu:
import torch_npu
from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm
@@ -367,57 +354,6 @@ def _forward_with_allreduce_fusion_quant_per_group(
return (bf16_out, fp8_out, scale_out), residual_out
def _fp8_static_input_scale(linear) -> Optional[torch.Tensor]:
"""Return the per-tensor static FP8 activation scale of ``linear`` if it is
an FP8 linear using static per-tensor activation scaling that can consume a
pre-quantized ``(fp8, scale)`` input; otherwise ``None``.
Recognizes both the native ``Fp8LinearMethod`` (non block/mxfp8/marlin) and
the compressed-tensors W8A8-FP8 scheme with a static per-tensor input scale
(e.g. RedHatAI ``*-FP8`` checkpoints). The flashinfer fused kernel only
supports per-tensor quant, hence the ``numel() == 1`` requirement.
"""
if linear is None:
return None
quant_method = getattr(linear, "quant_method", None)
if quant_method is None:
return None
if not _is_static_per_tensor_fp8_linear(quant_method, linear):
return None
input_scale = getattr(linear, "input_scale", None)
if input_scale is None or input_scale.numel() != 1:
return None
return input_scale
def _is_static_per_tensor_fp8_linear(quant_method, linear) -> bool:
try:
from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod
except ImportError:
Fp8LinearMethod = ()
if isinstance(quant_method, Fp8LinearMethod):
return not (
getattr(quant_method, "block_quant", False)
or getattr(quant_method, "use_mxfp8", False)
or getattr(quant_method, "use_marlin", False)
)
try:
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
CompressedTensorsLinearMethod,
)
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
CompressedTensorsW8A8Fp8,
)
except ImportError:
return False
if isinstance(quant_method, CompressedTensorsLinearMethod):
scheme = getattr(linear, "scheme", None)
return isinstance(scheme, CompressedTensorsW8A8Fp8) and getattr(
scheme, "is_static_input_scheme", False
)
return False
class RMSNorm(MultiPlatformOp):
def __init__(
self,
@@ -471,7 +407,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if x.numel() == 0:
if residual is not None:
@@ -501,20 +436,6 @@ class RMSNorm(MultiPlatformOp):
if needs_reshape:
out = out.reshape(original_shape)
return out
# Fuse the downstream FP8 static per-tensor activation quant into the
# norm when supported. Placed after the empty / variance-override /
# batch-invariant guards above (all incompatible with the fused kernel)
# and gated on not-HF-cast, so it only runs on the standard RMSNorm path.
if (
quant_linear is not None
and not self.cast_x_before_out_mul
and _flashinfer_rmsnorm_quant_available
):
scale = _fp8_static_input_scale(quant_linear)
if scale is not None:
return self.forward_with_per_tensor_quant_fusion(
x, scale, residual, post_residual_addition
)
if self.cast_x_before_out_mul and residual is None:
# Use HF-semantics kernel (cast to dtype before weight multiply).
if (
@@ -572,7 +493,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if residual is not None:
if post_residual_addition is not None:
@@ -588,7 +508,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
# Fix dsv4 dp attenton issue
# the symptom is torch.AcceleratorError: HIP error: invalid configuration argument
@@ -665,7 +584,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
# Fallback to native implementation if vllm is not available
if not _has_vllm_rms_norm:
@@ -705,7 +623,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
return self.forward_native(x, residual, post_residual_addition)
@@ -729,7 +646,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if not x.is_contiguous():
x = x.contiguous()
@@ -780,7 +696,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if _is_cpu_amx_available:
if residual is not None:
@@ -801,7 +716,6 @@ class RMSNorm(MultiPlatformOp):
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
quant_linear: Optional[nn.Module] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
if self.variance_size_override is not None:
return self.forward_native(x, residual, post_residual_addition)
@@ -855,65 +769,6 @@ class RMSNorm(MultiPlatformOp):
self, x, residual, self.weight, group_size, use_attn_tp_group, keep_bf16
)
def forward_with_per_tensor_quant_fusion(
self,
x: torch.Tensor,
scale: torch.Tensor,
residual: Optional[torch.Tensor] = None,
post_residual_addition: Optional[torch.Tensor] = None,
fp8_dtype: torch.dtype = torch.float8_e4m3fn,
) -> Union[
Tuple[torch.Tensor, torch.Tensor, torch.dtype],
Tuple[Tuple[torch.Tensor, torch.Tensor, torch.dtype], torch.Tensor],
]:
"""Fused RMSNorm + static per-tensor FP8 quantization.
The normed activation is quantized to ``fp8_dtype`` using the per-tensor
reciprocal ``scale`` (same convention as ``static_quant_fp8``:
``q = normed / scale``), so a downstream FP8 linear carrying a matching
static ``input_scale`` can skip its own activation quant.
The quantized activation is emitted as a ``(fp8_out, scale, orig_dtype)``
tuple; ``orig_dtype`` (the un-quantized activation dtype) is carried so
the downstream FP8 GEMM produces its output in the model's dtype rather
than defaulting to bf16.
Return contract mirrors ``forward``:
* no residual -> ``(fp8_out, scale, orig_dtype)``
* w/ residual -> ``((fp8_out, scale, orig_dtype), residual_out)``
"""
orig_dtype = x.dtype
needs_reshape = x.dim() != 2
if needs_reshape:
original_shape = x.shape
x = x.contiguous().reshape(-1, original_shape[-1])
elif not x.is_contiguous():
x = x.contiguous()
out = torch.empty_like(x, dtype=fp8_dtype)
if residual is not None:
if post_residual_addition is not None:
residual = residual + post_residual_addition
if residual.dim() != 2:
residual = residual.contiguous().reshape(-1, residual.shape[-1])
elif not residual.is_contiguous():
residual = residual.contiguous()
# In-place: residual += x, then out = quant(rmsnorm(residual) * w).
_flashinfer_fused_add_rmsnorm_quant(
out, x, residual, self.weight.data, scale, self.variance_epsilon
)
if needs_reshape:
out = out.reshape(original_shape)
residual = residual.reshape(original_shape)
return (out, scale, orig_dtype), residual
_flashinfer_rmsnorm_quant(
out, x, self.weight.data, scale, self.variance_epsilon
)
if needs_reshape:
out = out.reshape(original_shape)
return out, scale, orig_dtype
class LayerNorm(MultiPlatformOp):
def __init__(
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
-17
View File
@@ -231,23 +231,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsLinearScheme):
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if isinstance(x, tuple):
# Pre-quantized activation from a fused RMSNorm+FP8 quant kernel:
# x = (fp8_input, per_tensor_input_scale[, orig_dtype]).
# apply_fp8_linear detects the fp8 dtype and skips re-quantizing;
# orig_dtype (when present) sets the GEMM output dtype.
qx, x_scale = x[0], x[1]
out_dtype = x[2] if len(x) > 2 else None
return apply_fp8_linear(
input=qx,
weight=layer.weight,
weight_scale=layer.weight_scale,
input_scale=x_scale,
bias=bias,
use_per_token_if_dynamic=True,
compressed_tensor_quant=True,
pre_quant_output_dtype=out_dtype,
)
if self.weight_block_size is not None:
return self.w8a8_block_fp8_linear(
input=x,
python/sglang/srt/layers/quantization/fp8.py
+7 -20
View File
@@ -139,7 +139,10 @@ def _require_fp4_dtype():
if _use_aiter or _use_hip_int4:
from aiter.ops.shuffle import shuffle_scale, shuffle_weight
from aiter.ops.shuffle import (
shuffle_scale,
shuffle_weight,
)
if _use_aiter:
from sglang.srt.layers.quantization.fp8_utils import (
@@ -1025,24 +1028,6 @@ class Fp8LinearMethod(LinearMethodBase):
bias=bias,
)
if isinstance(x, tuple):
# Pre-quantized activation from a fused RMSNorm+FP8 quant kernel:
# x = (fp8_input, per_tensor_input_scale[, orig_dtype]).
# apply_fp8_linear detects the fp8 dtype and skips re-quantizing;
# orig_dtype (when present) sets the GEMM output dtype.
qx, x_scale = x[0], x[1]
out_dtype = x[2] if len(x) > 2 else None
return apply_fp8_linear(
input=qx,
weight=layer.weight,
weight_scale=layer.weight_scale,
input_scale=x_scale,
bias=bias,
cutlass_fp8_supported=self.cutlass_fp8_supported,
use_per_token_if_dynamic=self.use_per_token_if_dynamic,
pre_quant_output_dtype=out_dtype,
)
return apply_fp8_linear(
input=x,
weight=layer.weight,
@@ -1839,7 +1824,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
)
return qweight.view_as(weight), scale_u8
from sglang.srt.layers.quantization.mxfp8_block_convert import _ue8m0_to_fp32
from sglang.srt.layers.quantization.mxfp8_block_convert import (
_ue8m0_to_fp32,
)
def _quantize_for_deepgemm(weight: torch.Tensor):
weight = weight.contiguous()
python/sglang/srt/layers/quantization/fp8_utils.py
+18 -57
View File
@@ -59,7 +59,6 @@ logger = logging.getLogger(__name__)
_is_hip = is_hip()
_is_cuda = is_cuda()
_is_fp8_fnuz = is_fp8_fnuz()
_is_sm90_supported = is_sm90_supported()
_is_sm100_supported = is_sm100_supported()
_is_sm120_supported = is_sm120_supported()
_is_gfx95_supported = is_gfx95_supported()
@@ -313,9 +312,7 @@ FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None
if is_blackwell_supported() and is_flashinfer_available():
from flashinfer import (
SfLayout,
)
from flashinfer import SfLayout
from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8
from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8
from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize
@@ -1553,7 +1550,9 @@ def requant_block_scale_ue8m0_for_deepgemm(
scales are not already UE8M0, and DeepGEMM can run the layer (bf16 output,
aligned shape). Returns True when it requantizes.
"""
from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0
from sglang.srt.model_loader.utils import (
should_deepgemm_weight_requant_ue8m0,
)
if (
not use_deepgemm_runner
@@ -1842,7 +1841,6 @@ def apply_fp8_linear(
use_per_token_if_dynamic: bool = False,
pad_output: Optional[bool] = None,
compressed_tensor_quant: bool = False,
pre_quant_output_dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
# Note: we pad the input because torch._scaled_mm is more performant
# for matrices with batch dimension > 16.
@@ -1859,42 +1857,10 @@ def apply_fp8_linear(
input_2d = input.view(-1, input.shape[-1])
output_shape = [*input.shape[:-1], weight.shape[1]]
# A pre-quantized fp8 activation (e.g. from a fused RMSNorm+quant kernel)
# carries no original dtype: skip re-quant, reuse the supplied per-tensor
# input_scale, and emit ``pre_quant_output_dtype`` (the model's activation
# dtype, propagated by the producer) or bf16 if it was not provided.
input_prequantized = input_2d.dtype in (
torch.float8_e4m3fn,
torch.float8_e4m3fnuz,
)
if input_prequantized:
output_dtype = pre_quant_output_dtype or torch.bfloat16
else:
output_dtype = input.dtype
channelwise_cutlass = (
cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]
)
cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
use_cutlass_channelwise_gemm = (
channelwise_cutlass and cutlass_compatible_b and not use_triton_w8a8_fp8_kernel
)
native_scalar_a_scale = use_cutlass_channelwise_gemm and (
_is_sm90_supported or _is_sm100_supported or _is_sm120_supported
)
if input_prequantized:
assert input_scale is not None and input_scale.numel() == 1
qinput = input_2d
if channelwise_cutlass and not native_scalar_a_scale:
# Unsupported CUTLASS epilogues require one A scale per row.
x_scale = input_scale.repeat(input_2d.shape[0]).view(-1, 1)
else:
x_scale = input_scale
elif compressed_tensor_quant:
if compressed_tensor_quant:
# Maybe apply padding to output, see comment in __init__
num_token_padding = output_padding
if channelwise_cutlass:
if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
num_token_padding = None
# For static per-tensor activation scales when using inductor compiler,
# use pure PyTorch ops instead of the opaque sgl_kernel quant kernel.
@@ -1923,19 +1889,13 @@ def apply_fp8_linear(
num_token_padding=num_token_padding,
use_per_token_if_dynamic=use_per_token_if_dynamic,
)
if (
input_scale is not None
and channelwise_cutlass
and not native_scalar_a_scale
):
x_scale = input_scale.repeat(input_2d.shape[0]).view(-1, 1)
else:
# cutlass w8a8 fp8 sgl-kernel only supports per-token scale
if input_scale is not None:
assert input_scale.numel() == 1
# broadcast per-tensor scale to per-token scale when supporting cutlass
qinput, x_scale = static_quant_fp8(
input_2d,
input_scale,
repeat_scale=channelwise_cutlass and not native_scalar_a_scale,
input_2d, input_scale, repeat_scale=cutlass_fp8_supported
)
else:
# default use per-token quantization if dynamic
@@ -1956,12 +1916,13 @@ def apply_fp8_linear(
input_2d, group_size=input_2d.shape[1]
)
if channelwise_cutlass:
if not use_cutlass_channelwise_gemm:
if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
# Massage the input to be 2D
qinput = qinput.view(-1, qinput.shape[-1])
output = triton_scaled_mm(
qinput, weight, x_scale, weight_scale, output_dtype, bias
qinput, weight, x_scale, weight_scale, input.dtype, bias
)
else:
output = fp8_scaled_mm(
@@ -1969,7 +1930,7 @@ def apply_fp8_linear(
weight,
x_scale,
weight_scale,
out_dtype=output_dtype,
out_dtype=input.dtype,
bias=bias,
)
return output.view(*output_shape)
@@ -2002,7 +1963,7 @@ def apply_fp8_linear(
WQ=weight.T,
x_scale=x_scale,
w_scale=weight_scale,
dtype=output_dtype,
dtype=input.dtype,
)
if bias is not None:
output += bias
@@ -2018,7 +1979,7 @@ def apply_fp8_linear(
output = torch._scaled_mm(
qinput,
weight,
out_dtype=output_dtype,
out_dtype=input.dtype,
scale_a=x_scale,
scale_b=weight_scale.t(),
bias=bias,
@@ -2032,7 +1993,7 @@ def apply_fp8_linear(
output = torch._scaled_mm(
qinput,
weight,
out_dtype=output_dtype,
out_dtype=input.dtype,
scale_a=x_scale,
scale_b=weight_scale,
bias=bias,
@@ -2061,7 +2022,7 @@ def apply_fp8_linear(
input_2d.shape,
output_shape,
bias,
output_dtype,
input.dtype,
)
python/sglang/srt/models/llama.py
+3 -9
View File
@@ -348,13 +348,9 @@ class LlamaDecoderLayer(nn.Module):
# Self Attention
if residual is None:
residual = hidden_states
hidden_states = self.input_layernorm(
hidden_states, quant_linear=self.self_attn.qkv_proj
)
hidden_states = self.input_layernorm(hidden_states)
else:
hidden_states, residual = self.input_layernorm(
hidden_states, residual, quant_linear=self.self_attn.qkv_proj
)
hidden_states, residual = self.input_layernorm(hidden_states, residual)
hidden_states = self.self_attn(
positions=positions,
hidden_states=hidden_states,
@@ -362,9 +358,7 @@ class LlamaDecoderLayer(nn.Module):
)
# Fully Connected
hidden_states, residual = self.post_attention_layernorm(
hidden_states, residual, quant_linear=self.mlp.gate_up_proj
)
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
hidden_states = self.mlp(hidden_states)
return hidden_states, residual
python/sglang/srt/models/llama_eagle.py
+1 -1
View File
@@ -50,7 +50,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
# https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
if layer_id == 0:
del self.input_layernorm
setattr(self, "input_layernorm", lambda x, quant_linear=None: x)
setattr(self, "input_layernorm", lambda x: x)
class LlamaModel(nn.Module):
python/sglang/srt/models/qwen2.py
+3 -9
View File
@@ -291,13 +291,9 @@ class Qwen2DecoderLayer(nn.Module):
# Self Attention
if residual is None:
residual = hidden_states
hidden_states = self.input_layernorm(
hidden_states, quant_linear=self.self_attn.qkv_proj
)
hidden_states = self.input_layernorm(hidden_states)
else:
hidden_states, residual = self.input_layernorm(
hidden_states, residual, quant_linear=self.self_attn.qkv_proj
)
hidden_states, residual = self.input_layernorm(hidden_states, residual)
hidden_states = self.self_attn(
positions=positions,
hidden_states=hidden_states,
@@ -305,9 +301,7 @@ class Qwen2DecoderLayer(nn.Module):
)
# Fully Connected
hidden_states, residual = self.post_attention_layernorm(
hidden_states, residual, quant_linear=self.mlp.gate_up_proj
)
hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
hidden_states = self.mlp(hidden_states)
return hidden_states, residual
python/sglang/srt/models/qwen2_eagle.py
+1 -1
View File
@@ -51,7 +51,7 @@ class Qwen2DecoderLayer(Qwen2DecoderLayer):
# https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
if layer_id == 0:
del self.input_layernorm
setattr(self, "input_layernorm", lambda x, quant_linear=None: x)
setattr(self, "input_layernorm", lambda x: x)
class Qwen2Model(nn.Module):
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()