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, 12:15 UTC
[vLLM IR] rework gemma_rms_norm (#39014)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Jiangyun Zhu committed
Apr 7, 2026 at 16:37 UTC
8060bb0333855739d5bdee0e67f7eb227fa04ebb
8 files changed
+106
-75
tests/kernels/core/test_layernorm.py
+29
-1
@@ -6,7 +6,7 @@ import torch
6
7
from tests.kernels.quant_utils import FP8_DTYPE
8
from tests.kernels.utils import opcheck
9
-from vllm.model_executor.layers.layernorm import RMSNorm
9
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
10
from vllm.platforms import current_platform
11
from vllm.utils.torch_utils import set_random_seed
12
@@ -162,3 +162,31 @@ def test_fused_rms_norm_quant(
162
atol=1e-3,
163
rtol=1e-3,
164
)
165
+
166
+
167
+@torch.inference_mode()
168
+def test_gemma_rms_norm_mixed_input_weight_dtype(default_vllm_config) -> None:
169
+ if not torch.cuda.is_available():
170
+ pytest.skip("CUDA required")
171
+
172
+ device = CUDA_DEVICES[0]
173
+ torch.set_default_device(device)
174
+
175
+ num_tokens, hidden_size = 32, 1024
176
+ x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
177
+ layer = GemmaRMSNorm(hidden_size, eps=1e-6).to(device=device)
178
+ layer.weight.data.normal_(mean=0.0, std=0.1)
179
+
180
+ # Gemma uses fp32 weight parameter while activations can be bf16.
181
+ assert layer.weight.dtype == torch.float32
182
+ out = layer(x)
183
+
184
+ x_fp32 = x.float()
185
+ weight_fp32 = layer.weight.data.float() + 1.0
186
+ variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
187
+ ref = (x_fp32 * torch.rsqrt(variance + layer.variance_epsilon) * weight_fp32).to(
188
+ x.dtype
189
+ )
190
+
191
+ assert out.dtype == x.dtype
192
+ torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+21
-3
@@ -12,6 +12,9 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
12
from torch._inductor.pattern_matcher import PatternMatcherPass
13
14
import vllm.ir.ops
15
+from vllm.compilation.passes.fusion.rms_quant_fusion import (
16
+ _rms_input_weight_dtype_match,
17
+)
18
from vllm.config import VllmConfig
19
from vllm.config.utils import Range
20
from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
@@ -320,7 +323,12 @@ class AllReduceRMSNormPattern(BasePattern):
323
return allreduce[3], allreduce[1]
324
325
pm.register_replacement(
323
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
326
+ pattern,
327
+ replacement,
328
+ self.get_inputs(),
329
+ pm.fwd_only,
330
+ pm_pass,
331
+ extra_check=_rms_input_weight_dtype_match,
332
)
333
334
@@ -459,7 +467,12 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
467
return allreduce[4], allreduce[1]
468
469
pm.register_replacement(
462
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
470
+ pattern,
471
+ replacement,
472
+ self.get_inputs(),
473
+ pm.fwd_only,
474
+ pm_pass,
475
+ extra_check=_rms_input_weight_dtype_match,
476
)
477
478
@@ -621,7 +634,12 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
634
return allreduce[4], allreduce[1], allreduce[5]
635
636
pm.register_replacement(
624
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
637
+ pattern,
638
+ replacement,
639
+ self.get_inputs(),
640
+ pm.fwd_only,
641
+ pm_pass,
642
+ extra_check=_rms_input_weight_dtype_match,
643
)
644
645
vllm/compilation/passes/fusion/rms_quant_fusion.py
+29
-1
@@ -38,6 +38,22 @@ FP8_DTYPE = current_platform.fp8_dtype()
38
FP4_DTYPE = torch.uint8
39
40
41
+_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default
42
+
43
+
44
+# TODO: extend rmsnorm quant kernels to support mixed input/weight dtypes,
45
+# and remove this check.
46
+def _rms_input_weight_dtype_match(match: pm.Match) -> bool:
47
+ """Prevent fusion when rms_norm input and weight dtypes differ."""
48
+ for node in match.nodes:
49
+ if node.target == _RMS_NORM_OP:
50
+ # rms_norm(x, weight, epsilon, variance_size)
51
+ x, weight = node.args[0], node.args[1]
52
+ if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
53
+ return x.meta["val"].dtype == weight.meta["val"].dtype
54
+ return True
55
+
56
+
57
def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor:
58
return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda")
59
@@ -186,7 +202,14 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
202
]
203
pattern(*inputs)
204
189
- pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
205
+ pm.register_replacement(
206
+ pattern,
207
+ replacement,
208
+ inputs,
209
+ pm.fwd_only,
210
+ pm_pass,
211
+ extra_check=_rms_input_weight_dtype_match,
212
+ )
213
214
215
class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
@@ -249,6 +272,7 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
272
inputs,
273
pm.fwd_only,
274
pm_pass,
275
+ extra_check=_rms_input_weight_dtype_match,
276
)
277
278
@@ -350,6 +374,7 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
374
self.rmsnorm_matcher.inputs() + [scale],
375
pm.fwd_only,
376
pm_pass,
377
+ extra_check=_rms_input_weight_dtype_match,
378
)
379
380
@@ -445,6 +470,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
470
],
471
pm.fwd_only,
472
pm_pass,
473
+ extra_check=_rms_input_weight_dtype_match,
474
)
475
476
@@ -503,6 +529,7 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
529
],
530
pm.fwd_only,
531
pm_pass,
532
+ extra_check=_rms_input_weight_dtype_match,
533
)
534
535
@@ -559,6 +586,7 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
586
self.rmsnorm_matcher.inputs(),
587
pm.fwd_only,
588
pm_pass,
589
+ extra_check=_rms_input_weight_dtype_match,
590
)
591
592
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 requires no variance_size override and matching input/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
+13
-58
@@ -376,77 +376,32 @@ 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,
382
residual: torch.Tensor | None = None,
383
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
384
"""PyTorch-native implementation equivalent to forward()."""
425
- if residual is None:
426
- return self._forward_static_no_residual(
427
- self.weight.data, self.variance_epsilon, x
428
- )
429
- else:
430
- return self._forward_static_with_residual(
431
- self.weight.data, self.variance_epsilon, x, residual
385
+ orig_dtype = x.dtype
386
+ weight = self.weight.data.float() + 1.0
387
+ if residual is not None:
388
+ x = (
389
+ x.float() + residual.float()
390
+ if orig_dtype == torch.float16
391
+ else x + residual
392
)
393
+ residual = x
394
+ # ir.ops.rms_norm handles fp32 upcast internally
395
+ out = ir.ops.rms_norm(x, weight, self.variance_epsilon)
396
+ return (
397
+ out.to(orig_dtype) if residual is None else (out.to(orig_dtype), residual)
398
+ )
399
400
def forward_cuda(
401
self,
402
x: torch.Tensor,
403
residual: torch.Tensor | None = None,
404
) -> 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
405
return self.forward_native(x, residual)
406
407