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
[vLLM IR] gemma_rms_norm (#38780)
Signed-off-by: Icey <1790571317@qq.com>
Xiaoshuang Wang committed
Apr 5, 2026 at 01:55 UTC
156405d243924fbede2d4360494a81dea7203334
5 files changed
+25
-67
vllm/ir/ops/layernorm.py
+2
-3
@@ -16,7 +16,6 @@ def rms_norm(
16
x_var = x if variance_size is None else x[..., :variance_size]
17
variance = x_var.pow(2).mean(dim=-1, keepdim=True)
18
x = x * torch.rsqrt(variance + epsilon)
19
- x = x.to(orig_dtype)
19
if weight is not None:
21
- x = x * weight
22
- return x
20
+ x = x.to(weight.dtype) * weight
21
+ return x.to(orig_dtype)
vllm/kernels/aiter_ops.py
+4
-6
@@ -36,13 +36,11 @@ AITER_SUPPORTED = is_aiter_found()
36
37
rms_no_var_16bit_only = (
38
lambda x, weight, epsilon, variance_size=None: variance_size is None
39
- and x.dtype
40
- in (
41
- torch.float16,
42
- torch.bfloat16,
43
- )
39
+ and x.dtype in (torch.float16, torch.bfloat16)
40
+ and (weight is None or weight.dtype == x.dtype)
41
)
45
-"""AITER rms_norm only supports float16 and bfloat16 acts and no var_size override."""
42
+"""AITER rms_norm only supports float16 and bfloat16 acts, no var_size override,
43
+and requires weight dtype to match x dtype."""
44
45
46
@ir.ops.rms_norm.register_impl(
vllm/kernels/vllm_c.py
+5
-2
@@ -11,8 +11,11 @@ current_platform.import_kernels()
11
CUDA_ALIKE = current_platform.is_cuda_alike()
12
"""Most kernels in this file are supported on all CUDA-alike platforms."""
13
14
-rms_no_var_size = lambda x, weight, epsilon, variance_size=None: variance_size is None
15
-"""vLLM kernel does not support variance_size parameter."""
14
+rms_no_var_size = (
15
+ lambda x, weight, epsilon, variance_size=None: variance_size is None
16
+ and (weight is None or weight.dtype == x.dtype)
17
+)
18
+"""vLLM kernel does not support variance_size parameter or mismatched weight dtype."""
19
20
21
@ir.ops.rms_norm.register_impl(
vllm/kernels/xpu_ops.py
+3
-1
@@ -18,7 +18,9 @@ def is_xpu_kernels_found() -> bool:
18
XPU_KERNELS_SUPPORTED = is_xpu_kernels_found()
19
"""Kernels in this file are supported if vLLM XPU kernels are installed."""
20
21
-rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
21
+rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None and (
22
+ weight is None or weight.dtype == x.dtype
23
+)
24
25
26
@ir.ops.rms_norm.register_impl(
vllm/model_executor/layers/layernorm.py
+11
-55
@@ -376,46 +376,6 @@ class GemmaRMSNorm(CustomOp):
376
self.weight = nn.Parameter(torch.zeros(hidden_size))
377
self.variance_epsilon = eps
378
379
- @staticmethod
380
- def _forward_static_no_residual(
381
- weight: torch.Tensor,
382
- variance_epsilon: float,
383
- x: torch.Tensor,
384
- ) -> torch.Tensor:
385
- """PyTorch-native implementation equivalent to forward() without residual."""
386
- orig_dtype = x.dtype
387
- x = x.float()
388
- variance = x.pow(2).mean(dim=-1, keepdim=True)
389
- x = x * torch.rsqrt(variance + variance_epsilon)
390
- x = x * (1.0 + weight.float())
391
- x = x.to(orig_dtype)
392
- return x
393
-
394
- @staticmethod
395
- def _forward_static_with_residual(
396
- weight: torch.Tensor,
397
- variance_epsilon: float,
398
- x: torch.Tensor,
399
- residual: torch.Tensor,
400
- ) -> tuple[torch.Tensor, torch.Tensor]:
401
- """PyTorch-native implementation equivalent to forward() with residual."""
402
- orig_dtype = x.dtype
403
- x = (
404
- x.float() + residual.float()
405
- if orig_dtype == torch.float16
406
- else x + residual
407
- )
408
- residual = x
409
-
410
- x = x.float()
411
- variance = x.pow(2).mean(dim=-1, keepdim=True)
412
- x = x * torch.rsqrt(variance + variance_epsilon)
413
- # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16)
414
- # See https://github.com/huggingface/transformers/pull/29402
415
- x = x * (1.0 + weight.float())
416
- x = x.to(orig_dtype)
417
- return x, residual
418
-
379
def forward_native(
380
self,
381
x: torch.Tensor,
@@ -423,30 +383,26 @@ class GemmaRMSNorm(CustomOp):
383
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
384
"""PyTorch-native implementation equivalent to forward()."""
385
if residual is None:
426
- return self._forward_static_no_residual(
427
- self.weight.data, self.variance_epsilon, x
386
+ return ir.ops.rms_norm(
387
+ x, self.weight.data.float() + 1.0, self.variance_epsilon
388
)
389
else:
430
- return self._forward_static_with_residual(
431
- self.weight.data, self.variance_epsilon, x, residual
390
+ orig_dtype = x.dtype
391
+ x = (
392
+ x.float() + residual.float()
393
+ if orig_dtype == torch.float16
394
+ else x + residual
395
)
396
+ residual = x
397
+ return ir.ops.rms_norm(
398
+ x, self.weight.data.float() + 1.0, self.variance_epsilon
399
+ ).to(orig_dtype), residual
400
401
def forward_cuda(
402
self,
403
x: torch.Tensor,
404
residual: torch.Tensor | None = None,
405
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
439
- if torch.compiler.is_compiling():
440
- return self.forward_native(x, residual)
441
-
442
- if not getattr(self, "_is_compiled", False):
443
- self._forward_static_no_residual = torch.compile( # type: ignore
444
- self._forward_static_no_residual
445
- )
446
- self._forward_static_with_residual = torch.compile( # type: ignore
447
- self._forward_static_with_residual
448
- )
449
- self._is_compiled = True
406
return self.forward_native(x, residual)
407
408