[BugFix] Don't apply weight in batch-invariant RMSNorm when has_weight=False (#48741)
Signed-off-by: Josephasafg <ajgard7@gmail.com> Co-authored-by: Michael Gokhman <michael.gokhman@yahoo.com>
Asaf Gardin committed
Jul 16, 2026 at 06:59 UTC
df8a0900df22bd787806f13e449ba751a458d334
2 files changed
+21
-13
vllm/model_executor/layers/batch_invariant.py
+17
-12
@@ -781,6 +781,7 @@ def _rms_norm_kernel(
781
n_cols,
782
eps,
783
BLOCK_SIZE: tl.constexpr,
784
+ HAS_WEIGHT: tl.constexpr,
785
):
786
"""
787
Compute RMS normalization along the last dimension of a 2D tensor.
@@ -813,18 +814,19 @@ def _rms_norm_kernel(
814
col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
815
mask = col_idx < n_cols
816
vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
816
- weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0)
817
# Compute in float32 then convert back to input dtype
818
vals_f32 = vals.to(tl.float32)
819
- weight_f32 = weight.to(tl.float32)
820
- output_f32 = vals_f32 * inv_rms * weight_f32
819
+ output_f32 = vals_f32 * inv_rms
820
+ if HAS_WEIGHT:
821
+ weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0)
822
+ output_f32 = output_f32 * weight.to(tl.float32)
823
output = output_f32.to(vals.dtype)
824
tl.store(output_row_start_ptr + col_idx, output, mask=mask)
825
826
827
def rms_norm_batch_invariant(
828
input: torch.Tensor,
827
- weight: torch.Tensor,
829
+ weight: torch.Tensor | None,
830
eps: float = 1e-6,
831
residual: torch.Tensor | None = None,
832
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
@@ -834,7 +836,8 @@ def rms_norm_batch_invariant(
836
837
Args:
838
input: Input tensor of shape (..., hidden_size)
837
- weight: Weight tensor of shape (hidden_size,)
839
+ weight: Weight tensor of shape (hidden_size,), or None to skip the
840
+ per-channel multiply (``RMSNorm(has_weight=False)``)
841
eps: Small constant for numerical stability
842
residual: Optional residual tensor fused into the normalization path
843
@@ -851,17 +854,18 @@ def rms_norm_batch_invariant(
854
ops.fused_add_rms_norm(input, residual, weight, eps)
855
return input, residual
856
854
- assert weight.dim() == 1, "Weight must be 1-dimensional"
855
- assert input.shape[-1] == weight.shape[0], (
856
- f"Input last dimension ({input.shape[-1]}) must match "
857
- f"weight dimension ({weight.shape[0]})"
858
- )
857
+ if weight is not None:
858
+ assert weight.dim() == 1, "Weight must be 1-dimensional"
859
+ assert input.shape[-1] == weight.shape[0], (
860
+ f"Input last dimension ({input.shape[-1]}) must match "
861
+ f"weight dimension ({weight.shape[0]})"
862
+ )
863
+ weight = weight.contiguous()
864
865
# Flatten all dimensions except the last one
866
original_shape = input.shape
867
input_2d = input.reshape(-1, input.shape[-1])
868
input_2d = input_2d.contiguous()
864
- weight = weight.contiguous()
869
870
n_rows, n_cols = input_2d.shape
871
@@ -870,13 +874,14 @@ def rms_norm_batch_invariant(
874
grid = (n_rows,)
875
_rms_norm_kernel[grid](
876
input_2d,
873
- weight,
877
+ weight if weight is not None else input_2d,
878
output,
879
input_2d.stride(0),
880
output.stride(0),
881
n_cols,
882
eps,
883
BLOCK_SIZE=BLOCK_SIZE,
884
+ HAS_WEIGHT=weight is not None,
885
)
886
return output.reshape(original_shape)
887
vllm/model_executor/layers/layernorm.py
+4
-1
@@ -102,9 +102,12 @@ class RMSNorm(CustomOp):
102
assert self.variance_size_override is None, (
103
"Batch invariance is not supported for variance_size_override"
104
)
105
+ pass_weight = (
106
+ self.pass_weight_add if residual is not None else self.pass_weight
107
+ )
108
return rms_norm_batch_invariant(
109
x,
107
- self.weight.data,
110
+ self.weight.data if pass_weight else None,
111
self.variance_epsilon,
112
residual=residual,
113
)