A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 13:30 UTC
[IR][RmsNorm] pass None if not has_weight (#38961)
Signed-off-by: Linkun Chen <github@lkchen.net> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Linkun committed
Apr 4, 2026 at 08:02 UTC
a88ce94bbbe2ccd3fb7d654c7ba551a55da7005b
1 file changed
+5
-1
vllm/model_executor/layers/layernorm.py
+5
-1
@@ -241,8 +241,12 @@ class RMSNorm(CustomOp):
241
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
242
"""PyTorch-native implementation equivalent to forward()."""
243
if residual is None:
244
+ # TODO(luka): address the weight=None passing issue more generally
245
return ir.ops.rms_norm(
245
- x, self.weight.data, self.variance_epsilon, self.variance_size_override
246
+ x,
247
+ self.weight.data if self.has_weight else None,
248
+ self.variance_epsilon,
249
+ self.variance_size_override,
250
)
251
252
return self.forward_static(