@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 529c671e80

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 4, 08:45 UTC

[ROCm][FEAT] AITER Fused Allreduce + RMSNorm (#37646)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com> Signed-off-by: junkang1991 <junkangchow@gmail.com> Co-authored-by: Rita Brugarolas <Rita.BrugarolasBrufau@amd.com> Co-authored-by: junkang1991 <junkangchow@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

vllmellm committed May 1, 2026 at 23:07 UTC 529c671e8075d265a48b72e0eaaeb5e30d2f1630

9 files changed +454 -20

.buildkite/test-amd.yaml

@@ -1108,6 +1108,7 @@ steps:

1108 - export VLLM_TEST_CLEAN_GPU_MEMORY=1

1109 - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py

1110 - pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py

1111 + - pytest -v -s tests/compile/passes/distributed/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions

1112

1113 #----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------#

1114

tests/compile/fusions_e2e/test_tp2_ar_rms.py

+19 -3

@@ -19,6 +19,8 @@ from .models import (

19 FLASHINFER_ATTN,

20 FLASHINFER_MLA_ATTN,

21 FLASHMLA_SPARSE_ATTN,

22 + ROCM_AITER_UNIFIED_ATTN,

23 + ROCM_ATTN,

24 TRITON_ATTN,

25 deepseek_coder_v2_lite_fp8,

26 deepseek_r1_fp4,

@@ -34,7 +36,9 @@ from .models import (

36 qwen3_a3b_fp8,

37 )

37 -pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")

39 +pytestmark = pytest.mark.skipif(

40 + not current_platform.is_cuda_alike(), reason="Only test CUDA/ROCm"

41 +)

44 @multi_gpu_test(num_gpus=2)

@@ -55,6 +59,7 @@ pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only tes

59 @pytest.mark.parametrize("n_layers", [4])

60 @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))

61 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)

62 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")

63 def test_tp2_ar_rms_fp8_fusions(

64 model_name: str,

65 matches_fn: Callable[[int], Matches],

@@ -124,6 +129,7 @@ def test_tp2_ar_rms_fp8_fusions(

129 @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))

130 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)

131 @pytest.mark.skipif(not is_blackwell(), reason="Blackwell required for fp4")

132 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")

133 def test_tp2_ar_rms_fp4_fusions(

134 model_name: str,

135 matches_fn: Callable[[int], Matches],

@@ -176,10 +182,19 @@ def test_tp2_ar_rms_fp4_fusions(

182 "model_name, matches_fn, model_kwargs, hf_overrides",

183 [llama3_8b, qwen3_a3b, gpt_oss_20b],

184 )

179 -@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])

185 +@pytest.mark.parametrize(

186 + "attn_backend",

187 + [

188 + TRITON_ATTN,

189 + FLASHINFER_ATTN,

190 + ROCM_ATTN,

191 + ROCM_AITER_UNIFIED_ATTN,

192 + ],

193 +)

194 @pytest.mark.parametrize("n_layers", [4])

181 -@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))

195 +@pytest.mark.parametrize("custom_ops", tuple(custom_ops_combos("rms_norm")))

196 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)

197 +@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="Only test CUDA/ROCm")

