| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | import threading |
| 5 | from collections.abc import Iterator |
| 6 | from contextlib import contextmanager |
| 7 | from dataclasses import dataclass |
| 8 | from typing import Any |
| 9 | |
| 10 | import torch |
| 11 | import torch.distributed as dist |
| 12 | |
| 13 | import vllm.envs as envs |
| 14 | from vllm.config import get_current_vllm_config |
| 15 | from vllm.distributed import get_dp_group, get_ep_group, get_pcp_group |
| 16 | from vllm.distributed.utils import StatelessProcessGroup |
| 17 | from vllm.forward_context import get_forward_context |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.utils.flashinfer import ( |
| 20 | has_flashinfer_nvlink_one_sided, |
| 21 | has_flashinfer_nvlink_two_sided, |
| 22 | ) |
| 23 | from vllm.utils.func_utils import supports_kw |
| 24 | from vllm.utils.import_utils import ( |
| 25 | check_moonep_system_support, |
| 26 | has_deep_ep, |
| 27 | has_deep_ep_v2, |
| 28 | has_mori, |
| 29 | ) |
| 30 | |
| 31 | from .base_device_communicator import All2AllManagerBase, Cache |
| 32 | |
| 33 | if has_flashinfer_nvlink_two_sided(): |
| 34 | from flashinfer.comm import Mapping # type: ignore[import-not-found] |
| 35 | from flashinfer.comm.mnnvl import MnnvlConfig # type: ignore[import-not-found] |
| 36 | from flashinfer.comm.trtllm_alltoall import ( |
| 37 | MnnvlMoe, # type: ignore[import-not-found] |
| 38 | ) |
| 39 | |
| 40 | if has_flashinfer_nvlink_one_sided(): |
| 41 | from flashinfer.comm import Mapping # type: ignore[import-not-found] |
| 42 | from flashinfer.comm.mnnvl import MnnvlConfig # type: ignore[import-not-found] |
| 43 | from flashinfer.comm.trtllm_moe_alltoall import ( |
| 44 | MoeAlltoAll, # type: ignore[import-not-found] |
| 45 | moe_a2a_get_workspace_size_per_rank, |
| 46 | ) |
| 47 | |
| 48 | |
| 49 | logger = init_logger(__name__) |
| 50 | |
| 51 | |
| 52 | class AgRsAll2AllManager(All2AllManagerBase): |
| 53 | """An implementation of all2all communication based on |
| 54 | all-gather (dispatch) and reduce-scatter (combine). |
| 55 | """ |
| 56 | |
| 57 | def __init__(self, cpu_group, tcp_store_group=None): |
| 58 | super().__init__(cpu_group, tcp_store_group) |
| 59 | self.use_ep = get_current_vllm_config().parallel_config.enable_expert_parallel |
| 60 | |
| 61 | def _get_comm_group(self, is_sequence_parallel: bool) -> Any: |
| 62 | if is_sequence_parallel: |
| 63 | return get_ep_group() |
| 64 | if self.dp_world_size > 1: |
| 65 | if self.use_ep and get_pcp_group().world_size > 1: |
| 66 | assert self.tp_group.world_size == 1, ( |
| 67 | "DP+PCP with TP>1 requires sequence-parallel MoE inputs" |
| 68 | ) |
| 69 | return get_ep_group() |
| 70 | return get_dp_group() |
| 71 | return get_pcp_group() |
| 72 | |
| 73 | def _get_sizes(self, num_local_tokens: int, comm_group: Any) -> list[int]: |
| 74 | if self.dp_world_size == 1: |
| 75 | return [num_local_tokens] * comm_group.world_size |
| 76 | |
| 77 | dp_metadata = get_forward_context().dp_metadata |
| 78 | assert dp_metadata is not None |
| 79 | sizes = dp_metadata.get_chunk_sizes_across_dp_rank() |
| 80 | assert sizes is not None |
| 81 | assert len(sizes) == comm_group.world_size |
| 82 | return sizes |
| 83 | |
| 84 | def dispatch_router_logits( |
| 85 | self, |
| 86 | hidden_states: torch.Tensor, |
| 87 | router_logits: torch.Tensor, |
| 88 | is_sequence_parallel: bool = False, |
| 89 | extra_tensors: list[torch.Tensor] | None = None, |
| 90 | ) -> ( |
| 91 | tuple[torch.Tensor, torch.Tensor] |
| 92 | | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 93 | ): |
| 94 | """Gather hidden_states and router_logits from all dp ranks.""" |
| 95 | dist_group = self._get_comm_group(is_sequence_parallel) |
| 96 | sizes = self._get_sizes(hidden_states.shape[0], dist_group) |
| 97 | assert sizes[dist_group.rank_in_group] == hidden_states.shape[0] |
| 98 | |
| 99 | tensors_to_gather = [hidden_states, router_logits] |
| 100 | if extra_tensors is not None: |
| 101 | tensors_to_gather.extend(extra_tensors) |
| 102 | |
| 103 | gathered_tensors = dist_group.all_gatherv( |
| 104 | tensors_to_gather, |
| 105 | dim=0, |
| 106 | sizes=sizes, |
| 107 | ) |
| 108 | |
| 109 | if extra_tensors is not None: |
| 110 | return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:]) |
| 111 | return gathered_tensors[0], gathered_tensors[1] |
| 112 | |
| 113 | def dispatch( |
| 114 | self, |
| 115 | hidden_states: torch.Tensor, |
| 116 | topk_weights: torch.Tensor, |
| 117 | topk_ids: torch.Tensor, |
| 118 | is_sequence_parallel: bool = False, |
| 119 | extra_tensors: list[torch.Tensor] | None = None, |
| 120 | ) -> ( |
| 121 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 122 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 123 | ): |
| 124 | """Gather hidden_states and router_logits from all dp ranks.""" |
| 125 | dist_group = self._get_comm_group(is_sequence_parallel) |
| 126 | sizes = self._get_sizes(hidden_states.shape[0], dist_group) |
| 127 | assert sizes[dist_group.rank_in_group] == hidden_states.shape[0] |
| 128 | |
| 129 | tensors_to_gather = [hidden_states, topk_weights, topk_ids] |
| 130 | if extra_tensors is not None: |
| 131 | tensors_to_gather.extend(extra_tensors) |
| 132 | |
| 133 | gathered_tensors = dist_group.all_gatherv( |
| 134 | tensors_to_gather, |
| 135 | dim=0, |
| 136 | sizes=sizes, |
| 137 | ) |
| 138 | |
| 139 | hidden_states = gathered_tensors[0] |
| 140 | topk_weights = gathered_tensors[1] |
| 141 | topk_ids = gathered_tensors[2] |
| 142 | |
| 143 | if extra_tensors is None: |
| 144 | return hidden_states, topk_weights, topk_ids |
| 145 | |
| 146 | return hidden_states, topk_weights, topk_ids, gathered_tensors[3:] |
| 147 | |
| 148 | def combine( |
| 149 | self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False |
| 150 | ) -> torch.Tensor: |
| 151 | """Reduce-scatter hidden_states across all dp ranks.""" |
| 152 | dist_group = self._get_comm_group(is_sequence_parallel) |
| 153 | sizes = self._get_sizes( |
| 154 | hidden_states.shape[0] // dist_group.world_size, |
| 155 | dist_group, |
| 156 | ) |
| 157 | hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes) |
| 158 | return hidden_states |
| 159 | |
| 160 | def destroy(self): |
| 161 | pass |
| 162 | |
| 163 | |
| 164 | class DeepEPAll2AllManagerBase(All2AllManagerBase): |
| 165 | """All2All communication based on DeepEP High-Throughput kernels.""" |
| 166 | |
| 167 | def __init__(self, cpu_group, tcp_store_group=None): |
| 168 | assert has_deep_ep(), ( |
| 169 | "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md" |
| 170 | " to install DeepEP kernels." |
| 171 | ) # noqa |
| 172 | super().__init__(cpu_group, tcp_store_group) |
| 173 | self.handle_cache = Cache() |
| 174 | |
| 175 | # This is the DeepEP default. Stick to it till we can establish |
| 176 | # reasonable defaults based on profiling. |
| 177 | self.num_sms = 20 |
| 178 | |
| 179 | def get_handle(self, kwargs): |
| 180 | raise NotImplementedError |
| 181 | |
| 182 | def dispatch_router_logits( |
| 183 | self, |
| 184 | hidden_states: torch.Tensor, |
| 185 | router_logits: torch.Tensor, |
| 186 | is_sequence_parallel: bool = False, |
| 187 | extra_tensors: list[torch.Tensor] | None = None, |
| 188 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 189 | raise NotImplementedError |
| 190 | |
| 191 | def dispatch( |
| 192 | self, |
| 193 | hidden_states: torch.Tensor, |
| 194 | topk_weights: torch.Tensor, |
| 195 | topk_ids: torch.Tensor, |
| 196 | is_sequence_parallel: bool = False, |
| 197 | extra_tensors: list[torch.Tensor] | None = None, |
| 198 | ) -> ( |
| 199 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 200 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 201 | ): |
| 202 | raise NotImplementedError |
| 203 | |
| 204 | def combine( |
| 205 | self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False |
| 206 | ) -> torch.Tensor: |
| 207 | raise NotImplementedError |
| 208 | |
| 209 | def destroy(self): |
| 210 | with self.handle_cache._lock: |
| 211 | for _, handle in self.handle_cache._cache.items(): |
| 212 | handle.destroy() |
| 213 | self.handle_cache._cache.clear() |
| 214 | |
| 215 | |
| 216 | class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase): |
| 217 | """All2All communication based on DeepEP High-Throughput kernels.""" |
| 218 | |
| 219 | def __init__(self, cpu_group, tcp_store_group=None): |
| 220 | super().__init__(cpu_group, tcp_store_group) |
| 221 | |
| 222 | def _make_all2all_kwargs(self) -> dict[Any, Any]: |
| 223 | # Defaults for internode and intranode are taken from DeepEP tests. |
| 224 | num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024 |
| 225 | num_rdma_bytes = None |
| 226 | num_qps_per_rank = None |
| 227 | |
| 228 | if self.internode and not envs.VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: |
| 229 | num_rdma_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024 |
| 230 | num_qps_per_rank = self.num_sms // 2 |
| 231 | else: |
| 232 | num_rdma_bytes = 0 |
| 233 | num_qps_per_rank = 1 |
| 234 | |
| 235 | assert num_rdma_bytes is not None |
| 236 | assert num_qps_per_rank is not None |
| 237 | # TODO: remove platform-specific logic |
| 238 | # once ROCm DeepEP is updated with the latest APIs. |
| 239 | kwargs = dict( |
| 240 | group=self.cpu_group, |
| 241 | num_nvl_bytes=num_nvl_bytes, |
| 242 | num_rdma_bytes=num_rdma_bytes, |
| 243 | low_latency_mode=False, |
| 244 | num_qps_per_rank=num_qps_per_rank, |
| 245 | explicitly_destroy=True, |
| 246 | ) |
| 247 | return kwargs |
| 248 | |
| 249 | def get_handle(self, kwargs): |
| 250 | assert len(kwargs) == 0, ( |
| 251 | "DeepEPHTAll2AllManager expects no arguments. All the required " |
| 252 | "args are computed in the Manager itself." |
| 253 | ) |
| 254 | |
| 255 | import deep_ep # type: ignore[import-not-found] |
| 256 | |
| 257 | buffer_kwargs = self._make_all2all_kwargs() |
| 258 | logger.debug("DeepEP all2all args %s", buffer_kwargs) |
| 259 | handle: deep_ep.Buffer = self.handle_cache.get_or_create( |
| 260 | buffer_kwargs, deep_ep.Buffer |
| 261 | ) |
| 262 | return handle |
| 263 | |
| 264 | def set_num_sms(self, num_sms: int): |
| 265 | import deep_ep # type: ignore[import-not-found] |
| 266 | |
| 267 | # Right now the buffers are sized for only what the kernels were |
| 268 | # created with. So we can only reduce the number of SMS used |
| 269 | # but not increase it. |
| 270 | if num_sms > self.num_sms: |
| 271 | num_sms = self.num_sms |
| 272 | deep_ep.Buffer.set_num_sms(num_sms) |
| 273 | |
| 274 | |
| 275 | class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase): |
| 276 | """All2All communication based on DeepEP Low-Latency kernels.""" |
| 277 | |
| 278 | _buffer: Any = None |
| 279 | _mask: torch.Tensor | None = None |
| 280 | _last_mask: torch.Tensor | None = None |
| 281 | |
| 282 | def __init__(self, cpu_group, tcp_store_group=None): |
| 283 | super().__init__(cpu_group, tcp_store_group) |
| 284 | self.support_fault_tolerance = ( |
| 285 | get_current_vllm_config().parallel_config.enable_fault_tolerance |
| 286 | ) |
| 287 | |
| 288 | def _make_all2all_kwargs( |
| 289 | self, |
| 290 | max_num_tokens_per_dp_rank: int, |
| 291 | token_hidden_size: int, |
| 292 | num_ep_ranks: int, |
| 293 | num_global_experts: int, |
| 294 | num_local_experts: int, |
| 295 | ) -> dict[Any, Any]: |
| 296 | """max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank |
| 297 | can dispatch all the ranks must hold the same value. |
| 298 | token_hidden_size: the hidden dimension of each token. |
| 299 | num_ep_ranks: the number of EP group ranks. |
| 300 | num_global_experts: Number of experts in the model. |
| 301 | num_local_experts: Number of experts in an EP rank. |
| 302 | """ |
| 303 | import deep_ep # type: ignore[import-not-found] |
| 304 | |
| 305 | # Defaults for internode and intranode are taken from DeepEP tests. |
| 306 | num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024 |
| 307 | num_qps_per_rank = num_local_experts |
| 308 | num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( |
| 309 | num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank, |
| 310 | hidden=token_hidden_size, |
| 311 | num_ranks=num_ep_ranks, |
| 312 | num_experts=num_global_experts, |
| 313 | ) |
| 314 | |
| 315 | assert num_rdma_bytes is not None |
| 316 | # TODO: remove platform-specific logic |
| 317 | # once ROCm DeepEP is updated with the latest APIs. |
| 318 | kwargs = dict( |
| 319 | group=self.cpu_group, |
| 320 | num_nvl_bytes=num_nvl_bytes, |
| 321 | num_rdma_bytes=num_rdma_bytes, |
| 322 | low_latency_mode=True, |
| 323 | num_qps_per_rank=num_qps_per_rank, |
| 324 | allow_nvlink_for_low_latency_mode=True, |
| 325 | allow_mnnvl=envs.VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL, |
| 326 | explicitly_destroy=True, |
| 327 | enable_shrink=self.support_fault_tolerance, |
| 328 | ) |
| 329 | return kwargs |
| 330 | |
| 331 | def get_handle(self, kwargs): |
| 332 | """The kwargs for DeepEPLLAll2AllManager is dictated by |
| 333 | _make_all2all_kwargs. |
| 334 | """ |
| 335 | import deep_ep # type: ignore[import-not-found] |
| 336 | |
| 337 | buffer_kwargs = self._make_all2all_kwargs(**kwargs) |
| 338 | logger.debug("DeepEP all2all args %s", buffer_kwargs) |
| 339 | handle: deep_ep.Buffer = self.handle_cache.get_or_create( |
| 340 | buffer_kwargs, deep_ep.Buffer |
| 341 | ) |
| 342 | DeepEPLLAll2AllManager._buffer = handle |
| 343 | return handle |
| 344 | |
| 345 | # DeepEP LL uses RDMA so no SMs are used for communication |
| 346 | def max_sms_used(self) -> int | None: |
| 347 | return 0 |
| 348 | |
| 349 | def query_active_mask(self) -> torch.Tensor: |
| 350 | buf = DeepEPLLAll2AllManager._buffer |
| 351 | assert buf is not None |
| 352 | if DeepEPLLAll2AllManager._mask is None: |
| 353 | DeepEPLLAll2AllManager._mask = torch.zeros( |
| 354 | self.world_size, device="cuda", dtype=torch.int32 |
| 355 | ) |
| 356 | buf.low_latency_query_mask_buffer(DeepEPLLAll2AllManager._mask) |
| 357 | return DeepEPLLAll2AllManager._mask |
| 358 | |
| 359 | def query_fault(self) -> torch.Tensor: |
| 360 | current = self.query_active_mask() |
| 361 | if DeepEPLLAll2AllManager._last_mask is None: |
| 362 | DeepEPLLAll2AllManager._last_mask = torch.zeros_like(current) |
| 363 | has_fault = (current != DeepEPLLAll2AllManager._last_mask).any() |
| 364 | return has_fault |
| 365 | |
| 366 | def clean_buffers(self) -> None: |
| 367 | buf = DeepEPLLAll2AllManager._buffer |
| 368 | if buf is None: |
| 369 | return |
| 370 | buf.get_local_buffer_tensor(dtype=torch.int8, use_rdma_buffer=True).zero_() |
| 371 | torch.accelerator.synchronize() |
| 372 | buf.low_latency_clean_mask_buffer() |
| 373 | torch.accelerator.synchronize() |
| 374 | DeepEPLLAll2AllManager._last_mask = None |
| 375 | |
| 376 | |
| 377 | @dataclass |
| 378 | class _NixlEPBufferState: |
| 379 | buffer: Any |
| 380 | max_num_ep_ranks: int |
| 381 | connected_ep_size: int |
| 382 | active_ep_size: int |
| 383 | |
| 384 | |
| 385 | class NixlEPAll2AllManager(All2AllManagerBase): |
| 386 | """All2All communication based on NIXL EP kernels. |
| 387 | This backend supports elastic EP with dynamic rank connection/disconnection. |
| 388 | """ |
| 389 | |
| 390 | _buffer: _NixlEPBufferState | None = None |
| 391 | _lock = threading.RLock() |
| 392 | _mask: torch.Tensor | None = None |
| 393 | _last_mask: torch.Tensor | None = None |
| 394 | |
| 395 | def __init__(self, cpu_group, tcp_store_group=None): |
| 396 | if tcp_store_group is None: |
| 397 | tcp_store_group = StatelessProcessGroup( |
| 398 | rank=cpu_group.rank(), |
| 399 | world_size=cpu_group.size(), |
| 400 | store=dist.PrefixStore("nixl_ep", cpu_group.get_group_store()), |
| 401 | ) |
| 402 | super().__init__(cpu_group, tcp_store_group) |
| 403 | self.support_fault_tolerance = True |
| 404 | |
| 405 | @property |
| 406 | def max_num_ep_ranks(self) -> int: |
| 407 | assert NixlEPAll2AllManager._buffer is not None |
| 408 | return NixlEPAll2AllManager._buffer.max_num_ep_ranks |
| 409 | |
| 410 | def _init_buffer( |
| 411 | self, |
| 412 | max_num_tokens_per_dp_rank: int, |
| 413 | token_hidden_size: int, |
| 414 | num_experts_per_rank: int, |
| 415 | ) -> None: |
| 416 | from nixl_ep import Buffer # type: ignore[import-not-found] |
| 417 | |
| 418 | parallel_config = get_current_vllm_config().parallel_config |
| 419 | max_num_ep_ranks = ( |
| 420 | self.world_size |
| 421 | * parallel_config.elastic_ep_max_dp_size |
| 422 | // parallel_config.data_parallel_size |
| 423 | if parallel_config.enable_elastic_ep |
| 424 | else self.world_size |
| 425 | ) |
| 426 | max_num_global_experts = max_num_ep_ranks * num_experts_per_rank |
| 427 | num_rdma_bytes = Buffer.get_rdma_size_hint( |
| 428 | num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank, |
| 429 | hidden=token_hidden_size, |
| 430 | num_ranks=max_num_ep_ranks, |
| 431 | num_experts=max_num_global_experts, |
| 432 | ) |
| 433 | assert NixlEPAll2AllManager._buffer is None, ( |
| 434 | "NIXL EP buffer already initialized" |
| 435 | ) |
| 436 | buffer = Buffer( |
| 437 | rank=self.rank, |
| 438 | tcp_store_group=self.tcp_store_group.store, |
| 439 | ) |
| 440 | buffer.update_memory_buffers( |
| 441 | num_ranks=max_num_ep_ranks, |
| 442 | num_experts_per_rank=num_experts_per_rank, |
| 443 | num_rdma_bytes=num_rdma_bytes, |
| 444 | ) |
| 445 | ranks_to_connect = list(range(self.world_size)) |
| 446 | buffer.connect_ranks(ranks_to_connect) |
| 447 | NixlEPAll2AllManager._buffer = _NixlEPBufferState( |
| 448 | buffer=buffer, |
| 449 | max_num_ep_ranks=max_num_ep_ranks, |
| 450 | connected_ep_size=self.world_size, |
| 451 | active_ep_size=self.world_size, |
| 452 | ) |
| 453 | |
| 454 | def _connect_to_ep_size(self, ep_size: int, *, make_active: bool) -> None: |
| 455 | assert NixlEPAll2AllManager._buffer is not None |
| 456 | state = NixlEPAll2AllManager._buffer |
| 457 | if ep_size <= state.connected_ep_size: |
| 458 | return |
| 459 | |
| 460 | state.buffer.set_tcp_store_group(self.tcp_store_group.store) |
| 461 | ranks_to_connect = list(range(state.connected_ep_size, ep_size)) |
| 462 | state.buffer.connect_ranks(ranks_to_connect, activate=make_active) |
| 463 | state.connected_ep_size = ep_size |
| 464 | if make_active: |
| 465 | state.active_ep_size = ep_size |
| 466 | |
| 467 | def _disconnect_to_ep_size(self, ep_size: int) -> None: |
| 468 | assert NixlEPAll2AllManager._buffer is not None |
| 469 | state = NixlEPAll2AllManager._buffer |
| 470 | if ep_size >= state.connected_ep_size: |
| 471 | return |
| 472 | |
| 473 | state.buffer.set_tcp_store_group(self.tcp_store_group.store) |
| 474 | ranks_to_disconnect = list(range(ep_size, state.connected_ep_size)) |
| 475 | state.buffer.disconnect_ranks(ranks_to_disconnect) |
| 476 | state.connected_ep_size = ep_size |
| 477 | state.active_ep_size = min(state.active_ep_size, ep_size) |
| 478 | |
| 479 | def _unmask_connected_ranks(self, target_ep_size: int) -> None: |
| 480 | assert NixlEPAll2AllManager._buffer is not None |
| 481 | state = NixlEPAll2AllManager._buffer |
| 482 | state.buffer.set_tcp_store_group(self.tcp_store_group.store) |
| 483 | if target_ep_size <= state.active_ep_size: |
| 484 | return |
| 485 | assert state.connected_ep_size >= target_ep_size |
| 486 | |
| 487 | for rank in range(state.active_ep_size, target_ep_size): |
| 488 | state.buffer.update_mask_buffer(rank, mask=False) |
| 489 | torch.accelerator.synchronize() |
| 490 | state.active_ep_size = target_ep_size |
| 491 | |
| 492 | def _stage_ep_size(self) -> None: |
| 493 | assert NixlEPAll2AllManager._buffer is not None |
| 494 | state = NixlEPAll2AllManager._buffer |
| 495 | target_ep_size = self.world_size |
| 496 | |
| 497 | # Scale-up can safely connect standby ranks while leaving them masked. |
| 498 | # Scale-down must not disconnect active ranks until commit. |
| 499 | if target_ep_size > state.connected_ep_size: |
| 500 | self._connect_to_ep_size(target_ep_size, make_active=False) |
| 501 | |
| 502 | def stage_ep_size(self) -> None: |
| 503 | with NixlEPAll2AllManager._lock: |
| 504 | self._stage_ep_size() |
| 505 | |
| 506 | def commit_ep_size(self) -> None: |
| 507 | """Commit the staged EP size to the active communication set.""" |
| 508 | with NixlEPAll2AllManager._lock: |
| 509 | assert NixlEPAll2AllManager._buffer is not None |
| 510 | state = NixlEPAll2AllManager._buffer |
| 511 | target_ep_size = self.world_size |
| 512 | |
| 513 | if target_ep_size < state.connected_ep_size: |
| 514 | self._disconnect_to_ep_size(target_ep_size) |
| 515 | elif target_ep_size > state.connected_ep_size: |
| 516 | self._connect_to_ep_size(target_ep_size, make_active=True) |
| 517 | |
| 518 | self._unmask_connected_ranks(target_ep_size) |
| 519 | |
| 520 | @contextmanager |
| 521 | def mask_remote_ranks(self) -> Iterator[None]: |
| 522 | peers = [rank for rank in range(self.world_size) if rank != self.rank] |
| 523 | assert NixlEPAll2AllManager._buffer is not None |
| 524 | buffer = NixlEPAll2AllManager._buffer.buffer |
| 525 | for rank in peers: |
| 526 | buffer.update_mask_buffer(rank, mask=True) |
| 527 | torch.accelerator.synchronize() |
| 528 | try: |
| 529 | yield |
| 530 | finally: |
| 531 | for rank in peers: |
| 532 | buffer.update_mask_buffer(rank, mask=False) |
| 533 | torch.accelerator.synchronize() |
| 534 | |
| 535 | def get_handle(self, kwargs): |
| 536 | with NixlEPAll2AllManager._lock: |
| 537 | state = NixlEPAll2AllManager._buffer |
| 538 | if state is None: |
| 539 | max_num_tokens_per_dp_rank = kwargs["max_num_tokens_per_dp_rank"] |
| 540 | self._init_buffer( |
| 541 | max_num_tokens_per_dp_rank=max_num_tokens_per_dp_rank, |
| 542 | token_hidden_size=kwargs["token_hidden_size"], |
| 543 | num_experts_per_rank=kwargs["num_local_experts"], |
| 544 | ) |
| 545 | |
| 546 | assert NixlEPAll2AllManager._buffer is not None |
| 547 | handle = NixlEPAll2AllManager._buffer.buffer |
| 548 | return handle |
| 549 | |
| 550 | def dispatch( |
| 551 | self, |
| 552 | hidden_states: torch.Tensor, |
| 553 | topk_weights: torch.Tensor, |
| 554 | topk_ids: torch.Tensor, |
| 555 | is_sequence_parallel: bool = False, |
| 556 | extra_tensors: list[torch.Tensor] | None = None, |
| 557 | ) -> ( |
| 558 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 559 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 560 | ): |
| 561 | raise NotImplementedError |
| 562 | |
| 563 | def combine( |
| 564 | self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False |
| 565 | ) -> torch.Tensor: |
| 566 | raise NotImplementedError |
| 567 | |
| 568 | def destroy(self): |
| 569 | # NOTE(yongji): NIXLEPAll2AllManager instance is recreated during |
| 570 | # scale-up/down, so we cannot destroy the persistent buffer here. |
| 571 | assert NixlEPAll2AllManager._buffer is not None |
| 572 | buffer = NixlEPAll2AllManager._buffer.buffer |
| 573 | buffer.set_tcp_store_group(None) |
| 574 | |
| 575 | # NIXL EP uses RDMA so no SMs are used for communication |
| 576 | def max_sms_used(self) -> int | None: |
| 577 | return 0 |
| 578 | |
| 579 | def query_active_mask(self) -> torch.Tensor: |
| 580 | state = NixlEPAll2AllManager._buffer |
| 581 | assert state is not None |
| 582 | if NixlEPAll2AllManager._mask is None: |
| 583 | NixlEPAll2AllManager._mask = torch.zeros( |
| 584 | state.max_num_ep_ranks, device="cuda", dtype=torch.int32 |
| 585 | ) |
| 586 | state.buffer.query_mask_buffer(NixlEPAll2AllManager._mask) |
| 587 | return NixlEPAll2AllManager._mask[: state.active_ep_size] |
| 588 | |
| 589 | def query_fault(self) -> torch.Tensor: |
| 590 | current = self.query_active_mask() |
| 591 | last = NixlEPAll2AllManager._last_mask |
| 592 | if last is None or last.shape != current.shape: |
| 593 | NixlEPAll2AllManager._last_mask = torch.zeros_like(current) |
| 594 | last = NixlEPAll2AllManager._last_mask |
| 595 | has_fault = (current != last).any() |
| 596 | return has_fault |
| 597 | |
| 598 | def clean_buffers(self) -> None: |
| 599 | if NixlEPAll2AllManager._buffer is None: |
| 600 | return |
| 601 | state = NixlEPAll2AllManager._buffer |
| 602 | state.buffer.get_local_buffer_tensor( |
| 603 | dtype=torch.int8, use_rdma_buffer=True |
| 604 | ).zero_() |
| 605 | torch.accelerator.synchronize() |
| 606 | state.buffer.clean_mask_buffer() |
| 607 | torch.accelerator.synchronize() |
| 608 | NixlEPAll2AllManager._last_mask = None |
| 609 | |
| 610 | |
| 611 | class FlashInferNVLinkTwoSidedManager(All2AllManagerBase): |
| 612 | """All2All communication based on flashinfer all2allv/two-sided NVLink kernels.""" |
| 613 | |
| 614 | # This type lint could be removed after all of the work in |
| 615 | # https://github.com/vllm-project/vllm/issues/26533 done. |
| 616 | rank: int |
| 617 | world_size: int |
| 618 | |
| 619 | def __init__(self, cpu_group, tcp_store_group=None): |
| 620 | assert has_flashinfer_nvlink_two_sided(), ( |
| 621 | "flashinfer all2all module not found. Please install/check flashinfer" |
| 622 | ) # noqa |
| 623 | super().__init__(cpu_group, tcp_store_group) |
| 624 | logger.debug( |
| 625 | "Initialize for flashinfer All2All rank=%d, world size=%d", |
| 626 | self.rank, |
| 627 | self.world_size, |
| 628 | ) |
| 629 | self.initialized = False |
| 630 | self.alltoall_info = None |
| 631 | |
| 632 | def initialize( |
| 633 | self, |
| 634 | world_size: int, |
| 635 | rank: int, |
| 636 | gpus_per_node: int, |
| 637 | ): |
| 638 | """Initialize workspace.""" |
| 639 | if self.initialized: |
| 640 | return |
| 641 | |
| 642 | self.cleanup() |
| 643 | logger.debug("making map: rank=%d, world size=%d", rank, world_size) |
| 644 | self.mapping = Mapping( |
| 645 | world_size, |
| 646 | rank, |
| 647 | gpus_per_node, |
| 648 | tp_size=world_size, |
| 649 | ) |
| 650 | |
| 651 | from vllm.distributed.device_communicators.mnnvl_compat import ( |
| 652 | CustomCommunicator, |
| 653 | ) |
| 654 | |
| 655 | # MNNVL workspace is allocated per rank in the comm_backend's group; the |
| 656 | # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend |
| 657 | # must span the EP group (= DP*PCP*TP), not the DP group. |
| 658 | ep_config = MnnvlConfig( |
| 659 | comm_backend=CustomCommunicator(self.cpu_group), |
| 660 | fabric_page_size=1 << 29, # 512MB |
| 661 | allocation_granularity=0, # Auto-detect |
| 662 | ) |
| 663 | |
| 664 | self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, ep_config) |
| 665 | self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace( |
| 666 | self.mapping, ep_config |
| 667 | ) |
| 668 | |
| 669 | self.world_size = world_size |
| 670 | self.rank = rank |
| 671 | self.gpus_per_node = gpus_per_node |
| 672 | self.initialized = True |
| 673 | |
| 674 | logger.info( |
| 675 | "FlashInfer All2All initialized for rank %s, size %s", rank, world_size |
| 676 | ) |
| 677 | |
| 678 | def ensure_alltoall_workspace_initialized(self): |
| 679 | """Ensure workspace is initialized.""" |
| 680 | if not has_flashinfer_nvlink_two_sided(): |
| 681 | return False |
| 682 | |
| 683 | if self.world_size <= 1: |
| 684 | return False |
| 685 | |
| 686 | if not self.initialized: |
| 687 | self.initialize( |
| 688 | world_size=self.world_size, |
| 689 | rank=self.rank, |
| 690 | gpus_per_node=torch.accelerator.device_count, |
| 691 | ) |
| 692 | return self.initialized |
| 693 | |
| 694 | def get_handle(self, kwargs): |
| 695 | return self |
| 696 | |
| 697 | def cleanup(self): |
| 698 | """Clean up workspace.""" |
| 699 | if ( |
| 700 | self.initialized |
| 701 | and self.workspace_tensor is not None |
| 702 | and self.prepare_workspace_tensor is not None |
| 703 | ): |
| 704 | try: |
| 705 | del self.workspace_tensor |
| 706 | del self.prepare_workspace_tensor |
| 707 | except Exception as e: |
| 708 | logger.warning("Failed to cleanup FlashInfer workspace: %s", e) |
| 709 | finally: |
| 710 | self.workspace_tensor = None |
| 711 | self.prepare_workspace_tensor = None |
| 712 | self.mapping = None |
| 713 | self.initialized = False |
| 714 | |
| 715 | |
| 716 | class FlashInferNVLinkOneSidedManager(All2AllManagerBase): |
| 717 | """All2All communication based on FlashInfer's MoeAlltoAll/One-sided NVLink kernel. |
| 718 | This is a newer kernel from trtllm that should perform better than the kernel |
| 719 | used by flashinfer_nvlink_two_sided. |
| 720 | """ |
| 721 | |
| 722 | rank: int |
| 723 | world_size: int |
| 724 | |
| 725 | def __init__(self, cpu_group): |
| 726 | assert has_flashinfer_nvlink_one_sided(), ( |
| 727 | "flashinfer trtllm_moe_alltoall module not found. " |
| 728 | "Please install/check flashinfer" |
| 729 | ) |
| 730 | super().__init__(cpu_group) |
| 731 | logger.debug( |
| 732 | "Initialize FlashInfer One-sided NVLink rank=%d, world size=%d", |
| 733 | self.rank, |
| 734 | self.world_size, |
| 735 | ) |
| 736 | self.initialized = False |
| 737 | self.moe_alltoall: MoeAlltoAll | None = None |
| 738 | self.mapping = None |
| 739 | self.workspace_size = 0 |
| 740 | self.max_num_tokens = 0 |
| 741 | self.top_k = 0 |
| 742 | self.num_experts = 0 |
| 743 | self._combine_supports_output = False |
| 744 | |
| 745 | def initialize( |
| 746 | self, |
| 747 | max_num_tokens: int, |
| 748 | top_k: int, |
| 749 | num_experts: int, |
| 750 | hidden_size: int, |
| 751 | x_bytes_per_token: int, |
| 752 | x_sf_bytes_per_token: int, |
| 753 | ): |
| 754 | """Initialize (or grow) the MoeAlltoAll workspace.""" |
| 755 | total_dispatch_payload_size_per_token = ( |
| 756 | x_bytes_per_token |
| 757 | + x_sf_bytes_per_token |
| 758 | + top_k * 4 # int32 topks ids |
| 759 | + top_k * 4 # float32 topk weights |
| 760 | ) |
| 761 | combine_payload_size_per_token = hidden_size * 2 # bf16 hidden states |
| 762 | needed_workspace_size = moe_a2a_get_workspace_size_per_rank( |
| 763 | ep_size=self.world_size, |
| 764 | max_num_tokens=max_num_tokens, |
| 765 | total_dispatch_payload_size_per_token=total_dispatch_payload_size_per_token, |
| 766 | combine_payload_size_per_token=combine_payload_size_per_token, |
| 767 | ) |
| 768 | # workspace_size and max_num_tokens are kernel-side max-bounds, so |
| 769 | # heterogeneous MoE layers (e.g. NVFP4 base + bf16 MTP head) only |
| 770 | # need the shared workspace grown to the union. top_k and num_experts |
| 771 | # must match across layers: top_k is a strict-equality assert at |
| 772 | # dispatch (FlashInfer csrc/trtllm_moe_alltoall.cu), and num_experts |
| 773 | # feeds the expert-to-rank routing math, so any mismatch would crash |
| 774 | # or silently corrupt routing. All ranks see the same MoE layers in |
| 775 | # the same order with identical shapes, so the skip / rebuild |
| 776 | # branches are taken consistently across ranks. |
| 777 | if self.initialized: |
| 778 | assert top_k == self.top_k, ( |
| 779 | "FlashInfer one-sided MoeAlltoAll does not support " |
| 780 | f"heterogeneous top_k across MoE layers (got {top_k}, " |
| 781 | f"was built with {self.top_k})" |
| 782 | ) |
| 783 | assert num_experts == self.num_experts, ( |
| 784 | "FlashInfer one-sided MoeAlltoAll does not support " |
| 785 | f"heterogeneous num_experts across MoE layers (got " |
| 786 | f"{num_experts}, was built with {self.num_experts})" |
| 787 | ) |
| 788 | if ( |
| 789 | needed_workspace_size <= self.workspace_size |
| 790 | and max_num_tokens <= self.max_num_tokens |
| 791 | ): |
| 792 | return |
| 793 | |
| 794 | self.workspace_size = max(self.workspace_size, needed_workspace_size) |
| 795 | self.max_num_tokens = max(self.max_num_tokens, max_num_tokens) |
| 796 | self.top_k = top_k |
| 797 | self.num_experts = num_experts |
| 798 | |
| 799 | self.cleanup() |
| 800 | from vllm.platforms.interface import get_assigned_physical_gpu_ids |
| 801 | |
| 802 | assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() |
| 803 | gpus_per_node = ( |
| 804 | len(assigned_physical_gpu_ids) |
| 805 | if assigned_physical_gpu_ids is not None |
| 806 | else torch.accelerator.device_count() |
| 807 | ) |
| 808 | logger.debug( |
| 809 | "Making One-sided NVLink mapping: rank=%d, world size=%d", |
| 810 | self.rank, |
| 811 | self.world_size, |
| 812 | ) |
| 813 | self.mapping = Mapping( |
| 814 | self.world_size, |
| 815 | self.rank, |
| 816 | gpus_per_node, |
| 817 | tp_size=self.world_size, |
| 818 | moe_ep_size=self.world_size, |
| 819 | ) |
| 820 | |
| 821 | from vllm.distributed.device_communicators.mnnvl_compat import ( |
| 822 | CustomCommunicator, |
| 823 | ) |
| 824 | |
| 825 | # MNNVL workspace is allocated per rank in the comm_backend's group; the |
| 826 | # flashinfer kernel asserts workspace.size(0) == moe_ep_size, so the backend |
| 827 | # must span the EP group (= DP*PCP*TP), not the DP group. |
| 828 | ep_config = MnnvlConfig( |
| 829 | comm_backend=CustomCommunicator(self.cpu_group), |
| 830 | ) |
| 831 | |
| 832 | # Release cached allocator blocks before FlashInfer reserves the |
| 833 | # symmetric MNNVL fabric workspace. |
| 834 | torch.accelerator.empty_cache() |
| 835 | |
| 836 | self.moe_alltoall = MoeAlltoAll( |
| 837 | mapping=self.mapping, |
| 838 | max_num_tokens=self.max_num_tokens, |
| 839 | top_k=self.top_k, |
| 840 | num_experts=self.num_experts, |
| 841 | workspace_size_per_rank=self.workspace_size, |
| 842 | mnnvl_config=ep_config, |
| 843 | ) |
| 844 | try: |
| 845 | self._combine_supports_output = supports_kw( |
| 846 | self.moe_alltoall.combine, "output", allow_var_kwargs=False |
| 847 | ) |
| 848 | except (TypeError, ValueError): |
| 849 | self._combine_supports_output = False |
| 850 | |
| 851 | self.gpus_per_node = gpus_per_node |
| 852 | self.initialized = True |
| 853 | |
| 854 | logger.info( |
| 855 | "FlashInfer One-sided NVLink initialized for rank %s, size %s", |
| 856 | self.rank, |
| 857 | self.world_size, |
| 858 | ) |
| 859 | # Scope barrier to the EP group: with PP, different EP groups can |
| 860 | # rebuild a different number of times if their MoE layers have |
| 861 | # different shape sequences, so a world-level barrier would deadlock. |
| 862 | dist.barrier(group=self.cpu_group) |
| 863 | |
| 864 | def combine_into( |
| 865 | self, |
| 866 | payload: torch.Tensor, |
| 867 | runtime_max_tokens_per_rank: int, |
| 868 | output: torch.Tensor, |
| 869 | ) -> None: |
| 870 | """Combine into ``output``, with a fallback for older FlashInfer.""" |
| 871 | assert self.moe_alltoall is not None |
| 872 | if self._combine_supports_output: |
| 873 | self.moe_alltoall.combine( |
| 874 | payload=payload, |
| 875 | runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, |
| 876 | output=output, |
| 877 | ) |
| 878 | else: |
| 879 | combined_output = self.moe_alltoall.combine( |
| 880 | payload=payload, |
| 881 | runtime_max_tokens_per_rank=runtime_max_tokens_per_rank, |
| 882 | ) |
| 883 | output.copy_(combined_output) |
| 884 | |
| 885 | def get_handle(self, kwargs): |
| 886 | return self |
| 887 | |
| 888 | def cleanup(self): |
| 889 | """Clean up resources.""" |
| 890 | if self.initialized and self.moe_alltoall is not None: |
| 891 | try: |
| 892 | del self.moe_alltoall |
| 893 | except Exception as e: |
| 894 | logger.warning( |
| 895 | "Failed to cleanup FlashInfer One-sided NVLink workspace: %s", e |
| 896 | ) |
| 897 | finally: |
| 898 | self.moe_alltoall = None |
| 899 | self.mapping = None |
| 900 | self.initialized = False |
| 901 | |
| 902 | def checkpoint_prepare(self) -> None: |
| 903 | if self.initialized: |
| 904 | assert self.moe_alltoall is not None |
| 905 | self.moe_alltoall.checkpoint_prepare() |
| 906 | |
| 907 | def checkpoint_restore(self) -> None: |
| 908 | if self.initialized: |
| 909 | assert self.moe_alltoall is not None |
| 910 | from vllm.distributed.device_communicators.mnnvl_compat import ( |
| 911 | CustomCommunicator, |
| 912 | ) |
| 913 | |
| 914 | self.moe_alltoall.checkpoint_restore(CustomCommunicator(self.cpu_group)) |
| 915 | |
| 916 | |
| 917 | class MoriAll2AllManager(All2AllManagerBase): |
| 918 | def __init__(self, cpu_group, all2all_backend: str): |
| 919 | assert has_mori(), ( |
| 920 | "MoRI kernels not found. Please follow https://github.com/ROCm/mori/blob/main/README.md" |
| 921 | " to install MoRI kernels." |
| 922 | ) # noqa |
| 923 | assert all2all_backend in ( |
| 924 | "mori_high_throughput", |
| 925 | "mori_low_latency", |
| 926 | ), f"unsupported MoRI all2all backend: {all2all_backend!r}" |
| 927 | import mori |
| 928 | |
| 929 | super().__init__(cpu_group) |
| 930 | self._all2all_backend = all2all_backend |
| 931 | self.handle_cache = Cache() |
| 932 | |
| 933 | torch._C._distributed_c10d._register_process_group("mori", cpu_group) |
| 934 | mori.shmem.shmem_torch_process_group_init("mori") |
| 935 | |
| 936 | def _make_all2all_kwargs( |
| 937 | self, |
| 938 | rank: int, |
| 939 | num_ep_ranks: int, |
| 940 | input_dtype: torch.dtype, |
| 941 | quant_dtype: torch.dtype, |
| 942 | token_hidden_size: int, |
| 943 | scale_dim: int, |
| 944 | scale_type_size: int, |
| 945 | max_num_tokens_per_dp_rank: int, |
| 946 | num_local_experts: int, |
| 947 | num_experts_per_token: int, |
| 948 | ): |
| 949 | import mori # type: ignore[import-not-found] |
| 950 | |
| 951 | from vllm.platforms.rocm import on_gfx942, on_gfx950 |
| 952 | |
| 953 | assert on_gfx942() or on_gfx950(), ( |
| 954 | "mori currently only support arch gfx942 and gfx950" |
| 955 | ) |
| 956 | |
| 957 | if not self.internode: |
| 958 | # single node |
| 959 | kernel_type = mori.ops.EpDispatchCombineKernelType.IntraNode |
| 960 | rdma_block_num = 0 |
| 961 | warp_num_per_block = 16 |
| 962 | block_num = 80 |
| 963 | else: |
| 964 | # Multi-node: kernel follows --all2all-backend (mirrors deepep_* split). |
| 965 | # mori_low_latency → InterNodeV1LL; mori_high_throughput → V1. |
| 966 | if self._all2all_backend == "mori_low_latency": |
| 967 | kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1LL |
| 968 | else: |
| 969 | kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1 |
| 970 | if on_gfx942(): |
| 971 | warp_num_per_block = 16 |
| 972 | block_num = 32 |
| 973 | rdma_block_num = 16 |
| 974 | elif on_gfx950(): |
| 975 | warp_num_per_block = 8 |
| 976 | block_num = 64 |
| 977 | rdma_block_num = 32 |
| 978 | else: |
| 979 | raise NotImplementedError( |
| 980 | "mori currently only support arch gfx942 and gfx950" |
| 981 | ) |
| 982 | |
| 983 | return dict( |
| 984 | rank=rank, |
| 985 | world_size=num_ep_ranks, |
| 986 | data_type=quant_dtype, |
| 987 | hidden_dim=token_hidden_size, |
| 988 | scale_dim=scale_dim, |
| 989 | scale_type_size=scale_type_size, |
| 990 | max_token_type_size=input_dtype.itemsize, |
| 991 | max_num_inp_token_per_rank=max_num_tokens_per_dp_rank, |
| 992 | num_experts_per_rank=num_local_experts, |
| 993 | num_experts_per_token=num_experts_per_token, |
| 994 | warp_num_per_block=warp_num_per_block, |
| 995 | block_num=block_num, |
| 996 | kernel_type=kernel_type, |
| 997 | rdma_block_num=rdma_block_num, |
| 998 | gpu_per_node=min(8, num_ep_ranks), |
| 999 | ) |
| 1000 | |
| 1001 | def _make_handle(self, **kwargs): |
| 1002 | import mori # type: ignore[import-not-found] |
| 1003 | |
| 1004 | mori_config = mori.ops.EpDispatchCombineConfig(**kwargs) |
| 1005 | handle = mori.ops.EpDispatchCombineOp(mori_config) |
| 1006 | return handle |
| 1007 | |
| 1008 | def get_handle(self, kwargs): |
| 1009 | import mori # type: ignore[import-not-found] |
| 1010 | |
| 1011 | mori_kwargs = self._make_all2all_kwargs(**kwargs) |
| 1012 | logger.debug("MoRI all2all args %s", mori_kwargs) |
| 1013 | handle: mori.ops.EpDispatchCombineOp = self.handle_cache.get_or_create( |
| 1014 | mori_kwargs, self._make_handle |
| 1015 | ) |
| 1016 | return handle |
| 1017 | |
| 1018 | |
| 1019 | class DeepEPV2All2AllManager(All2AllManagerBase): |
| 1020 | """All2All communication based on DeepEP v2 ElasticBuffer (unified API). |
| 1021 | Uses NCCL Gin backend with analytical SM calculation. |
| 1022 | """ |
| 1023 | |
| 1024 | def __init__(self, cpu_group, tcp_store_group=None, device_group=None): |
| 1025 | assert has_deep_ep_v2(), ( |
| 1026 | "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 " |
| 1027 | "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4." |
| 1028 | ) |
| 1029 | super().__init__(cpu_group, tcp_store_group) |
| 1030 | self._device_group = device_group |
| 1031 | self.handle_cache = Cache() |
| 1032 | self._num_sms: int | None = None |
| 1033 | self._gin_checked = False |
| 1034 | |
| 1035 | def _make_all2all_kwargs( |
| 1036 | self, |
| 1037 | num_max_tokens_per_rank: int, |
| 1038 | hidden: int, |
| 1039 | num_topk: int, |
| 1040 | use_fp8_dispatch: bool, |
| 1041 | ) -> dict: |
| 1042 | return dict( |
| 1043 | group=self._device_group |
| 1044 | if self._device_group is not None |
| 1045 | else self.cpu_group, |
| 1046 | num_max_tokens_per_rank=num_max_tokens_per_rank, |
| 1047 | hidden=hidden, |
| 1048 | num_topk=num_topk, |
| 1049 | use_fp8_dispatch=use_fp8_dispatch, |
| 1050 | allow_hybrid_mode=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE, |
| 1051 | prefer_overlap_with_compute=envs.VLLM_DEEPEP_V2_PREFER_OVERLAP, |
| 1052 | allow_multiple_reduction=(envs.VLLM_DEEPEP_V2_ALLOW_MULTIPLE_REDUCTION), |
| 1053 | explicitly_destroy=True, |
| 1054 | ) |
| 1055 | |
| 1056 | def _check_gin_support(self, group) -> None: |
| 1057 | from vllm.utils.nccl import query_nccl_gin_type |
| 1058 | |
| 1059 | # ProcessGroupNCCL creates communicators lazily. Initialize this exact |
| 1060 | # group before querying so a null comm pointer is not mistaken for |
| 1061 | # missing GIN support. |
| 1062 | probe = torch.zeros(1, device="cuda") |
| 1063 | torch.distributed.all_reduce(probe, group=group) |
| 1064 | # DeepEPv2 respects EP_DISABLE_GIN, so skip the GIN requirement check. |
| 1065 | if os.environ.get("EP_DISABLE_GIN", "0") != "0": |
| 1066 | return |
| 1067 | |
| 1068 | gin_type = query_nccl_gin_type( |
| 1069 | group, railed=envs.VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE |
| 1070 | ) |
| 1071 | if gin_type is None: |
| 1072 | raise RuntimeError( |
| 1073 | "DeepEPv2 communicator properties query failed; " |
| 1074 | "networking capability could not be determined." |
| 1075 | ) |
| 1076 | if gin_type == 0: |
| 1077 | raise RuntimeError( |
| 1078 | "DeepEPv2 requires NCCL GIN (GPU-Initiated Networking). " |
| 1079 | "This usually means IBGDA-capable InfiniBand NICs or drivers " |
| 1080 | "are not available. See tools/ep_kernels/README.md for " |
| 1081 | "requirements." |
| 1082 | ) |
| 1083 | |
| 1084 | def get_handle(self, kwargs): |
| 1085 | import deep_ep # type: ignore[import-not-found] |
| 1086 | |
| 1087 | num_experts = kwargs.pop("num_experts", 256) |
| 1088 | buffer_kwargs = self._make_all2all_kwargs(**kwargs) |
| 1089 | if not self._gin_checked: |
| 1090 | self._check_gin_support(buffer_kwargs["group"]) |
| 1091 | self._gin_checked = True |
| 1092 | logger.debug("DeepEP v2 all2all args %s", buffer_kwargs) |
| 1093 | handle: deep_ep.ElasticBuffer = self.handle_cache.get_or_create( |
| 1094 | buffer_kwargs, deep_ep.ElasticBuffer |
| 1095 | ) |
| 1096 | if self._num_sms is None: |
| 1097 | self._num_sms = handle.get_theoretical_num_sms( |
| 1098 | num_experts=num_experts, |
| 1099 | num_topk=kwargs["num_topk"], |
| 1100 | ) |
| 1101 | return handle |
| 1102 | |
| 1103 | def max_sms_used(self) -> int | None: |
| 1104 | return self._num_sms |
| 1105 | |
| 1106 | def destroy(self): |
| 1107 | with self.handle_cache._lock: |
| 1108 | for _, handle in self.handle_cache._cache.items(): |
| 1109 | handle.destroy() |
| 1110 | self.handle_cache._cache.clear() |
| 1111 | |
| 1112 | |
| 1113 | class MoonEPAll2AllManager(All2AllManagerBase): |
| 1114 | """All2All communication based on MoonEP |
| 1115 | (https://github.com/MoonshotAI/MoonEP). |
| 1116 | |
| 1117 | MoonEP keeps token loads perfectly balanced across EP ranks by planning |
| 1118 | a small number of dynamically redundant experts online and prefetching |
| 1119 | their weights before expert compute. Every rank receives exactly |
| 1120 | S x K token slots regardless of router skew, so all communication and |
| 1121 | compute shapes are static. |
| 1122 | |
| 1123 | Requires NVLink symmetric-memory / multicast capable topologies |
| 1124 | (single node NVSwitch, e.g. H100/H200/B200/GB300 class). |
| 1125 | """ |
| 1126 | |
| 1127 | def __init__(self, cpu_group, tcp_store_group=None, device_group=None): |
| 1128 | check_moonep_system_support() |
| 1129 | super().__init__(cpu_group, tcp_store_group) |
| 1130 | self._device_group = device_group |
| 1131 | self.handle_cache = Cache() |
| 1132 | |
| 1133 | def _make_buffer_kwargs( |
| 1134 | self, |
| 1135 | max_num_tokens_per_dp_rank: int, |
| 1136 | token_hidden_size: int, |
| 1137 | num_topk: int, |
| 1138 | num_global_experts: int, |
| 1139 | num_prefetch_slots: int, |
| 1140 | token_padding: int, |
| 1141 | num_sms: int, |
| 1142 | ) -> dict: |
| 1143 | return dict( |
| 1144 | S=max_num_tokens_per_dp_rank, |
| 1145 | H=token_hidden_size, |
| 1146 | K=num_topk, |
| 1147 | E=num_global_experts, |
| 1148 | num_ep_ranks=self.world_size, |
| 1149 | num_sms=num_sms, |
| 1150 | token_padding=token_padding, |
| 1151 | B=num_prefetch_slots, |
| 1152 | group=self._device_group |
| 1153 | if self._device_group is not None |
| 1154 | else self.cpu_group, |
| 1155 | explicitly_destroy=True, |
| 1156 | ) |
| 1157 | |
| 1158 | def get_handle(self, kwargs): |
| 1159 | from vllm.model_executor.layers.fused_moe.prepare_finalize.moonep import ( |
| 1160 | MoonEPBufferPool, |
| 1161 | ) |
| 1162 | |
| 1163 | buffer_kwargs = self._make_buffer_kwargs(**kwargs) |
| 1164 | logger.debug("MoonEP all2all args %s", buffer_kwargs) |
| 1165 | |
| 1166 | def make_pool(**kw): |
| 1167 | return MoonEPBufferPool(kw, max_tokens_per_rank=kw["S"]) |
| 1168 | |
| 1169 | handle: MoonEPBufferPool = self.handle_cache.get_or_create( |
| 1170 | buffer_kwargs, make_pool |
| 1171 | ) |
| 1172 | return handle |
| 1173 | |
| 1174 | def destroy(self): |
| 1175 | with self.handle_cache._lock: |
| 1176 | for _, handle in self.handle_cache._cache.items(): |
| 1177 | handle.destroy() |
| 1178 | self.handle_cache._cache.clear() |