@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 740f379fae

[ROCm][AITER] Directly Implement AITER Custom All-reduce in CudaCommunicator (#46065)

Signed-off-by: BadrBasowid <badr.basowid@gmail.com>

BadrBasowid committed Jul 6, 2026 at 20:16 UTC 740f379faefd1ddeb47c85ebd600af3ccae77b7a
12 files changed +499 -148
.buildkite/test-amd.yaml
+5 -1
@@ -2782,15 +2782,19 @@ steps:
2782 - vllm/envs.py
2783 - examples/offline_inference/data_parallel.py
2784 - tests/distributed/test_context_parallel.py
2785 + - tests/distributed/test_rocm_aiter_custom_ar.py
2786 - tests/distributed/test_rocm_quick_reduce.py
2787 - tests/distributed/test_quick_all_reduce.py
2788 + - tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
2789 - tests/v1/distributed/test_dbo.py
2790 - tests/utils.py
2791 commands:
2792 - pytest -v -s tests/distributed/test_context_parallel.py
2791 - - pytest -v -s tests/v1/distributed/test_dbo.py
2793 + - pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
2794 + - pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
2795 - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
2796 - pytest -v -s tests/distributed/test_quick_all_reduce.py
2797 + - pytest -v -s tests/v1/distributed/test_dbo.py
2798
2799 #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
2800
tests/compile/fusions_e2e/conftest.py
+1
@@ -79,6 +79,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
79 ):
80 monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1" if use_deepgemm else "0")
81 monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_aiter else "0")
82 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1" if use_aiter else "0")
83 from vllm._aiter_ops import rocm_aiter_ops
84
85 rocm_aiter_ops.refresh_env_variables()
tests/compile/passes/distributed/test_fusion_all_reduce.py
+9 -1
@@ -30,6 +30,9 @@ from vllm.config import (
30 set_current_vllm_config,
31 )
32 from vllm.distributed import tensor_model_parallel_all_reduce
33 +from vllm.distributed.device_communicators.aiter_custom_all_reduce import (
34 + AiterCustomAllreduce,
35 +)
36 from vllm.distributed.parallel_state import (
37 init_distributed_environment,
38 initialize_model_parallel,
@@ -504,8 +507,12 @@ def all_reduce_fusion_pass_on_test_model(
507 "MASTER_ADDR": "localhost",
508 "MASTER_PORT": "12345",
509 "VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend,
510 + "VLLM_ROCM_USE_AITER": str(int(use_aiter)),
511 + "VLLM_ROCM_USE_AITER_CUSTOM_AR": str(int(use_aiter)),
512 }
513 )
514 + if use_aiter:
515 + rocm_aiter_ops.refresh_env_variables()
516
517 init_distributed_environment()
518
@@ -616,7 +623,7 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
623 m.setenv("VLLM_ROCM_USE_AITER", "1")
624 rocm_aiter_ops.refresh_env_variables()
625
619 - if not rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group():
626 + if not AiterCustomAllreduce.build_supports_per_group_quant():
627 pytest.skip(
628 "aiter build is missing 'fused_ar_rms_per_group_quant' (needs "
629 "ROCm/aiter PR #2823); the new patterns aren't registered."
@@ -671,6 +678,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model(
678 "MASTER_ADDR": "localhost",
679 "MASTER_PORT": "12345",
680 "VLLM_ROCM_USE_AITER": "1",
681 + "VLLM_ROCM_USE_AITER_CUSTOM_AR": "1",
682 }
683 )
684 rocm_aiter_ops.refresh_env_variables()
tests/distributed/test_rocm_aiter_custom_ar.py new
+134
@@ -0,0 +1,134 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import ray
6 +import torch
7 +import torch.distributed as dist
8 +
9 +from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops
10 +from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
11 +from vllm.distributed.parallel_state import get_tp_group, graph_capture
12 +from vllm.envs import disable_envs_cache
13 +from vllm.platforms import current_platform
14 +
15 +from ..utils import (
16 + assert_rocm_custom_allreduce_backend_state,
17 + ensure_model_parallel_initialized,
18 + init_test_distributed_environment,
19 + multi_gpu_test,
20 + multi_process_parallel,
21 +)
22 +
23 +pytestmark = pytest.mark.skipif(
24 + not current_platform.is_rocm(),
25 + reason="ROCm-only AITER custom allreduce tests",
26 +)
27 +
28 +test_cases = [
29 + ((2, 7168), torch.float16),
30 + ((2, 7168), torch.bfloat16),
31 + ((128, 8192), torch.float16),
32 + ((128, 8192), torch.bfloat16),
33 +]
34 +
35 +
36 +def _configure_aiter_custom_ar_env(monkeypatch: pytest.MonkeyPatch) -> None:
37 + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
38 + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False)
39 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
40 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1")
41 + monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE")
42 + disable_envs_cache()
43 + rocm_aiter_ops.refresh_env_variables()
44 +
45 +
46 +def _assert_aiter_handles_input(inp: torch.Tensor) -> None:
47 + aiter_ar_comm = get_tp_group().device_communicator.aiter_ar_comm
48 + assert aiter_ar_comm is not None
49 + assert aiter_ar_comm.should_custom_ar(inp), (
50 + f"AITER CustomAllreduce does not support input shape {inp.shape}."
51 + )
52 +
53 +
54 +@ray.remote(num_gpus=1, max_calls=1)
55 +def graph_allreduce(
56 + monkeypatch: pytest.MonkeyPatch,
57 + tp_size,
58 + pp_size,
59 + rank,
60 + distributed_init_port,
61 +) -> None:
62 + with monkeypatch.context() as m:
63 + _configure_aiter_custom_ar_env(m)
64 +
65 + device = torch.device(f"cuda:{rank}")
66 + torch.accelerator.set_device_index(device)
67 + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
68 + ensure_model_parallel_initialized(tp_size, pp_size)
69 + assert_rocm_custom_allreduce_backend_state(True, "NONE")
70 + group = get_tp_group().device_group
71 +
72 + # A small all_reduce for warmup.
73 + # this is needed because device communicators might be created lazily
74 + # (e.g. NCCL). This will ensure that the communicator is initialized
75 + # before any communication happens, so that this group can be used for
76 + # graph capture immediately.
77 + data = torch.zeros(1)
78 + data = data.to(device=device)
79 + dist.all_reduce(data, group=group)
80 + torch.accelerator.synchronize()
81 + del data
82 +
83 + for shape, dtype in test_cases:
84 + with graph_capture(device=device) as graph_capture_context:
85 + inp = torch.ones(shape, dtype=dtype, device=device)
86 + _assert_aiter_handles_input(inp)
87 + expected = inp * tp_size
88 +
89 + torch.accelerator.synchronize()
90 + graph = torch.cuda.CUDAGraph()
91 + with torch.cuda.graph(graph, stream=graph_capture_context.stream):
92 + out = tensor_model_parallel_all_reduce(inp)
93 +
94 + graph.replay()
95 + torch.testing.assert_close(out, expected)
96 +
97 +
98 +@ray.remote(num_gpus=1, max_calls=1)
99 +def eager_allreduce(
100 + monkeypatch: pytest.MonkeyPatch,
101 + tp_size,
102 + pp_size,
103 + rank,
104 + distributed_init_port,
105 +) -> None:
106 + with monkeypatch.context() as m:
107 + _configure_aiter_custom_ar_env(m)
108 +
109 + device = torch.device(f"cuda:{rank}")
110 + torch.accelerator.set_device_index(device)
111 + init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
112 + ensure_model_parallel_initialized(tp_size, pp_size)
113 + assert_rocm_custom_allreduce_backend_state(True, "NONE")
114 +
115 + for shape, dtype in test_cases:
116 + inp = torch.ones(shape, dtype=dtype, device=device)
117 + _assert_aiter_handles_input(inp)
118 + expected = inp * tp_size
119 + out = tensor_model_parallel_all_reduce(inp)
120 + torch.testing.assert_close(out, expected)
121 +
122 +
123 +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed")
124 +@multi_gpu_test(num_gpus=2)
125 +@pytest.mark.parametrize("tp_size", [2])
126 +@pytest.mark.parametrize("pipeline_parallel_size", [1])
127 +@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])
128 +def test_rocm_aiter_custom_allreduce(
129 + monkeypatch: pytest.MonkeyPatch,
130 + tp_size,
131 + pipeline_parallel_size,
132 + test_target,
133 +):
134 + multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
tests/utils.py
+40
@@ -1459,6 +1459,46 @@ def multi_process_parallel(
1459 ray.shutdown()
1460
1461
1462 +def assert_rocm_custom_allreduce_backend_state(
1463 + use_aiter_custom_ar: bool,
1464 + quick_reduce_quantization: str,
1465 +) -> None:
1466 + from vllm.distributed.parallel_state import get_tp_group
1467 +
1468 + device_communicator = get_tp_group().device_communicator
1469 + aiter_ar_comm = device_communicator.aiter_ar_comm
1470 + if use_aiter_custom_ar:
1471 + assert aiter_ar_comm is not None, "AITER CustomAllreduce was not initialized."
1472 + assert not aiter_ar_comm.disabled, "AITER CustomAllreduce is disabled."
1473 + assert device_communicator.ca_comm is None, (
1474 + "vLLM CustomAllreduce should not be initialized when AITER CA is used."
1475 + )
1476 + else:
1477 + assert aiter_ar_comm is None, (
1478 + "AITER CustomAllreduce should not be initialized when disabled."
1479 + )
1480 + assert device_communicator.ca_comm is not None, (
1481 + "vLLM CustomAllreduce should be initialized when AITER CA is disabled."
1482 + )
1483 +
1484 + qr_comm = device_communicator.qr_comm
1485 + assert qr_comm is not None, "QuickReduce communicator was not initialized."
1486 + if quick_reduce_quantization == "NONE":
1487 + assert qr_comm.disabled, "QuickReduce should be disabled."
1488 + else:
1489 + assert not qr_comm.disabled, "QuickReduce should be enabled."
1490 +
1491 +
1492 +def assert_rocm_custom_allreduce_backend_state_on_worker(
1493 + _worker,
1494 + use_aiter_custom_ar: bool,
1495 + quick_reduce_quantization: str,
1496 +) -> None:
1497 + assert_rocm_custom_allreduce_backend_state(
1498 + use_aiter_custom_ar, quick_reduce_quantization
1499 + )
1500 +
1501 +
1502 @contextmanager
1503 def error_on_warning(category: type[Warning] = Warning):
1504 """
tests/v1/e2e/general/test_rocm_aiter_custom_ar.py new
+118
@@ -0,0 +1,118 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops
7 +from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
8 +from vllm.envs import disable_envs_cache
9 +from vllm.platforms import current_platform
10 +
11 +from ....conftest import VllmRunner
12 +from ....utils import (
13 + assert_rocm_custom_allreduce_backend_state_on_worker,
14 + multi_gpu_test,
15 +)
16 +
17 +PROMPTS = ["Hello, my name is", "The capital of France is"]
18 +
19 +
20 +def _run_generation(
21 + vllm_runner: type[VllmRunner],
22 + monkeypatch: pytest.MonkeyPatch,
23 + compilation_config: CompilationConfig,
24 + *,
25 + model: str,
26 + max_tokens: int,
27 + use_aiter_custom_ar: bool,
28 + quick_reduce_quantization: str,
29 +) -> list[tuple[list[int], str]]:
30 + with monkeypatch.context() as m:
31 + m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
32 + m.setenv("VLLM_ROCM_USE_AITER", "1")
33 + m.setenv(
34 + "VLLM_ROCM_USE_AITER_CUSTOM_AR",
35 + "1" if use_aiter_custom_ar else "0",
36 + )
37 + m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quick_reduce_quantization)
38 + disable_envs_cache()
39 + rocm_aiter_ops.refresh_env_variables()
40 +
41 + with vllm_runner(
42 + model,
43 + dtype="half",
44 + tensor_parallel_size=2,
45 + compilation_config=compilation_config,
46 + max_model_len=256,
47 + max_num_seqs=len(PROMPTS),
48 + gpu_memory_utilization=0.7,
49 + ) as llm:
50 + llm.get_llm().collective_rpc(
51 + assert_rocm_custom_allreduce_backend_state_on_worker,
52 + args=(use_aiter_custom_ar, quick_reduce_quantization),
53 + )
54 +
55 + return llm.generate_greedy(PROMPTS, max_tokens)
56 +
57 +
58 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-only")
59 +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed")
60 +@multi_gpu_test(num_gpus=2)
61 +@pytest.mark.parametrize(
62 + "quick_reduce_quantization",
63 + [
64 + pytest.param("FP", id="quick-reduce-on"),
65 + pytest.param("NONE", id="quick-reduce-off"),
66 + ],
67 +)
68 +@pytest.mark.parametrize(
69 + "cudagraph_mode",
70 + [
71 + pytest.param(CUDAGraphMode.NONE, id="cudagraph-none"),
72 + pytest.param(CUDAGraphMode.FULL, id="cudagraph-full"),
73 + ],
74 +)
75 +@pytest.mark.parametrize(
76 + "model,max_tokens",
77 + [
78 + pytest.param("facebook/opt-125m", 8, id="opt-125m"),
79 + ],
80 +)
81 +def test_rocm_aiter_custom_ar_e2e(
82 + vllm_runner: type[VllmRunner],
83 + monkeypatch: pytest.MonkeyPatch,
84 + cudagraph_mode: CUDAGraphMode,
85 + quick_reduce_quantization: str,
86 + model: str,
87 + max_tokens: int,
88 +):
89 + compilation_mode = (
90 + CompilationMode.NONE
91 + if cudagraph_mode == CUDAGraphMode.NONE
92 + else CompilationMode.VLLM_COMPILE
93 + )
94 + compilation_config = CompilationConfig(
95 + mode=compilation_mode,
96 + cudagraph_mode=cudagraph_mode,
97 + )
98 +
99 + baseline_generations = _run_generation(
100 + vllm_runner,
101 + monkeypatch,
102 + compilation_config,
103 + model=model,
104 + max_tokens=max_tokens,
105 + use_aiter_custom_ar=False,
106 + quick_reduce_quantization=quick_reduce_quantization,
107 + )
108 + aiter_custom_ar_generations = _run_generation(
109 + vllm_runner,
110 + monkeypatch,
111 + compilation_config,
112 + model=model,
113 + max_tokens=max_tokens,
114 + use_aiter_custom_ar=True,
115 + quick_reduce_quantization=quick_reduce_quantization,
116 + )
117 +
118 + assert aiter_custom_ar_generations == baseline_generations
vllm/_aiter_ops.py
+34 -95
@@ -2,12 +2,9 @@
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
5
6 import torch
7 from torch._ops import OpOverload
10 -from torch.distributed import ProcessGroup
8
9 import vllm.envs as envs
10 from vllm.platforms import current_platform
@@ -52,42 +49,6 @@ def is_aiter_found() -> bool:
49 IS_AITER_FOUND = is_aiter_found()
50
51
55 -class AiterCustomAllreduceProto(Protocol):
56 - max_size: int
57 - world_size: int
58 - fully_connected: bool
59 -
60 - @contextmanager
61 - def capture(self): ...
62 - def close(self) -> None: ...
63 - def fused_ar_rms(
64 - self,
65 - inp: torch.Tensor,
66 - res_inp: torch.Tensor,
67 - *,
68 - w: torch.Tensor,
69 - eps: float,
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 -
52 def is_aiter_found_and_supported() -> bool:
53 """Check if AITER library is available and platform supports it.
54
@@ -830,6 +791,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
791 ) -> tuple[torch.Tensor, torch.Tensor]:
792 aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
793 assert aiter_ar is not None, "aiter allreduce must be initialized"
794 + ca = aiter_ar.aiter_ca
795
796 total_bytes = input_.numel() * input_.element_size()
797 hidden_dim = input_.shape[-1]
@@ -840,8 +802,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
802 else:
803 hidden_ok = False
804 token_ok = token_num <= 80
843 - world_size = aiter_ar.world_size
844 - full_nvlink = aiter_ar.fully_connected
805 + world_size = ca.world_size
806 + full_nvlink = ca.fully_connected
807
808 if world_size == 2:
809 size_ok = True
@@ -854,12 +816,11 @@ def _rocm_aiter_fused_allreduce_rmsnorm_impl(
816
817 use_1stage = hidden_ok and token_ok and size_ok
818
857 - result = aiter_ar.fused_ar_rms(
819 + result = ca.custom_fused_ar_rms(
820 input_,
821 residual,
860 - w=weight,
861 - eps=epsilon,
862 - registered=torch.cuda.is_current_stream_capturing(),
822 + weight,
823 + epsilon,
824 use_1stage=use_1stage,
825 )
826 assert result is not None
@@ -890,6 +851,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
851 """
852 aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
853 assert aiter_ar is not None, "aiter allreduce must be initialized"
854 + ca = aiter_ar.aiter_ca
855
856 total_bytes = input_.numel() * input_.element_size()
857 hidden_dim = input_.shape[-1]
@@ -900,8 +862,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
862 else:
863 hidden_ok = False
864 token_ok = token_num <= 80
903 - world_size = aiter_ar.world_size
904 - full_nvlink = aiter_ar.fully_connected
865 + world_size = ca.world_size
866 + full_nvlink = ca.fully_connected
867
868 if world_size == 2:
869 size_ok = True
@@ -914,7 +876,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl(
876
877 use_1stage = hidden_ok and token_ok and size_ok
878
917 - result = aiter_ar.fused_ar_rms_per_group_quant(
879 + result = ca.fused_ar_rms_per_group_quant(
880 input_,
881 residual,
882 w=weight,
@@ -962,6 +924,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
924 """
925 aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
926 assert aiter_ar is not None, "aiter allreduce must be initialized"
927 + ca = aiter_ar.aiter_ca
928
929 total_bytes = input_.numel() * input_.element_size()
930 hidden_dim = input_.shape[-1]
@@ -972,8 +935,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
935 else:
936 hidden_ok = False
937 token_ok = token_num <= 80
975 - world_size = aiter_ar.world_size
976 - full_nvlink = aiter_ar.fully_connected
938 + world_size = ca.world_size
939 + full_nvlink = ca.fully_connected
940
941 if world_size == 2:
942 size_ok = True
@@ -986,7 +949,7 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl(
949
950 use_1stage = hidden_ok and token_ok and size_ok
951
989 - result = aiter_ar.fused_ar_rms_per_group_quant(
952 + result = ca.fused_ar_rms_per_group_quant(
953 input_,
954 residual,
955 w=weight,
@@ -1577,6 +1540,7 @@ class rocm_aiter_ops:
1540
1541 # Check if the env variable is set
1542 _AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
1543 + _CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR
1544 _LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
1545 _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
1546 _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
@@ -1598,9 +1562,6 @@ class rocm_aiter_ops:
1562 # num_shared_experts / shared_expert_scoring_func args (7-arg form).
1563 _TOPK_SOFTMAX_FUSED_SIGMOID: bool | None = None
1564
1601 - _ALL_REDUCE_MAX_SIZE: int = 8192 * 1024 * 8 * 2
1602 - _CUSTOM_ALL_REDUCE: AiterCustomAllreduceProto | None = None
1603 -
1565 @classmethod
1566 def refresh_env_variables(cls):
1567 """
@@ -1611,6 +1572,7 @@ class rocm_aiter_ops:
1572 you can call this function to reload the env variables.
1573 """
1574 cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
1575 + cls._CUSTOM_ALL_REDUCE_ENABLED = envs.VLLM_ROCM_USE_AITER_CUSTOM_AR
1576 cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
1577 cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
1578 cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
@@ -1770,6 +1732,11 @@ class rocm_aiter_ops:
1732 def is_mha_enabled(cls) -> bool:
1733 return cls._AITER_ENABLED and cls._MHA_ENABLED
1734
1735 + @classmethod
1736 + @if_aiter_supported
1737 + def is_custom_all_reduce_enabled(cls) -> bool:
1738 + return cls._AITER_ENABLED and cls._CUSTOM_ALL_REDUCE_ENABLED
1739 +
1740 @classmethod
1741 @if_aiter_supported
1742 def is_shuffle_kv_cache_enabled(cls) -> bool:
@@ -1824,33 +1791,20 @@ class rocm_aiter_ops:
1791 return cls.is_linear_enabled() and on_gfx950()
1792
1793 @classmethod
1827 - def initialize_aiter_allreduce(
1828 - cls, group: ProcessGroup, device: torch.device
1829 - ) -> None:
1830 - try:
1831 - from aiter.dist.device_communicators.custom_all_reduce import (
1832 - CustomAllreduce as AiterCustomAllreduce,
1833 - )
1834 -
1835 - cls._CUSTOM_ALL_REDUCE = AiterCustomAllreduce(group, device)
1836 - except Exception:
1837 - cls._CUSTOM_ALL_REDUCE = None
1838 -
1839 - @classmethod
1840 - def get_aiter_allreduce(cls) -> AiterCustomAllreduceProto | None:
1841 - return cls._CUSTOM_ALL_REDUCE
1842 -
1843 - @classmethod
1844 - def destroy_aiter_allreduce(cls) -> None:
1845 - if cls._CUSTOM_ALL_REDUCE is not None:
1846 - cls._CUSTOM_ALL_REDUCE.close()
1847 - cls._CUSTOM_ALL_REDUCE = None
1794 + def get_aiter_allreduce(cls):
1795 + """Return the TP device communicator's AITER custom-allreduce if it has
1796 + one, return None otherwise
1797 + """
1798 + from vllm.distributed.device_communicators.aiter_custom_all_reduce import (
1799 + AiterCustomAllreduce,
1800 + )
1801 + from vllm.distributed.parallel_state import get_tp_group
1802
1849 - @classmethod
1850 - def get_aiter_allreduce_max_size(cls) -> int | None:
1851 - # effective max input size (based on upstream aiter version: v0.1.10.post3)
1852 - # https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/aiter/dist/device_communicators/custom_all_reduce.py#L272-L273
1853 - return int(cls._ALL_REDUCE_MAX_SIZE / 2)
1803 + device_comm = get_tp_group().device_communicator
1804 + aiter_ar_comm = getattr(device_comm, "aiter_ar_comm", None)
1805 + return (
1806 + aiter_ar_comm if isinstance(aiter_ar_comm, AiterCustomAllreduce) else None
1807 + )
1808
1809 @classmethod
1810 @if_aiter_supported
@@ -2165,21 +2119,6 @@ class rocm_aiter_ops:
2119 def get_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_op() -> OpOverload: # noqa: E501
2120 return torch.ops.vllm.rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm.default # noqa: E501
2121
2168 - # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823).
2169 - @classmethod
2170 - def has_fused_allreduce_rmsnorm_quant_per_group(cls) -> bool:
2171 - """True if the running AITER build exposes the per-group AR+RMS+quant
2172 - kernel (added in ROCm/aiter PR #2823).
2173 -
2174 - The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off
2175 - this so vLLM degrades to the AR+RMS-only fusion when run against an
2176 - older aiter that lacks the per-group launcher.
2177 - """
2178 - aiter_ar = cls.get_aiter_allreduce()
2179 - return aiter_ar is not None and hasattr(
2180 - aiter_ar, "fused_ar_rms_per_group_quant"
2181 - )
2182 -
2122 @staticmethod
2123 def get_fused_mla_dual_rms_norm_op() -> OpOverload:
2124 return torch.ops.vllm.fused_mla_dual_rms_norm.default
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+10 -39
@@ -19,7 +19,6 @@ from vllm.compilation.passes.fusion.rms_quant_fusion import (
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
22 from vllm.distributed.parallel_state import (
23 get_tensor_model_parallel_rank,
24 get_tensor_model_parallel_world_size,
@@ -1473,39 +1472,23 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1472 )
1473 return
1474
1476 - device_comm = get_tp_group().device_communicator
1477 - if device_comm is None:
1478 - logger.warning_once("Device communicator is required.")
1479 - return
1480 -
1481 - ca_comm = getattr(device_comm, "ca_comm", None)
1475 + ca_comm = rocm_aiter_ops.get_aiter_allreduce()
1476 if ca_comm is None:
1483 - logger.warning_once("Custom Allreduce is required.")
1477 + logger.warning_once(
1478 + "AITER allreduce fusions are disabled "
1479 + "because AITER Custom All Reduce is not enabled. "
1480 + "Set VLLM_ROCM_USE_AITER_CUSTOM_AR=1 "
1481 + "to enable it."
1482 + )
1483 return
1484 self.ca_comm = ca_comm
1485
1487 - assert isinstance(ca_comm, CustomAllreduce)
1488 -
1489 - group = get_tp_group().cpu_group
1490 - rocm_aiter_ops.initialize_aiter_allreduce(group, self.device)
1486 hidden_dim = config.model_config.get_hidden_size()
1487 element_size = torch.tensor([], dtype=self.model_dtype).element_size()
1493 - max_size = rocm_aiter_ops.get_aiter_allreduce_max_size()
1494 - if max_size is None:
1495 - logger.warning("AITER allreduce fusion must be initialized")
1496 - return
1497 -
1498 - # Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim.
1499 - # Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM
1500 - # and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12
1501 - # hidden_dim is a runtime argument. Detect the older API via the missing
1502 - # `_pool` attribute and skip fusion for unsupported sizes.
1503 - # Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590
1504 - aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
1488 + max_size = ca_comm.effective_max_size()
1489 _AITER_OLD_FUSED_AR_RMS_HIDDEN = (512, 1024, 2048, 4096)
1490 if (
1507 - aiter_ar is not None
1508 - and not hasattr(aiter_ar, "_pool")
1491 + not ca_comm.supports_dynamic_hidden_dim
1492 and hidden_dim not in _AITER_OLD_FUSED_AR_RMS_HIDDEN
1493 ):
1494 logger.warning_once(
@@ -1515,10 +1498,6 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1498 _AITER_OLD_FUSED_AR_RMS_HIDDEN,
1499 hidden_dim,
1500 )
1518 - # Tear down aiter's custom-allreduce so its IPC handles don't
1519 - # race with vllm's ca_comm on the unfused fallback path.
1520 - with contextlib.suppress(Exception):
1521 - rocm_aiter_ops.destroy_aiter_allreduce()
1501 return
1502
1503 max_token_num = max_size // (hidden_dim * element_size)
@@ -1532,9 +1511,7 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1511 # fall back to the AR+RMS-only fusion paired with PR #41825's
1512 # standalone RMS+quant fusion -- still correct, just leaves the
1513 # post-AR quant as a standalone kernel.
1535 - supports_per_group_quant = (
1536 - rocm_aiter_ops.has_fused_allreduce_rmsnorm_quant_per_group()
1537 - )
1514 + supports_per_group_quant = ca_comm.supports_per_group_quant
1515 if not supports_per_group_quant:
1516 logger.warning_once(
1517 "AITER AR+RMS+per-group-FP8-quant fusion disabled: aiter "
@@ -1609,9 +1586,3 @@ class RocmAiterAllReduceFusionPass(VllmFusionPatternMatcherPass):
1586 logger.debug(
1587 "%s Replaced %s patterns", self.__class__.__name__, self.matched_count
1588 )
1612 -
1613 - def __del__(self) -> None:
1614 - if getattr(self, "disabled", True):
1615 - return
1616 - with contextlib.suppress(Exception):
1617 - rocm_aiter_ops.destroy_aiter_allreduce()
vllm/config/vllm.py
+7 -2
@@ -1850,8 +1850,13 @@ class VllmConfig:
1850 tp_size = self.parallel_config.tensor_parallel_size
1851 from vllm._aiter_ops import rocm_aiter_ops
1852
1853 - if rocm_aiter_ops.is_enabled():
1854 - max_size = rocm_aiter_ops.get_aiter_allreduce_max_size()
1853 + max_size: int | None = None
1854 + if rocm_aiter_ops.is_custom_all_reduce_enabled():
1855 + from vllm.distributed.device_communicators.aiter_custom_all_reduce import ( # noqa: E501
1856 + AiterCustomAllreduce,
1857 + )
1858 +
1859 + max_size = AiterCustomAllreduce.effective_max_size()
1860 else:
1861 max_size = compilation_config.pass_config.flashinfer_max_size(tp_size)
1862 if max_size is not None and self.model_config is not None:
vllm/distributed/device_communicators/aiter_custom_all_reduce.py new
+95
@@ -0,0 +1,95 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""vLLM-owned wrapper over AITER's ``CustomAllreduce``.
4 +
5 +vLLM's ``CudaCommunicator`` stores one of these as ``aiter_ar_comm`` (when
6 +``VLLM_ROCM_USE_AITER_CUSTOM_AR`` is set) so the plain allreduce and
7 +the fused allreduce+RMSNorm path share a single AITER instance with its IPC buffers.
8 +
9 +"""
10 +
11 +import torch
12 +from torch.distributed import ProcessGroup
13 +
14 +from vllm.logger import init_logger
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +class AiterCustomAllreduce:
20 + # Default IPC buffer size for AITER's CustomAllreduce.
21 + MAX_SIZE: int = 8192 * 1024 * 8 * 2
22 +
23 + @classmethod
24 + def effective_max_size(cls) -> int:
25 + """
26 + Max input byte size eligible for AITER custom allreduce.
27 + """
28 + return cls.MAX_SIZE // 2
29 +
30 + def __init__(
31 + self,
32 + group: ProcessGroup,
33 + device: int | str | torch.device,
34 + max_size: int | None = None,
35 + ):
36 + from aiter.dist.device_communicators.custom_all_reduce import (
37 + CustomAllreduce as _AiterCustomAllreduce,
38 + )
39 +
40 + if max_size is None:
41 + max_size = self.MAX_SIZE
42 +
43 + self._impl = _AiterCustomAllreduce(group, device, max_size=max_size)
44 +
45 + @property
46 + def aiter_ca(self):
47 + return self._impl
48 +
49 + @property
50 + def disabled(self) -> bool:
51 + return self._impl.disabled
52 +
53 + def should_custom_ar(self, inp: torch.Tensor) -> bool:
54 + return self._impl.should_custom_ar(inp)
55 +
56 + def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None:
57 + return self._impl.custom_all_reduce(inp)
58 +
59 + def capture(self):
60 + return self._impl.capture()
61 +
62 + def close(self) -> None:
63 + self._impl.close()
64 +
65 + @property
66 + def supports_dynamic_hidden_dim(self) -> bool:
67 + """Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim.
68 + Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM
69 + and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12
70 + hidden_dim is a runtime argument. Older builds are detected via
71 + AiterCustomAllreduce.supports_dynamic_hidden_dim; This function is used to
72 + skip fusion for unsupported sizes on them.
73 + Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590
74 + """
75 + return hasattr(self._impl, "_pool")
76 +
77 + @staticmethod
78 + def build_supports_per_group_quant() -> bool:
79 + """True if the running AITER build exposes the per-group AR+RMS+quant
80 + kernel (added in ROCm/aiter PR #2823).
81 +
82 + The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off
83 + this so vLLM degrades to the AR+RMS-only fusion when run against an
84 + older aiter that lacks the per-group launcher.
85 + """
86 + from aiter.dist.device_communicators.custom_all_reduce import (
87 + CustomAllreduce as _AiterCustomAllreduce,
88 + )
89 +
90 + return hasattr(_AiterCustomAllreduce, "fused_ar_rms_per_group_quant")
91 +
92 + # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823).
93 + @property
94 + def supports_per_group_quant(self) -> bool:
95 + return self.build_supports_per_group_quant()
vllm/distributed/device_communicators/cuda_communicator.py
+40 -10
@@ -6,6 +6,7 @@ import torch
6 from torch.distributed import ProcessGroup
7
8 import vllm.envs as envs
9 +from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.distributed.device_communicators.all_reduce_utils import (
11 NCCL_SYMM_MEM_ALL_REDUCE_CONFIG,
12 should_nccl_symm_mem_ag_rs,
@@ -19,6 +20,7 @@ from vllm.logger import init_logger
20 from vllm.platforms import current_platform
21
22 from ..utils import StatelessProcessGroup
23 +from .aiter_custom_all_reduce import AiterCustomAllreduce
24 from .base_device_communicator import DeviceCommunicatorBase
25
26 logger = init_logger(__name__)
@@ -48,16 +50,21 @@ class CudaCommunicator(DeviceCommunicatorBase):
50 use_custom_allreduce = False
51 use_torch_symm_mem = False
52 use_flashinfer_allreduce = False
53 + use_aiter_allreduce = False
54 else:
55 from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE
56
57 use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
58 use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
59 use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER
60 + use_aiter_allreduce = use_custom_allreduce and bool(
61 + rocm_aiter_ops.is_custom_all_reduce_enabled()
62 + )
63
64 self.use_custom_allreduce = use_custom_allreduce
65 self.use_torch_symm_mem = use_torch_symm_mem
66 self.use_flashinfer_allreduce = use_flashinfer_allreduce
67 + self.use_aiter_allreduce = use_aiter_allreduce
68
69 # lazy import to avoid documentation build error
70 from vllm.distributed.device_communicators.custom_all_reduce import (
@@ -85,6 +92,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
92 self.qr_comm: QuickAllReduce | None = None
93 self.symm_mem_comm: SymmMemCommunicator | None = None
94 self.fi_ar_comm: FlashInferAllReduce | None = None
95 + self.aiter_ar_comm: AiterCustomAllreduce | None = None
96
97 if use_torch_symm_mem and current_platform.is_cuda():
98 self.symm_mem_comm = SymmMemCommunicator(
@@ -98,7 +106,13 @@ class CudaCommunicator(DeviceCommunicatorBase):
106 device=self.device,
107 )
108
101 - if use_custom_allreduce and self.world_size > 1:
109 + if self.use_aiter_allreduce and self.world_size > 1:
110 + self.aiter_ar_comm = AiterCustomAllreduce(
111 + group=self.cpu_group,
112 + device=self.device,
113 + )
114 +
115 + if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1:
116 # Initialize a custom fast all-reduce implementation.
117 self.ca_comm = CustomAllreduce(
118 group=self.cpu_group,
@@ -108,13 +122,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
122 ),
123 )
124
111 - if current_platform.is_rocm():
112 - # Initialize a custom quick all-reduce implementation for AMD.
113 - # Quick reduce is designed as a complement to custom allreduce.
114 - # Based on quickreduce (https://github.com/mk1-project/quickreduce).
115 - # If it's a rocm, 'use_custom_allreduce==True' means it must
116 - # currently be an MI300 series.
117 - self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device)
125 + if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm():
126 + # Initialize a custom quick all-reduce implementation for AMD.
127 + # Quick reduce is designed as a complement to custom allreduce
128 + # (vLLM's or AITER's), so it is initialized for either backend.
129 + # Based on quickreduce (https://github.com/mk1-project/quickreduce).
130 + # On ROCm, 'use_custom_allreduce==True' means it must currently be
131 + # an MI300 series.
132 + self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device)
133
134 if self.world_size > 1:
135 self._log_all_reduce_backend_selection()
@@ -203,6 +218,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
218 "NCCL_SYMM_MEM",
219 "QUICK_REDUCE",
220 "FLASHINFER",
221 + "AITER_CUSTOM",
222 "CUSTOM",
223 "SYMM_MEM",
224 "PYNCCL",
@@ -236,6 +252,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
252 enabled_ar_backends.append("QUICK_REDUCE")
253 if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
254 enabled_ar_backends.append("FLASHINFER")
255 + if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled:
256 + enabled_ar_backends.append("AITER_CUSTOM")
257 if self.ca_comm is not None and not self.ca_comm.disabled:
258 enabled_ar_backends.append("CUSTOM")
259 if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled:
@@ -261,8 +279,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
279 out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_)
280 if out is not None:
281 return out
264 - # always try quick reduce first, then flashinfer, then custom allreduce,
265 - # and then pynccl. (quick reduce just for ROCM MI3*)
282 + # always try quick reduce first, then flashinfer, then the AITER or vLLM
283 + # custom allreduce, and then pynccl. (quick reduce just for ROCM MI3*)
284 qr_comm = self.qr_comm
285 if (
286 qr_comm is not None
@@ -281,6 +299,15 @@ class CudaCommunicator(DeviceCommunicatorBase):
299 out = fi_ar_comm.all_reduce(input_)
300 assert out is not None
301 return out
302 + aiter_ar_comm = self.aiter_ar_comm
303 + if (
304 + aiter_ar_comm is not None
305 + and not aiter_ar_comm.disabled
306 + and aiter_ar_comm.should_custom_ar(input_)
307 + ):
308 + out = aiter_ar_comm.custom_all_reduce(input_)
309 + assert out is not None
310 + return out
311 ca_comm = self.ca_comm
312 if (
313 ca_comm is not None
@@ -509,6 +536,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
536 self.pynccl_comm = None
537 if self.ca_comm is not None:
538 self.ca_comm = None
539 + if self.aiter_ar_comm is not None:
540 + self.aiter_ar_comm.close()
541 + self.aiter_ar_comm = None
542 if self.fi_ar_comm is not None:
543 self.fi_ar_comm.destroy()
544 self.fi_ar_comm = None
vllm/envs.py
+6
@@ -119,6 +119,7 @@ if TYPE_CHECKING:
119 VLLM_USE_OINK_OPS: bool = False
120 VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: bool = True
121 VLLM_ROCM_USE_AITER: bool = False
122 + VLLM_ROCM_USE_AITER_CUSTOM_AR: bool = True
123 VLLM_ROCM_USE_AITER_PAGED_ATTN: bool = False
124 VLLM_ROCM_USE_AITER_LINEAR: bool = True
125 VLLM_ROCM_USE_AITER_LINEAR_HIPBMM: bool = False
@@ -1146,6 +1147,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1147 "VLLM_ROCM_USE_AITER": lambda: (
1148 os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1")
1149 ),
1150 + # Use AITER's CustomAllreduce as the custom-allreduce backend inside vLLM's
1151 + # CudaCommunicator on ROCm.
1152 + "VLLM_ROCM_USE_AITER_CUSTOM_AR": lambda: (
1153 + os.getenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1")
1154 + ),
1155 # Whether to use aiter paged attention.
1156 # By default is disabled.
1157 "VLLM_ROCM_USE_AITER_PAGED_ATTN": lambda: (