198 def test_tp2_ar_rms_fusions(

199 model_name: str,

200 matches_fn: Callable[[int], Matches],

@@ -221,4 +236,5 @@ def test_tp2_ar_rms_fusions(

236 compilation_config,

237 matches_check,

238 tp_size=2,

239 + use_aiter=current_platform.is_rocm(),

240 )

tests/compile/passes/distributed/test_fusion_all_reduce.py

+77 -12

@@ -8,8 +8,12 @@ import torch

8 import vllm.envs as envs

9 from tests.compile.backend import TestBackend

10 from tests.utils import TestFP8Layer, has_module_attribute, multi_gpu_test

11 +from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops

12 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant

12 -from vllm.compilation.passes.fusion.allreduce_rms_fusion import AllReduceFusionPass

13 +from vllm.compilation.passes.fusion.allreduce_rms_fusion import (

14 + AllReduceFusionPass,

15 + RocmAiterAllReduceFusionPass,

16 +)

17 from vllm.compilation.passes.utility.fix_functionalization import (

18 FixFunctionalizationPass,

19 )

@@ -42,13 +46,19 @@ DEVICE_TYPE = current_platform.device_type

47 class TestAllReduceRMSNormModel(torch.nn.Module):

48 def __init__(

45 - self, hidden_size=16, token_num=16, eps=1e-6, dtype: torch.dtype = torch.float16

49 + self,

50 + hidden_size=16,

51 + token_num=16,

52 + eps=1e-6,

53 + dtype: torch.dtype = torch.float16,

54 + use_aiter: bool = False,

55 ):

56 super().__init__()

57 self.hidden_size = hidden_size

58 self.eps = eps

59 self.norm = [RMSNorm(hidden_size, eps) for i in range(4)]

60 self.w = [torch.rand(hidden_size, hidden_size) for _ in range(3)]

61 + self.use_aiter = use_aiter

63 def forward(self, x):

64 # avoid having graph input be an arg to a pattern directly

@@ -76,6 +86,8 @@ class TestAllReduceRMSNormModel(torch.nn.Module):

86 return [torch.ops.vllm.all_reduce.default]

88 def ops_in_model_after(self):

89 + if self.use_aiter:

90 + return [rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()]

91 return [torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default]

@@ -194,12 +206,36 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):

206

207 @multi_gpu_test(num_gpus=2)

208 @pytest.mark.parametrize(

197 - "test_model, enable_quant_fp8_custom_op",

209 + "test_model, enable_quant_fp8_custom_op, use_aiter",

210 [

199 - (TestAllReduceRMSNormModel, False),

200 - (TestAllReduceRMSNormStaticQuantFP8Model, True),

201 - (TestAllReduceRMSNormStaticQuantFP8Model, False),

202 - (TestAllReduceFusedAddRMSNormStaticQuantFP4Model, False),

211 + (TestAllReduceRMSNormModel, False, IS_AITER_FOUND),

212 + pytest.param(

213 + TestAllReduceRMSNormStaticQuantFP8Model,

214 + True,

215 + False,

216 + marks=pytest.mark.skipif(

217 + current_platform.is_rocm(),

218 + reason="Not supported on ROCm platform",

219 + ),

220 + ),

221 + pytest.param(

222 + TestAllReduceRMSNormStaticQuantFP8Model,

223 + False,

224 + False,

225 + marks=pytest.mark.skipif(

226 + current_platform.is_rocm(),

227 + reason="Not supported on ROCm platform",

228 + ),

229 + ),

230 + pytest.param(

231 + TestAllReduceFusedAddRMSNormStaticQuantFP4Model,

232 + False,

233 + False,

234 + marks=pytest.mark.skipif(

235 + current_platform.is_rocm(),

236 + reason="Not supported on ROCm platform",

237 + ),

238 + ),

239 ],

240 )

241 @pytest.mark.parametrize("batch_size", [8])

@@ -210,9 +246,18 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):

246 @pytest.mark.parametrize("flashinfer_allreduce_backend", ["trtllm", "mnnvl"])

247 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")

248 @pytest.mark.skipif(

213 - not find_spec("flashinfer")

214 - or not has_module_attribute("flashinfer.comm", "allreduce_fusion")

215 - or not has_module_attribute("flashinfer.comm", "create_allreduce_fusion_workspace"),

249 + current_platform.is_rocm() and not IS_AITER_FOUND,

250 + reason="aiter is not found",

251 +)

