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

[bugfix] Honor cast_x_before_out_mul in RMSNorm.forward_cuda residual path (#25920)

Browse Source 
Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com>
This commit is contained in:
Xingyu Liu
2026-05-28 01:22:46 -07:00
committed by GitHub
parent 12e28bdf0c
commit 578d27e56a
4 changed files with 108 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh
+26 -11
View File
@@ -41,15 +41,19 @@ struct VecTypeTrait<fp16_t, 32> {
using vec_t = device::AlignedVector<packed_t, 8>;
};
template <typename packed_t>
SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) {
float2 valf = device::cast<fp32x2_t, packed_t>(val);
template <bool kCastXBeforeOutMul, typename packed_t>
SGL_DEVICE packed_t rms(float2 valf, packed_t& weight, float rsqrt_square_sum) {
float2 weightf = device::cast<fp32x2_t, packed_t>(weight);
if constexpr (kCastXBeforeOutMul) {
auto rounded = device::cast<packed_t, fp32x2_t>(make_float2(valf.x * rsqrt_square_sum, valf.y * rsqrt_square_sum));
valf = device::cast<fp32x2_t, packed_t>(rounded);
return device::cast<packed_t, fp32x2_t>(make_float2(valf.x * weightf.x, valf.y * weightf.y));
}
return device::cast<packed_t, fp32x2_t>(
make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum));
}
template <typename T, int VEC_SIZE_IN_BYTE>
template <bool kCastXBeforeOutMul, typename T, int VEC_SIZE_IN_BYTE>
__global__ void fused_add_rmsnorm_reg_kernel(
T* __restrict__ input, T* __restrict__ residual, const T* __restrict__ weight, int vec_hidden_size, float eps) {
constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8;
@@ -58,10 +62,11 @@ __global__ void fused_add_rmsnorm_reg_kernel(
using vec_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::vec_t;
using packed_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::packed_t;
vec_t v; // Save input
vec_t v_res; // Save residual
vec_t v_weight; // Save weight
vec_t v_out; // Save output
vec_t v; // Save input
vec_t v_res; // Save residual
vec_t v_weight; // Save weight
vec_t v_out; // Save output
float2 inp_res_cache[inner_loop]; // fp32 sum cache; only read when kCastXBeforeOutMul=true
auto token_id = blockIdx.x;
float2 acc_square = make_float2(0.0f, 0.0f); // Sum of squares for each thread
@@ -84,6 +89,9 @@ __global__ void fused_add_rmsnorm_reg_kernel(
acc_square.x += inp_res.x * inp_res.x;
acc_square.y += inp_res.y * inp_res.y;
v[i] = device::cast<packed_t, fp32x2_t>(inp_res);
if constexpr (kCastXBeforeOutMul) {
inp_res_cache[i] = inp_res;
}
}
// Store inp+res to residual
@@ -114,14 +122,21 @@ __global__ void fused_add_rmsnorm_reg_kernel(
if (threadIdx.x < vec_hidden_size) {
float rsqrt_square_sum = buffer[threadIdx.x / 32]; // Read rsqrt from Shared Memory(Broadcast)
for (int i = 0; i < inner_loop; i++) {
v_out[i] = rms(v[i], v_weight[i], rsqrt_square_sum);
// HF parity needs the full fp32 sum (not the DType-rounded v[i]).
float2 valf;
if constexpr (kCastXBeforeOutMul) {
valf = inp_res_cache[i];
} else {
valf = device::cast<fp32x2_t, packed_t>(v[i]);
}
v_out[i] = rms<kCastXBeforeOutMul>(valf, v_weight[i], rsqrt_square_sum);
}
vec_t* p_out = reinterpret_cast<vec_t*>(input) + token_id * vec_hidden_size;
p_out[threadIdx.x] = v_out;
}
}
template <typename DType>
template <bool kCastXBeforeOutMul, typename DType>
struct FusedAddRMSNormKernel {
static void
run(const tvm::ffi::TensorView input,
@@ -164,7 +179,7 @@ struct FusedAddRMSNormKernel {
elements_in_vec);
// Launch kernel
auto kernel = fused_add_rmsnorm_reg_kernel<DType, device::kMaxVecBytes>;
auto kernel = fused_add_rmsnorm_reg_kernel<kCastXBeforeOutMul, DType, device::kMaxVecBytes>;
LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
.enable_pdl(false)(
kernel,
python/sglang/jit_kernel/norm.py
+11 -3
View File
@@ -66,9 +66,15 @@ def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module:
)
def is_supported_jit_fused_add_rmsnorm_hidden_size(hidden_size: int) -> bool:
return hidden_size > 0 and hidden_size % 16 == 0 and hidden_size <= 8192
@cache_once
def _jit_fused_add_rmsnorm_module(dtype: torch.dtype) -> Module:
args = make_cpp_args(dtype)
def _jit_fused_add_rmsnorm_module(
dtype: torch.dtype, cast_x_before_out_mul: bool
) -> Module:
args = make_cpp_args(cast_x_before_out_mul, dtype)
return load_jit(
"fused_add_rmsnorm",
*args,
@@ -144,8 +150,10 @@ def fused_add_rmsnorm(
residual: torch.Tensor,
weight: torch.Tensor,
eps: float = 1e-6,
*,
cast_x_before_out_mul: bool = False,
) -> None:
module = _jit_fused_add_rmsnorm_module(input.dtype)
module = _jit_fused_add_rmsnorm_module(input.dtype, cast_x_before_out_mul)
module.fused_add_rmsnorm(input, residual, weight, eps)
python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py
+45 -15
View File
@@ -7,16 +7,23 @@ import torch
from sglang.jit_kernel.utils import get_ci_test_range
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
def sglang_jit_fused_add_rmsnorm(
input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
input: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
eps: float,
*,
cast_x_before_out_mul: bool = False,
) -> None:
from sglang.jit_kernel.norm import fused_add_rmsnorm
fused_add_rmsnorm(input, residual, weight, eps)
fused_add_rmsnorm(
input, residual, weight, eps, cast_x_before_out_mul=cast_x_before_out_mul
)
def flashinfer_fused_add_rmsnorm(
@@ -27,6 +34,16 @@ def flashinfer_fused_add_rmsnorm(
fused_add_rmsnorm(input, residual, weight, eps=eps)
def forward_native_hf_reference(
x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor, eps: float
) -> tuple[torch.Tensor, torch.Tensor]:
sum_fp32 = x.to(torch.float32) + residual.to(torch.float32)
residual_out = sum_fp32.to(x.dtype)
variance = sum_fp32.pow(2).mean(-1, keepdim=True)
out = w * (sum_fp32 * torch.rsqrt(variance + eps)).to(x.dtype)
return out, residual_out
BS_LIST = [2**n for n in range(0, 14)]
BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
@@ -36,31 +53,44 @@ HIDDEN_SIZE_LIST = get_ci_test_range(
)
DEVICE = "cuda"
DTYPE = torch.bfloat16
EPS = torch.finfo(torch.bfloat16).eps
@pytest.mark.parametrize(
"batch_size,hidden_size", list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST))
"batch_size,hidden_size,cast_x_before_out_mul",
list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST, [False, True])),
)
def test_fused_add_rmsnorm(batch_size: int, hidden_size: int) -> None:
def test_fused_add_rmsnorm(
batch_size: int, hidden_size: int, cast_x_before_out_mul: bool
) -> None:
torch.manual_seed(0)
input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
residual = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE)
input_sglang = input.clone()
residual_sglang = residual.clone()
input_flashinfer = input.clone()
residual_flashinfer = residual.clone()
sglang_jit_fused_add_rmsnorm(
input_sglang, residual_sglang, weight, torch.finfo(torch.bfloat16).eps
)
flashinfer_fused_add_rmsnorm(
input_flashinfer, residual_flashinfer, weight, torch.finfo(torch.bfloat16).eps
)
torch.testing.assert_close(input_sglang, input_flashinfer, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(
residual_sglang, residual_flashinfer, atol=1e-2, rtol=1e-2
input_sglang,
residual_sglang,
weight,
EPS,
cast_x_before_out_mul=cast_x_before_out_mul,
)
if cast_x_before_out_mul:
out_ref, residual_ref = forward_native_hf_reference(
input, residual, weight, EPS
)
else:
input_ref = input.clone()
residual_ref_buf = residual.clone()
flashinfer_fused_add_rmsnorm(input_ref, residual_ref_buf, weight, EPS)
out_ref, residual_ref = input_ref, residual_ref_buf
torch.testing.assert_close(input_sglang, out_ref, atol=1e-2, rtol=1e-2)
torch.testing.assert_close(residual_sglang, residual_ref, atol=1e-2, rtol=1e-2)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
python/sglang/srt/layers/layernorm.py
+26
View File
@@ -124,6 +124,11 @@ if _is_cuda:
_jit_rmsnorm_hf = None
from sglang.jit_kernel.norm import fused_add_rmsnorm as _jit_fused_add_rmsnorm
from sglang.jit_kernel.norm import (
is_supported_jit_fused_add_rmsnorm_hidden_size,
)
logger = logging.getLogger(__name__)
@@ -270,6 +275,27 @@ class RMSNorm(MultiPlatformOp):
out = out.reshape(original_shape)
return out
if residual is not None:
if self.cast_x_before_out_mul:
if (
x.dtype in (torch.float16, torch.bfloat16)
and self.weight.data.dtype == x.dtype
and (
post_residual_addition is None
or post_residual_addition.dtype == x.dtype
)
and is_supported_jit_fused_add_rmsnorm_hidden_size(x.shape[-1])
):
if post_residual_addition is not None:
residual = residual + post_residual_addition
_jit_fused_add_rmsnorm(
x,
residual,
self.weight.data,
self.variance_epsilon,
cast_x_before_out_mul=self.cast_x_before_out_mul,
)
return x, residual
return self.forward_native(x, residual, post_residual_addition)
# TODO: Ideally we want to have (hidden_states+residual)+post_residual_addition.
# but right now we can only have hidden_states+(residual+post_residual_addition).
# (hidden_states+residual)+post_residual_addition != hidden_states+(residual+post_residual_addition),