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 7, 15:15 UTC
[Bugfix] Add missing activation attr to RMSNormGated (#35423)
Signed-off-by: tibG <naps@qubes.milou> Co-authored-by: tibG <naps@qubes.milou>
Tib committed
Feb 27, 2026 at 13:53 UTC
6467b635b6acfd5b30dd31804c79ef70ad4bf834
1 file changed
+3
vllm/model_executor/layers/layernorm.py
+3
@@ -510,6 +510,7 @@ class RMSNormGated(CustomOp):
510
norm_before_gate: bool = False,
511
device: torch.device | None = None,
512
dtype: torch.dtype | None = None,
513
+ activation: str = "swish",
514
):
515
"""Initialize RMSNormGated.
516
@@ -524,10 +525,12 @@ class RMSNormGated(CustomOp):
525
If False and z is provided: out = norm(x * silu(z))
526
device: Device to create parameters on
527
dtype: Data type for parameters
528
+ activation: Activation function name for gating
529
"""
530
factory_kwargs = {"device": device, "dtype": dtype}
531
super().__init__()
532
self.eps = eps
533
+ self.activation = activation
534
self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
535
self.register_parameter("bias", None)
536
self.group_size = group_size