@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d0975a4b50
Fork 0

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 1, 23:45 UTC

[perf] Add gemma RMS AR fusion (#42646)

Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

Jiahan Chang (Cyrus) committed Jun 4, 2026 at 16:33 UTC d0975a4b50140a9d953f00955a1cbb2a4945edef
3 files changed +225 -16
tests/compile/passes/distributed/test_fusion_all_reduce.py
+59 -1
@@ -14,6 +14,7 @@ from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
14 AllReduceFusionPass,
15 RocmAiterAllReduceFusionPass,
16 )
17 +from vllm.compilation.passes.fx_utils import find_op_nodes
18 from vllm.compilation.passes.utility.fix_functionalization import (
19 FixFunctionalizationPass,
20 )
@@ -33,7 +34,7 @@ from vllm.distributed.parallel_state import (
34 init_distributed_environment,
35 initialize_model_parallel,
36 )
36 -from vllm.model_executor.layers.layernorm import RMSNorm
37 +from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
38 from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 kFp8StaticTensorSym,
40 )
@@ -91,6 +92,49 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
92 return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
93
94
95 +class TestAllReduceGemmaRMSNormModel(torch.nn.Module):
96 + def __init__(
97 + self,
98 + hidden_size=16,
99 + token_num=16,
100 + eps=1e-6,
101 + dtype: torch.dtype = torch.float16,
102 + ):
103 + super().__init__()
104 + self.hidden_size = hidden_size
105 + self.eps = eps
106 + self.norm = [GemmaRMSNorm(hidden_size, eps) for _ in range(4)]
107 + # Non-trivial weight (~Gemma range) so (1 + w) exercises the scale path.
108 + for n in self.norm:
109 + n.weight.data.normal_(mean=0.0, std=0.1)
110 + self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)]
111 +
112 + def forward(self, x):
113 + # avoid having graph input be an arg to a pattern directly
114 + z = torch.relu(x)
115 + x = resid = tensor_model_parallel_all_reduce(z)
116 + y = self.norm[0](x)
117 +
118 + z2 = torch.mm(y, self.w[0])
119 + x2 = tensor_model_parallel_all_reduce(z2)
120 + y2, resid = self.norm[1](x2, resid)
121 +
122 + z3 = torch.mm(y2, self.w[1])
123 + x3 = tensor_model_parallel_all_reduce(z3)
124 + y3, resid = self.norm[2](x3, resid)
125 +
126 + z4 = torch.mm(y3, self.w[2])
127 + x4 = tensor_model_parallel_all_reduce(z4)
128 + y4, resid = self.norm[3](x4, resid)
129 + return y4
130 +
131 + def ops_in_model_before(self):
132 + return [torch.ops.vllm.all_reduce.default]
133 +
134 + def ops_in_model_after(self):
135 + return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]
136 +
137 +
138 class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
139 quant_key = kFp8StaticTensorSym
140
@@ -209,6 +253,15 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
253 "test_model, enable_quant_fp8_custom_op, use_aiter",
254 [
255 (TestAllReduceRMSNormModel, False, IS_AITER_FOUND),
256 + pytest.param(
257 + TestAllReduceGemmaRMSNormModel,
258 + False,
259 + False,
260 + marks=pytest.mark.skipif(
261 + current_platform.is_rocm(),
262 + reason="Not supported on ROCm platform",
263 + ),
264 + ),
265 pytest.param(
266 TestAllReduceRMSNormStaticQuantFP8Model,
267 True,
@@ -404,4 +457,9 @@ def all_reduce_fusion_pass_on_test_model(
457 )
458 backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
459 backend.check_after_ops(model.ops_in_model_after())
460 + if test_model_cls is TestAllReduceGemmaRMSNormModel:
461 + fused_op = torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default
462 + fused_nodes = list(find_op_nodes(fused_op, backend.graph_post_pass))
463 + assert fused_nodes
464 + assert all(n.kwargs.get("weight_bias") == 1.0 for n in fused_nodes)
465 del all_reduce_fusion_pass
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+162 -1
@@ -44,6 +44,24 @@ from .matcher_utils import MatcherQuantFP8
44
45 FP8_DTYPE = current_platform.fp8_dtype()
46
47 +_IR_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default
48 +_IR_FUSED_ADD_RMS_NORM_OP = torch.ops.vllm_ir.fused_add_rms_norm.default
49 +
50 +
51 +def _norm_input_weight_dtype_match(match: pm.Match) -> bool:
52 + """Prevent fusion when the norm input and weight dtypes differ (e.g. a Gemma
53 + fp32 weight.float()+1 gamma), covering rms_norm and fused_add_rms_norm."""
54 + for node in match.nodes:
55 + if node.target == _IR_RMS_NORM_OP:
56 + x, weight = node.args[0], node.args[1]
57 + elif node.target == _IR_FUSED_ADD_RMS_NORM_OP:
58 + x, weight = node.args[0], node.args[2]
59 + else:
60 + continue
61 + if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
62 + return x.meta["val"].dtype == weight.meta["val"].dtype
63 + return True
64 +
65
66 # The empirical value for small batch
67 PDL_ADVANCE_LAUNCH_TOKENS = 16
@@ -132,6 +150,7 @@ if flashinfer_comm is not None:
150 quant_out: torch.Tensor | None = None,
151 scale_out: torch.Tensor | None = None,
152 scale_factor: torch.Tensor | None = None,
153 + weight_bias: float = 0.0,
154 ) -> None:
155 num_tokens, hidden_size = allreduce_in.shape
156 element_size = allreduce_in.element_size()
@@ -208,6 +227,7 @@ if flashinfer_comm is not None:
227 layout_code=layout_code,
228 use_oneshot=use_oneshot,
229 fp32_acc=fp32_acc,
230 + weight_bias=weight_bias,
231 trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
232 )
233
@@ -225,6 +245,7 @@ if flashinfer_comm is not None:
245 quant_out: torch.Tensor | None = None,
246 scale_out: torch.Tensor | None = None,
247 scale_factor: torch.Tensor | None = None,
248 + weight_bias: float = 0.0,
249 ) -> None:
250 pass
251
@@ -399,14 +420,142 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):
420 # allreduce_in, residual
421 return allreduce[1], allreduce[2]
422
423 + # extra_check routes a Gemma fp32 gamma to AllReduceFusedAddGemmaRMSNormPattern.
424 pm.register_replacement(
403 - pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
425 + pattern,
426 + replacement,
427 + self.get_inputs(),
428 + pm.fwd_only,
429 + pm_pass,
430 + extra_check=_norm_input_weight_dtype_match,
431 )
432
433 # Same pattern, but only return the output and not residual
434 # (helpful for end of graph where residual is not used again)
435 first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0]
436
437 + pm.register_replacement(
438 + first_return_only(pattern), # type: ignore[no-untyped-call]
439 + first_return_only(replacement), # type: ignore[no-untyped-call]
440 + self.get_inputs(),
441 + pm.fwd_only,
442 + pm_pass,
443 + extra_check=_norm_input_weight_dtype_match,
444 + )
445 +
446 +
447 +class AllReduceGemmaRMSNormPattern(BasePattern):
448 + """Gemma-style variant of AllReduceRMSNormPattern (no residual)."""
449 +
450 + def __init__(
451 + self,
452 + epsilon: float,
453 + dtype: torch.dtype,
454 + device: str | None,
455 + allreduce_params: FlashInferFusedAllReduceParams,
456 + ) -> None:
457 + super().__init__(dtype, device)
458 + self.epsilon = epsilon
459 + self.allreduce_params = allreduce_params
460 +
461 + def get_inputs(self) -> list[torch.Tensor]:
462 + return [self.empty(5, 16), self.empty(16)]
463 +
464 + def register(self, pm_pass: PatternMatcherPass) -> None:
465 + def pattern(
466 + input: torch.Tensor, weight: torch.Tensor
467 + ) -> tuple[torch.Tensor, torch.Tensor]:
468 + allreduce_output = tensor_model_parallel_all_reduce(input)
469 + rms = vllm.ir.ops.rms_norm(
470 + allreduce_output, weight.float() + 1.0, self.epsilon
471 + )
472 + return rms, allreduce_output
473 +
474 + def replacement(
475 + input: torch.Tensor, weight: torch.Tensor
476 + ) -> tuple[torch.Tensor, torch.Tensor]:
477 + residual = torch.zeros_like(input)
478 + rms_result = torch.empty_like(input)
479 + assert flashinfer_comm is not None, "FlashInfer must be enabled"
480 + allreduce = auto_functionalized(
481 + flashinfer_trtllm_fused_allreduce_norm,
482 + allreduce_in=input,
483 + residual=residual,
484 + norm_out=rms_result,
485 + quant_out=None,
486 + scale_out=None,
487 + rms_gamma=weight,
488 + rms_eps=self.epsilon,
489 + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
490 + weight_bias=1.0,
491 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
492 + )
493 + return allreduce[3], allreduce[1]
494 +
495 + pm.register_replacement(
496 + pattern,
497 + replacement,
498 + self.get_inputs(),
499 + pm.fwd_only,
500 + pm_pass,
501 + )
502 +
503 +
504 +class AllReduceFusedAddGemmaRMSNormPattern(BasePattern):
505 + """Gemma-style variant of AllReduceFusedAddRMSNormPattern (with residual)."""
506 +
507 + def __init__(
508 + self,
509 + epsilon: float,
510 + dtype: torch.dtype,
511 + device: str | None,
512 + allreduce_params: FlashInferFusedAllReduceParams,
513 + ) -> None:
514 + super().__init__(dtype, device)
515 + self.epsilon = epsilon
516 + self.allreduce_params = allreduce_params
517 +
518 + def get_inputs(self) -> list[torch.Tensor]:
519 + input = self.empty(5, 16)
520 + residual = self.empty(5, 16)
521 + weight = self.empty(16)
522 + return [residual, input.to(self.dtype), weight]
523 +
524 + def register(self, pm_pass: PatternMatcherPass) -> None:
525 + def pattern(
526 + residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
527 + ) -> tuple[torch.Tensor, torch.Tensor]:
528 + allreduce_output = tensor_model_parallel_all_reduce(input)
529 + rms, residual = vllm.ir.ops.fused_add_rms_norm(
530 + allreduce_output, residual, weight.float() + 1.0, self.epsilon
531 + )
532 + return rms, residual
533 +
534 + def replacement(
535 + residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor
536 + ) -> tuple[torch.Tensor, torch.Tensor]:
537 + assert flashinfer_comm is not None, "FlashInfer must be enabled"
538 + allreduce = auto_functionalized(
539 + flashinfer_trtllm_fused_allreduce_norm,
540 + allreduce_in=input,
541 + residual=residual,
542 + norm_out=None,
543 + quant_out=None,
544 + scale_out=None,
545 + rms_gamma=weight,
546 + rms_eps=self.epsilon,
547 + pattern_code=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
548 + weight_bias=1.0,
549 + **self.allreduce_params.get_trtllm_fused_allreduce_kwargs(),
550 + )
551 + return allreduce[1], allreduce[2]
552 +
553 + pm.register_replacement(
554 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
555 + )
556 +
557 + first_return_only = lambda fn: lambda a, b, c: fn(a, b, c)[0]
558 +
559 pm.register_replacement(
560 first_return_only(pattern), # type: ignore[no-untyped-call]
561 first_return_only(replacement), # type: ignore[no-untyped-call]
@@ -881,6 +1030,18 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1030 self.device,
1031 self.allreduce_params,
1032 ).register(self.patterns)
1033 + AllReduceGemmaRMSNormPattern(
1034 + epsilon,
1035 + self.model_dtype,
1036 + self.device,
1037 + self.allreduce_params,
1038 + ).register(self.patterns)
1039 + AllReduceFusedAddGemmaRMSNormPattern(
1040 + epsilon,
1041 + self.model_dtype,
1042 + self.device,
1043 + self.allreduce_params,
1044 + ).register(self.patterns)
1045
1046 # WARNING: This is a hack to clear the pattern matcher cache
1047 # and allow multiple values of epsilon.
vllm/model_executor/layers/layernorm.py
+4 -14
@@ -159,20 +159,10 @@ class GemmaRMSNorm(CustomOp):
159 residual: torch.Tensor | None = None,
160 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
161 """PyTorch-native implementation equivalent to forward()."""
162 - orig_dtype = x.dtype
163 - weight = self.weight.data.float() + 1.0
164 - if residual is not None:
165 - x = (
166 - x.float() + residual.float()
167 - if orig_dtype == torch.float16
168 - else x + residual
169 - )
170 - residual = x
171 - # ir.ops.rms_norm handles fp32 upcast internally
172 - out = ir.ops.rms_norm(x, weight, self.variance_epsilon)
173 - return (
174 - out.to(orig_dtype) if residual is None else (out.to(orig_dtype), residual)
175 - )
162 + weight = self.weight.float() + 1.0
163 + if residual is None:
164 + return ir.ops.rms_norm(x, weight, self.variance_epsilon)
165 + return ir.ops.fused_add_rms_norm(x, residual, weight, self.variance_epsilon)
166
167 def forward_cuda(
168 self,