252 +@pytest.mark.skipif(

253 + current_platform.is_cuda()

254 + and (

255 + not find_spec("flashinfer")

256 + or not has_module_attribute("flashinfer.comm", "allreduce_fusion")

257 + or not has_module_attribute(

258 + "flashinfer.comm", "create_allreduce_fusion_workspace"

259 + )

260 + ),

261 reason="flashinfer is not found or flashinfer "

262 "is not compiled with allreduce_fusion",

263 )

@@ -225,7 +270,14 @@ def test_all_reduce_fusion_pass_replace(

270 enable_rms_norm_custom_op,

271 enable_quant_fp8_custom_op,

272 flashinfer_allreduce_backend,

273 + use_aiter: bool,

274 + monkeypatch: pytest.MonkeyPatch,

275 ):

276 + if use_aiter:

277 + with monkeypatch.context() as m:

278 + m.setenv("VLLM_ROCM_USE_AITER", str(use_aiter))

279 + rocm_aiter_ops.refresh_env_variables()

280 +

281 num_processes = 2

282 if (

283 test_model == TestAllReduceFusedAddRMSNormStaticQuantFP4Model

@@ -249,6 +301,8 @@ def test_all_reduce_fusion_pass_replace(

301 enable_rms_norm_custom_op,

302 enable_quant_fp8_custom_op,

303 flashinfer_allreduce_backend,

304 + use_aiter,

305 + monkeypatch,

306 ),

307 nprocs=nprocs,

308 )

@@ -267,6 +321,8 @@ def all_reduce_fusion_pass_on_test_model(

321 enable_rms_norm_custom_op,

322 enable_quant_fp8_custom_op,

323 flashinfer_allreduce_backend,

324 + use_aiter: bool,

325 + monkeypatch: pytest.MonkeyPatch,

326 ):

327 set_random_seed(0)

328

@@ -313,7 +369,11 @@ def all_reduce_fusion_pass_on_test_model(

369 )

370 with set_current_vllm_config(vllm_config):

371 initialize_model_parallel(tensor_model_parallel_size=world_size)

316 - all_reduce_fusion_pass = AllReduceFusionPass(vllm_config)

372 + all_reduce_fusion_pass = (

373 + RocmAiterAllReduceFusionPass(vllm_config)

374 + if use_aiter

375 + else AllReduceFusionPass(vllm_config)

376 + )

377 noop_pass = NoOpEliminationPass(vllm_config)

378 func_pass = FixFunctionalizationPass(vllm_config)

379 cleanup_pass = PostCleanupPass(vllm_config)

@@ -323,7 +383,12 @@ def all_reduce_fusion_pass_on_test_model(

383 )

384

385 token_num = batch_size * seq_len

326 - model = test_model_cls(hidden_size, token_num, dtype=dtype)

386 + if test_model_cls is TestAllReduceRMSNormModel:

387 + model = test_model_cls(

388 + hidden_size, token_num, dtype=dtype, use_aiter=use_aiter

389 + )

390 + else:

391 + model = test_model_cls(hidden_size, token_num, dtype=dtype)

392

393 hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)

394

vllm/_aiter_ops.py

+115

@@ -2,9 +2,12 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import functools

4 from collections.abc import Callable

5 +from contextlib import contextmanager

6 +from typing import Protocol

8 import torch

9 from torch._ops import OpOverload

10 +from torch.distributed import ProcessGroup

12 import vllm.envs as envs

13 from vllm.platforms import current_platform

@@ -39,6 +42,27 @@ def is_aiter_found() -> bool:

42 IS_AITER_FOUND = is_aiter_found()

45 +class AiterCustomAllreduceProto(Protocol):

46 + max_size: int

47 + world_size: int

48 + fully_connected: bool

49 +

50 + @contextmanager

51 + def capture(self): ...

52 + def close(self) -> None: ...

53 + def fused_ar_rms(

54 + self,

55 + inp: torch.Tensor,

56 + res_inp: torch.Tensor,

57 + *,

58 + w: torch.Tensor,

59 + eps: float,

60 + registered: bool = False,

61 + use_1stage: bool = False,

62 + ) -> tuple[torch.Tensor, torch.Tensor]: ...

