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

A high-throughput and memory-efficient inference and serving engine for LLMs

[ROCm][Compile] Fuse AR + RMSNorm + per-group FP8 quant (+ DSv3.2 indexer fan-out) (#42864)

Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Frida Andersson <fanderss@amd.com>

Markus Hartikainen committed Jun 9, 2026 at 15:06 UTC 80e2c4462dd9a94252b4b0061d00ed66908131c1
3 files changed +795 -2
tests/compile/passes/distributed/test_fusion_all_reduce.py
+276
@@ -189,6 +189,118 @@ class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
189 ]
190
191
192 +class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
193 + """Exercises the new ROCm AITER AR+RMS+per-group-FP8-quant patterns.
194 +
195 + Four ``rms_norm`` sites that together hit every pattern registered by
196 + ``RocmAiterAllReduceFusionPass`` for the per-group FP8 quant path:
197 +
198 + * ``norm[0]``: ``all_reduce -> rms_norm -> group_fp8_quant`` (no residual)
199 + -> ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern``
200 + * ``norm[1]``: ``all_reduce -> fused_add_rms_norm -> group_fp8_quant``
201 + (single ``rms`` consumer)
202 + -> ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern``
203 + * ``norm[2..3]``: ``all_reduce -> fused_add_rms_norm
204 + -> (group_fp8_quant + rocm_unquantized_gemm)`` (two ``rms`` consumers,
205 + modeling the DSv3.2 indexer fan-out)
206 + -> ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern``
207 +
208 + The chain feeds the next AllReduce by dequantizing the FP8 output (FP8
209 + cast back to bf16 multiplied by the per-group scale), which is enough to
210 + keep the matmul chain bf16 without depending on a real FP8 block-scaled
211 + GEMM kernel.
212 + """
213 +
214 + quant_group_size = 128
215 + indexer_out_dim = 8
216 +
217 + def __init__(
218 + self,
219 + hidden_size=128,
220 + token_num=16,
221 + eps=1e-6,
222 + dtype: torch.dtype = torch.bfloat16,
223 + use_triton_quant: bool = False,
224 + ):
225 + super().__init__()
226 + self.hidden_size = hidden_size
227 + self.eps = eps
228 + self.use_triton_quant = use_triton_quant
229 + assert hidden_size % self.quant_group_size == 0, (
230 + f"hidden_size ({hidden_size}) must be a multiple of "
231 + f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant"
232 + )
233 + self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)]
234 + self.w = [torch.rand(hidden_size, hidden_size, dtype=dtype) for _ in range(3)]
235 + self.indexer_w = [
236 + torch.rand(self.indexer_out_dim, hidden_size, dtype=dtype) for _ in range(2)
237 + ]
238 +
239 + def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
240 + if self.use_triton_quant:
241 + return torch.ops.vllm.triton_per_token_group_quant_fp8(
242 + rms, self.quant_group_size
243 + )
244 + return torch.ops.vllm.rocm_aiter_group_fp8_quant.default(
245 + rms, self.quant_group_size
246 + )
247 +
248 + def _dequantize_to_bf16(
249 + self, q: torch.Tensor, s: torch.Tensor, ref: torch.Tensor
250 + ) -> torch.Tensor:
251 + # Broadcast the per-group scale across each group of `quant_group_size`
252 + # so we can chain the FP8 output back into a bf16 matmul. This avoids
253 + # depending on a real FP8 block-scaled GEMM kernel in the test.
254 + s_full = s.repeat_interleave(self.quant_group_size, dim=-1).to(ref.dtype)
255 + return q.to(ref.dtype) * s_full
256 +
257 + def forward(self, hidden_states):
258 + z = torch.relu(hidden_states)
259 + x = resid = tensor_model_parallel_all_reduce(z)
260 + rms = self.norm[0](x)
261 + q0, s0 = self._group_quant(rms)
262 + y = self._dequantize_to_bf16(q0, s0, rms)
263 +
264 + z2 = torch.mm(y, self.w[0])
265 + x2 = tensor_model_parallel_all_reduce(z2)
266 + rms2, resid = self.norm[1](x2, resid)
267 + q1, s1 = self._group_quant(rms2)
268 + y2 = self._dequantize_to_bf16(q1, s1, rms2)
269 +
270 + z3 = torch.mm(y2, self.w[1])
271 + x3 = tensor_model_parallel_all_reduce(z3)
272 + rms3, resid = self.norm[2](x3, resid)
273 + q2, s2 = self._group_quant(rms3)
274 + # Second consumer of ``rms3``: forces the with-indexer pattern.
275 + idx2 = torch.ops.vllm.rocm_unquantized_gemm(rms3, self.indexer_w[0], None)
276 + y3 = self._dequantize_to_bf16(q2, s2, rms3)
277 +
278 + z4 = torch.mm(y3, self.w[2])
279 + x4 = tensor_model_parallel_all_reduce(z4)
280 + rms4, resid = self.norm[3](x4, resid)
281 + q3, s3 = self._group_quant(rms4)
282 + # Second consumer of ``rms4``: forces the with-indexer pattern.
283 + idx3 = torch.ops.vllm.rocm_unquantized_gemm(rms4, self.indexer_w[1], None)
284 + y4 = self._dequantize_to_bf16(q3, s3, rms4)
285 + return y4, idx2, idx3
286 +
287 + def ops_in_model_before(self):
288 + return [
289 + torch.ops.vllm.all_reduce.default,
290 + (
291 + torch.ops.vllm.triton_per_token_group_quant_fp8.default
292 + if self.use_triton_quant
293 + else torch.ops.vllm.rocm_aiter_group_fp8_quant.default
294 + ),
295 + ]
296 +
297 + def ops_in_model_after(self):
298 + return [
299 + rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_op(),
300 + rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op(), # noqa: E501
301 + ]
302 +
303 +
304 class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
305 def __init__(
306 self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16
@@ -463,3 +575,167 @@ def all_reduce_fusion_pass_on_test_model(
575 assert fused_nodes
576 assert all(n.kwargs.get("weight_bias") == 1.0 for n in fused_nodes)
577 del all_reduce_fusion_pass
578 +
579 +
580 +@multi_gpu_test(num_gpus=2)
581 +@pytest.mark.parametrize("use_triton_quant", [True, False])
582 +@pytest.mark.parametrize("batch_size", [8])
583 +@pytest.mark.parametrize("seq_len", [8])
584 +@pytest.mark.parametrize("hidden_size", [128])
585 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
586 +@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
587 +@pytest.mark.skipif(
588 + not current_platform.is_rocm(),
589 + reason="ROCm AITER AR+RMS+per-group-FP8-quant fusion is ROCm-only",
590 +)
591 +@pytest.mark.skipif(not IS_AITER_FOUND, reason="aiter is not found")
592 +def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
593 + batch_size: int,
594 + seq_len: int,
595 + hidden_size: int,
596 + dtype: torch.dtype,
597 + enable_rms_norm_custom_op: bool,
598 + use_triton_quant: bool,
599 + monkeypatch: pytest.MonkeyPatch,
600 +):
601 + """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new
602 + ROCm AITER AR+RMS+per-group-FP8-quant fusion patterns.
603 +
604 + Validates the three new ``VllmPatternReplacement`` patterns added to
605 + ``RocmAiterAllReduceFusionPass``:
606 +
607 + * ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern`` (no-residual)
608 + * ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual,
609 + single ``rms`` consumer)
610 + * ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with-
611 + residual, DSv3.2 indexer fan-out; parametrized over both
612 + ``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant``
613 + producers).
614 + """
615 + with monkeypatch.context() as m:
616 + m.setenv("VLLM_ROCM_USE_AITER", "1")
617 + rocm_aiter_ops.refresh_env_variables()
618 +
619 + if not rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group():
620 + pytest.skip(
621 + "aiter build is missing 'fused_ar_rms_per_group_quant' (needs "
622 + "ROCm/aiter PR #2823); the new patterns aren't registered."
623 + )
624 +
625 + num_processes = 2
626 +
627 + def run_torch_spawn(fn, nprocs):
628 + torch.multiprocessing.spawn(
629 + fn,
630 + args=(
631 + num_processes,
632 + TestAiterAllReduceRMSNormGroupQuantFP8Model,
633 + batch_size,
634 + seq_len,
635 + hidden_size,
636 + dtype,
637 + enable_rms_norm_custom_op,
638 + use_triton_quant,
639 + monkeypatch,
640 + ),
641 + nprocs=nprocs,
642 + )
643 +
644 + run_torch_spawn(rocm_aiter_group_quant_fusion_pass_on_test_model, num_processes)
645 +
646 +
647 +def rocm_aiter_group_quant_fusion_pass_on_test_model(
648 + local_rank: int,
649 + world_size: int,
650 + test_model_cls: torch.nn.Module,
651 + batch_size: int,
652 + seq_len: int,
653 + hidden_size: int,
654 + dtype: torch.dtype,
655 + enable_rms_norm_custom_op: bool,
656 + use_triton_quant: bool,
657 + monkeypatch: pytest.MonkeyPatch,
658 +):
659 + set_random_seed(0)
660 +
661 + device = torch.device(f"{DEVICE_TYPE}:{local_rank}")
662 + torch.accelerator.set_device_index(device)
663 + torch.set_default_device(device)
664 + torch.set_default_dtype(dtype)
665 +
666 + update_environment_variables(
667 + {
668 + "RANK": str(local_rank),
669 + "LOCAL_RANK": str(local_rank),
670 + "WORLD_SIZE": str(world_size),
671 + "MASTER_ADDR": "localhost",
672 + "MASTER_PORT": "12345",
673 + "VLLM_ROCM_USE_AITER": "1",
674 + }
675 + )
676 + rocm_aiter_ops.refresh_env_variables()
677 +
678 + init_distributed_environment()
679 +
680 + custom_ops = []
681 + if enable_rms_norm_custom_op:
682 + custom_ops.append("+rms_norm")
683 + # ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip``
684 + # only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at
685 + # the call site). The patterns in this PR are robust to both Triton and
686 + # rocm_aiter forms; we always enable +quant_fp8 so the matcher's example
687 + # trace finds the same form the test model uses.
688 + custom_ops.append("+quant_fp8")
689 +
690 + vllm_config = VllmConfig(
691 + compilation_config=CompilationConfig(
692 + mode=CompilationMode.VLLM_COMPILE, custom_ops=custom_ops
693 + )
694 + )
695 + vllm_config.compilation_config.pass_config = PassConfig(
696 + fuse_allreduce_rms=True, eliminate_noops=True
697 + )
698 + vllm_config.device_config = DeviceConfig(device=torch.device(DEVICE_TYPE))
699 + vllm_config.parallel_config.rank = local_rank
700 +
701 + model_name = "RedHatAI/Llama-3.2-1B-Instruct-FP8"
702 + vllm_config.model_config = ModelConfig(
703 + model=model_name, trust_remote_code=True, dtype=dtype, seed=42
704 + )
705 + with set_current_vllm_config(vllm_config):
706 + initialize_model_parallel(tensor_model_parallel_size=world_size)
707 + all_reduce_fusion_pass = RocmAiterAllReduceFusionPass(vllm_config)
708 + noop_pass = NoOpEliminationPass(vllm_config)
709 + func_pass = FixFunctionalizationPass(vllm_config)
710 + cleanup_pass = PostCleanupPass(vllm_config)
711 +
712 + backend = TestBackend(
713 + noop_pass, all_reduce_fusion_pass, func_pass, cleanup_pass
714 + )
715 +
716 + token_num = batch_size * seq_len
717 + model = test_model_cls(
718 + hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant
719 + )
720 +
721 + hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)
722 +
723 + compiled_model = torch.compile(model, backend=backend)
724 + compiled_model(hidden_states)
725 +
726 + results_unfused = model(hidden_states)
727 + results_fused = compiled_model(hidden_states)
728 + # The fused per-group AR+RMS+QUANT op is bit-equivalent to the unfused
729 + # chain modulo the small AllReduce + RMSNorm reordering inside aiter.
730 + # Per-group FP8 quant introduces step noise <=1 per group; use the
731 + # same tolerance as the sibling FP8 static test.
732 + torch.testing.assert_close(results_unfused, results_fused, atol=1e-2, rtol=1e-2)
733 +
734 + # Four pattern firings: norm[0] (no-add quant), norm[1] (add quant,
735 + # single ``rms`` consumer), norm[2..3] (add quant + indexer fan-out).
736 + assert all_reduce_fusion_pass.matched_count == 4, (
737 + f"{all_reduce_fusion_pass.matched_count=}"
738 + )
739 + backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)
740 + backend.check_after_ops(model.ops_in_model_after())
741 + del all_reduce_fusion_pass
vllm/_aiter_ops.py
+196
@@ -70,6 +70,21 @@ class AiterCustomAllreduceProto(Protocol):
70 registered: bool = False,
71 use_1stage: bool = False,
72 ) -> tuple[torch.Tensor, torch.Tensor]: ...
73 + def fused_ar_rms_per_group_quant(
74 + self,
75 + inp: torch.Tensor,
76 + res_inp: torch.Tensor,
77 + *,
78 + w: torch.Tensor,
79 + eps: float,
80 + group_size: int = 128,
81 + registered: bool = False,
82 + use_1stage: bool = False,
83 + emit_bf16: bool = False,
84 + ) -> (
85 + tuple[torch.Tensor, torch.Tensor, torch.Tensor]
86 + | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
87 + ): ...
88 def should_custom_ar(self, inp: torch.Tensor) -> bool: ...
89
90
@@ -850,6 +865,152 @@ def _rocm_aiter_fused_allreduce_rmsnorm_fake(
865 return torch.empty_like(input_), torch.empty_like(residual)
866
867
868 +def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
869 + input_: torch.Tensor,
870 + residual: torch.Tensor,
871 + weight: torch.Tensor,
872 + epsilon: float,
873 + group_size: int,
874 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
875 + """Fused AllReduce + RMSNorm + per-group FP8 quant.
876 +
877 + Mirrors the eligibility logic of ``_rocm_aiter_fused_allreduce_rmsnorm_impl``
878 + for the 1-stage vs 2-stage AITER kernel dispatch (both variants run inside
879 + AITER, the only choice we make here is the launcher to call into).
880 + """
881 + aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
882 + assert aiter_ar is not None, "aiter allreduce must be initialized"
883 +
884 + total_bytes = input_.numel() * input_.element_size()
885 + hidden_dim = input_.shape[-1]
886 + token_num = input_.shape[0]
887 + if input_.dtype in (torch.bfloat16, torch.float16):
888 + pack_size = 16 // input_.element_size()
889 + hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
890 + else:
891 + hidden_ok = False
892 + token_ok = token_num <= 80
893 + world_size = aiter_ar.world_size
894 + full_nvlink = aiter_ar.fully_connected
895 +
896 + if world_size == 2:
897 + size_ok = True
898 + elif full_nvlink and world_size <= 4:
899 + size_ok = total_bytes < 256 * 1024
900 + elif full_nvlink and world_size <= 8:
901 + size_ok = total_bytes < 128 * 1024
902 + else:
903 + size_ok = False
904 +
905 + use_1stage = hidden_ok and token_ok and size_ok
906 +
907 + result = aiter_ar.fused_ar_rms_per_group_quant(
908 + input_,
909 + residual,
910 + w=weight,
911 + eps=epsilon,
912 + group_size=group_size,
913 + registered=torch.cuda.is_current_stream_capturing(),
914 + use_1stage=use_1stage,
915 + )
916 + assert result is not None
917 + return result[0], result[1], result[2]
918 +
919 +
920 +def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_fake(
921 + input_: torch.Tensor,
922 + residual: torch.Tensor,
923 + weight: torch.Tensor,
924 + epsilon: float,
925 + group_size: int,
926 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
927 + hidden_dim = input_.shape[-1]
928 + num_groups = hidden_dim // group_size
929 + quant_out = torch.empty(input_.shape, dtype=FP8_DTYPE, device=input_.device)
930 + residual_out = torch.empty_like(residual)
931 + scale_out = torch.empty(
932 + input_.shape[:-1] + (num_groups,),
933 + dtype=torch.float32,
934 + device=input_.device,
935 + )
936 + return quant_out, residual_out, scale_out
937 +
938 +
939 +def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
940 + input_: torch.Tensor,
941 + residual: torch.Tensor,
942 + weight: torch.Tensor,
943 + epsilon: float,
944 + group_size: int,
945 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
946 + """Fused AllReduce + add-RMSNorm + per-group FP8 quant + bf16 normed act.
947 +
948 + Wraps the same AITER launcher as ``_rocm_aiter_fused_allreduce_rmsnorm_
949 + quant_per_group_impl`` with ``emit_bf16=True``, emitting the pre-quant
950 + bf16/fp16 normed activation for a parallel consumer (DeepSeek V3.2 sparse
951 + indexer ``wk_weights_proj``).
952 + """
953 + aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
954 + assert aiter_ar is not None, "aiter allreduce must be initialized"
955 +
956 + total_bytes = input_.numel() * input_.element_size()
957 + hidden_dim = input_.shape[-1]
958 + token_num = input_.shape[0]
959 + if input_.dtype in (torch.bfloat16, torch.float16):
960 + pack_size = 16 // input_.element_size()
961 + hidden_ok = hidden_dim % pack_size == 0 and hidden_dim // pack_size <= 1024
962 + else:
963 + hidden_ok = False
964 + token_ok = token_num <= 80
965 + world_size = aiter_ar.world_size
966 + full_nvlink = aiter_ar.fully_connected
967 +
968 + if world_size == 2:
969 + size_ok = True
970 + elif full_nvlink and world_size <= 4:
971 + size_ok = total_bytes < 256 * 1024
972 + elif full_nvlink and world_size <= 8:
973 + size_ok = total_bytes < 128 * 1024
974 + else:
975 + size_ok = False
976 +
977 + use_1stage = hidden_ok and token_ok and size_ok
978 +
979 + result = aiter_ar.fused_ar_rms_per_group_quant(
980 + input_,
981 + residual,
982 + w=weight,
983 + eps=epsilon,
984 + group_size=group_size,
985 + registered=torch.cuda.is_current_stream_capturing(),
986 + use_1stage=use_1stage,
987 + emit_bf16=True,
988 + )
989 + assert result is not None
990 + assert len(result) == 4, "emit_bf16=True must return four tensors from aiter"
991 + return result[0], result[1], result[2], result[3]
992 +
993 +
994 +def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake(
995 + input_: torch.Tensor,
996 + residual: torch.Tensor,
997 + weight: torch.Tensor,
998 + epsilon: float,
999 + group_size: int,
1000 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1001 + hidden_dim = input_.shape[-1]
1002 + num_groups = hidden_dim // group_size
1003 + quant_out = torch.empty(input_.shape, dtype=FP8_DTYPE, device=input_.device)
1004 + residual_out = torch.empty_like(residual)
1005 + scale_out = torch.empty(
1006 + input_.shape[:-1] + (num_groups,),
1007 + dtype=torch.float32,
1008 + device=input_.device,
1009 + )
1010 + bf16_norm_out = torch.empty_like(input_)
1011 + return quant_out, residual_out, scale_out, bf16_norm_out
1012 +
1013 +
1014 def _rocm_aiter_per_tensor_quant_impl(
1015 x: torch.Tensor,
1016 quant_dtype: torch.dtype,
@@ -1838,6 +1999,18 @@ class rocm_aiter_ops:
1999 fake_impl=_rocm_aiter_fused_allreduce_rmsnorm_fake,
2000 )
2001
2002 + direct_register_custom_op(
2003 + op_name="rocm_aiter_fused_allreduce_rmsnorm_quant_per_group",
2004 + op_func=(_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl),
2005 + fake_impl=(_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_fake),
2006 + )
2007 +
2008 + direct_register_custom_op(
2009 + op_name="rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm", # noqa: E501
2010 + op_func=_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl, # noqa: E501
2011 + fake_impl=_rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_fake, # noqa: E501
2012 + )
2013 +
2014 direct_register_custom_op(
2015 op_name="fused_mla_dual_rms_norm",
2016 op_func=_fused_mla_dual_rms_norm_impl,
@@ -1892,6 +2065,29 @@ class rocm_aiter_ops:
2065 def get_fused_allreduce_rmsnorm_op() -> OpOverload:
2066 return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm.default
2067
2068 + @staticmethod
2069 + def get_fused_allreduce_rmsnorm_quant_per_group_op() -> OpOverload:
2070 + return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group.default
2071 +
2072 + @staticmethod
2073 + def get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() -> OpOverload: # noqa: E501
2074 + return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm.default # noqa: E501
2075 +
2076 + # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823).
2077 + @classmethod
2078 + def has_fused_allreduce_rmsnorm_quant_per_group(cls) -> bool:
2079 + """True if the running AITER build exposes the per-group AR+RMS+quant
2080 + kernel (added in ROCm/aiter PR #2823).
2081 +
2082 + The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off
2083 + this so vLLM degrades to the AR+RMS-only fusion when run against an
2084 + older aiter that lacks the per-group launcher.
2085 + """
2086 + aiter_ar = cls.get_aiter_allreduce()
2087 + return aiter_ar is not None and hasattr(
2088 + aiter_ar, "fused_ar_rms_per_group_quant"
2089 + )
2090 +
2091 @staticmethod
2092 def get_fused_mla_dual_rms_norm_op() -> OpOverload:
2093 return torch.ops.vllm.fused_mla_dual_rms_norm.default
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+323 -2
@@ -26,6 +26,9 @@ from vllm.distributed.parallel_state import (
26 )
27 from vllm.logger import init_logger
28 from vllm.model_executor.layers.quantization.utils.quant_utils import (
29 + GroupShape,
30 + QuantKey,
31 + ScaleDesc,
32 kFp8StaticTensorSym,
33 )
34 from vllm.platforms import current_platform
@@ -1109,7 +1112,7 @@ class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement):
1112 allreduce = self.FUSED_AR_RMSNORM_OP(
1113 input_=input,
1114 residual=residual,
1112 - weight=weight,
1115 + weight=weight.to(input.dtype),
1116 epsilon=self.epsilon,
1117 )
1118 return allreduce[0], allreduce[1]
@@ -1155,7 +1158,7 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
1158 allreduce = self.FUSED_AR_RMSNORM_OP(
1159 input_=input,
1160 residual=residual,
1158 - weight=weight,
1161 + weight=weight.to(input.dtype),
1162 epsilon=self.epsilon,
1163 )
1164 return allreduce[0], allreduce[1]
@@ -1163,6 +1166,279 @@ class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):
1166 return _replacement
1167
1168
1169 +class AiterAllreduceFusedRMSNormGroupQuantFP8Pattern(
1170 + BasePattern, VllmPatternReplacement
1171 +):
1172 + """Fuse AllReduce + RMSNorm + per-group FP8 quant into a single AITER
1173 + custom op.
1174 +
1175 + Matches the AR-side analogue of ``AiterRMSFp8GroupQuantPattern`` in
1176 + ``rocm_aiter_fusion.py``: ``all_reduce -> rms_norm -> group_fp8_quant``
1177 + fans out into ``rocm_aiter_fused_allreduce_rmsnorm_quant_per_group``.
1178 +
1179 + Without this pattern, ``RocmAiterAllReduceFusionPass`` would fuse the
1180 + ``all_reduce + rms_norm`` half (PR #41825 wires that), but the trailing
1181 + ``rocm_aiter_group_fp8_quant`` would still launch as a separate kernel.
1182 + That standalone quant accounts for ~535us / decode step on DSv3.2 MI355X
1183 + TP4 -- this pattern eliminates it by absorbing the quant into the AR
1184 + epilogue. Group size 128 matches the FP8 block-scaled MM kernel used by
1185 + DSv3.2's linear weights.
1186 + """
1187 +
1188 + def __init__(
1189 + self,
1190 + epsilon: float,
1191 + dtype: torch.dtype,
1192 + device: str | None,
1193 + group_size: int = 128,
1194 + ) -> None:
1195 + super().__init__(dtype, device)
1196 + self.epsilon = epsilon
1197 + self.dtype = dtype
1198 + self.group_size = group_size
1199 + self.FUSED_AR_RMS_QUANT_OP = (
1200 + rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_op()
1201 + )
1202 + self.quant_dtype = current_platform.fp8_dtype()
1203 + self.quant_matcher = MatcherQuantFP8(
1204 + QuantKey(
1205 + dtype=self.quant_dtype,
1206 + scale=ScaleDesc(torch.float32, False, GroupShape(1, group_size)),
1207 + symmetric=True,
1208 + ),
1209 + match_rocm_aiter=True,
1210 + )
1211 +
1212 + def get_inputs(self) -> list[torch.Tensor]:
1213 + # input, weight; hidden dim must be a group_size multiple so the
1214 + # group quant matcher's example trace is well-defined.
1215 + return [self.empty(5, self.group_size), self.empty(self.group_size)]
1216 +
1217 + @property
1218 + def pattern(self):
1219 + def _pattern(
1220 + input: torch.Tensor, weight: torch.Tensor
1221 + ) -> tuple[torch.Tensor, torch.Tensor]:
1222 + allreduce_output = tensor_model_parallel_all_reduce(input)
1223 + rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon)
1224 + quant, scale = self.quant_matcher(rms)
1225 + return quant, scale
1226 +
1227 + return _pattern
1228 +
1229 + @property
1230 + def replacement(self):
1231 + def _replacement(
1232 + input: torch.Tensor, weight: torch.Tensor
1233 + ) -> tuple[torch.Tensor, torch.Tensor]:
1234 + residual = torch.zeros_like(input)
1235 + result = self.FUSED_AR_RMS_QUANT_OP(
1236 + input_=input,
1237 + residual=residual,
1238 + weight=weight.to(input.dtype),
1239 + epsilon=self.epsilon,
1240 + group_size=self.group_size,
1241 + )
1242 + # quant_out, scale_out (residual is unused on the no-add path,
1243 + # mirroring how AiterAllreduceFusedRMSNormPattern drops the
1244 + # residual output)
1245 + return result[0], result[2]
1246 +
1247 + return _replacement
1248 +
1249 +
1250 +class AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern(
1251 + BasePattern, VllmPatternReplacement
1252 +):
1253 + """``fused_add`` variant of ``AiterAllreduceFusedRMSNormGroupQuantFP8Pattern``.
1254 +
1255 + Targets the dominant DSv3.2-style post-attention / post-MLP path:
1256 + ``all_reduce -> fused_add_rms_norm -> group_fp8_quant``. Returns the
1257 + FP8 quant output, the residual carry-over, and the per-group scale.
1258 + """
1259 +
1260 + def __init__(
1261 + self,
1262 + epsilon: float,
1263 + dtype: torch.dtype,
1264 + device: str | None,
1265 + group_size: int = 128,
1266 + ) -> None:
1267 + super().__init__(dtype, device)
1268 + self.epsilon = epsilon
1269 + self.dtype = dtype
1270 + self.group_size = group_size
1271 + self.FUSED_AR_RMS_QUANT_OP = (
1272 + rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_op()
1273 + )
1274 + self.quant_dtype = current_platform.fp8_dtype()
1275 + self.quant_matcher = MatcherQuantFP8(
1276 + QuantKey(
1277 + dtype=self.quant_dtype,
1278 + scale=ScaleDesc(torch.float32, False, GroupShape(1, group_size)),
1279 + symmetric=True,
1280 + ),
1281 + match_rocm_aiter=True,
1282 + )
1283 +
1284 + def get_inputs(self) -> list[torch.Tensor]:
1285 + # residual, input, weight
1286 + return [
1287 + self.empty(5, self.group_size),
1288 + self.empty(5, self.group_size),
1289 + self.empty(self.group_size),
1290 + ]
1291 +
1292 + @property
1293 + def pattern(self):
1294 + def _pattern(
1295 + residual: torch.Tensor,
1296 + input: torch.Tensor,
1297 + weight: torch.Tensor,
1298 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1299 + allreduce_output = tensor_model_parallel_all_reduce(input)
1300 + rms, residual_out = vllm.ir.ops.fused_add_rms_norm(
1301 + allreduce_output, residual, weight, self.epsilon
1302 + )
1303 + quant, scale = self.quant_matcher(rms)
1304 + return quant, scale, residual_out
1305 +
1306 + return _pattern
1307 +
1308 + @property
1309 + def replacement(self):
1310 + def _replacement(
1311 + residual: torch.Tensor,
1312 + input: torch.Tensor,
1313 + weight: torch.Tensor,
1314 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1315 + result = self.FUSED_AR_RMS_QUANT_OP(
1316 + input_=input,
1317 + residual=residual,
1318 + weight=weight.to(input.dtype),
1319 + epsilon=self.epsilon,
1320 + group_size=self.group_size,
1321 + )
1322 + # quant_out, scale_out, residual_out
1323 + return result[0], result[2], result[1]
1324 +
1325 + return _replacement
1326 +
1327 +
1328 +class AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern(
1329 + BasePattern, VllmPatternReplacement
1330 +):
1331 + """Indexer-fan-out variant of ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern``.
1332 +
1333 + Targets the DSv3.2 post-attention / post-MLP path where the post-AR normed
1334 + activation has two consumers: a per-group FP8 quant for ``fused_qkv_a_proj``
1335 + *and* a bf16 ``rocm_unquantized_gemm`` for the indexer ``wk_weights_proj``.
1336 + The single-consumer pattern above cannot fire when this fan-out is present,
1337 + so without this pattern the standalone FP8 quant kernel survives unfused
1338 + (~535us / decode step on DSv3.2 MI355X TP4).
1339 +
1340 + Lowers to ``rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm``
1341 + (the ``emit_bf16=True`` variant of the AR+RMS+QUANT launcher, which returns
1342 + FP8 quant + scales + bf16 normed activations in one kernel) and rewires the
1343 + indexer GEMM onto the emitted bf16 norm output. The RMS output is also a
1344 + graph output in DSv3.2's residual carry; it is returned as a pattern output
1345 + so the matcher can substitute the bf16 norm in its place.
1346 +
1347 + The trailing FP8 group-quant is matched via ``MatcherQuantFP8`` (consistent
1348 + with the sibling patterns above), which traces both ``QuantFP8.forward_hip``
1349 + and ``forward_native`` paths and so matches whichever op the call site
1350 + lowers to (``vllm.triton_per_token_group_quant_fp8`` or
1351 + ``vllm.rocm_aiter_group_fp8_quant``).
1352 + """
1353 +
1354 + def __init__(
1355 + self,
1356 + epsilon: float,
1357 + dtype: torch.dtype,
1358 + device: str | None,
1359 + group_size: int = 128,
1360 + ) -> None:
1361 + super().__init__(dtype, device)
1362 + self.epsilon = epsilon
1363 + self.dtype = dtype
1364 + self.group_size = group_size
1365 + self.FUSED_AR_RMS_QUANT_BF16_OP = (
1366 + rocm_aiter_ops.get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() # noqa: E501
1367 + )
1368 + self.quant_dtype = current_platform.fp8_dtype()
1369 + self.quant_matcher = MatcherQuantFP8(
1370 + QuantKey(
1371 + dtype=self.quant_dtype,
1372 + scale=ScaleDesc(torch.float32, False, GroupShape(1, group_size)),
1373 + symmetric=True,
1374 + ),
1375 + match_rocm_aiter=True,
1376 + )
1377 +
1378 + def get_inputs(self) -> list[torch.Tensor]:
1379 + h = self.group_size
1380 + indexer_out = 8
1381 + return [
1382 + self.empty(5, h),
1383 + self.empty(5, h),
1384 + self.empty(h),
1385 + self.empty(indexer_out, h),
1386 + ]
1387 +
1388 + @property
1389 + def pattern(self):
1390 + eps = self.epsilon
1391 +
1392 + def _pattern(
1393 + residual: torch.Tensor,
1394 + input_: torch.Tensor,
1395 + norm_weight: torch.Tensor,
1396 + indexer_weight: torch.Tensor,
1397 + ) -> tuple[
1398 + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
1399 + ]:
1400 + ar_out = tensor_model_parallel_all_reduce(input_)
1401 + rms, res_out = vllm.ir.ops.fused_add_rms_norm(
1402 + ar_out, residual, norm_weight, eps
1403 + )
1404 + q, s = self.quant_matcher(rms)
1405 + idx = torch.ops.vllm.rocm_unquantized_gemm(rms, indexer_weight)
1406 + return q, s, res_out, idx, rms
1407 +
1408 + return _pattern
1409 +
1410 + @property
1411 + def replacement(self):
1412 + gs = self.group_size
1413 + eps = self.epsilon
1414 +
1415 + def _replacement(
1416 + residual: torch.Tensor,
1417 + input_: torch.Tensor,
1418 + norm_weight: torch.Tensor,
1419 + indexer_weight: torch.Tensor,
1420 + ) -> tuple[
1421 + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
1422 + ]:
1423 + fused = self.FUSED_AR_RMS_QUANT_BF16_OP(
1424 + input_=input_,
1425 + residual=residual,
1426 + weight=norm_weight.to(input_.dtype),
1427 + epsilon=eps,
1428 + group_size=gs,
1429 + )
1430 + quant_out, residual_out, scale_out, bf16_norm = (
1431 + fused[0],
1432 + fused[1],
1433 + fused[2],
1434 + fused[3],
1435 + )
1436 + idx = torch.ops.vllm.rocm_unquantized_gemm(bf16_norm, indexer_weight)
1437 + return quant_out, scale_out, residual_out, idx, bf16_norm
1438 +
1439 + return _replacement
1440 +
1441 +
1442 class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1443 def __init__(self, config: VllmConfig) -> None:
1444 super().__init__(config, "rocm_aiter_allreduce_fusion_pass")
@@ -1232,7 +1508,52 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1508 config.scheduler_config.max_num_batched_tokens,
1509 )
1510
1511 + # Only register the AR+RMS+per-group-FP8-quant patterns when the
1512 + # running aiter exposes the kernel. Older aiter builds (pre PR #2823)
1513 + # fall back to the AR+RMS-only fusion paired with PR #41825's
1514 + # standalone RMS+quant fusion -- still correct, just leaves the
1515 + # post-AR quant as a standalone kernel.
1516 + supports_per_group_quant = (
1517 + rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group()
1518 + )
1519 + if not supports_per_group_quant:
1520 + logger.warning_once(
1521 + "AITER AR+RMS+per-group-FP8-quant fusion disabled: aiter "
1522 + "build is missing 'fused_ar_rms_per_group_quant'. Upgrade "
1523 + "aiter past PR #2823 to enable the trailing per-group "
1524 + "FP8 quant fusion."
1525 + )
1526 +
1527 for epsilon in [1e-5, 1e-6]:
1528 + # Quant-fused variants must register first so the pattern matcher
1529 + # tries them before the AR+RMS-only variants. Otherwise the
1530 + # AR+RMS-only fusion runs first and consumes the all_reduce node,
1531 + # leaving the trailing quant op stranded as an unfused kernel.
1532 + # Register larger subgraphs first (DeepSeek indexer fan-out, then
1533 + # quant-only AR+RMS+quant, then AR+RMS-only).
1534 + if supports_per_group_quant:
1535 + self.register(
1536 + AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern(
1537 + epsilon,
1538 + self.model_dtype,
1539 + self.device,
1540 + )
1541 + )
1542 + self.register(
1543 + AiterAllreduceFusedRMSNormGroupQuantFP8Pattern(
1544 + epsilon,
1545 + self.model_dtype,
1546 + self.device,
1547 + )
1548 + )
1549 + self.register(
1550 + AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern(
1551 + epsilon,
1552 + self.model_dtype,
1553 + self.device,
1554 + )
1555 + )
1556 +
1557 self.register(
1558 AiterAllreduceFusedRMSNormPattern(
1559 epsilon,