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

Preserve RMSNorm shape in batch-invariant mode (#30956)

Browse Source 
Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
This commit is contained in:
Lianmin Zheng
2026-07-12 17:51:53 -07:00
committed by GitHub
co-authored by Lu Fang
parent 7da30f4e55
commit 22c08a9bee
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/layernorm.py
+4 -1
View File
@@ -287,11 +287,14 @@ class RMSNorm(MultiPlatformOp):
or get_server_args().rl_on_policy_target == "fsdp" or get_server_args().rl_on_policy_target == "fsdp"
): ):
return self.forward_native(x, residual, post_residual_addition) return self.forward_native(x, residual, post_residual_addition)
return rms_norm_batch_invariant( out = rms_norm_batch_invariant(
x, x,
self.weight.data, self.weight.data,
self.variance_epsilon, self.variance_epsilon,
) )
if needs_reshape:
out = out.reshape(original_shape)
return out
if self.cast_x_before_out_mul and residual is None: if self.cast_x_before_out_mul and residual is None:
# Use HF-semantics kernel (cast to dtype before weight multiply). # Use HF-semantics kernel (cast to dtype before weight multiply).
if ( if (