63 + def should_custom_ar(self, inp: torch.Tensor) -> bool: ...

64 +

65 +

66 def is_aiter_found_and_supported() -> bool:

67 """Check if AITER library is available and platform supports it.

@@ -750,6 +774,55 @@ def _rocm_aiter_rmsnorm_fused_dynamic_quant_fake(

774 return out, y_scale

775

776

777 +def _rocm_aiter_fused_allreduce_rmsnorm_impl(

778 + input_: torch.Tensor,

779 + residual: torch.Tensor,

780 + weight: torch.Tensor,

781 + epsilon: float,

782 +) -> tuple[torch.Tensor, torch.Tensor]:

783 + aiter_ar = rocm_aiter_ops.get_aiter_allreduce()

784 + assert aiter_ar is not None, "aiter allreduce must be initialized"

785 +

786 + total_bytes = input_.numel() * input_.element_size()

787 + hidden_dim = input_.shape[-1]

788 + token_num = input_.shape[0]

789 + hidden_ok = hidden_dim in (512, 1024, 2048, 4096, 7168)

790 + token_ok = token_num <= 80

791 + world_size = aiter_ar.world_size

792 + full_nvlink = aiter_ar.fully_connected

793 +

794 + if world_size == 2:

795 + size_ok = True

796 + elif full_nvlink and world_size <= 4:

797 + size_ok = total_bytes < 256 * 1024

798 + elif full_nvlink and world_size <= 8:

799 + size_ok = total_bytes < 128 * 1024

800 + else:

801 + size_ok = False

802 +

803 + use_1stage = hidden_ok and token_ok and size_ok

804 +

805 + result = aiter_ar.fused_ar_rms(

806 + input_,

807 + residual,

808 + w=weight,

809 + eps=epsilon,

810 + registered=torch.cuda.is_current_stream_capturing(),

811 + use_1stage=use_1stage,

812 + )

813 + assert result is not None

814 + return result[0], result[1]

815 +

816 +

817 +def _rocm_aiter_fused_allreduce_rmsnorm_fake(

818 + input_: torch.Tensor,

819 + residual: torch.Tensor,

820 + weight: torch.Tensor,

821 + epsilon: float,

822 +) -> tuple[torch.Tensor, torch.Tensor]:

823 + return torch.empty_like(input_), torch.empty_like(residual)

824 +

825 +

826 def _rocm_aiter_per_tensor_quant_impl(

827 x: torch.Tensor,

828 quant_dtype: torch.dtype,

@@ -1188,6 +1261,9 @@ class rocm_aiter_ops:

1261 # TODO: Consolidate under _LINEAR_ENABLED

1262 _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

1263

1264 + _ALL_REDUCE_MAX_SIZE: int = 8192 * 1024 * 8 * 2

1265 + _CUSTOM_ALL_REDUCE: AiterCustomAllreduceProto | None = None

1266 +

1267 @classmethod

1268 def refresh_env_variables(cls):

1269 """

@@ -1362,6 +1438,35 @@ class rocm_aiter_ops:

1438

1439 return cls.is_linear_enabled() and on_gfx950()

1440

1441 + @classmethod

1442 + def initialize_aiter_allreduce(

1443 + cls, group: ProcessGroup, device: torch.device

1444 + ) -> None:

1445 + try:

1446 + from aiter.dist.device_communicators.custom_all_reduce import (

1447 + CustomAllreduce as AiterCustomAllreduce,

1448 + )

1449 +

1450 + cls._CUSTOM_ALL_REDUCE = AiterCustomAllreduce(group, device)

1451 + except Exception:

1452 + cls._CUSTOM_ALL_REDUCE = None

1453 +

1454 + @classmethod

1455 + def get_aiter_allreduce(cls) -> AiterCustomAllreduceProto | None:

