From 4f423bd5bcf38cb2fc09db5dc9e0bc8a6bfb1638 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Thu, 4 Jun 2026 05:40:34 +0200 Subject: [PATCH] [EPLB] Nixl communicator optimization. Zero-copy transfers (#41633) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: ilmarkov Signed-off-by: Markov Ilya Signed-off-by: Tyler Michael Smith Co-authored-by: Markov Ilya Co-authored-by: Nicolò Lucchesi Co-authored-by: Tyler Michael Smith Co-authored-by: Tyler Michael Smith --- tests/distributed/test_eplb_execute.py | 30 +- .../distributed/test_eplb_fused_moe_layer.py | 12 +- .../test_eplb_fused_moe_layer_dep_nvfp4.py | 12 +- tests/kernels/moe/test_moe_layer.py | 5 +- .../device_communicators/custom_all_reduce.py | 6 +- .../device_communicators/quick_all_reduce.py | 6 +- .../distributed/elastic_ep/elastic_execute.py | 6 +- vllm/distributed/eplb/async_worker.py | 1 + vllm/distributed/eplb/eplb_communicator.py | 385 +++++++++--------- vllm/distributed/eplb/eplb_state.py | 4 +- vllm/distributed/eplb/rebalance_execute.py | 25 +- vllm/distributed/utils.py | 14 + 12 files changed, 287 insertions(+), 219 deletions(-) diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index d9e6a739b01..0b87477950f 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -277,12 +277,15 @@ def assert_verification_synced(local_ok: bool, msg: str) -> None: assert bool(ok_tensor.item()), msg -def create_eplb_communicator_or_raise(*, group_coordinator, backend, expert_weights): +def create_eplb_communicator_or_raise( + *, group_coordinator, backend, expert_weights, expert_buffer +): try: return create_eplb_communicator( group_coordinator=group_coordinator, backend=backend, expert_weights=expert_weights, + expert_buffer=expert_buffer, ) except Exception as exc: raise RuntimeError( @@ -355,7 +358,8 @@ def _test_async_transfer_layer_without_mtp_worker( communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, backend=eplb_communicator, - expert_weights=expert_weights[0], + expert_weights=expert_weights, + expert_buffer=expert_buffer, ) communicator.set_stream(cuda_stream) @@ -368,6 +372,7 @@ def _test_async_transfer_layer_without_mtp_worker( ep_group=ep_group, communicator=communicator, cuda_stream=cuda_stream, + layer_idx=layer_idx, ) cuda_stream.synchronize() move_from_buffer( @@ -460,10 +465,12 @@ def _test_rearrange_expert_weights_with_redundancy( num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices ) + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, backend=eplb_communicator, - expert_weights=expert_weights[0], + expert_weights=expert_weights, + expert_buffer=expert_buffer, ) # Execute weight rearrangement @@ -471,9 +478,9 @@ def _test_rearrange_expert_weights_with_redundancy( old_indices, new_indices, expert_weights, + expert_buffer, ep_group, - is_profile=False, - communicator=communicator, + communicator, ) # Verify the rearrangement result @@ -593,10 +600,12 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: layer_copy.append(weight.clone()) original_weights.append(layer_copy) + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, backend="torch_nccl", - expert_weights=expert_weights[0], + expert_weights=expert_weights, + expert_buffer=expert_buffer, ) # Execute rearrangement (should be no change) @@ -604,9 +613,9 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: indices, indices, # Same indices expert_weights, + expert_buffer, ep_group, communicator, - is_profile=False, ) # Verify that the weights have not changed @@ -726,10 +735,12 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: layer_copy.append(weight.clone()) original_weights.append(layer_copy) + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] communicator = create_eplb_communicator_or_raise( group_coordinator=ep_group_coordinator, backend="torch_nccl", - expert_weights=expert_weights[0], + expert_weights=expert_weights, + expert_buffer=expert_buffer, ) # Execute profile mode rearrangement @@ -737,9 +748,10 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None: old_indices, new_indices, expert_weights, + expert_buffer, ep_group, communicator, - is_profile=True, # Profile mode + is_profile=True, ) # In profile mode, the weights should remain unchanged diff --git a/tests/distributed/test_eplb_fused_moe_layer.py b/tests/distributed/test_eplb_fused_moe_layer.py index eacdb3abc36..87ed4485d3d 100644 --- a/tests/distributed/test_eplb_fused_moe_layer.py +++ b/tests/distributed/test_eplb_fused_moe_layer.py @@ -9,9 +9,11 @@ import pytest import torch from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace from vllm.distributed.parallel_state import ( ensure_model_parallel_initialized, + get_eplb_group, get_tp_group, ) from vllm.model_executor.layers.fused_moe.layer import FusedMoE @@ -213,12 +215,20 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): for lidx in range(test_config.num_layers): shuffled_indices[lidx] = torch.randperm(test_config.num_experts) + expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]] + communicator = create_eplb_communicator( + group_coordinator=get_eplb_group(), + backend="torch_nccl", + expert_weights=rank_expert_weights, + expert_buffer=expert_buffer, + ) rearrange_expert_weights_inplace( indices, shuffled_indices, rank_expert_weights, + expert_buffer, ep_group, - is_profile=False, + communicator, ) num_local_experts = test_config.num_local_experts diff --git a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py index 9ab785af313..4818f8a2c8c 100644 --- a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py +++ b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py @@ -10,11 +10,13 @@ import torch from tests.kernels.moe.utils import make_test_quant_config from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace from vllm.distributed.parallel_state import ( ensure_model_parallel_initialized, get_dp_group, + get_eplb_group, ) from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe.layer import FusedMoE @@ -171,12 +173,20 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): for lidx in range(test_config.num_layers): shuffled_indices[lidx] = torch.randperm(test_config.num_experts) + expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]] + communicator = create_eplb_communicator( + group_coordinator=get_eplb_group(), + backend="torch_nccl", + expert_weights=rank_expert_weights, + expert_buffer=expert_buffer, + ) rearrange_expert_weights_inplace( indices, shuffled_indices, rank_expert_weights, + expert_buffer, ep_group, - is_profile=False, + communicator, ) num_global_experts = test_config.num_experts diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index e0f73cd657e..188f4448137 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -1287,10 +1287,12 @@ def _test_body_eplb( expert_weights = [list(eplb_moe_layer.get_expert_weights())] + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]] communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), backend=vllm_config.parallel_config.eplb_config.communicator, - expert_weights=expert_weights[0], + expert_weights=expert_weights, + expert_buffer=expert_buffer, ) # Rearrange expert weights across EP ranks @@ -1298,6 +1300,7 @@ def _test_body_eplb( old_global_expert_indices=initial_indices.unsqueeze(0), new_global_expert_indices=shuffled_indices.unsqueeze(0), expert_weights=expert_weights, + expert_buffer=expert_buffer, ep_group=cpu_group, communicator=communicator, ) diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index 65a19626468..c57cc74fc06 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -40,11 +40,7 @@ def _can_p2p(rank: int, world_size: int) -> bool: return True -def is_weak_contiguous(inp: torch.Tensor): - return inp.is_contiguous() or ( - inp.storage().nbytes() - inp.storage_offset() * inp.element_size() - == inp.numel() * inp.element_size() - ) +from vllm.distributed.utils import is_weak_contiguous # noqa: E402 class CustomAllreduce: diff --git a/vllm/distributed/device_communicators/quick_all_reduce.py b/vllm/distributed/device_communicators/quick_all_reduce.py index 9c9d39a91a9..8c7ee7452f1 100644 --- a/vllm/distributed/device_communicators/quick_all_reduce.py +++ b/vllm/distributed/device_communicators/quick_all_reduce.py @@ -24,11 +24,7 @@ except Exception: quick_ar = False -def is_weak_contiguous(inp: torch.Tensor): - return inp.is_contiguous() or ( - inp.storage().nbytes() - inp.storage_offset() * inp.element_size() - == inp.numel() * inp.element_size() - ) +from vllm.distributed.utils import is_weak_contiguous # noqa: E402, F401 class QuickReduceRegime(Enum): diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 2cd6decb3a5..b08c17b141d 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -470,10 +470,14 @@ class ElasticEPScalingExecutor: module._replace_quant_method(module.quant_method.old_quant_method) prepare_communication_buffer_for_model(self.worker.model_runner.model) + eplb_model_state.expert_buffer = [ + torch.empty_like(w) for w in model.expert_weights[0] + ] eplb_model_state.communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), backend=parallel_config.eplb_config.communicator, - expert_weights=model.expert_weights[0], + expert_weights=model.expert_weights, + expert_buffer=eplb_model_state.expert_buffer, ) if ( diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py index 542606fe741..eb2ec260907 100644 --- a/vllm/distributed/eplb/async_worker.py +++ b/vllm/distributed/eplb/async_worker.py @@ -120,6 +120,7 @@ def transfer_run_periodically( ep_group=eplb_group, is_profile=is_profile, cuda_stream=cuda_stream, + layer_idx=layer_idx, ) # Wait until all writes to expert_buffer have finished before making the diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index f8ee90b934f..9cccc05b2ce 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -30,6 +30,7 @@ from vllm.distributed.parallel_state import ( is_local_first_rank, ) from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator +from vllm.distributed.utils import is_weak_contiguous from vllm.logger import init_logger from vllm.platforms import current_platform @@ -63,8 +64,22 @@ class EplbCommunicator(ABC): pass @abstractmethod - def execute(self, old_indices: np.ndarray | None = None) -> None: - pass + def execute(self) -> None: + """Complete all enqueued transfers. + + Some backends perform communication here; others (e.g. NIXL) + issue transfers eagerly in add_recv and only wait here. + On return, all data is available in the destination buffers. + """ + + def set_transfer_context( # noqa: B027 + self, old_indices: np.ndarray, layer_idx: int + ) -> None: + """Pre-set layer context before add_recv calls. + + Default is a no-op; overridden by backends (e.g. NIXL) that need + layer-level context to issue transfers inside add_recv. + """ @property def needs_profile_buffer_reservation(self) -> bool: @@ -125,7 +140,7 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator): ) ) - def execute(self, old_indices: np.ndarray | None = None) -> None: + def execute(self) -> None: if not self._p2p_ops: return try: @@ -168,7 +183,7 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): for tensor in tensors: self._ops.append(("recv", tensor, src_rank)) - def execute(self, old_indices: np.ndarray | None = None) -> None: + def execute(self) -> None: if not self._ops: return @@ -229,29 +244,47 @@ class NixlEplbCommunicator(EplbCommunicator): def __init__( self, cpu_group: ProcessGroup, - expert_weights: Sequence[torch.Tensor], - cuda_stream: torch.cuda.Stream | None = None, + all_expert_weights: Sequence[Sequence[torch.Tensor]], + expert_buffer: Sequence[torch.Tensor], ) -> None: - assert expert_weights, "NixlEplbCommunicator requires non-empty expert_weights." + assert all_expert_weights, ( + "NixlEplbCommunicator requires non-empty all_expert_weights." + ) + assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer." nixl_wrapper_cls = nixl_utils.NixlWrapper if nixl_wrapper_cls is None: raise RuntimeError("NIXL/ RIXL is unavailable.") + self._cpu_group = cpu_group - self._cuda_stream = cuda_stream self._world_size = cpu_group.size() self._rank = cpu_group.rank() - # expert_id -> weight tensors to pack into the send buffer. - self._expert_send_map: dict[int, list[torch.Tensor]] = {} - # src_rank -> expert_id -> weight tensors to unpack after transfer. - self._recv_map: dict[int, dict[int, list[torch.Tensor]]] = {} - self._num_local_experts: int = expert_weights[0].shape[0] - self._device = expert_weights[0].device - for tensor in expert_weights: - assert tensor.device == self._device, ( - "All local EPLB tensors are expected to be on the same device: " - f"expected={self._device}, got={tensor.device}" + + self._all_expert_weights = all_expert_weights + self._expert_buffer = expert_buffer + self._num_local_experts: int = all_expert_weights[0][0].shape[0] + self._device = all_expert_weights[0][0].device + + for layer_tensors in all_expert_weights: + for tensor in layer_tensors: + assert is_weak_contiguous(tensor), ( + "Expert weight tensors must be contiguous in memory" + ) + assert tensor.device == self._device, ( + "All local EPLB tensors are expected to be on the same " + f"device: expected={self._device}, got={tensor.device}" + ) + for tensor in expert_buffer: + assert is_weak_contiguous(tensor), ( + "expert_buffer tensors must be contiguous in memory" ) + # (local_dlist, remote_dlist, xfer_handle) for in-flight READs; + # accumulated by add_recv, drained by execute. + self._xfer_entries: list[tuple[int, int, int]] = [] + # Per-rank expert_id -> physical row; set by set_transfer_context. + self._expert_to_src_row: list[dict[int, int]] | None = None + self._layer_idx: int | None = None + nixl_agent_config = nixl_utils.nixl_agent_config config = ( nixl_agent_config(capture_telemetry=False) @@ -260,15 +293,16 @@ class NixlEplbCommunicator(EplbCommunicator): ) self._nixl_wrapper = nixl_wrapper_cls(self._make_agent_name(), config) self._nixl_memory_type = "VRAM" - self._registered_desc: object | None = None + # NIXL registration handles; deregistered in __del__. + self._registered_descs: list[object] = [] self._remote_agents: dict[int, str] = {} - self._remote_send_meta: dict[int, tuple[int, int]] = {} - self._send_buffer: torch.Tensor = torch.empty(0) - self._recv_buffer: torch.Tensor = torch.empty(0) - self._expert_bytes: int = 0 + # peer -> (layer, tensor) -> (base_ptr, bytes_per_expert, dev_id). + self._remote_send_meta: dict[ + int, dict[tuple[int, int], tuple[int, int, int]] + ] = {} self._cuda_device_id = int(self._device.index or 0) - self._init_step("buffers", self._init_registered_buffers, expert_weights) + self._init_step("buffers", self._init_registered_buffers) self._init_step("agents", self._init_remote_agents) self._init_step("send meta", self._exchange_remote_send_meta) self._log_initialized() @@ -291,19 +325,34 @@ class NixlEplbCommunicator(EplbCommunicator): uid = uuid.uuid4().hex[:8] return f"eplb-{self._rank}{pp_suffix}-{uid}" + def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: + pass + def add_send( self, tensors: list[torch.Tensor], dst_rank: int, expert_id: int, ) -> None: - assert dst_rank != self._rank, ( - "EPLB communicator should not enqueue same-rank sends: " - f"rank={self._rank}, dst_rank={dst_rank}" + # No-op: NIXL READ is receiver-initiated. The sender's expert + # weights are pre-registered and always readable in-place. + pass + + def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None: + # Pre-compute expert_id -> src_row mapping for every rank so that + # add_recv can immediately issue NIXL READs. + assert not self._xfer_entries, ( + f"set_transfer_context() called with {len(self._xfer_entries)} " + f"pending transfers from layer {self._layer_idx}; " + f"execute() was not called after previous add_recv() calls" ) - # An expert sent to multiple peers is packed only once; skip duplicates. - if expert_id not in self._expert_send_map: - self._expert_send_map[expert_id] = tensors + self._layer_idx = layer_idx + n = self._num_local_experts + rank_experts = old_indices[: self._world_size * n].reshape(self._world_size, n) + self._expert_to_src_row = [ + {int(eid): i for i, eid in enumerate(row) if eid != -1} + for row in rank_experts + ] def add_recv( self, @@ -311,13 +360,44 @@ class NixlEplbCommunicator(EplbCommunicator): src_rank: int, expert_id: int, ) -> None: - assert src_rank != self._rank, ( - "EPLB communicator should not enqueue same-rank recvs: " - f"rank={self._rank}, src_rank={src_rank}" + # Build NIXL descriptors and issue the RDMA READ immediately, + # overlapping the transfer with the remaining Python loop in + # move_to_buffer. + assert self._expert_to_src_row is not None and self._layer_idx is not None, ( + "set_transfer_context() must be called before add_recv()" ) - recv_experts = self._recv_map.setdefault(src_rank, {}) - if expert_id not in recv_experts: - recv_experts[expert_id] = tensors + src_row = self._expert_to_src_row[src_rank][expert_id] + layer_idx = self._layer_idx + + local_descs: list[tuple[int, int, int]] = [] + remote_descs: list[tuple[int, int, int]] = [] + for t_idx, t in enumerate(tensors): + send_base, send_stride, remote_dev = self._remote_send_meta[src_rank][ + (layer_idx, t_idx) + ] + assert t.nbytes == send_stride, ( + f"tensor {t_idx} size {t.nbytes} != remote stride {send_stride}" + ) + local_descs.append( + ( + t.data_ptr(), + t.nbytes, + self._cuda_device_id, + ) + ) + remote_descs.append( + ( + send_base + src_row * send_stride, + send_stride, + remote_dev, + ) + ) + + local_h, remote_h, xfer_h = self._create_peer_xfer( + src_rank, local_descs, remote_descs + ) + self._nixl_wrapper.transfer(xfer_h) + self._xfer_entries.append((local_h, remote_h, xfer_h)) def _init_remote_agents(self) -> None: local_metadata = self._nixl_wrapper.get_agent_metadata() @@ -334,73 +414,60 @@ class NixlEplbCommunicator(EplbCommunicator): peer_metadata ) - def _init_registered_buffers(self, expert_weights: Sequence[torch.Tensor]) -> None: - total_bytes = max(sum(t.nbytes for t in expert_weights), 1) - assert total_bytes % self._num_local_experts == 0, ( - f"Number of bytes in moe layer {total_bytes} is not divisible " - f"by number of local experts {self._num_local_experts}" - ) - self._expert_bytes = total_bytes // self._num_local_experts + def _init_registered_buffers(self) -> None: + all_tensors: list[torch.Tensor] = [] + for layer_tensors in self._all_expert_weights: + all_tensors.extend(layer_tensors) + all_tensors.extend(self._expert_buffer) - self._send_buffer = torch.empty( - total_bytes, device=self._device, dtype=torch.uint8 - ) - self._recv_buffer = torch.empty( - total_bytes, device=self._device, dtype=torch.uint8 - ) - - descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer]) + descs = self._nixl_wrapper.get_reg_descs(all_tensors) self._nixl_wrapper.register_memory(descs) - self._registered_desc = descs + self._registered_descs.append(descs) def _exchange_remote_send_meta(self) -> None: - """Exchange send-buffer metadata so each rank can build dynamic - descriptors at execute time.""" - local_meta: tuple[int, int] = ( - self._send_buffer.data_ptr(), - self._cuda_device_id, - ) - gathered_meta: list[tuple[int, int] | None] = [None] * self._world_size + """Exchange per-layer per-tensor metadata so receivers can compute + remote RDMA addresses at transfer time.""" + local_meta: dict[tuple[int, int], tuple[int, int, int]] = {} + for layer_idx, layer_tensors in enumerate(self._all_expert_weights): + for t_idx, t in enumerate(layer_tensors): + nbytes_per_expert = t.nbytes // self._num_local_experts + local_meta[(layer_idx, t_idx)] = ( + t.data_ptr(), + nbytes_per_expert, + self._cuda_device_id, + ) + + # Per-rank map: (layer_idx, tensor_idx) -> (base_ptr, bytes_per_expert, dev_id). + # add_recv uses base_ptr + src_row * bytes_per_expert to compute + # the remote RDMA address for each expert. + gathered_meta: list[dict[tuple[int, int], tuple[int, int, int]] | None] = [ + None + ] * self._world_size torch.distributed.all_gather_object( gathered_meta, local_meta, group=self._cpu_group ) + local_keys = set(local_meta.keys()) for peer in self._remote_agents: peer_meta = gathered_meta[peer] assert peer_meta is not None + peer_keys = set(peer_meta.keys()) + if peer_keys != local_keys: + raise RuntimeError( + f"NIXL EPLB metadata key mismatch with rank {peer}: " + f"local={sorted(local_keys)}, peer={sorted(peer_keys)}" + ) + for key in local_keys: + _, local_stride, _ = local_meta[key] + _, peer_stride, _ = peer_meta[key] + if local_stride != peer_stride: + raise RuntimeError( + f"NIXL EPLB nbytes_per_expert mismatch for {key} " + f"with rank {peer}: " + f"local={local_stride}, peer={peer_stride}" + ) self._remote_send_meta[peer] = peer_meta - @staticmethod - def _pack_send_buffer( - in_tensors: list[torch.Tensor], - send_buffer: torch.Tensor, - byte_offset: int, - ) -> None: - for tensor in in_tensors: - raw = tensor.reshape(-1).view(torch.uint8) - if raw.numel() == 0: - continue - send_buffer[byte_offset : byte_offset + raw.numel()].copy_( - raw, non_blocking=True - ) - byte_offset += raw.numel() - - @staticmethod - def _unpack_recv_buffer( - recv_buffer: torch.Tensor, - out_tensors: list[torch.Tensor], - byte_offset: int, - ) -> None: - for tensor in out_tensors: - num_bytes = tensor.numel() * tensor.element_size() - if num_bytes == 0: - continue - tensor.reshape(-1).view(torch.uint8).copy_( - recv_buffer[byte_offset : byte_offset + num_bytes], - non_blocking=True, - ) - byte_offset += num_bytes - def _wait_for_all_transfers(self, handles: list[int]) -> None: pending = set(handles) while pending: @@ -456,110 +523,52 @@ class NixlEplbCommunicator(EplbCommunicator): ) return (local_handle, remote_handle, xfer_handle) - def execute(self, old_indices: np.ndarray | None = None) -> None: - assert old_indices is not None, ( - "NixlEplbCommunicator.execute requires old_indices" + def execute(self) -> None: + assert self._layer_idx is not None or not self._xfer_entries, ( + "set_transfer_context() must be called before execute() " + "if any add_recv() calls were made" ) - - xfer_entries: list[tuple[int, int, int]] = [] try: - n = self._num_local_experts - rank_experts = old_indices[: self._world_size * n].reshape( - self._world_size, n - ) - # Build expert_id -> send slot mapping per rank. - expert_to_send_slot: list[dict[int, int]] = [ - {int(eid): i for i, eid in enumerate(row) if eid != -1} - for row in rank_experts - ] + self._wait_for_all_transfers([x[2] for x in self._xfer_entries]) - # Phase 1: pack each expert at its slot offset in the send buffer. - with torch.cuda.stream(self._cuda_stream): - for expert_id, tensors in self._expert_send_map.items(): - slot = expert_to_send_slot[self._rank][expert_id] - byte_offset = slot * self._expert_bytes - self._pack_send_buffer(tensors, self._send_buffer, byte_offset) - - # Ensure all packed data is visible in device memory before pulls. - if self._cuda_stream is not None: - self._cuda_stream.synchronize() - else: - torch.cuda.current_stream().synchronize() - # READ is receiver-initiated; synchronize all ranks before transfer. - # We use monitored_barrier so a rank that crashes or exits early - # produces a diagnostic timeout instead of a silent hang. + # Post-READ barrier. + # Correctness fence for zero-copy: prevents overwrite-while- + # remote-read race. torch.distributed.monitored_barrier( group=self._cpu_group, timeout=timedelta(minutes=5), ) - - # Phase 2: issue one batched READ per peer. - recv_offsets: dict[tuple[int, int], int] = {} - recv_offset = 0 - recv_base = self._recv_buffer.data_ptr() - for src in range(self._world_size): - if src == self._rank: - continue - recv_experts = self._recv_map.get(src) - if not recv_experts: - continue - expert_ids = list(recv_experts.keys()) - remote_base, remote_dev = self._remote_send_meta[src] - local_descs: list[tuple[int, int, int]] = [] - remote_descs: list[tuple[int, int, int]] = [] - for expert_id in expert_ids: - slot = expert_to_send_slot[src][expert_id] - remote_off = slot * self._expert_bytes - recv_offsets[(src, expert_id)] = recv_offset - local_descs.append( - ( - recv_base + recv_offset, - self._expert_bytes, - self._cuda_device_id, - ) - ) - remote_descs.append( - (remote_base + remote_off, self._expert_bytes, remote_dev) - ) - recv_offset += self._expert_bytes - assert recv_offset <= self._recv_buffer.nbytes - local_h, remote_h, xfer_h = self._create_peer_xfer( - src, local_descs, remote_descs - ) - self._nixl_wrapper.transfer(xfer_h) - xfer_entries.append((local_h, remote_h, xfer_h)) - - # Phase 3: wait for all in-flight transfers, then unpack. - self._wait_for_all_transfers([x[2] for x in xfer_entries]) - - with torch.cuda.stream(self._cuda_stream): - for (src, expert_id), offset in recv_offsets.items(): - self._unpack_recv_buffer( - self._recv_buffer, - self._recv_map[src][expert_id], - offset, - ) finally: - for local_h, remote_h, xfer_h in xfer_entries: + for local_h, remote_h, xfer_h in self._xfer_entries: with contextlib.suppress(Exception): self._nixl_wrapper.release_xfer_handle(xfer_h) with contextlib.suppress(Exception): self._nixl_wrapper.release_dlist_handle(local_h) with contextlib.suppress(Exception): self._nixl_wrapper.release_dlist_handle(remote_h) - self._expert_send_map.clear() - self._recv_map.clear() + self._xfer_entries.clear() + self._expert_to_src_row = None + self._layer_idx = None def __del__(self) -> None: - try: - if self._registered_desc is not None: - self._nixl_wrapper.deregister_memory(self._registered_desc) - self._registered_desc = None + with contextlib.suppress(Exception): + for local_h, remote_h, xfer_h in self._xfer_entries: + with contextlib.suppress(Exception): + self._nixl_wrapper.release_xfer_handle(xfer_h) + with contextlib.suppress(Exception): + self._nixl_wrapper.release_dlist_handle(local_h) + with contextlib.suppress(Exception): + self._nixl_wrapper.release_dlist_handle(remote_h) + with contextlib.suppress(Exception): + for descs in self._registered_descs: + with contextlib.suppress(Exception): + self._nixl_wrapper.deregister_memory(descs) + self._registered_descs.clear() + with contextlib.suppress(Exception): for agent_name in self._remote_agents.values(): - self._nixl_wrapper.remove_remote_agent(agent_name) + with contextlib.suppress(Exception): + self._nixl_wrapper.remove_remote_agent(agent_name) self._remote_agents.clear() - except Exception as e: - logger.warning("Error during NixlEplbCommunicator cleanup: %s", e) class PyNcclEplbCommunicator(EplbCommunicator): @@ -600,7 +609,7 @@ class PyNcclEplbCommunicator(EplbCommunicator): for tensor in tensors: self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream) - def execute(self, old_indices: np.ndarray | None = None) -> None: + def execute(self) -> None: if self._group_started: self._pynccl_comm.group_end() self._group_started = False @@ -609,7 +618,8 @@ class PyNcclEplbCommunicator(EplbCommunicator): def create_eplb_communicator( group_coordinator: GroupCoordinator, backend: str | None, - expert_weights: Sequence[torch.Tensor], + expert_weights: Sequence[Sequence[torch.Tensor]], + expert_buffer: Sequence[torch.Tensor], ) -> EplbCommunicator: """Create an EPLB communicator for the given backend. @@ -624,16 +634,18 @@ def create_eplb_communicator( ``"pynccl"`` in that case. When tensors reside on CPU, ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU process group. - expert_weights: Expert weight tensors from *one* MoE layer. - NixlEplbCommunicator pre-allocates send/recv buffers sized - to this layer, so all other MoE layers must have the same - tensor count, shapes, and dtypes. + expert_weights: Expert weight tensors for *all* MoE layers. + Shape ``(num_layers)(num_tensors_per_layer)``. + NixlEplbCommunicator registers all layers with NIXL for + zero-copy RDMA reads. + expert_buffer: Pre-allocated receive buffer tensors (one per + weight tensor in a single layer). """ - # Keep a safe default for callers that have not resolved communicator yet. if backend is None: backend = "torch_nccl" - tensor_device_type = expert_weights[0].device.type if expert_weights else "cpu" + first_layer = expert_weights[0] if expert_weights else [] + tensor_device_type = first_layer[0].device.type if first_layer else "cpu" torch_group = ( group_coordinator.cpu_group if tensor_device_type == "cpu" @@ -649,7 +661,7 @@ def create_eplb_communicator( unsupported_dtypes = sorted( { tensor.dtype - for tensor in expert_weights + for tensor in first_layer if not ncclDataTypeEnum.supports_torch_dtype(tensor.dtype) }, key=str, @@ -704,7 +716,8 @@ def create_eplb_communicator( try: return NixlEplbCommunicator( cpu_group=group_coordinator.cpu_group, - expert_weights=expert_weights, + all_expert_weights=expert_weights, + expert_buffer=expert_buffer, ) except Exception as exc: raise RuntimeError( diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 319a5f22c92..2ad4aa9e779 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -450,7 +450,8 @@ class EplbState: communicator = create_eplb_communicator( group_coordinator=get_eplb_group(), backend=self.parallel_config.eplb_config.communicator, - expert_weights=model.expert_weights[0], + expert_weights=model.expert_weights, + expert_buffer=expert_buffer, ) model_state = EplbModelState( @@ -766,6 +767,7 @@ class EplbState: eplb_model_state.physical_to_logical_map, new_physical_to_logical_map, eplb_model_state.model.expert_weights, + eplb_model_state.expert_buffer, ep_group, eplb_model_state.communicator, is_profile, diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 50b7013295c..53b0356dcd8 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -178,6 +178,7 @@ def move_to_buffer( cuda_stream: torch.cuda.Stream | None, ep_rank: int, communicator: EplbCommunicator, + layer_idx: int = 0, ) -> TransferMetadata: """ Rearranges expert weights during EPLB rebalancing. @@ -193,6 +194,7 @@ def move_to_buffer( cuda_stream: CUDA stream for async copies (can be None for sync mode). ep_rank: Rank of this process in expert parallel group. communicator: EplbCommunicator instance for P2P communication. + layer_idx: Index of the MoE layer being transferred. Returns: TransferMetadata: Metadata needed for completing remote weight transfers. @@ -265,6 +267,8 @@ def move_to_buffer( for w, b in zip(expert_weights, expert_weights_buffers): b[dst].copy_(w[src_local], non_blocking=True) + communicator.set_transfer_context(old_indices, layer_idx) + # 2. Post sends if send_count > 0: experts = send_expert_ids[:send_count] @@ -331,9 +335,8 @@ def move_to_buffer( expert_id=int(expert), ) - # 4. Execute the P2P operations. The real communication happens here. - communicator.execute(old_indices=old_indices) - # wait for the communication to finish + # 4. Execute transfers and wait for completion. + communicator.execute() return TransferMetadata( is_unchanged=is_unchanged, is_received_locally=is_received_locally, @@ -431,6 +434,7 @@ def transfer_layer( is_profile: bool = False, cuda_stream: torch.cuda.Stream | None = None, rank_mapping: dict[int, int] | None = None, + layer_idx: int = 0, ) -> TransferMetadata: """ Rearranges the expert weights in place according to the new expert indices. @@ -452,6 +456,7 @@ def transfer_layer( communications to reserve enough memory for the buffers. cuda_stream: CUDA stream for async copies (can be None for sync mode). rank_mapping: Optional rank mapping for elastic expert parallelism. + layer_idx: Index of the MoE layer being transferred. Returns: TransferMetadata: Metadata needed for completing remote weight transfers, @@ -499,6 +504,7 @@ def transfer_layer( cuda_stream=cuda_stream, ep_rank=ep_group.rank(), communicator=communicator, + layer_idx=layer_idx, ) @@ -506,6 +512,7 @@ def rearrange_expert_weights_inplace( old_global_expert_indices: torch.Tensor, new_global_expert_indices: torch.Tensor, expert_weights: Sequence[Sequence[torch.Tensor]], + expert_buffer: Sequence[torch.Tensor], ep_group: ProcessGroup, communicator: EplbCommunicator, is_profile: bool = False, @@ -524,6 +531,8 @@ def rearrange_expert_weights_inplace( of tensors of shape (num_local_physical_experts, hidden_size_i). For example, a linear layer may have up and down projection, so weight_count = 2. Each weight's hidden size can be different. + expert_buffer: Pre-allocated receive buffer tensors (one per + weight tensor in a single layer). ep_group: The device process group for expert parallelism. communicator: EplbCommunicator instance for P2P communication. is_profile (bool): If `True`, do not perform any actual weight copy. @@ -566,10 +575,10 @@ def rearrange_expert_weights_inplace( # Reserve NCCL communication buffers via a dummy all_gather. # Backends that pre-allocate their own transfer buffers # skip this to avoid the extra memory spike during profiling. - weights_buffer: list[torch.Tensor] = [ + profile_buffer: list[torch.Tensor] = [ torch.empty_like(w) for w in first_layer_weights ] - for weight, buffer in zip(expert_weights[0], weights_buffer): + for weight, buffer in zip(expert_weights[0], profile_buffer): dummy_recv_buffer = [buffer for _ in range(ep_size)] torch.distributed.barrier() all_gather( @@ -579,10 +588,7 @@ def rearrange_expert_weights_inplace( ) return - # Buffers to hold the expert weights during the exchange. - # NOTE: Currently we assume the same weights across different layers - # have the same shape. - weights_buffer = [torch.empty_like(w) for w in first_layer_weights] + weights_buffer = list(expert_buffer) old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy() new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy() @@ -597,6 +603,7 @@ def rearrange_expert_weights_inplace( cuda_stream=None, ep_rank=ep_rank, communicator=communicator, + layer_idx=layer_idx, ) move_from_buffer( diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index ba0419a2800..1e38794603c 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -64,6 +64,20 @@ def divide(numerator, denominator): return numerator // denominator +def is_weak_contiguous(inp: torch.Tensor) -> bool: + """Check that *inp* occupies a single contiguous block of memory. + + Unlike ``torch.Tensor.is_contiguous()``, this also accepts tensors + whose strides are not strictly C-contiguous (e.g. column-major) as + long as the underlying storage from the tensor's offset onward is + exactly ``numel * element_size`` bytes. + """ + return inp.is_contiguous() or ( + inp.storage().nbytes() - inp.storage_offset() * inp.element_size() + == inp.numel() * inp.element_size() + ) + + def split_tensor_along_last_dim( tensor: torch.Tensor, num_partitions: int,