| 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 | """Max input byte size eligible for AITER custom allreduce.""" |
| 26 | return cls.MAX_SIZE // 2 |
| 27 | |
| 28 | def __init__( |
| 29 | self, |
| 30 | group: ProcessGroup, |
| 31 | device: int | str | torch.device, |
| 32 | max_size: int | None = None, |
| 33 | ): |
| 34 | from aiter.dist.device_communicators.custom_all_reduce import ( |
| 35 | CustomAllreduce as _AiterCustomAllreduce, |
| 36 | ) |
| 37 | |
| 38 | if max_size is None: |
| 39 | max_size = self.MAX_SIZE |
| 40 | |
| 41 | self._impl = _AiterCustomAllreduce(group, device, max_size=max_size) |
| 42 | |
| 43 | @property |
| 44 | def aiter_ca(self): |
| 45 | return self._impl |
| 46 | |
| 47 | @property |
| 48 | def disabled(self) -> bool: |
| 49 | return self._impl.disabled |
| 50 | |
| 51 | def should_custom_ar(self, inp: torch.Tensor) -> bool: |
| 52 | return self._impl.should_custom_ar(inp) |
| 53 | |
| 54 | def custom_all_reduce(self, inp: torch.Tensor) -> torch.Tensor | None: |
| 55 | return self._impl.custom_all_reduce(inp) |
| 56 | |
| 57 | def should_custom_ag(self, inp: torch.Tensor) -> bool: |
| 58 | return self._impl.should_custom_ag(inp) |
| 59 | |
| 60 | def should_custom_rs(self, inp: torch.Tensor, dim: int) -> bool: |
| 61 | return self._impl.should_custom_rs(inp, dim) |
| 62 | |
| 63 | def custom_all_gather(self, inp: torch.Tensor, dim: int = 0) -> torch.Tensor | None: |
| 64 | return self._impl.custom_all_gather(inp, dim=dim) |
| 65 | |
| 66 | def custom_reduce_scatter( |
| 67 | self, inp: torch.Tensor, out: torch.Tensor, dim: int = 0 |
| 68 | ) -> torch.Tensor | None: |
| 69 | return self._impl.custom_reduce_scatter(inp, out, dim=dim) |
| 70 | |
| 71 | def use_1stage_fused_ar_rms(self, inp: torch.Tensor) -> bool: |
| 72 | """Whether AITER's fused allreduce+RMSNorm runs as its one-stage kernel. |
| 73 | |
| 74 | Mirrors the launcher contract of aiter's ``fused_allreduce_rmsnorm`` |
| 75 | (csrc/include/custom_all_reduce.cuh): rows of 16-byte packs, at most |
| 76 | 1024 packs per row, at most 80 tokens, and the byte cap of the |
| 77 | one-stage custom allreduce for this TP size and topology. Outside it |
| 78 | the fused op runs the two-stage variant (cross-device reduce-scatter |
| 79 | + local norm), which is slower than an explicit ``all_reduce`` + norm, |
| 80 | so callers that can fall back should require this. Capture-static: |
| 81 | depends only on shape, dtype, TP size and topology. |
| 82 | """ |
| 83 | hidden_dim = inp.shape[-1] |
| 84 | # Token cap first: prefill-sized inputs leave here with one comparison. |
| 85 | if inp.numel() // hidden_dim > 80: |
| 86 | return False |
| 87 | if inp.dtype not in (torch.bfloat16, torch.float16): |
| 88 | return False |
| 89 | pack_size = 16 // inp.element_size() |
| 90 | if hidden_dim % pack_size != 0 or hidden_dim // pack_size > 1024: |
| 91 | return False |
| 92 | ca = self._impl |
| 93 | world_size = ca.world_size |
| 94 | if world_size == 2: |
| 95 | return True |
| 96 | if not ca.fully_connected: |
| 97 | return False |
| 98 | total_bytes = inp.numel() * inp.element_size() |
| 99 | if world_size <= 4: |
| 100 | return total_bytes < 256 * 1024 |
| 101 | if world_size <= 8: |
| 102 | return total_bytes < 128 * 1024 |
| 103 | return False |
| 104 | |
| 105 | def capture(self): |
| 106 | return self._impl.capture() |
| 107 | |
| 108 | def close(self) -> None: |
| 109 | self._impl.close() |
| 110 | |
| 111 | @property |
| 112 | def supports_dynamic_hidden_dim(self) -> bool: |
| 113 | """Aiter's fused_allreduce_rmsnorm kernel dispatches on hidden_dim. |
| 114 | Before aiter v0.1.12 the launcher was template-specialized on HIDDEN_DIM |
| 115 | and silently no-op'd for sizes outside {512, 1024, 2048, 4096}. From v0.1.12 |
| 116 | hidden_dim is a runtime argument. Older builds are detected via |
| 117 | AiterCustomAllreduce.supports_dynamic_hidden_dim; This function is used to |
| 118 | skip fusion for unsupported sizes on them. |
| 119 | Ref (old kernel): https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/csrc/include/custom_all_reduce.cuh#L2590 |
| 120 | """ |
| 121 | return hasattr(self._impl, "_pool") |
| 122 | |
| 123 | @staticmethod |
| 124 | def build_supports_per_group_quant() -> bool: |
| 125 | """True if the running AITER build exposes the per-group AR+RMS+quant |
| 126 | kernel (added in ROCm/aiter PR #2823). |
| 127 | |
| 128 | The pattern registration in ``RocmAiterAllReduceFusionPass`` keys off |
| 129 | this so vLLM degrades to the AR+RMS-only fusion when run against an |
| 130 | older aiter that lacks the per-group launcher. |
| 131 | """ |
| 132 | from aiter.dist.device_communicators.custom_all_reduce import ( |
| 133 | CustomAllreduce as _AiterCustomAllreduce, |
| 134 | ) |
| 135 | |
| 136 | return hasattr(_AiterCustomAllreduce, "fused_ar_rms_per_group_quant") |
| 137 | |
| 138 | # TODO(frida-andersson): drop once vLLM pins AITER >= 0.1.14 (ROCm/aiter#2823). |
| 139 | @property |
| 140 | def supports_per_group_quant(self) -> bool: |
| 141 | return self.build_supports_per_group_quant() |