1456 + return cls._CUSTOM_ALL_REDUCE

1457 +

1458 + @classmethod

1459 + def destroy_aiter_allreduce(cls) -> None:

1460 + if cls._CUSTOM_ALL_REDUCE is not None:

1461 + cls._CUSTOM_ALL_REDUCE.close()

1462 + cls._CUSTOM_ALL_REDUCE = None

1463 +

1464 + @classmethod

1465 + def get_aiter_allreduce_max_size(cls) -> int | None:

1466 + # effective max input size (based on upstream aiter version: v0.1.10.post3)

1467 + # https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/aiter/dist/device_communicators/custom_all_reduce.py#L272-L273

1468 + return int(cls._ALL_REDUCE_MAX_SIZE / 2)

1469 +

1470 @staticmethod

1471 @if_aiter_supported

1472 def register_ops_once() -> None:

@@ -1552,6 +1657,12 @@ class rocm_aiter_ops:

1657 fake_impl=_triton_rotary_embedding_fake,

1658 )

1659

1660 + direct_register_custom_op(

1661 + op_name="rocm_aiter_fused_allreduce_rmsnorm",

1662 + op_func=_rocm_aiter_fused_allreduce_rmsnorm_impl,

1663 + fake_impl=_rocm_aiter_fused_allreduce_rmsnorm_fake,

1664 + )

1665 +

1666 direct_register_custom_op(

1667 op_name="fused_mla_dual_rms_norm",

1668 op_func=_fused_mla_dual_rms_norm_impl,

@@ -1605,6 +1716,10 @@ class rocm_aiter_ops:

1716 def get_triton_rotary_embedding_op() -> OpOverload:

1717 return torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default

1718

1719 + @staticmethod

1720 + def get_fused_allreduce_rmsnorm_op() -> OpOverload:

1721 + return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm.default

1722 +

1723 @staticmethod

1724 def get_fused_mla_dual_rms_norm_op() -> OpOverload:

1725 return torch.ops.vllm.fused_mla_dual_rms_norm.default

vllm/compilation/passes/fusion/act_quant_fusion.py

@@ -190,6 +190,7 @@ class SiluMulBlockQuantPattern(ActivationQuantPattern):

190 is_scale_transposed: bool = False,

191 is_e8m0: bool = False,

192 is_tma_aligned: bool = False,

193 + match_aiter: bool = False,

194 ) -> None:

195 super().__init__(quant_key)

