[codex] Optimize hidden-size 512 RMSNorm dispatch (#24710)

Co-authored-by: Codex <codex@example.com>

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

@@ -85,17 +85,22 @@ __global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_double(const RMSNormPar
   }
 
   sum_of_squares = warp::reduce_sum(sum_of_squares);
-  const auto warp_id = threadIdx.x / kWarpThreads;
-  smem[warp_id] = sum_of_squares;
-  __syncthreads();
-  if (warp_id == 0) {
-    const auto tx = threadIdx.x;
-    const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
-    sum_of_squares = warp::reduce_sum(local_sum);
-    smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+  float norm_factor;
+  if constexpr (kNumWarps == 1) {
+    norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
+  } else {
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    smem[warp_id] = sum_of_squares;
+    __syncthreads();
+    if (warp_id == 0) {
+      const auto tx = threadIdx.x;
+      const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
+      sum_of_squares = warp::reduce_sum(local_sum);
+      smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+    }
+    __syncthreads();
+    norm_factor = smem[warp_id];
   }
-  __syncthreads();
-  const float norm_factor = smem[warp_id];
 
   Storage output_first, output_second;
 #pragma unroll
@@ -147,17 +152,22 @@ __global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_wide(const RMSNormParam
   }
 
   sum_of_squares = warp::reduce_sum(sum_of_squares);
-  const auto warp_id = threadIdx.x / kWarpThreads;
-  smem[warp_id] = sum_of_squares;
-  __syncthreads();
-  if (warp_id == 0) {
-    const auto tx = threadIdx.x;
-    const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
-    sum_of_squares = warp::reduce_sum(local_sum);
-    smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+  float norm_factor;
+  if constexpr (kNumWarps == 1) {
+    norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
+  } else {
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    smem[warp_id] = sum_of_squares;
+    __syncthreads();
+    if (warp_id == 0) {
+      const auto tx = threadIdx.x;
+      const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
+      sum_of_squares = warp::reduce_sum(local_sum);
+      smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+    }
+    __syncthreads();
+    norm_factor = smem[warp_id];
   }
-  __syncthreads();
-  const float norm_factor = smem[warp_id];
 
   Storage output_vec;
 #pragma unroll

python/sglang/jit_kernel/norm.py

@@ -46,6 +46,8 @@ def _is_supported_rmsnorm_hidden_size(d: int) -> bool:
 def _rmsnorm_kernel_class(hidden_size: int) -> str:
     if hidden_size in _RMSNORM_WARP_SIZES:
         return "RMSNormWarpKernel"
+    if hidden_size == 512:
+        return "RMSNormHalfKernel"
     if hidden_size >= _RMSNORM_HALF_BLOCK_MIN_SIZE:
         if hidden_size % 512 == 0:
             return "RMSNormHalfKernel"

python/sglang/jit_kernel/tests/test_rmsnorm.py

@@ -88,7 +88,7 @@ def test_rmsnorm_hidden_size_support(hidden_size: int) -> None:
         (64, "RMSNormWarpKernel"),
         (128, "RMSNormWarpKernel"),
         (256, "RMSNormWarpKernel"),
-        (512, "RMSNormKernel"),
+        (512, "RMSNormHalfKernel"),
         (1536, "RMSNormKernel"),
         (2048, "RMSNormHalfKernel"),
         (2304, "RMSNormKernel"),  # NOTE: not 512 aligned