@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f887aa1a53
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 3, 15:15 UTC

[Aiter][ROCm] RMSNormGated+GroupedQuantFP8 fusion (#40710)

Signed-off-by: Tres Popp <tres.popp@amd.com> Signed-off-by: Tres Popp <trespopp@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Tres committed May 14, 2026 at 21:37 UTC f887aa1a53e273d90ac537fcd399504f70aff2c7
6 files changed +555 -32
tests/compile/passes/test_fusion.py
+240 -1
@@ -40,7 +40,7 @@ from vllm.model_executor.kernels.linear import (
40 TritonFp8BlockScaledMMKernel,
41 _KernelT,
42 )
43 -from vllm.model_executor.layers.layernorm import RMSNorm
43 +from vllm.model_executor.layers.layernorm import RMSNorm, RMSNormGated
44 from vllm.model_executor.layers.quantization.utils.quant_utils import (
45 GroupShape,
46 create_fp8_quant_key,
@@ -441,3 +441,242 @@ def test_aiter_fusion_rmsnorm_quant(
441 _run_fusion_test(
442 model, fusion_pass, vllm_config, dtype, hidden_size, num_tokens
443 )
444 +
445 +
446 +class TestGatedModel(torch.nn.Module):
447 + """Model that uses RMSNormGated + reshape + group FP8 quant + linear.
448 +
449 + Mimics GatedDeltaNetAttention's output projection path where:
450 + - RMSNormGated operates on per-head tensors (N*H, D)
451 + - Output is reshaped to (N, H*D) before group quantization + linear
452 + """
453 +
454 + def __init__(
455 + self,
456 + num_heads: int,
457 + head_dim: int,
458 + eps: float,
459 + force_kernel: type[_KernelT],
460 + group_shape: GroupShape,
461 + dtype: torch.dtype,
462 + use_aiter_quant: bool = True,
463 + ):
464 + super().__init__()
465 + self.num_heads = num_heads
466 + self.head_dim = head_dim
467 + hidden_dim = num_heads * head_dim
468 +
469 + self.norm = RMSNormGated(
470 + head_dim,
471 + eps=eps,
472 + group_size=None,
473 + norm_before_gate=True,
474 + )
475 +
476 + self.activation_quant_key = create_fp8_quant_key(
477 + static=False, group_shape=group_shape
478 + )
479 + self.weight_quant_key = create_fp8_quant_key(
480 + static=True, group_shape=GroupShape(group_shape.col, group_shape.col)
481 + )
482 +
483 + self.fp8_linear = TestFP8Layer(
484 + weight_shape=(hidden_dim, hidden_dim),
485 + activation_quant_key=self.activation_quant_key,
486 + weight_quant_key=self.weight_quant_key,
487 + force_kernel=force_kernel,
488 + transpose_weights=True,
489 + input_dtype=dtype,
490 + )
491 + self.fp8_linear.kernel.quant_fp8.use_aiter = use_aiter_quant
492 +
493 + def forward(self, x, z):
494 + num_heads = self.num_heads
495 + head_dim = self.head_dim
496 + hidden_dim = num_heads * head_dim
497 + x = torch.relu(x)
498 + z = torch.relu(z)
499 + x_heads = x.reshape(-1, num_heads, head_dim).reshape(-1, head_dim)
500 + z_heads = z.reshape(-1, num_heads, head_dim).reshape(-1, head_dim)
501 + normed = self.norm(x_heads, z_heads)
502 + merged = normed.reshape(-1, hidden_dim)
503 + out = self.fp8_linear(merged)
504 + return out
505 +
506 + def ops_in_model_after(self):
507 + from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
508 + AiterRMSNormGatedFp8GroupQuantPattern,
509 + )
510 +
511 + return [AiterRMSNormGatedFp8GroupQuantPattern.FUSED_OP]
512 +
513 +
514 +class _MockGDNLayer:
515 + """Minimal mock to populate static_forward_context for pass discovery.
516 +
517 + Uses __class__ assignment to pass isinstance checks against
518 + GatedDeltaNetAttention without requiring a full config-based init.
519 + """
520 +
521 + def __init__(self, num_v_heads: int, head_v_dim: int, tp_size: int = 1):
522 + self.num_v_heads = num_v_heads
523 + self.head_v_dim = head_v_dim
524 + self.tp_size = tp_size
525 +
526 + from vllm.model_executor.layers.mamba.gdn_linear_attn import (
527 + GatedDeltaNetAttention,
528 + )
529 +
530 + self.__class__ = GatedDeltaNetAttention
531 +
532 +
533 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
534 +@pytest.mark.parametrize("num_heads", [2])
535 +@pytest.mark.parametrize("head_dim", [128])
536 +@pytest.mark.parametrize("num_tokens", [8])
537 +@pytest.mark.parametrize("eps", [1e-5, 1e-6])
538 +@pytest.mark.skipif(
539 + (not current_platform.is_rocm() or not IS_AITER_FOUND),
540 + reason="Only test on ROCm with aiter package installed",
541 +)
542 +def test_aiter_fusion_rmsnorm_gated_quant(
543 + dtype: torch.dtype,
544 + num_heads: int,
545 + head_dim: int,
546 + num_tokens: int,
547 + eps: float,
548 + monkeypatch: pytest.MonkeyPatch,
549 +):
550 + group_shape = GroupShape(1, 128)
551 + vllm_config = VllmConfig(
552 + model_config=ModelConfig(dtype=dtype),
553 + compilation_config=CompilationConfig(
554 + mode=CompilationMode.VLLM_COMPILE,
555 + custom_ops=["-rms_norm", "-silu_and_mul", "-quant_fp8"],
556 + pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True),
557 + ),
558 + )
559 +
560 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
561 + from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
562 + RocmAiterRMSNormQuantFusionPass,
563 + )
564 +
565 + m.setenv("VLLM_ROCM_USE_AITER", "1")
566 + rocm_aiter_ops.refresh_env_variables()
567 +
568 + # Register a mock GDN layer so the pass discovers num_heads/head_dim
569 + mock_gdn = _MockGDNLayer(num_v_heads=num_heads, head_v_dim=head_dim, tp_size=1)
570 + vllm_config.compilation_config.static_forward_context["mock_gdn_layer"] = (
571 + mock_gdn
572 + )
573 +
574 + torch.set_default_device("cuda")
575 + torch.set_default_dtype(dtype)
576 + torch.manual_seed(1)
577 +
578 + fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config)
579 +
580 + model = TestGatedModel(
581 + num_heads=num_heads,
582 + head_dim=head_dim,
583 + eps=eps,
584 + force_kernel=AiterFp8BlockScaledMMKernel,
585 + group_shape=group_shape,
586 + dtype=dtype,
587 + use_aiter_quant=True,
588 + )
589 +
590 + noop_pass = NoOpEliminationPass(vllm_config)
591 + cleanup_pass = PostCleanupPass(vllm_config)
592 +
593 + backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
594 + backend2 = TestBackend(noop_pass, cleanup_pass)
595 +
596 + hidden_dim = num_heads * head_dim
597 + x = torch.rand(num_tokens, hidden_dim)
598 + z = torch.rand(num_tokens, hidden_dim)
599 + torch._dynamo.mark_dynamic(x, 0)
600 + torch._dynamo.mark_dynamic(z, 0)
601 +
602 + model_fused = torch.compile(model, backend=backend)
603 + result_fused = model_fused(x, z)
604 +
605 + model_unfused = torch.compile(model, backend=backend2)
606 + result_unfused = model_unfused(x, z)
607 +
608 + torch.testing.assert_close(result_fused, result_unfused, atol=1e-2, rtol=1e-2)
609 +
610 + assert fusion_pass.matched_count == 1
611 + backend.check_after_ops(model.ops_in_model_after())
612 +
613 +
614 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
615 +@pytest.mark.parametrize("num_heads", [2])
616 +@pytest.mark.parametrize("head_dim", [128])
617 +@pytest.mark.parametrize("num_tokens", [8])
618 +@pytest.mark.parametrize("eps", [1e-6])
619 +@pytest.mark.skipif(
620 + (not current_platform.is_rocm() or not IS_AITER_FOUND),
621 + reason="Only test on ROCm with aiter package installed",
622 +)
623 +def test_aiter_fusion_rmsnorm_gated_quant_no_gdn_layers(
624 + dtype: torch.dtype,
625 + num_heads: int,
626 + head_dim: int,
627 + num_tokens: int,
628 + eps: float,
629 + monkeypatch: pytest.MonkeyPatch,
630 +):
631 + """Verify that without GDN layers in static_forward_context,
632 + the gated pattern is not registered and no matches occur."""
633 + group_shape = GroupShape(1, 128)
634 + vllm_config = VllmConfig(
635 + model_config=ModelConfig(dtype=dtype),
636 + compilation_config=CompilationConfig(
637 + mode=CompilationMode.VLLM_COMPILE,
638 + custom_ops=["-rms_norm", "-silu_and_mul", "-quant_fp8"],
639 + pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True),
640 + ),
641 + )
642 +
643 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
644 + from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
645 + RocmAiterRMSNormQuantFusionPass,
646 + )
647 +
648 + m.setenv("VLLM_ROCM_USE_AITER", "1")
649 + rocm_aiter_ops.refresh_env_variables()
650 +
651 + torch.set_default_device("cuda")
652 + torch.set_default_dtype(dtype)
653 + torch.manual_seed(1)
654 +
655 + # No mock GDN layer registered -- pass should not register gated pattern
656 + fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config)
657 +
658 + model = TestGatedModel(
659 + num_heads=num_heads,
660 + head_dim=head_dim,
661 + eps=eps,
662 + force_kernel=AiterFp8BlockScaledMMKernel,
663 + group_shape=group_shape,
664 + dtype=dtype,
665 + use_aiter_quant=True,
666 + )
667 +
668 + noop_pass = NoOpEliminationPass(vllm_config)
669 + cleanup_pass = PostCleanupPass(vllm_config)
670 +
671 + backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
672 +
673 + hidden_dim = num_heads * head_dim
674 + x = torch.rand(num_tokens, hidden_dim)
675 + z = torch.rand(num_tokens, hidden_dim)
676 + torch._dynamo.mark_dynamic(x, 0)
677 + torch._dynamo.mark_dynamic(z, 0)
678 +
679 + model_fused = torch.compile(model, backend=backend)
680 + model_fused(x, z)
681 +
682 + assert fusion_pass.matched_count == 0
vllm/_aiter_ops.py
+58
@@ -913,6 +913,50 @@ def _rocm_aiter_rmsnorm_fp8_group_quant_fake(
913 )
914
915
916 +def _rocm_aiter_fused_rms_gated_fp8_group_quant_impl(
917 + x: torch.Tensor,
918 + weight: torch.Tensor,
919 + bias: torch.Tensor | None,
920 + z: torch.Tensor,
921 + eps: float,
922 + norm_before_gate: bool,
923 + activation: str,
924 + group_size: int,
925 +) -> tuple[torch.Tensor, torch.Tensor]:
926 + """Fused gated-RMSNorm + FP8 group quantization via aiter Triton kernel."""
927 + from aiter.ops.triton.quant import fused_rms_gated_fp8_group_quant
928 +
929 + return fused_rms_gated_fp8_group_quant(
930 + x,
931 + weight,
932 + bias,
933 + z,
934 + eps,
935 + norm_before_gate=norm_before_gate,
936 + activation=activation,
937 + out_dtype=FP8_DTYPE,
938 + group_size=group_size,
939 + )
940 +
941 +
942 +def _rocm_aiter_fused_rms_gated_fp8_group_quant_fake(
943 + x: torch.Tensor,
944 + weight: torch.Tensor,
945 + bias: torch.Tensor | None,
946 + z: torch.Tensor,
947 + eps: float,
948 + norm_before_gate: bool,
949 + activation: str,
950 + group_size: int,
951 +) -> tuple[torch.Tensor, torch.Tensor]:
952 + M, N = x.shape
953 + scale_shape = (M, (N + group_size - 1) // group_size)
954 + return (
955 + torch.empty_like(x, dtype=FP8_DTYPE, device=x.device),
956 + torch.empty(scale_shape, dtype=torch.float32, device=x.device),
957 + )
958 +
959 +
960 def _rocm_aiter_group_fp8_quant_impl(
961 x: torch.Tensor,
962 group_size: int,
@@ -1480,6 +1524,9 @@ class rocm_aiter_ops:
1524 try:
1525 import aiter.ops.triton.causal_conv1d_update_single_token # noqa: F401
1526 import aiter.ops.triton.gated_delta_net # noqa: F401
1527 + from aiter.ops.triton.quant import ( # noqa: F401
1528 + fused_rms_gated_fp8_group_quant,
1529 + )
1530
1531 return True
1532 except (ImportError, ModuleNotFoundError):
@@ -1598,6 +1645,12 @@ class rocm_aiter_ops:
1645 fake_impl=_rocm_aiter_rmsnorm_fp8_group_quant_fake,
1646 )
1647
1648 + direct_register_custom_op(
1649 + op_name="rocm_aiter_fused_rms_gated_fp8_group_quant",
1650 + op_func=_rocm_aiter_fused_rms_gated_fp8_group_quant_impl,
1651 + fake_impl=_rocm_aiter_fused_rms_gated_fp8_group_quant_fake,
1652 + )
1653 +
1654 direct_register_custom_op(
1655 op_name="rocm_aiter_rmsnorm_with_add_fp8_group_quant",
1656 op_func=_rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl,
@@ -1689,6 +1742,11 @@ class rocm_aiter_ops:
1742 def get_rmsnorm_group_fused_quant_op() -> OpOverload:
1743 return torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default
1744
1745 + @staticmethod
1746 + def get_fused_rms_gated_fp8_group_quant_op() -> OpOverload:
1747 + """Return the fused gated-RMSNorm + FP8 group quant custom op."""
1748 + return torch.ops.vllm.rocm_aiter_fused_rms_gated_fp8_group_quant.default
1749 +
1750 @staticmethod
1751 def get_rmsnorm_group_add_fused_quant_op() -> OpOverload:
1752 return torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant.default
vllm/compilation/passes/fusion/matcher_utils.py
+62 -1
@@ -10,6 +10,7 @@ from torch._ops import OpOverload
10 from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config import get_current_vllm_config
12 from vllm.model_executor.layers.activation import SiluAndMul
13 +from vllm.model_executor.layers.layernorm import RMSNormGated
14 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
15 from vllm.model_executor.layers.quantization.utils.quant_utils import (
16 GroupShape,
@@ -28,7 +29,6 @@ from vllm.model_executor.layers.rotary_embedding.deepseek_scaling_rope import (
29 )
30 from vllm.platforms import current_platform
31
31 -RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
32 ROTARY_OP = torch.ops._C.rotary_embedding.default
33 FLASHINFER_ROTARY_OP = torch.ops.vllm.flashinfer_rotary_embedding.default
34
@@ -161,6 +161,67 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
161 return result
162
163
164 +class MatcherRMSNormGated(MatcherCustomOp):
165 + """Matches RMSNormGated with norm_before_gate=True and group_size=None."""
166 +
167 + def __init__(
168 + self,
169 + epsilon: float,
170 + enabled: bool | None = None,
171 + norm_before_gate: bool = True,
172 + group_size: int | None = None,
173 + ) -> None:
174 + if enabled is None:
175 + enabled = RMSNormGated.enabled()
176 +
177 + super().__init__(enabled)
178 + self.epsilon = epsilon
179 + self.norm_before_gate = norm_before_gate
180 + self.group_size = group_size
181 +
182 + def inputs(self) -> list[torch.Tensor]:
183 + x = self.empty(5, 16)
184 + z = self.empty(5, 16)
185 + weight = self.empty(16)
186 + return [x, z, weight]
187 +
188 + def forward_custom(
189 + self,
190 + x: torch.Tensor,
191 + z: torch.Tensor,
192 + weight: torch.Tensor,
193 + ) -> torch.Tensor:
194 + from vllm.model_executor.layers.fla.ops.layernorm_guard import (
195 + rmsnorm_fn,
196 + )
197 +
198 + return rmsnorm_fn(
199 + x,
200 + weight,
201 + bias=None,
202 + z=z,
203 + eps=self.epsilon,
204 + group_size=self.group_size,
205 + norm_before_gate=self.norm_before_gate,
206 + )
207 +
208 + def forward_native(
209 + self,
210 + x: torch.Tensor,
211 + z: torch.Tensor,
212 + weight: torch.Tensor,
213 + ) -> torch.Tensor:
214 + return RMSNormGated.forward_static(
215 + x,
216 + z,
217 + weight,
218 + self.epsilon,
219 + self.model_dtype,
220 + group_size=self.group_size,
221 + norm_before_gate=self.norm_before_gate,
222 + )
223 +
224 +
225 class MatcherDeepseekScalingRotaryEmbedding(MatcherCustomOp):
226 def __init__(
227 self,
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+124 -1
@@ -12,7 +12,7 @@ from torch._inductor.pattern_matcher import PatternMatcherPass
12 import vllm.ir.ops
13 import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401
14 from vllm._aiter_ops import rocm_aiter_ops
15 -from vllm.config import VllmConfig
15 +from vllm.config import VllmConfig, get_layers_from_vllm_config
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.quantization.utils.quant_utils import (
18 GroupShape,
@@ -28,9 +28,12 @@ from ..vllm_inductor_pass import (
28 VllmInductorPass,
29 VllmPatternMatcherPass,
30 VllmPatternReplacement,
31 + _fx_view_to_reshape,
32 + fold_consecutive_reshapes,
33 )
34 from .matcher_utils import (
35 MatcherQuantFP8,
36 + MatcherRMSNormGated,
37 MatcherSiluAndMul,
38 )
39 from .rms_quant_fusion import (
@@ -449,6 +452,97 @@ class DoubleAiterRMSFp8GroupQuantViewPattern(AiterRMSNormQuantPattern):
452 )
453
454
455 +class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern):
456 + """
457 + Matches decomposed RMSNormGated + reshape + group FP8 quant and replaces
458 + with rocm_aiter_fused_rms_gated_fp8_group_quant.
459 +
460 + The norm operates per-head on (N*H, D) tensors. The compiler folds the
461 + reshape chain so after norm the result goes through reshape->merge->quant.
462 + The pattern reshapes from (N*H, D) to (N, H*D) before calling
463 + MatcherQuantFP8 so that _quantize_group_native sees the full hidden dim
464 + and computes the correct num_groups.
465 + """
466 +
467 + FUSED_OP = rocm_aiter_ops.get_fused_rms_gated_fp8_group_quant_op()
468 +
469 + def __init__(
470 + self,
471 + epsilon: float,
472 + quant_dtype: torch.dtype,
473 + group_shape: GroupShape,
474 + num_heads: int,
475 + head_dim: int,
476 + match_aiter_quant: bool = True,
477 + symmetric: bool = True,
478 + ) -> None:
479 + scale = ScaleDesc(torch.float32, False, group_shape)
480 + key = FusedRMSQuantKey(
481 + fused_add=False,
482 + quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
483 + )
484 + super().__init__(epsilon, key, match_aiter_quant)
485 + self.rmsnorm_gated_matcher = MatcherRMSNormGated(epsilon)
486 + self.num_heads = num_heads
487 + self.head_dim = head_dim
488 +
489 + def register(self, pm_pass: PatternMatcherPass) -> None:
490 + num_heads = self.num_heads
491 + head_dim = self.head_dim
492 + hidden_dim = num_heads * head_dim
493 + quant_matcher = self.quant_matcher
494 +
495 + def pattern(
496 + x: torch.Tensor,
497 + z: torch.Tensor,
498 + weight: torch.Tensor,
499 + ) -> tuple[torch.Tensor, torch.Tensor]:
500 + normed = self.rmsnorm_gated_matcher(x, z, weight)
501 + merged = normed.reshape(-1, hidden_dim)
502 + quant_out, scales_out = quant_matcher(merged)
503 + return quant_out, scales_out
504 +
505 + def replacement(
506 + x: torch.Tensor,
507 + z: torch.Tensor,
508 + weight: torch.Tensor,
509 + ) -> tuple[torch.Tensor, torch.Tensor]:
510 + fused = self.FUSED_OP(
511 + x=x,
512 + weight=weight,
513 + bias=None,
514 + z=z,
515 + eps=self.epsilon,
516 + norm_before_gate=True,
517 + activation="silu",
518 + group_size=head_dim,
519 + )
520 + fp8_out = fused[0]
521 + scales_out = fused[1]
522 + fp8_reshaped = fp8_out.reshape(-1, hidden_dim)
523 + scales_reshaped = scales_out.reshape(-1, num_heads)
524 + return fp8_reshaped, scales_reshaped
525 +
526 + n_tokens = 2
527 + x = self.empty(n_tokens * num_heads, head_dim)
528 + z = self.empty(n_tokens * num_heads, head_dim)
529 + w = self.empty(head_dim)
530 +
531 + def trace_fn(*args, **kwargs):
532 + gm = pm.fwd_only(*args, **kwargs)
533 + _fx_view_to_reshape(gm)
534 + fold_consecutive_reshapes(gm)
535 + return gm
536 +
537 + pm.register_replacement(
538 + pattern,
539 + replacement,
540 + [x, z, w],
541 + trace_fn,
542 + pm_pass,
543 + )
544 +
545 +
546 class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
547 """
548 This pass fuses aiter rms_norm & vllm/aiter quant custom ops
@@ -464,6 +558,19 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
558 pass_name="rocm_aiter_rms_norm_quant_fusion_pass"
559 )
560
561 + # Discover (num_heads, head_dim) pairs for gated RMSNorm patterns
562 + # from GatedDeltaNetAttention layers in static_forward_context.
563 + from vllm.model_executor.layers.mamba.gdn_linear_attn import (
564 + GatedDeltaNetAttention,
565 + )
566 +
567 + gdn_layers = get_layers_from_vllm_config(config, GatedDeltaNetAttention)
568 + gated_norm_shapes: set[tuple[int, int]] = set()
569 + for layer in gdn_layers.values():
570 + gated_norm_shapes.add(
571 + (layer.num_v_heads // layer.tp_size, layer.head_v_dim)
572 + )
573 +
574 # Make sure fused add patterns are before simple rms norm,
575 # as the latter is a subset of the former in torch ops.
576 # The DoubleQuant patterns handle 1 rms_norm -> 2 group_fp8_quant
@@ -517,6 +624,21 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
624 epsilon, FP8_DTYPE, match_aiter_quant=match_aiter_quant
625 ).register(self.patterns)
626
627 + # Fuse decomposed RMSNormGated + group fp8 quant.
628 + # The replacement op (fused_rms_gated_fp8_group_quant) requires
629 + # an aiter version that includes the GDN triton kernel renames.
630 + if gated_norm_shapes and rocm_aiter_ops.are_gdn_triton_kernels_available():
631 + for num_heads, head_dim in gated_norm_shapes:
632 + if head_dim != 128:
633 + continue
634 + AiterRMSNormGatedFp8GroupQuantPattern(
635 + epsilon,
636 + FP8_DTYPE,
637 + GroupShape(1, 128),
638 + num_heads=num_heads,
639 + head_dim=head_dim,
640 + ).register(self.patterns)
641 +
642 self.dump_patterns(config, self.patterns)
643
644 @VllmInductorPass.time_and_log
@@ -534,6 +656,7 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
656 AiterFusedAddRMSFp8GroupQuantPattern,
657 DoubleAiterRMSFp8GroupQuantPattern,
658 DoubleAiterRMSFp8GroupQuantViewPattern,
659 + AiterRMSNormGatedFp8GroupQuantPattern,
660 ]
661 return self.hash_source(self, *fusion_patterns)
662
vllm/compilation/passes/vllm_inductor_pass.py
+27
@@ -252,6 +252,33 @@ def _fx_view_to_reshape(gm: fx.GraphModule) -> None:
252 view_to_reshape(gm)
253
254
255 +def fold_consecutive_reshapes(gm: fx.GraphModule) -> None:
256 + """Fold consecutive reshape ops into a single reshape.
257 +
258 + ``make_fx`` faithfully records every view/reshape the Python code performs,
259 + so patterns like ``x.reshape(a, b).reshape(c, d)`` produce two reshape
260 + nodes. Inductor's own optimisation would fold these, but
261 + ``pm.register_replacement``'s ``trace_fn`` runs before Inductor, so we
262 + must fold them ourselves for the pattern to match the compiled graph.
263 +
264 + When reshape(A, shape1) feeds only into reshape(result, shape2),
265 + the first reshape is redundant -- replace with reshape(A, shape2).
266 + """
267 + aten_reshape = torch.ops.aten.reshape.default
268 + for node in list(gm.graph.nodes):
269 + if not is_func(node, aten_reshape):
270 + continue
271 + inp = node.args[0]
272 + if not isinstance(inp, fx.Node) or not is_func(inp, aten_reshape):
273 + continue
274 + if len(inp.users) != 1:
275 + continue
276 + original_input = inp.args[0]
277 + node.args = (original_input, node.args[1])
278 + inp.replace_all_uses_with(original_input)
279 + gm.graph.erase_node(inp)
280 +
281 +
282 def _remove_noop_permutes(gm: fx.GraphModule) -> None:
283 for node in gm.graph.nodes:
284 if not is_func(node, torch.ops.aten.permute.default):
vllm/model_executor/layers/layernorm.py
+44 -29
@@ -229,56 +229,71 @@ class RMSNormGated(CustomOp):
229 def reset_parameters(self):
230 torch.nn.init.ones_(self.weight)
231
232 - def forward_native(
233 - self, x: torch.Tensor, z: torch.Tensor | None = None
232 + @staticmethod
233 + def forward_static(
234 + x: torch.Tensor,
235 + z: torch.Tensor | None,
236 + weight: torch.Tensor,
237 + epsilon: float,
238 + orig_dtype: torch.dtype,
239 + group_size: int | None = None,
240 + norm_before_gate: bool = True,
241 + activation: str = "swish",
242 ) -> torch.Tensor:
235 - """
236 - Native PyTorch implementation of RMS normalization with gating.
237 -
238 - Args:
239 - x: Input tensor
240 - z: Optional gating tensor
243 + """Pure-PyTorch RMS normalization with optional gating.
244
242 - Returns:
243 - Normalized (and optionally gated) tensor
245 + This static method contains the full native logic so that both
246 + ``forward_native`` and ``MatcherRMSNormGated`` (used by the
247 + compilation pattern matcher) can share the same implementation.
248
245 - If z is not None:
246 - - norm_before_gate=True: out = norm(x) * silu(z)
247 - - norm_before_gate=False: out = norm(x * silu(z))
249 + If *z* is not None and *norm_before_gate* is True:
250 + ``out = rms_norm(x) * act(z)``
251 + If *z* is not None and *norm_before_gate* is False:
252 + ``out = rms_norm(x * act(z))``
253 """
249 - orig_dtype = x.dtype
254 x = x.float()
251 - weight = self.weight.float()
252 - z = z.float() if z is not None else None
255 + weight = weight.float()
256 + if z is not None:
257 + z = z.float()
258
254 - assert self.activation in ["silu", "sigmoid", "swish"]
255 - act_fn = F.sigmoid if self.activation == "sigmoid" else F.silu
259 + assert activation in ["silu", "sigmoid", "swish"]
260 + act_fn = F.sigmoid if activation == "sigmoid" else F.silu
261
257 - # Apply gating before normalization if needed
258 - if z is not None and not self.norm_before_gate:
262 + if z is not None and not norm_before_gate:
263 x = x * act_fn(z)
264
261 - # RMS Normalization
262 - if self.group_size is None:
263 - # Standard RMS norm across the last dimension
265 + if group_size is None:
266 variance = x.pow(2).mean(dim=-1, keepdim=True)
265 - x_normed = x * torch.rsqrt(variance + self.eps)
267 + x_normed = x * torch.rsqrt(variance + epsilon)
268 out = x_normed * weight
269 else:
268 - # Group RMS norm
270 from einops import rearrange
271
271 - x_group = rearrange(x, "... (g d) -> ... g d", d=self.group_size)
272 + x_group = rearrange(x, "... (g d) -> ... g d", d=group_size)
273 variance = x_group.pow(2).mean(dim=-1, keepdim=True)
273 - x_normed = x_group * torch.rsqrt(variance + self.eps)
274 + x_normed = x_group * torch.rsqrt(variance + epsilon)
275 out = rearrange(x_normed, "... g d -> ... (g d)") * weight
276
276 - # Apply gating after normalization if needed
277 - if z is not None and self.norm_before_gate:
277 + if z is not None and norm_before_gate:
278 out = out * act_fn(z)
279
280 return out.to(orig_dtype)
281
282 + def forward_native(
283 + self, x: torch.Tensor, z: torch.Tensor | None = None
284 + ) -> torch.Tensor:
285 + """PyTorch-native implementation equivalent to forward()."""
286 + return self.forward_static(
287 + x,
288 + z,
289 + self.weight,
290 + self.eps,
291 + x.dtype,
292 + group_size=self.group_size,
293 + norm_before_gate=self.norm_before_gate,
294 + activation=self.activation,
295 + )
296 +
297 def forward_cuda(
298 self, x: torch.Tensor, z: torch.Tensor | None = None
299 ) -> torch.Tensor: