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

Signed-off-by: Xingyu Liu <charlotteliu12x@gmail.com>

python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh

@@ -41,15 +41,19 @@ struct VecTypeTrait<fp16_t, 32> {
   using vec_t = device::AlignedVector<packed_t, 8>;
 };
 
-template <typename packed_t>
+template <bool kCastXBeforeOutMul, typename packed_t>
-SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) {
+SGL_DEVICE packed_t rms(float2 valf, packed_t& weight, float rsqrt_square_sum) {
-  float2 valf = device::cast<fp32x2_t, packed_t>(val);
   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;                           // Save input
-  vec_t v_res;     // Save residual
+  vec_t v_res;                       // Save residual
-  vec_t v_weight;  // Save weight
+  vec_t v_weight;                    // Save weight
-  vec_t v_out;     // Save output
+  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

@@ -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:
+def _jit_fused_add_rmsnorm_module(
-    args = make_cpp_args(dtype)
+    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

@@ -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
+        input_sglang,
-    )
+        residual_sglang,
-    flashinfer_fused_add_rmsnorm(
+        weight,
-        input_flashinfer, residual_flashinfer, weight, torch.finfo(torch.bfloat16).eps
+        EPS,
-    )
+        cast_x_before_out_mul=cast_x_before_out_mul,
-    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
     )
 
+    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

@@ -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),