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
Revert "[vLLM IR] gemma_rms_norm" (#38998)
Robert Shaw committed
Apr 4, 2026 at 17:48 UTC
7d266abb2254a5f8afad0901be11284b7c922e84
5 files changed
+67
-25
vllm/ir/ops/layernorm.py
+3
-2
@@ -16,6 +16,7 @@ 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)
20
if weight is not None:
20
- x = x.to(weight.dtype) * weight
21
- return x.to(orig_dtype)
21
+ x = x * weight
22
+ return x
vllm/kernels/aiter_ops.py
+6
-4
@@ -36,11 +36,13 @@ 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 in (torch.float16, torch.bfloat16)
40
- and (weight is None or weight.dtype == x.dtype)
39
+ and x.dtype
40
+ in (
41
+ torch.float16,
42
+ torch.bfloat16,
43
+ )
44
)
42
-"""AITER rms_norm only supports float16 and bfloat16 acts, no var_size override,
43
-and requires weight dtype to match x dtype."""
45
+"""AITER rms_norm only supports float16 and bfloat16 acts and no var_size override."""
46
47
48
@ir.ops.rms_norm.register_impl(
vllm/kernels/vllm_c.py
+2
-5
@@ -11,11 +11,8 @@ 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 = (
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."""
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."""
16
17
18
@ir.ops.rms_norm.register_impl(
vllm/kernels/xpu_ops.py
+1
-3
@@ -18,9 +18,7 @@ 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 and (
22
- weight is None or weight.dtype == x.dtype
23
-)
21
+rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
22
23
24
@ir.ops.rms_norm.register_impl(
vllm/model_executor/layers/layernorm.py
+55
-11
@@ -376,6 +376,46 @@ 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
+
419
def forward_native(
420
self,
421
x: torch.Tensor,
@@ -383,26 +423,30 @@ class GemmaRMSNorm(CustomOp):
423
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
424
"""PyTorch-native implementation equivalent to forward()."""
425
if residual is None:
386
- return ir.ops.rms_norm(
387
- x, self.weight.data.float() + 1.0, self.variance_epsilon
426
+ return self._forward_static_no_residual(
427
+ self.weight.data, self.variance_epsilon, x
428
)
429
else:
390
- orig_dtype = x.dtype
391
- x = (
392
- x.float() + residual.float()
393
- if orig_dtype == torch.float16
394
- else x + residual
430
+ return self._forward_static_with_residual(
431
+ self.weight.data, self.variance_epsilon, x, residual
432
)
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
433
434
def forward_cuda(
435
self,
436
x: torch.Tensor,
437
residual: torch.Tensor | None = None,
438
) -> 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
450
return self.forward_native(x, residual)
451
452