minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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);
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
+2
View File
@@ -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
+1 -1
View File
@@ -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