From 4aa03d91fd86da1dbd0272e7e2cee7f66d093383 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=99=88=E4=B8=80=E6=B6=B5?= <yingluosanqian@gmail.com>
Date: Thu, 5 Feb 2026 20:36:18 +0800
Subject: [PATCH] [diffusion] fix: fix accuracy bug caused by #14717 (#18296)

---
 python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py | 6 +++---
 .../sglang/jit_kernel/tests/test_fused_norm_scale_shift.py  | 4 ++--
 python/sglang/multimodal_gen/runtime/layers/layernorm.py    | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py b/python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py
index 246bc1856..978d9bd6b 100644
--- a/python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py
+++ b/python/sglang/jit_kernel/diffusion/cutedsl/common/reduce.py
@@ -17,7 +17,7 @@ def cta_reduce_sum(
     val: cute.Numeric, num_warps: cutlass.Constexpr, tidx: cutlass.Int32
 ) -> cute.Numeric:
     smem = cutlass.utils.SmemAllocator()
-    acc = smem.allocate_tensor(cutlass.Float32, num_warps)
+    acc = smem.allocate_tensor(cutlass.Float32, num_warps + 1)
     warp_id = tidx >> 5
     lane_id = tidx & 31
     if lane_id == 0:
@@ -27,7 +27,7 @@ def cta_reduce_sum(
         val = acc[lane_id] if lane_id < num_warps else cutlass.Float32(0)
         val = warp_reduce_sum(val)
         if lane_id == 0:
-            acc[0] = val
+            acc[num_warps] = val
     cute.arch.sync_threads()
-    val = acc[0]
+    val = acc[num_warps]
     return val
diff --git a/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py b/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py
index e7c45041e..73a815f5e 100644
--- a/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py
+++ b/python/sglang/jit_kernel/tests/test_fused_norm_scale_shift.py
@@ -24,11 +24,11 @@ SHAPE_MAP = {
 }
 SHAPES = [
     # (B, S, F, D)
-    (1, 1024, 8, 3072),
-    (4, 512, 16, 3072),
     (1, 115200, 1, 3072),  # Hunyuan
     (1, 32760, 1, 1536),  # Wan
     (1, 6, 1, 3072),  # Qwen
+    (1, 1024, 8, 3072),
+    (4, 512, 16, 3072),
 ]
 DTYPES = [torch.float16, torch.bfloat16, torch.float32]
 NORM_TYPES = ["layer", "rms"]
diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
index eaf87151f..6998bffdd 100644
--- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py
+++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py
@@ -51,7 +51,7 @@ class RMSNorm(CustomOp):
         var_hidden_size: Optional[int] = None,
     ) -> None:
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
+        self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
         self.hidden_size = hidden_size
         self.variance_size_override = (