196 self.quant_matcher = MatcherQuantFP8(

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

+209 -1

@@ -12,12 +12,14 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized

12 from torch._inductor.pattern_matcher import PatternMatcherPass

14 import vllm.ir.ops

15 +from vllm._aiter_ops import rocm_aiter_ops

16 from vllm.compilation.passes.fusion.rms_quant_fusion import (

17 _rms_input_weight_dtype_match,

18 )

19 from vllm.config import VllmConfig

20 from vllm.config.utils import Range

21 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce

22 +from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce

23 from vllm.distributed.parallel_state import (

24 get_tensor_model_parallel_rank,

25 get_tensor_model_parallel_world_size,

@@ -32,7 +34,12 @@ from vllm.utils.torch_utils import (

34 )

36 from ..inductor_pass import enable_fake_mode

35 -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

37 +from ..vllm_inductor_pass import (

38 + VllmFusionPatternMatcherPass,

39 + VllmInductorPass,

40 + VllmPatternMatcherPass,

41 + VllmPatternReplacement,

42 +)

43 from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8

45 FP8_DTYPE = current_platform.fp8_dtype()

@@ -889,3 +896,204 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

896 return

897 with contextlib.suppress(Exception):

898 destroy_fi_ar_workspace()

899 +

900 +

901 +# TODO: make BasePattern to inherit from VllmPatternReplacement

902 +class AiterAllreduceFusedRMSNormPattern(BasePattern, VllmPatternReplacement):

903 + def __init__(

904 + self,

905 + epsilon: float,

906 + dtype: torch.dtype,

907 + device: str | None,

908 + use_aiter_rmsnorm: bool = True,

909 + ) -> None:

910 + super().__init__(dtype, device)

911 + self.dtype = dtype

912 + self.epsilon = epsilon

913 + self.FUSED_AR_RMSNORM_OP = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()

914 +

915 + def get_inputs(self) -> list[torch.Tensor]:

916 + return [self.empty(5, 16), self.empty(16)]

917 +

918 + @property

919 + def pattern(self):

920 + def _pattern(

921 + input: torch.Tensor, weight: torch.Tensor

922 + ) -> tuple[torch.Tensor, torch.Tensor]:

923 + allreduce_output = tensor_model_parallel_all_reduce(input)

924 + rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon)

925 +

926 + return rms, allreduce_output

927 +

928 + return _pattern

929 +

930 + @property

931 + def replacement(self):

932 + def _replacement(

933 + input: torch.Tensor, weight: torch.Tensor

934 + ) -> tuple[torch.Tensor, torch.Tensor]:

935 + residual = torch.empty_like(input)

936 + allreduce = self.FUSED_AR_RMSNORM_OP(

937 + input_=input,

938 + residual=residual,

939 + weight=weight,

940 + epsilon=self.epsilon,

941 + )

942 + return allreduce[0], allreduce[1]

943 +

944 + return _replacement

945 +

946 +

947 +class AiterAllreduceFusedAddRMSNormPattern(BasePattern, VllmPatternReplacement):

948 + def __init__(

949 + self,

950 + epsilon: float,

951 + dtype: torch.dtype,

952 + device: str | None,

953 + use_aiter_rmsnorm: bool = True,

954 + ) -> None:

955 + super().__init__(dtype, device)

956 + self.epsilon = epsilon

957 + self.dtype = dtype

958 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(

959 + epsilon, match_rocm_aiter=use_aiter_rmsnorm

960 + )

961 + self.FUSED_AR_RMSNORM_OP = rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()

962 +

963 + def get_inputs(self) -> list[torch.Tensor]:

964 + input, residual, weight = self.rmsnorm_matcher.inputs()

965 +

966 + return [residual, input.to(self.dtype), weight]

967 +

968 + @property

969 + def pattern(self):

970 + def _pattern(

971 + residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor

972 + ) -> tuple[torch.Tensor, torch.Tensor]:

973 + allreduce_output = tensor_model_parallel_all_reduce(input)

974 + rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)

975 +

976 + return rms, residual

977 +

978 + return _pattern

979 +

980 + @property

981 + def replacement(self):

982 + def _replacement(

983 + residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor

984 + ) -> tuple[torch.Tensor, torch.Tensor]:

985 + allreduce = self.FUSED_AR_RMSNORM_OP(

986 + input_=input,

987 + residual=residual,

988 + weight=weight,

989 + epsilon=self.epsilon,

990 + )

991 + return allreduce[0], allreduce[1]

992 +

993 + return _replacement

994 +

995 +

996 +class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):

997 + def __init__(self, config: VllmConfig) -> None:

998 + super().__init__(config, "rocm_aiter_allreduce_fusion_pass")

999 + self.disabled = True

1000 + self.tp_size = get_tensor_model_parallel_world_size()

1001 + if self.tp_size <= 1:

1002 + logger.warning_once("AllReduce fusion pass is disabled for tp_size <= 1.")

1003 + return

1004 +

1005 + if config.model_config is None:

1006 + logger.warning_once(

1007 + "AllReduce fusion pass is disabled for missing model_config."

1008 + )

1009 + return

1010 +

1011 + device_comm = get_tp_group().device_communicator

1012 + if device_comm is None:

1013 + logger.warning_once("Device communicator is required.")

1014 + return

1015 +

