diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh b/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh index db1cef119..c9830da6d 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh +++ b/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh @@ -41,15 +41,19 @@ struct VecTypeTrait { using vec_t = device::AlignedVector; }; -template -SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) { - float2 valf = device::cast(val); +template +SGL_DEVICE packed_t rms(float2 valf, packed_t& weight, float rsqrt_square_sum) { float2 weightf = device::cast(weight); + if constexpr (kCastXBeforeOutMul) { + auto rounded = device::cast(make_float2(valf.x * rsqrt_square_sum, valf.y * rsqrt_square_sum)); + valf = device::cast(rounded); + return device::cast(make_float2(valf.x * weightf.x, valf.y * weightf.y)); + } return device::cast( make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum)); } -template +template __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::vec_t; using packed_t = typename VecTypeTrait::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(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(v[i]); + } + v_out[i] = rms(valf, v_weight[i], rsqrt_square_sum); } vec_t* p_out = reinterpret_cast(input) + token_id * vec_hidden_size; p_out[threadIdx.x] = v_out; } } -template +template 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; + auto kernel = fused_add_rmsnorm_reg_kernel; LaunchKernel(static_cast(N.unwrap()), threads, device.unwrap()) .enable_pdl(false)( kernel, diff --git a/python/sglang/jit_kernel/norm.py b/python/sglang/jit_kernel/norm.py index f1e19fea0..4fb3c9451 100644 --- a/python/sglang/jit_kernel/norm.py +++ b/python/sglang/jit_kernel/norm.py @@ -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) diff --git a/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py b/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py index 2caf7b0dd..2d8ec7a07 100644 --- a/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py +++ b/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py @@ -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"])) diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 84d78c2f3..7db57684a 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -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),