[diffusion] fix: fix accuracy bug caused by #14717 (#18296)

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

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"]

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 = (