minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
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):