@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ae48466cf3

[Perf][ROCm] Add AITER custom AG/RS (DP only) (#48247)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>

Simon Danielsson committed Sep 11, 2026 at 07:19 UTC ae48466cf33b4d7cfbe0f32370c425a3aa9fde61
8 files changed +552 -77
benchmarks/kernels/benchmark_device_communicators.py
+204 -51
@@ -29,6 +29,10 @@ import torch
29 import torch.distributed as dist
30 from torch.distributed import ProcessGroup
31
32 +from vllm._aiter_ops import rocm_aiter_ops
33 +from vllm.distributed.device_communicators.aiter_custom_all_reduce import (
34 + AiterCustomAllreduce,
35 +)
36 from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
37 from vllm.distributed.device_communicators.flashinfer_all_reduce import (
38 FlashInferAllReduce,
@@ -42,6 +46,7 @@ from vllm.distributed.device_communicators.pynccl_allocator import (
46 )
47 from vllm.distributed.device_communicators.symm_mem import SymmMemCommunicator
48 from vllm.logger import init_logger
49 +from vllm.platforms import current_platform
50 from vllm.utils.argparse_utils import FlexibleArgumentParser
51
52 logger = init_logger(__name__)
@@ -56,6 +61,30 @@ BENCHMARK_DTYPE = torch.bfloat16
61 # CUDA graph settings
62 CUDA_GRAPH_CAPTURE_CYCLES = 10
63
64 +# All-gather and reduce-scatter are distinct collectives from all-reduce, so
65 +# results are compared within a collective (each against its own NCCL/RCCL
66 +# baseline) rather than across collectives.
67 +COMM_COLLECTIVE = {
68 + "ca_1stage": "all-reduce",
69 + "ca_2stage": "all-reduce",
70 + "pynccl": "all-reduce",
71 + "pynccl-symm": "all-reduce",
72 + "symm_mem_multimem": "all-reduce",
73 + "symm_mem_two_shot": "all-reduce",
74 + "flashinfer_trtllm": "all-reduce",
75 + "flashinfer_mnnvl": "all-reduce",
76 + "aiter_ag": "all-gather",
77 + "pynccl_ag": "all-gather",
78 + "aiter_rs": "reduce-scatter",
79 + "pynccl_rs": "reduce-scatter",
80 +}
81 +COLLECTIVE_ORDER = ["all-reduce", "all-gather", "reduce-scatter"]
82 +COLLECTIVE_BASELINE = {
83 + "all-reduce": "pynccl",
84 + "all-gather": "pynccl_ag",
85 + "reduce-scatter": "pynccl_rs",
86 +}
87 +
88
89 class CommunicatorBenchmark:
90 """Benchmark class for testing device communicators."""
@@ -77,6 +106,11 @@ class CommunicatorBenchmark:
106 max_seq_len = max(sequence_lengths)
107 max_tensor_elements = max_seq_len * HIDDEN_SIZE
108 self.max_size_override = max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1
109 + # AITER has effective max size of half the max size, so we double it, and
110 + # account for AG output size (which is world_size larger)
111 + self.aiter_max_size = (
112 + 2 * world_size * max_tensor_elements * BENCHMARK_DTYPE.itemsize + 1
113 + )
114
115 # Initialize communicators
116 self.custom_allreduce = None
@@ -85,6 +119,7 @@ class CommunicatorBenchmark:
119 self.symm_mem_comm_multimem = None
120 self.symm_mem_comm_two_shot = None
121 self.fi_ar_comm = None
122 + self.aiter_comm = None
123
124 self._init_communicators()
125
@@ -181,6 +216,27 @@ class CommunicatorBenchmark:
216 )
217 self.fi_ar_comm = None
218
219 + if current_platform.is_rocm() and rocm_aiter_ops.is_custom_all_reduce_enabled():
220 + try:
221 + # Also used for AG/RS for DP
222 + self.aiter_comm = AiterCustomAllreduce(
223 + group=self.cpu_group,
224 + device=self.device,
225 + max_size=self.aiter_max_size,
226 + )
227 + if not self.aiter_comm.disabled:
228 + logger.info("Rank %s: AITER custom AG/RS initialized", self.rank)
229 + else:
230 + logger.info("Rank %s: AITER custom AG/RS disabled", self.rank)
231 + self.aiter_comm = None
232 + except Exception as e:
233 + logger.warning(
234 + "Rank %s: Failed to initialize AITER custom AG/RS: %s",
235 + self.rank,
236 + e,
237 + )
238 + self.aiter_comm = None
239 +
240 def benchmark_allreduce(
241 self, sequence_length: int, num_warmup: int, num_trials: int
242 ) -> dict[str, float]:
@@ -385,29 +441,116 @@ class CommunicatorBenchmark:
441 f"CUDA graph benchmark failed for communicator: {e}"
442 ) from e
443
444 + def benchmark_ag_rs(
445 + self, sequence_length: int, num_warmup: int, num_trials: int
446 + ) -> dict[str, float]:
447 + """Benchmark all-gather and reduce-scatter: custom (AITER) vs NCCL/RCCL.
448 +
449 + Each collective is measured for both the AITER custom kernel and the
450 + PyNccl baseline so they compare like-for-like (print_results groups
451 + columns by collective). All-gather maps ``(seq, hidden)`` ->
452 + ``(world_size * seq, hidden)``; reduce-scatter maps ``(seq, hidden)`` ->
453 + ``(seq // world_size, hidden)`` and so only runs when ``seq`` divides
454 + evenly across ranks. Both reuse ``benchmark_allreduce_single`` timed on a
455 + ``(seq, hidden)`` input; baselines write into a pre-allocated buffer.
456 + """
457 + results: dict[str, float] = {}
458 + entries = []
459 +
460 + # ---- all-gather: (seq, hidden) -> (world_size * seq, hidden) ----
461 + if self.aiter_comm is not None:
462 + c = self.aiter_comm
463 + entries.append(
464 + (
465 + "aiter_ag",
466 + lambda t, c=c: c.custom_all_gather(t, dim=0),
467 + lambda t, c=c: c.should_custom_ag(t),
468 + c.capture(),
469 + )
470 + )
471 + if self.pynccl_comm is not None:
472 + c = self.pynccl_comm
473 + ag_out = torch.empty(
474 + sequence_length * self.world_size,
475 + HIDDEN_SIZE,
476 + dtype=BENCHMARK_DTYPE,
477 + device=self.device,
478 + )
479 + entries.append(
480 + (
481 + "pynccl_ag",
482 + lambda t, c=c, o=ag_out: c.all_gather(o, t),
483 + lambda t: True,
484 + nullcontext(),
485 + )
486 + )
487 +
488 + # ---- reduce-scatter: (seq, hidden) -> (seq // world_size, hidden) ----
489 + if sequence_length % self.world_size == 0:
490 + rs_shape = (sequence_length // self.world_size, HIDDEN_SIZE)
491 + if self.aiter_comm is not None:
492 + c = self.aiter_comm
493 + rs_out = torch.empty(
494 + rs_shape, dtype=BENCHMARK_DTYPE, device=self.device
495 + )
496 + entries.append(
497 + (
498 + "aiter_rs",
499 + lambda t, c=c, o=rs_out: c.custom_reduce_scatter(t, o, dim=0),
500 + lambda t, c=c: c.should_custom_rs(t, dim=0),
501 + c.capture(),
502 + )
503 + )
504 + if self.pynccl_comm is not None:
505 + c = self.pynccl_comm
506 + rs_out = torch.empty(
507 + rs_shape, dtype=BENCHMARK_DTYPE, device=self.device
508 + )
509 + entries.append(
510 + (
511 + "pynccl_rs",
512 + lambda t, c=c, o=rs_out: c.reduce_scatter(o, t),
513 + lambda t: True,
514 + nullcontext(),
515 + )
516 + )
517
389 -def _calculate_speedup_info(comm_results: dict[str, float]) -> str:
390 - """Calculate speedup information for a single tensor size."""
518 + for name, fn, should_use_fn, context in entries:
519 + latency = self.benchmark_allreduce_single(
520 + sequence_length, fn, should_use_fn, context, num_warmup, num_trials
521 + )
522 + if latency is not None:
523 + results[name] = latency
524 +
525 + return results
526 +
527 +
528 +def _calculate_speedup_info(comm_results: dict[str, float], baseline: str) -> str:
529 + """Fastest comm in a collective group and its speedup vs ``baseline``."""
530 if not comm_results:
531 return "N/A"
532
394 - # Find the fastest communicator
533 fastest_comm = min(comm_results.keys(), key=lambda k: comm_results[k])
534 fastest_time = comm_results[fastest_comm]
535
398 - # Calculate speedup vs PyNccl if available
399 - if "pynccl" in comm_results:
400 - pynccl_time = comm_results["pynccl"]
401 - speedup = pynccl_time / fastest_time
536 + if baseline in comm_results and fastest_time > 0:
537 + speedup = comm_results[baseline] / fastest_time
538 return f"{fastest_comm} ({speedup:.2f}x)"
403 - else:
404 - return f"{fastest_comm} (N/A)"
539 + return f"{fastest_comm} (N/A)"
540
541
542 def print_results(
408 - results: dict[str, dict[str, float]], sequence_lengths: list[int], world_size: int
543 + results: dict[int, dict[str, float]], sequence_lengths: list[int], world_size: int
544 ):
410 - """Print benchmark results in a formatted table."""
545 + """Print benchmark results as one table per collective.
546 +
547 + All-reduce, all-gather and reduce-scatter are different collectives, so each
548 + gets its own table and its speedup is computed against that collective's
549 + NCCL/RCCL baseline rather than across collectives.
550 + """
551 + present: set[str] = set()
552 + for size_results in results.values():
553 + present.update(size_results.keys())
554
555 print(f"\n{'=' * 130}")
556 print("Device Communicator Benchmark Results")
@@ -415,53 +558,45 @@ def print_results(
558 f"World Size: {world_size}, Data Type: {BENCHMARK_DTYPE}, "
559 f"Hidden Size: {HIDDEN_SIZE}"
560 )
561 + print("Tensor Shape is the per-rank input; times are ms per operation.")
562 print(f"{'=' * 130}")
563
420 - # Get all communicator names
421 - all_comms = set()
422 - for size_results in results.values():
423 - all_comms.update(size_results.keys())
424 -
425 - all_comms = sorted(list(all_comms))
426 -
427 - # Print header
428 - header = f"{'Tensor Shape':<20}{'Tensor Size':<15}"
429 - for comm in all_comms:
430 - header += f"{comm:<20}"
431 - header += f"{'Best (Speedup vs PyNccl)':<30}"
432 - print(header)
433 - print("-" * len(header))
434 -
435 - # Print results for each sequence length
436 - for seq_len in sequence_lengths:
437 - if seq_len in results:
438 - # Calculate tensor size in elements and bytes
439 - tensor_elements = seq_len * HIDDEN_SIZE
440 - tensor_bytes = tensor_elements * BENCHMARK_DTYPE.itemsize
441 -
442 - # Format tensor size (MB)
443 - tensor_size_mb = tensor_bytes / (1024 * 1024)
444 - tensor_size_str = f"{tensor_size_mb:.2f} MB"
445 -
446 - # Format tensor shape
564 + for collective in COLLECTIVE_ORDER:
565 + comms = sorted(c for c in present if COMM_COLLECTIVE.get(c) == collective)
566 + if not comms:
567 + continue
568 + baseline = COLLECTIVE_BASELINE[collective]
569 +
570 + print(f"\n{collective} (speedup vs {baseline})")
571 + header = f"{'Tensor Shape':<20}{'Tensor Size':<15}"
572 + for comm in comms:
573 + header += f"{comm:<20}"
574 + header += f"{'Best':<30}"
575 + print(header)
576 + print("-" * len(header))
577 +
578 + for seq_len in sequence_lengths:
579 + if seq_len not in results:
580 + continue
581 + row_results = {
582 + c: results[seq_len][c] for c in comms if c in results[seq_len]
583 + }
584 +
585 + tensor_bytes = seq_len * HIDDEN_SIZE * BENCHMARK_DTYPE.itemsize
586 + tensor_size_str = f"{tensor_bytes / (1024 * 1024):.2f} MB"
587 tensor_shape = f"({seq_len}, {HIDDEN_SIZE})"
588
589 row = f"{tensor_shape:<20}{tensor_size_str:<15}"
450 - for comm in all_comms:
451 - if comm in results[seq_len]:
452 - row += f"{results[seq_len][comm]:<20.3f}"
590 + for comm in comms:
591 + if comm in row_results:
592 + row += f"{row_results[comm]:<20.3f}"
593 else:
594 row += f"{'N/A':<20}"
455 -
456 - # Calculate speedup information
457 - speedup_info = _calculate_speedup_info(results[seq_len])
458 - row += f"{speedup_info:<30}"
459 -
595 + row += f"{_calculate_speedup_info(row_results, baseline):<30}"
596 print(row)
597
462 - print(f"{'=' * 130}")
463 - print("All times are in milliseconds (ms) per allreduce operation")
464 - print("Speedup column shows: fastest_algorithm (speedup_vs_pynccl)")
598 + print(f"\n{'=' * 130}")
599 + print("All times are in milliseconds (ms) per operation")
600
601
602 def main():
@@ -526,6 +661,13 @@ def main():
661 num_warmup=args.num_warmup,
662 num_trials=args.num_trials,
663 )
664 + results.update(
665 + benchmark.benchmark_ag_rs(
666 + sequence_length=seq_len,
667 + num_warmup=args.num_warmup,
668 + num_trials=args.num_trials,
669 + )
670 + )
671
672 all_results[seq_len] = results
673
@@ -538,12 +680,23 @@ def main():
680
681 # Save to JSON if requested
682 if args.output_json:
541 - # Add speedup information to results
683 + # Add per-collective speedup information to results.
684 enhanced_results = {}
685 for seq_len, comm_results in all_results.items():
686 + speedup_info = {}
687 + for collective in COLLECTIVE_ORDER:
688 + group = {
689 + k: v
690 + for k, v in comm_results.items()
691 + if COMM_COLLECTIVE.get(k) == collective
692 + }
693 + if group:
694 + speedup_info[collective] = _calculate_speedup_info(
695 + group, COLLECTIVE_BASELINE[collective]
696 + )
697 enhanced_results[seq_len] = {
698 "timings": comm_results,
546 - "speedup_info": _calculate_speedup_info(comm_results),
699 + "speedup_info": speedup_info,
700 }
701
702 output_data = {
tests/distributed/test_comm_ops.py
+15
@@ -410,6 +410,21 @@ def test_flashinfer_all_reduce_precedes_nccl(monkeypatch: pytest.MonkeyPatch) ->
410 nccl_selector.assert_not_called()
411
412
413 +def test_aiter_all_gather_precedes_pynccl(monkeypatch: pytest.MonkeyPatch) -> None:
414 + """Ensure can use aiter all_gather if available even if pynccl is not available."""
415 + output = torch.empty(2)
416 + aiter_comm = Mock()
417 + aiter_comm.should_custom_ag.return_value = True
418 + aiter_comm.custom_all_gather.return_value = output
419 + communicator = CudaCommunicator.__new__(CudaCommunicator)
420 + communicator.world_size = 2
421 + communicator.aiter_ar_comm = aiter_comm
422 + communicator.pynccl_comm = Mock(disabled=True)
423 + monkeypatch.setattr(communicator, "_can_use_aiter_ag_rs", Mock(return_value=True))
424 +
425 + assert communicator.all_gatherv(torch.empty(1)) is output
426 +
427 +
428 def test_isend_object_posts_size_then_object_and_releases_on_wait(
429 monkeypatch: pytest.MonkeyPatch,
430 ) -> None:
tests/distributed/test_rocm_aiter_custom_ar.py
+171 -1
@@ -8,7 +8,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
11 +from vllm.distributed.parallel_state import get_dp_group, get_tp_group, graph_capture
12 from vllm.envs import disable_envs_cache
13 from vllm.platforms import current_platform
14
@@ -132,3 +132,173 @@ def test_rocm_aiter_custom_allreduce(
132 test_target,
133 ):
134 multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
135 +
136 +
137 +def _get_aiter_ag_rs_comm():
138 + device_communicator = get_dp_group().device_communicator
139 + assert device_communicator.use_aiter_ag_rs, (
140 + "AITER custom AG/RS was not enabled on the DP group."
141 + )
142 + aiter_comm = device_communicator.aiter_ar_comm
143 + assert aiter_comm is not None, "AITER custom AG/RS was not initialized."
144 + assert not aiter_comm.disabled, "AITER custom AG/RS is disabled."
145 + return aiter_comm
146 +
147 +
148 +def _assert_aiter_handles_ag(aiter_comm, inp: torch.Tensor) -> None:
149 + assert aiter_comm.should_custom_ag(inp), (
150 + f"AITER custom all-gather does not support input shape {inp.shape}."
151 + )
152 +
153 +
154 +def _assert_aiter_handles_rs(aiter_comm, inp: torch.Tensor) -> None:
155 + assert aiter_comm.should_custom_rs(inp, dim=0), (
156 + f"AITER custom reduce-scatter does not support input shape {inp.shape}."
157 + )
158 +
159 +
160 +@ray.remote(num_gpus=1, max_calls=1)
161 +def eager_ag_rs(
162 + monkeypatch: pytest.MonkeyPatch,
163 + tp_size,
164 + pp_size,
165 + rank,
166 + distributed_init_port,
167 + data_parallel_size,
168 + data_parallel_master_port,
169 +) -> None:
170 + with monkeypatch.context() as m:
171 + _configure_aiter_custom_ar_env(m)
172 +
173 + device = torch.device(f"cuda:{rank}")
174 + torch.accelerator.set_device_index(device)
175 + init_test_distributed_environment(
176 + tp_size,
177 + pp_size,
178 + rank,
179 + distributed_init_port,
180 + data_parallel_size=data_parallel_size,
181 + data_parallel_master_port=data_parallel_master_port,
182 + )
183 +
184 + dp_group = get_dp_group()
185 + group = dp_group.device_group
186 + dp_world = dp_group.world_size
187 + aiter_comm = _get_aiter_ag_rs_comm()
188 +
189 + for shape, dtype in test_cases:
190 + num_tokens, hidden = shape
191 +
192 + # all-gather: each rank contributes (num_tokens, hidden).
193 + inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1)
194 + _assert_aiter_handles_ag(aiter_comm, inp)
195 + expected = torch.empty(
196 + (num_tokens * dp_world, hidden), dtype=dtype, device=device
197 + )
198 + dist.all_gather_into_tensor(expected, inp, group=group)
199 + out = aiter_comm.custom_all_gather(inp, dim=0)
200 + assert out is not None
201 + torch.testing.assert_close(out, expected)
202 +
203 + # reduce-scatter: each rank contributes (num_tokens * dp, hidden).
204 + rs_in = torch.ones(
205 + (num_tokens * dp_world, hidden), dtype=dtype, device=device
206 + ) * (rank + 1)
207 + _assert_aiter_handles_rs(aiter_comm, rs_in)
208 + rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
209 + dist.reduce_scatter_tensor(rs_expected, rs_in, group=group)
210 + rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
211 + aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0)
212 + torch.testing.assert_close(rs_out, rs_expected)
213 +
214 +
215 +@ray.remote(num_gpus=1, max_calls=1)
216 +def graph_ag_rs(
217 + monkeypatch: pytest.MonkeyPatch,
218 + tp_size,
219 + pp_size,
220 + rank,
221 + distributed_init_port,
222 + data_parallel_size,
223 + data_parallel_master_port,
224 +) -> None:
225 + with monkeypatch.context() as m:
226 + _configure_aiter_custom_ar_env(m)
227 +
228 + device = torch.device(f"cuda:{rank}")
229 + torch.accelerator.set_device_index(device)
230 + init_test_distributed_environment(
231 + tp_size,
232 + pp_size,
233 + rank,
234 + distributed_init_port,
235 + data_parallel_size=data_parallel_size,
236 + data_parallel_master_port=data_parallel_master_port,
237 + )
238 +
239 + dp_group = get_dp_group()
240 + group = dp_group.device_group
241 + dp_world = dp_group.world_size
242 + aiter_comm = _get_aiter_ag_rs_comm()
243 +
244 + # Warmup so DP comms is initialized before graph capture
245 + data = torch.zeros(1, device=device)
246 + dist.all_reduce(data, group=group)
247 + torch.accelerator.synchronize()
248 + del data
249 +
250 + for shape, dtype in test_cases:
251 + num_tokens, hidden = shape
252 +
253 + # all-gather under graph capture.
254 + inp = torch.ones(shape, dtype=dtype, device=device) * (rank + 1)
255 + _assert_aiter_handles_ag(aiter_comm, inp)
256 + ag_expected = torch.empty(
257 + (num_tokens * dp_world, hidden), dtype=dtype, device=device
258 + )
259 + dist.all_gather_into_tensor(ag_expected, inp, group=group)
260 + with graph_capture(device=device) as graph_capture_context:
261 + torch.accelerator.synchronize()
262 + graph = torch.cuda.CUDAGraph()
263 + with torch.cuda.graph(graph, stream=graph_capture_context.stream):
264 + ag_out = aiter_comm.custom_all_gather(inp, dim=0)
265 + graph.replay()
266 + torch.testing.assert_close(ag_out, ag_expected)
267 +
268 + # reduce-scatter under graph capture.
269 + rs_in = torch.ones(
270 + (num_tokens * dp_world, hidden), dtype=dtype, device=device
271 + ) * (rank + 1)
272 + _assert_aiter_handles_rs(aiter_comm, rs_in)
273 + rs_expected = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
274 + dist.reduce_scatter_tensor(rs_expected, rs_in, group=group)
275 + rs_out = torch.empty((num_tokens, hidden), dtype=dtype, device=device)
276 + with graph_capture(device=device) as graph_capture_context:
277 + torch.accelerator.synchronize()
278 + graph = torch.cuda.CUDAGraph()
279 + with torch.cuda.graph(graph, stream=graph_capture_context.stream):
280 + aiter_comm.custom_reduce_scatter(rs_in, rs_out, dim=0)
281 + graph.replay()
282 + torch.testing.assert_close(rs_out, rs_expected)
283 +
284 +
285 +@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed")
286 +@multi_gpu_test(num_gpus=2)
287 +@pytest.mark.parametrize("tp_size", [1])
288 +@pytest.mark.parametrize("pipeline_parallel_size", [1])
289 +@pytest.mark.parametrize("data_parallel_size", [2])
290 +@pytest.mark.parametrize("test_target", [eager_ag_rs, graph_ag_rs])
291 +def test_rocm_aiter_custom_ag_rs(
292 + monkeypatch: pytest.MonkeyPatch,
293 + tp_size,
294 + pipeline_parallel_size,
295 + data_parallel_size,
296 + test_target,
297 +):
298 + multi_process_parallel(
299 + monkeypatch,
300 + tp_size,
301 + pipeline_parallel_size,
302 + test_target,
303 + data_parallel_size=data_parallel_size,
304 + )
tests/utils.py
+52 -10
@@ -1416,6 +1416,8 @@ def init_test_distributed_environment(
1416 rank: int,
1417 distributed_init_port: str,
1418 local_rank: int = -1,
1419 + data_parallel_size: int = 1,
1420 + data_parallel_master_port: int | None = None,
1421 ) -> None:
1422 # Note: This function is often called from Ray worker processes, so we
1423 # can't rely on pytest fixtures to set the config. We check if the config
@@ -1428,6 +1430,29 @@ def init_test_distributed_environment(
1430
1431 distributed_init_method = f"tcp://localhost:{distributed_init_port}"
1432
1433 + if data_parallel_size > 1:
1434 + # For DP we need to set a common DP master port
1435 + from vllm.config.parallel import ParallelConfig
1436 +
1437 + assert data_parallel_master_port is not None, (
1438 + "data_parallel_master_port is required when data_parallel_size > 1"
1439 + )
1440 + tp_pp_world = tp_size * pp_size
1441 + parallel_config = ParallelConfig(
1442 + data_parallel_size=data_parallel_size,
1443 + data_parallel_rank=rank // tp_pp_world,
1444 + _data_parallel_master_port_list=[int(data_parallel_master_port)],
1445 + )
1446 + with set_current_vllm_config(VllmConfig(parallel_config=parallel_config)):
1447 + init_distributed_environment(
1448 + world_size=tp_pp_world,
1449 + rank=rank % tp_pp_world,
1450 + distributed_init_method=distributed_init_method,
1451 + local_rank=local_rank if local_rank >= 0 else rank,
1452 + )
1453 + ensure_model_parallel_initialized(tp_size, pp_size)
1454 + return
1455 +
1456 if get_current_vllm_config_or_none() is not None:
1457 # Config already set, use it directly
1458 init_distributed_environment(
@@ -1454,6 +1479,7 @@ def multi_process_parallel(
1479 tp_size: int,
1480 pp_size: int,
1481 test_target: Any,
1482 + data_parallel_size: int = 1,
1483 ) -> None:
1484 import ray
1485
@@ -1474,18 +1500,34 @@ def multi_process_parallel(
1500 )
1501
1502 distributed_init_port = get_open_port()
1503 + # Separate port for the DP master group; only used when data_parallel_size > 1.
1504 + data_parallel_master_port = get_open_port() if data_parallel_size > 1 else None
1505 + world_size = data_parallel_size * tp_size * pp_size
1506 try:
1507 refs = []
1479 - for rank in range(tp_size * pp_size):
1480 - refs.append(
1481 - test_target.remote(
1482 - monkeypatch,
1483 - tp_size,
1484 - pp_size,
1485 - rank,
1486 - distributed_init_port,
1487 - ),
1488 - )
1508 + for rank in range(world_size):
1509 + if data_parallel_size > 1:
1510 + refs.append(
1511 + test_target.remote(
1512 + monkeypatch,
1513 + tp_size,
1514 + pp_size,
1515 + rank,
1516 + distributed_init_port,
1517 + data_parallel_size,
1518 + data_parallel_master_port,
1519 + ),
1520 + )
1521 + else:
1522 + refs.append(
1523 + test_target.remote(
1524 + monkeypatch,
1525 + tp_size,
1526 + pp_size,
1527 + rank,
1528 + distributed_init_port,
1529 + ),
1530 + )
1531 ray.get(refs)
1532 finally:
1533 ray.shutdown()
vllm/distributed/device_communicators/aiter_custom_all_reduce.py
+14
@@ -56,6 +56,20 @@ class AiterCustomAllreduce:
56 def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None:
57 return self._impl.custom_all_reduce(inp)
58
59 + def should_custom_ag(self, inp: torch.Tensor) -> bool:
60 + return self._impl.should_custom_ag(inp)
61 +
62 + def should_custom_rs(self, inp: torch.Tensor, dim: int) -> bool:
63 + return self._impl.should_custom_rs(inp, dim)
64 +
65 + def custom_all_gather(self, inp: torch.Tensor, dim: int = 0) -> torch.Tensor | None:
66 + return self._impl.custom_all_gather(inp, dim=dim)
67 +
68 + def custom_reduce_scatter(
69 + self, inp: torch.Tensor, out: torch.Tensor, dim: int = 0
70 + ) -> torch.Tensor | None:
71 + return self._impl.custom_reduce_scatter(inp, out, dim=dim)
72 +
73 def use_1stage_fused_ar_rms(self, inp: torch.Tensor) -> bool:
74 """Whether AITER's fused allreduce+RMSNorm runs as its one-stage kernel.
75
vllm/distributed/device_communicators/cuda_communicator.py
+83 -4
@@ -109,6 +109,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
109 self.fi_ar_comm: FlashInferAllReduce | None = None
110 self.fi_pcie_ipc_ar_comm: FlashInferPcieIpcAllReduce | None = None
111 self.aiter_ar_comm: AiterCustomAllreduce | None = None
112 + self.use_aiter_ag_rs: bool = False
113
114 if use_torch_symm_mem and current_platform.is_cuda():
115 self.symm_mem_comm = SymmMemCommunicator(
@@ -149,6 +150,22 @@ class CudaCommunicator(DeviceCommunicatorBase):
150 ),
151 )
152
153 + # AITER custom all-gather/reduce-scatter DP-attention dispatch/combine
154 + if (
155 + "dp" in unique_name
156 + and self.world_size in (2, 4, 8)
157 + and current_platform.is_rocm()
158 + and rocm_aiter_ops.is_custom_all_reduce_enabled()
159 + ):
160 + self.aiter_ar_comm = AiterCustomAllreduce(
161 + group=self.cpu_group,
162 + device=self.device,
163 + )
164 + if self.aiter_ar_comm.disabled:
165 + self.aiter_ar_comm = None
166 + else:
167 + self.use_aiter_ag_rs = True
168 +
169 if use_custom_allreduce and self.world_size > 1 and current_platform.is_rocm():
170 # Initialize a custom quick all-reduce implementation for AMD.
171 # Quick reduce is designed as a complement to custom allreduce
@@ -285,7 +302,11 @@ class CudaCommunicator(DeviceCommunicatorBase):
302 enabled_ar_backends.append("NCCL_SYMM_MEM")
303 if self.qr_comm is not None and not self.qr_comm.disabled:
304 enabled_ar_backends.append("QUICK_REDUCE")
288 - if self.aiter_ar_comm is not None and not self.aiter_ar_comm.disabled:
305 + if (
306 + self.use_aiter_allreduce
307 + and self.aiter_ar_comm is not None
308 + and not self.aiter_ar_comm.disabled
309 + ):
310 enabled_ar_backends.append("AITER_CUSTOM")
311 if self.ca_comm is not None and not self.ca_comm.disabled:
312 enabled_ar_backends.append("CUSTOM")
@@ -340,7 +361,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
361 return out
362 aiter_ar_comm = self.aiter_ar_comm
363 if (
343 - aiter_ar_comm is not None
364 + self.use_aiter_allreduce
365 + and aiter_ar_comm is not None
366 and not aiter_ar_comm.disabled
367 and aiter_ar_comm.should_custom_ar(input_)
368 ):
@@ -462,6 +484,11 @@ class CudaCommunicator(DeviceCommunicatorBase):
484 # Convert negative dim to positive.
485 dim += input_.dim()
486
487 + # 'sizes' is not needed if all inputs in the same group have the same
488 + # shape
489 + if sizes is not None and all(s == sizes[0] for s in sizes):
490 + sizes = None
491 +
492 # Note: This will produce an incorrect answer if we don't make
493 # the input_tensor contiguous. Possible bug in reduce_scatter_tensor?
494 input_tensor = input_.movedim(0, dim).contiguous()
@@ -475,6 +502,16 @@ class CudaCommunicator(DeviceCommunicatorBase):
502 chunk_size = input_tensor.shape[0] // world_size
503 output_shape = (chunk_size,) + input_tensor.shape[1:]
504
505 + if self._can_use_aiter_ag_rs(sizes):
506 + aiter_comm = self.aiter_ar_comm
507 + assert aiter_comm is not None
508 + if aiter_comm.should_custom_rs(input_tensor, dim=0):
509 + output = torch.empty(
510 + output_shape, dtype=input_tensor.dtype, device=input_tensor.device
511 + )
512 + aiter_comm.custom_reduce_scatter(input_tensor, output, dim=0)
513 + return output.movedim(0, dim).contiguous()
514 +
515 # Symmetric memory is only used when all ranks have uniform sizes.
516 # ncclCommWindowRegister is collective: asymmetric pool allocations
517 # from variable per-rank sizes cause deadlocks.
@@ -632,6 +669,34 @@ class CudaCommunicator(DeviceCommunicatorBase):
669 self.all2all_manager.destroy()
670 self.all2all_manager = None # type: ignore[assignment]
671
672 + def _can_use_aiter_ag_rs(self, sizes: list[int] | None) -> bool:
673 + """Whether the AITER custom AG/RS fast path may run for this collective.
674 +
675 + Requires:
676 + - uniform batches
677 + - FULL CUDAgraphs
678 + """
679 + if (
680 + not self.use_aiter_ag_rs
681 + or self.aiter_ar_comm is None
682 + or self.aiter_ar_comm.disabled
683 + ):
684 + return False
685 + if sizes is not None:
686 + return False
687 +
688 + from vllm.config.compilation import CUDAGraphMode
689 + from vllm.forward_context import get_forward_context
690 +
691 + try:
692 + ctx = get_forward_context()
693 + except AssertionError:
694 + return False
695 + if ctx.cudagraph_runtime_mode != CUDAGraphMode.FULL:
696 + return False
697 + bd = ctx.batch_descriptor
698 + return bd is not None and bd.uniform
699 +
700 def suspend(self) -> None:
701 if self.pynccl_comm is not None:
702 self.pynccl_comm.suspend()
@@ -665,14 +730,28 @@ class CudaCommunicator(DeviceCommunicatorBase):
730 if dim != 0:
731 raise NotImplementedError("only dim 0 all-gatherv is supported")
732 world_size = self.world_size
668 - pynccl_comm = self.pynccl_comm
669 - assert pynccl_comm is not None and not pynccl_comm.disabled
733
734 # 'sizes' is not needed if all inputs in the same group have the same
735 # shape
736 if sizes is not None and all(s == sizes[0] for s in sizes):
737 sizes = None
738
739 + if self._can_use_aiter_ag_rs(sizes):
740 + aiter_comm = self.aiter_ar_comm
741 + assert aiter_comm is not None
742 + if isinstance(input_, torch.Tensor):
743 + if aiter_comm.should_custom_ag(input_):
744 + out = aiter_comm.custom_all_gather(input_, dim=0)
745 + if out is not None:
746 + return out
747 + elif all(aiter_comm.should_custom_ag(inp) for inp in input_):
748 + outs = [aiter_comm.custom_all_gather(inp, dim=0) for inp in input_]
749 + if all(o is not None for o in outs):
750 + return outs
751 +
752 + pynccl_comm = self.pynccl_comm
753 + assert pynccl_comm is not None and not pynccl_comm.disabled
754 +
755 # Symmetric memory is only used when all ranks have uniform sizes.
756 # ncclCommWindowRegister is collective: asymmetric pool allocations
757 # from variable per-rank sizes cause deadlocks.
vllm/distributed/parallel_state.py
+12 -10
@@ -673,11 +673,11 @@ class GroupCoordinator:
673 else:
674 stream = graph_capture_context.stream
675
676 - # only cuda uses this function,
676 + # only cuda/rocm uses this function,
677 # so we don't abstract it into the base class
678 maybe_ca_context = nullcontext()
679 maybe_fi_pcie_ipc_context: AbstractContextManager[Any] = nullcontext()
680 - maybe_aiter_context = nullcontext()
680 + maybe_aiter_ar_context = nullcontext()
681 from vllm.distributed.device_communicators.cuda_communicator import (
682 CudaCommunicator,
683 )
@@ -698,12 +698,10 @@ class GroupCoordinator:
698 if fi_pcie_ipc_ar_comm is not None:
699 maybe_fi_pcie_ipc_context = fi_pcie_ipc_ar_comm.capture()
700
701 - from vllm._aiter_ops import rocm_aiter_ops
702 -
703 - if rocm_aiter_ops.is_enabled():
704 - aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
705 - if aiter_ar is not None:
706 - maybe_aiter_context = aiter_ar.capture() # type: ignore
701 + # Capture each group's own comm. A global lookup would double-capture
702 + aiter_ar_comm = getattr(self.device_communicator, "aiter_ar_comm", None)
703 + if aiter_ar_comm is not None:
704 + maybe_aiter_ar_context = aiter_ar_comm.capture() # type: ignore
705
706 # ensure all initialization operations complete before attempting to
707 # capture the graph on another stream
@@ -715,7 +713,7 @@ class GroupCoordinator:
713 torch.cuda.stream(stream),
714 maybe_ca_context,
715 maybe_fi_pcie_ipc_context,
718 - maybe_aiter_context,
716 + maybe_aiter_ar_context,
717 ):
718 yield graph_capture_context
719
@@ -1640,7 +1638,11 @@ def graph_capture(
1638 context = graph_capture_context or GraphCaptureContext(
1639 torch.cuda.Stream(device=device)
1640 )
1643 - with get_tp_group().graph_capture(context), get_pp_group().graph_capture(context):
1641 + with (
1642 + get_tp_group().graph_capture(context),
1643 + get_pp_group().graph_capture(context),
1644 + get_dp_group().graph_capture(context),
1645 + ):
1646 yield context
1647
1648
vllm/envs.py
+1 -1
@@ -1261,7 +1261,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
1261 os.getenv("VLLM_ROCM_USE_AITER", "False").lower() in ("true", "1")
1262 ),
1263 # Use AITER's CustomAllreduce as the custom-allreduce backend inside vLLM's
1264 - # CudaCommunicator on ROCm.
1264 + # CudaCommunicator on ROCm. Also enables AITER AG/RS for DP communication.
1265 "VLLM_ROCM_USE_AITER_CUSTOM_AR": lambda: (
1266 os.getenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "True").lower() in ("true", "1")
1267 ),