minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Codex <codex@example.com>
This commit is contained in:
Xiaoyu Zhang
2026-05-19 09:26:10 +08:00
committed by GitHub
co-authored by Codex
parent 170fe57cf0
commit 2424303dfb
3 changed files with 33 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh
+30 -20
View File
@@ -85,17 +85,22 @@ __global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_double(const RMSNormPar
} }
sum_of_squares = warp::reduce_sum(sum_of_squares); sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto warp_id = threadIdx.x / kWarpThreads; float norm_factor;
smem[warp_id] = sum_of_squares; if constexpr (kNumWarps == 1) {
__syncthreads(); norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
if (warp_id == 0) { } else {
const auto tx = threadIdx.x; const auto warp_id = threadIdx.x / kWarpThreads;
const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f; smem[warp_id] = sum_of_squares;
sum_of_squares = warp::reduce_sum(local_sum); __syncthreads();
smem[tx] = math::rsqrt(sum_of_squares / kDim + eps); 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; Storage output_first, output_second;
#pragma unroll #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); sum_of_squares = warp::reduce_sum(sum_of_squares);
const auto warp_id = threadIdx.x / kWarpThreads; float norm_factor;
smem[warp_id] = sum_of_squares; if constexpr (kNumWarps == 1) {
__syncthreads(); norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
if (warp_id == 0) { } else {
const auto tx = threadIdx.x; const auto warp_id = threadIdx.x / kWarpThreads;
const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f; smem[warp_id] = sum_of_squares;
sum_of_squares = warp::reduce_sum(local_sum); __syncthreads();
smem[tx] = math::rsqrt(sum_of_squares / kDim + eps); 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; Storage output_vec;
#pragma unroll #pragma unroll
python/sglang/jit_kernel/norm.py
+2
View File
@@ -46,6 +46,8 @@ def _is_supported_rmsnorm_hidden_size(d: int) -> bool:
def _rmsnorm_kernel_class(hidden_size: int) -> str: def _rmsnorm_kernel_class(hidden_size: int) -> str:
if hidden_size in _RMSNORM_WARP_SIZES: if hidden_size in _RMSNORM_WARP_SIZES:
return "RMSNormWarpKernel" return "RMSNormWarpKernel"
if hidden_size == 512:
return "RMSNormHalfKernel"
if hidden_size >= _RMSNORM_HALF_BLOCK_MIN_SIZE: if hidden_size >= _RMSNORM_HALF_BLOCK_MIN_SIZE:
if hidden_size % 512 == 0: if hidden_size % 512 == 0:
return "RMSNormHalfKernel" return "RMSNormHalfKernel"
python/sglang/jit_kernel/tests/test_rmsnorm.py
+1 -1
View File
@@ -88,7 +88,7 @@ def test_rmsnorm_hidden_size_support(hidden_size: int) -> None:
(64, "RMSNormWarpKernel"), (64, "RMSNormWarpKernel"),
(128, "RMSNormWarpKernel"), (128, "RMSNormWarpKernel"),
(256, "RMSNormWarpKernel"), (256, "RMSNormWarpKernel"),
(512, "RMSNormKernel"), (512, "RMSNormHalfKernel"),
(1536, "RMSNormKernel"), (1536, "RMSNormKernel"),
(2048, "RMSNormHalfKernel"), (2048, "RMSNormHalfKernel"),
(2304, "RMSNormKernel"), # NOTE: not 512 aligned (2304, "RMSNormKernel"), # NOTE: not 512 aligned