1016 + ca_comm = getattr(device_comm, "ca_comm", None)

1017 + if ca_comm is None:

1018 + logger.warning_once("Custom Allreduce is required.")

1019 + return

1020 + self.ca_comm = ca_comm

1021 +

1022 + assert isinstance(ca_comm, CustomAllreduce)

1023 +

1024 + group = get_tp_group().cpu_group

1025 + rocm_aiter_ops.initialize_aiter_allreduce(group, self.device)

1026 + hidden_dim = config.model_config.get_hidden_size()

1027 + element_size = torch.tensor([], dtype=self.model_dtype).element_size()

1028 + max_size = rocm_aiter_ops.get_aiter_allreduce_max_size()

1029 + if max_size is None:

1030 + logger.warning("AITER allreduce fusion must be initialized")

1031 + return

1032 +

1033 + # Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim.

1034 + # Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM

1035 + # and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12

1036 + # hidden_dim is a runtime argument. Detect the older API via the missing

1037 + # `_pool` attribute and skip fusion for unsupported sizes.

1038 + # Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590

1039 + aiter_ar = rocm_aiter_ops.get_aiter_allreduce()

1040 + _AITER_OLD_FUSED_AR_RMS_HIDDEN = (512, 1024, 2048, 4096)

1041 + if (

1042 + aiter_ar is not None

1043 + and not hasattr(aiter_ar, "_pool")

1044 + and hidden_dim not in _AITER_OLD_FUSED_AR_RMS_HIDDEN

1045 + ):

1046 + logger.warning_once(

1047 + "AITER allreduce-rmsnorm fusion disabled: aiter<0.1.12 "

1048 + "only supports hidden_dim in %s; got %d. Upgrade aiter to "

1049 + ">=0.1.12 to enable fusion for this model.",

1050 + _AITER_OLD_FUSED_AR_RMS_HIDDEN,

1051 + hidden_dim,

1052 + )

1053 + # Tear down aiter's custom-allreduce so its IPC handles don't

1054 + # race with vllm's ca_comm on the unfused fallback path.

1055 + with contextlib.suppress(Exception):

1056 + rocm_aiter_ops.destroy_aiter_allreduce()

1057 + return

1058 +

1059 + max_token_num = max_size // (hidden_dim * element_size)

1060 + self.max_token_num = min(

1061 + max_token_num,

1062 + config.scheduler_config.max_num_batched_tokens,

1063 + )

1064 +

1065 + for epsilon in [1e-5, 1e-6]:

1066 + self.register(

1067 + AiterAllreduceFusedRMSNormPattern(

1068 + epsilon,

1069 + self.model_dtype,

1070 + self.device,

1071 + )

1072 + )

1073 + self.register(

1074 + AiterAllreduceFusedAddRMSNormPattern(

1075 + epsilon,

1076 + self.model_dtype,

1077 + self.device,

1078 + )

1079 + )

1080 +

1081 + # WARNING: This is a hack to clear the pattern matcher cache

1082 + # and allow multiple values of epsilon.

1083 + torch._inductor.pattern_matcher._seen_patterns.clear()

1084 +

1085 + self.disabled = False

1086 +

1087 + self.dump_patterns(config, self.pm_pass)

1088 +

1089 + def is_applicable_for_range(self, compile_range: Range) -> bool:

1090 + if self.disabled:

1091 + logger.warning_once("AllReduce fusion pass is disabled.")

1092 + return False

1093 + return bool(compile_range.end <= self.max_token_num)

1094 +

1095 + def __del__(self) -> None:

1096 + if getattr(self, "disabled", True):

1097 + return

1098 + with contextlib.suppress(Exception):

1099 + rocm_aiter_ops.destroy_aiter_allreduce()

vllm/compilation/passes/pass_manager.py

+7 -1

@@ -18,6 +18,9 @@ from .ir.lowering_pass import VllmIRLoweringPass

18 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

20 if rocm_aiter_ops.is_enabled():

21 + from .fusion.allreduce_rms_fusion import (

22 + RocmAiterAllReduceFusionPass,

23 + )

24 from .fusion.rocm_aiter_fusion import (

25 MLADualRMSNormFusionPass,

26 RocmAiterRMSNormQuantFusionPass,

@@ -137,7 +140,10 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]

140 self.passes += [AsyncTPPass(config)]

141

142 if self.pass_config.fuse_allreduce_rms:

140 - self.passes += [AllReduceFusionPass(config)]

143 + if rocm_aiter_ops.is_enabled():

144 + self.passes += [RocmAiterAllReduceFusionPass(config)]

145 + else:

146 + self.passes += [AllReduceFusionPass(config)]

147

148 if self.pass_config.fuse_minimax_qk_norm:

149 self.passes += [MiniMaxQKNormPass(config)]

vllm/config/vllm.py

+16 -2

@@ -121,6 +121,15 @@ def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:

121 from vllm.platforms import current_platform

122 from vllm.utils.flashinfer import has_flashinfer

123

124 + if current_platform.is_rocm():

125 + from vllm._aiter_ops import rocm_aiter_ops

126 +

127 + return (

128 + rocm_aiter_ops.is_enabled()

129 + and rocm_aiter_ops.is_rmsnorm_enabled()

130 + and cfg.parallel_config.tensor_parallel_size > 1

131 + )

132 +

133 return (

134 cfg.parallel_config.tensor_parallel_size > 1

135 and current_platform.is_cuda()

@@ -1604,10 +1613,15 @@ class VllmConfig:

1613 if compile_range_end is not None:

1614 computed_compile_ranges_endpoints.append(compile_range_end)

1615

1607 - # Add the compile ranges for flashinfer

1616 + # Add the compile ranges for flashinfer/aiter.

1617 if compilation_config.pass_config.fuse_allreduce_rms:

1618 tp_size = self.parallel_config.tensor_parallel_size

1610 - max_size = compilation_config.pass_config.flashinfer_max_size(tp_size)

1619 + from vllm._aiter_ops import rocm_aiter_ops

1620 +

1621 + if rocm_aiter_ops.is_enabled():

1622 + max_size = rocm_aiter_ops.get_aiter_allreduce_max_size()

1623 + else:

1624 + max_size = compilation_config.pass_config.flashinfer_max_size(tp_size)

1625 if max_size is not None:

1626 assert isinstance(self.model_config.dtype, torch.dtype)

1627 max_token_num = max_size // (

vllm/distributed/parallel_state.py

+9 -1

@@ -472,6 +472,7 @@ class GroupCoordinator:

472 # only cuda uses this function,

473 # so we don't abstract it into the base class

474 maybe_ca_context = nullcontext()

475 + maybe_aiter_context = nullcontext()

476 from vllm.distributed.device_communicators.cuda_communicator import (

477 CudaCommunicator,

478 )

@@ -482,13 +483,20 @@ class GroupCoordinator:

483 if ca_comm is not None:

484 maybe_ca_context = ca_comm.capture() # type: ignore

485

486 + from vllm._aiter_ops import rocm_aiter_ops

487 +

488 + if rocm_aiter_ops.is_enabled():

489 + aiter_ar = rocm_aiter_ops.get_aiter_allreduce()

490 + if aiter_ar is not None:

491 + maybe_aiter_context = aiter_ar.capture() # type: ignore

492 +

493 # ensure all initialization operations complete before attempting to

494 # capture the graph on another stream

495 curr_stream = torch.cuda.current_stream()

496 if curr_stream != stream:

497 stream.wait_stream(curr_stream)

498

491 - with torch.cuda.stream(stream), maybe_ca_context:

499 + with torch.cuda.stream(stream), maybe_ca_context, maybe_aiter_context:

500 yield graph_capture_context

501

502 def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: