From 50dd4cb42726777635acda9ed4f5440ae4a2e281 Mon Sep 17 00:00:00 2001 From: Ilya Markov Date: Mon, 20 Apr 2026 12:24:23 +0200 Subject: [PATCH] [EPLB] Add nixl-based eplb communicator (#36276) Signed-off-by: ilmarkov Signed-off-by: Markov Ilya --- docs/serving/expert_parallel_deployment.md | 1 + tests/distributed/test_eplb_execute.py | 17 +- vllm/config/parallel.py | 3 +- vllm/distributed/eplb/eplb_communicator.py | 446 +++++++++++++++++- vllm/distributed/eplb/rebalance_execute.py | 34 +- .../kv_transfer/kv_connector/v1/nixl/stats.py | 4 +- .../kv_transfer/kv_connector/v1/nixl/utils.py | 46 -- .../kv_connector/v1/nixl/worker.py | 3 +- vllm/distributed/nixl_utils.py | 54 +++ vllm/v1/worker/gpu_model_runner.py | 28 +- 10 files changed, 556 insertions(+), 80 deletions(-) create mode 100644 vllm/distributed/nixl_utils.py diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index d75ae7feb49..fef4df770fa 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -153,6 +153,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` | | `policy` | The policy type for expert parallel load balancing | `"default"` | +| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` | For example: diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 7256ad56a3f..8a46087e991 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -9,7 +9,10 @@ import torch import torch.distributed from vllm.config import VllmConfig, set_current_vllm_config -from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator +from vllm.distributed.eplb.eplb_communicator import ( + create_eplb_communicator, + has_nixl, +) from vllm.distributed.eplb.rebalance_execute import ( move_from_buffer, rearrange_expert_weights_inplace, @@ -527,7 +530,9 @@ def _test_rearrange_expert_weights_with_redundancy( (4, 8, 8, 16), ], ) -@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"]) +@pytest.mark.parametrize( + "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"] +) def test_rearrange_expert_weights_with_redundancy( world_size, num_layers, @@ -537,6 +542,8 @@ def test_rearrange_expert_weights_with_redundancy( ): """Test the functionality of rearranging expert weights with redundancy.""" + if eplb_communicator == "nixl" and not has_nixl(): + pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") distributed_run( @@ -633,7 +640,9 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None: (2, 2, 2, 3), ], ) -@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"]) +@pytest.mark.parametrize( + "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"] +) def test_async_transfer_layer_without_mtp( world_size: int, num_layers: int, @@ -643,6 +652,8 @@ def test_async_transfer_layer_without_mtp( ): """Exercise async EPLB transfer path without MTP/spec decode.""" + if eplb_communicator == "nixl" and not has_nixl(): + pytest.skip("NIXL is not available") if torch.accelerator.device_count() < world_size: pytest.skip(f"Need at least {world_size} GPUs to run the test") diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index a42b8422ef3..afd0d1dd501 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -36,7 +36,7 @@ DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"] DataParallelBackend = Literal["ray", "mp"] EPLBPolicyOption = Literal["default"] DCPCommBackend = Literal["ag_rs", "a2a"] -EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "pynccl"] +EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"] All2AllBackend = Literal[ "naive", "pplx", @@ -90,6 +90,7 @@ class EPLBConfig: Backend for EPLB expert weight communication: - "torch_nccl": Use torch.distributed on the device process group - "torch_gloo": Use torch.distributed gloo with CPU staging + - "nixl": Use NIXL/ RIXL with staged send/recv buffers - "pynccl": Use PyNccl send/recv - None: Auto-select backend ("torch_gloo" for async, "torch_nccl" for sync) """ diff --git a/vllm/distributed/eplb/eplb_communicator.py b/vllm/distributed/eplb/eplb_communicator.py index 059d6b7cffb..6ff41272fce 100644 --- a/vllm/distributed/eplb/eplb_communicator.py +++ b/vllm/distributed/eplb/eplb_communicator.py @@ -4,8 +4,12 @@ EPLB communicator implementations and factory. """ +import contextlib +import time +import uuid from abc import ABC, abstractmethod from collections.abc import Sequence +from datetime import timedelta import torch from torch.distributed import ( @@ -18,13 +22,27 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.device_communicators.pynccl_wrapper import ( ncclDataTypeEnum, ) -from vllm.distributed.parallel_state import GroupCoordinator, is_local_first_rank +from vllm.distributed.nixl_utils import ( + NixlWrapper, + nixl_agent_config, +) +from vllm.distributed.parallel_state import ( + GroupCoordinator, + get_pp_group, + is_local_first_rank, +) from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) +def has_nixl() -> bool: + """Whether the optional NIXL / RIXL package is available.""" + return NixlWrapper is not None + + class EplbCommunicator(ABC): """Abstract EPLB communicator for expert weight transfers.""" @@ -40,6 +58,12 @@ class EplbCommunicator(ABC): def execute(self) -> None: pass + @property + def needs_profile_buffer_reservation(self) -> bool: + """Whether the profile path must run a dummy collective operation to reserve + communication buffers.""" + return True + def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None: self._cuda_stream = cuda_stream @@ -167,6 +191,385 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator): dst_tensor.copy_(cpu_tensor, non_blocking=True) +class NixlEplbCommunicator(EplbCommunicator): + """EPLB communicator backed by NIXL READ transfers.""" + + def __init__( + self, + cpu_group: ProcessGroup, + expert_weights: Sequence[torch.Tensor], + cuda_stream: torch.cuda.Stream | None = None, + ) -> None: + assert expert_weights, "NixlEplbCommunicator requires non-empty expert_weights." + if NixlWrapper 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() + self._send_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {} + self._recv_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {} + self._dtypes: list[torch.dtype] = [] + 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}" + ) + if tensor.dtype not in self._dtypes: + self._dtypes.append(tensor.dtype) + + config = ( + nixl_agent_config(capture_telemetry=False) + if nixl_agent_config is not None + else None + ) + self._nixl_wrapper = NixlWrapper(self._make_agent_name(), config) + self._nixl_memory_type = "VRAM" + self._registered_desc: object | None = None + self._remote_agents: dict[int, str] = {} + self._remote_send_meta: dict[int, tuple[int, int, int]] = {} + self._send_buffer: torch.Tensor = torch.empty(0) + self._recv_buffer: torch.Tensor = torch.empty(0) + self._peer_partition_bytes: int = 0 + self._dtype_max_bytes: dict[torch.dtype, int] = {} + self._cuda_device_id = int(self._device.index or 0) + self._xfer_cache: dict[tuple[int, int, int], tuple[int, int, int]] = {} + self._init_step("buffers", self._init_registered_buffers, expert_weights) + self._init_step("agents", self._init_remote_agents) + self._init_step("send meta", self._exchange_remote_send_meta) + self._log_initialized() + + @property + def needs_profile_buffer_reservation(self) -> bool: + return False + + @staticmethod + def _init_step(name: str, fn: object, *args: object, **kwargs: object) -> None: + try: + fn(*args, **kwargs) # type: ignore[operator] + except Exception as exc: + raise RuntimeError(f"NIXL EPLB init failed: {name}") from exc + + def _make_agent_name(self) -> str: + """Build a deployment-unique nixl agent name.""" + pp_size = get_pp_group().world_size + pp_suffix = f"-pp{get_pp_group().rank_in_group}" if pp_size > 1 else "" + uid = uuid.uuid4().hex[:8] + return f"eplb-{self._rank}{pp_suffix}-{uid}" + + def _get_peer_buckets( + self, + bucket_map: dict[torch.dtype, list[list[torch.Tensor]]], + dtype: torch.dtype, + ) -> list[list[torch.Tensor]]: + peer_buckets = bucket_map.get(dtype) + if peer_buckets is None: + peer_buckets = [[] for _ in range(self._world_size)] + bucket_map[dtype] = peer_buckets + return peer_buckets + + def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None: + assert dst_rank != self._rank, ( + "EPLB communicator should not enqueue same-rank sends: " + f"rank={self._rank}, dst_rank={dst_rank}" + ) + self._get_peer_buckets(self._send_tensors, tensor.dtype)[dst_rank].append( + tensor + ) + + def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None: + assert src_rank != self._rank, ( + "EPLB communicator should not enqueue same-rank recvs: " + f"rank={self._rank}, src_rank={src_rank}" + ) + self._get_peer_buckets(self._recv_tensors, tensor.dtype)[src_rank].append( + tensor + ) + + def _init_remote_agents(self) -> None: + local_metadata = self._nixl_wrapper.get_agent_metadata() + gathered_metadata: list[bytes | None] = [None] * self._world_size + torch.distributed.all_gather_object( + gathered_metadata, local_metadata, group=self._cpu_group + ) + for peer in range(self._world_size): + if peer == self._rank: + continue + peer_metadata = gathered_metadata[peer] + assert peer_metadata is not None + self._remote_agents[peer] = self._nixl_wrapper.add_remote_agent( + peer_metadata + ) + + def _init_registered_buffers(self, expert_weights: Sequence[torch.Tensor]) -> None: + total_max_bytes = 0 + for dtype in self._dtypes: + max_numel = max( + sum(t.numel() for t in expert_weights if t.dtype == dtype), 1 + ) + max_bytes = max_numel * dtype.itemsize + self._dtype_max_bytes[dtype] = max_bytes + total_max_bytes += max_bytes + + self._peer_partition_bytes = total_max_bytes + + # The send buffer needs world_size partitions because remote peers + # READ from fixed offsets (rank * partition_bytes). + # This allocates world_size * partition_bytes + # which can cause OOM on large models. + # TODO(ilmarkov): shrink to const * partition_bytes and execute + # communication in multiple steps dealing with the worst case. + send_total_bytes = self._peer_partition_bytes * self._world_size + + self._send_buffer = torch.empty( + send_total_bytes, device=self._device, dtype=torch.uint8 + ) + self._recv_buffer = torch.empty( + self._peer_partition_bytes, device=self._device, dtype=torch.uint8 + ) + + descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer]) + self._nixl_wrapper.register_memory(descs) + self._registered_desc = 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, int] = ( + self._send_buffer.data_ptr(), + self._peer_partition_bytes, + self._cuda_device_id, + ) + gathered_meta: list[tuple[int, int, int] | None] = [None] * self._world_size + torch.distributed.all_gather_object( + gathered_meta, local_meta, group=self._cpu_group + ) + + for peer in self._remote_agents: + peer_meta = gathered_meta[peer] + assert peer_meta is not None + self._remote_send_meta[peer] = peer_meta + + @staticmethod + def _pack_send_buffer( + peer_tensors: list[torch.Tensor], + send_buffer: torch.Tensor, + byte_offset: int, + ) -> int: + """ + Returns the byte offset after the last written byte. + """ + for tensor in peer_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() + return byte_offset + + @staticmethod + def _unpack_recv_buffer( + recv_buffer: torch.Tensor, + peer_tensors: list[torch.Tensor], + byte_offset: int, + ) -> int: + """ + Returns the byte offset after the last read byte. + """ + for tensor in peer_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 + return byte_offset + + def _release_all_cached_handles(self) -> None: + """Best-effort release of every cached dlist and xfer handle.""" + for local_dlist, remote_dlist, xfer in self._xfer_cache.values(): + for release_fn, handle in ( + (self._nixl_wrapper.release_xfer_handle, xfer), + (self._nixl_wrapper.release_dlist_handle, local_dlist), + (self._nixl_wrapper.release_dlist_handle, remote_dlist), + ): + with contextlib.suppress(Exception): + release_fn(handle) + self._xfer_cache.clear() + + def _wait_for_all_transfers(self, handles: list[int]) -> None: + pending = set(handles) + while pending: + completed: list[int] = [] + for handle in pending: + state = self._nixl_wrapper.check_xfer_state(handle) + if state == "DONE": + completed.append(handle) + continue + if state != "PROC": + raise RuntimeError(f"NIXL transfer failed with state={state}") + for handle in completed: + pending.remove(handle) + if pending: + time.sleep(0.0005) + + def _get_or_create_xfer(self, src: int, total_bytes: int, recv_offset: int) -> int: + """Return a cached xfer handle or create and cache a new one.""" + key = (src, total_bytes, recv_offset) + cached = self._xfer_cache.get(key) + if cached is not None: + return cached[2] + + recv_base = self._recv_buffer.data_ptr() + local_desc = self._nixl_wrapper.get_xfer_descs( + [ + ( + recv_base + recv_offset, + total_bytes, + self._cuda_device_id, + ) + ], + self._nixl_memory_type, + ) + local_handle = self._nixl_wrapper.prep_xfer_dlist( + "NIXL_INIT_AGENT", + local_desc, + ) + + remote_base, remote_part_bytes, remote_dev = self._remote_send_meta[src] + agent_name = self._remote_agents[src] + remote_desc = self._nixl_wrapper.get_xfer_descs( + [ + ( + remote_base + self._rank * remote_part_bytes, + total_bytes, + remote_dev, + ) + ], + self._nixl_memory_type, + ) + remote_handle = self._nixl_wrapper.prep_xfer_dlist( + agent_name, + remote_desc, + ) + + xfer_handle = self._nixl_wrapper.make_prepped_xfer( + "READ", + local_handle, + [0], + remote_handle, + [0], + ) + self._xfer_cache[key] = (local_handle, remote_handle, xfer_handle) + return xfer_handle + + def execute(self) -> None: + xfer_handles: list[int] = [] + try: + # Phase 1: pack send buffers. + with torch.cuda.stream(self._cuda_stream): + for dst in range(self._world_size): + byte_offset = dst * self._peer_partition_bytes + for dtype in self._dtypes: + peer_tensors = self._send_tensors.get( + dtype, [[] for _ in range(self._world_size)] + )[dst] + actual_bytes = sum( + t.numel() * t.element_size() for t in peer_tensors + ) + if actual_bytes > self._dtype_max_bytes[dtype]: + raise RuntimeError( + "NIXL EPLB send overflow for dtype " + f"{dtype}: peer={dst}, " + f"required={actual_bytes}, " + f"capacity={self._dtype_max_bytes[dtype]}" + ) + byte_offset = self._pack_send_buffer( + peer_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. + torch.distributed.monitored_barrier( + group=self._cpu_group, + timeout=timedelta(minutes=5), + ) + + # Phase 2: look up or create descriptors and issue all READs. + # Data from all peers is packed sequentially into the single + # partition-sized recv buffer at running offsets. + recv_offsets: dict[int, int] = {} + recv_offset = 0 + for src in range(self._world_size): + if src == self._rank: + continue + actual_total_bytes = 0 + for dtype in self._dtypes: + peer_tensors = self._recv_tensors.get( + dtype, [[] for _ in range(self._world_size)] + )[src] + actual_total_bytes += sum( + t.numel() * t.element_size() for t in peer_tensors + ) + if actual_total_bytes == 0: + continue + + recv_offsets[src] = recv_offset + xfer_handle = self._get_or_create_xfer( + src, actual_total_bytes, recv_offset + ) + self._nixl_wrapper.transfer(xfer_handle) + xfer_handles.append(xfer_handle) + recv_offset += actual_total_bytes + + # Phase 3: single wait for all in-flight transfers, then unpack. + self._wait_for_all_transfers(xfer_handles) + + with torch.cuda.stream(self._cuda_stream): + for src, offset in recv_offsets.items(): + byte_offset = offset + for dtype in self._dtypes: + peer_tensors = self._recv_tensors.get( + dtype, [[] for _ in range(self._world_size)] + )[src] + byte_offset = self._unpack_recv_buffer( + self._recv_buffer, + peer_tensors, + byte_offset, + ) + except Exception: + self._release_all_cached_handles() + raise + finally: + self._send_tensors.clear() + self._recv_tensors.clear() + + def __del__(self) -> None: + try: + self._release_all_cached_handles() + if self._registered_desc is not None: + self._nixl_wrapper.deregister_memory(self._registered_desc) + self._registered_desc = None + for agent_name in self._remote_agents.values(): + 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): """EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv.""" @@ -204,6 +607,24 @@ def create_eplb_communicator( backend: str | None, expert_weights: Sequence[torch.Tensor], ) -> EplbCommunicator: + """Create an EPLB communicator for the given backend. + + Args: + group_coordinator: Process-group coordinator that provides the + device and CPU communication groups. + backend: Communicator backend name (``"torch_nccl"``, + ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``). + Falls back to ``"torch_nccl"`` when *None*. + Stateless (elastic EP) groups only support ``"torch_nccl"`` + and ``"pynccl"``; ``"torch_nccl"`` is silently promoted to + ``"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. + """ # Keep a safe default for callers that have not resolved communicator yet. if backend is None: backend = "torch_nccl" @@ -256,7 +677,7 @@ def create_eplb_communicator( if backend not in ("torch_nccl", "pynccl"): raise ValueError( f"Elastic EP requires 'torch_nccl' or 'pynccl' EPLB communicator " - f"(got '{backend}'). torch_gloo is not supported with stateless groups." + f"(got '{backend}')." ) if backend == "torch_nccl": logger.warning( @@ -266,7 +687,26 @@ def create_eplb_communicator( backend = "pynccl" return _create_pynccl() - if backend == "torch_gloo": + if backend == "nixl": + if not has_nixl(): + raise RuntimeError( + "EPLB communicator 'nixl' requested but NIXL is unavailable." + ) + if not (current_platform.is_cuda_alike() and tensor_device_type != "cpu"): + raise RuntimeError( + "EPLB communicator 'nixl' supports only cuda-like devices " + f"(got {tensor_device_type})." + ) + try: + return NixlEplbCommunicator( + cpu_group=group_coordinator.cpu_group, + expert_weights=expert_weights, + ) + except Exception as exc: + raise RuntimeError( + f"Failed to initialize NixlEplbCommunicator ({exc})." + ) from exc + elif backend == "torch_gloo": return TorchDistGlooStagedEplbCommunicator( cpu_group=group_coordinator.cpu_group, ) diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 6670ccc6ee6..6c8a2dcc319 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -541,23 +541,29 @@ def rearrange_expert_weights_inplace( assert num_physical_experts == ep_size * num_local_physical_experts first_layer_weights = list(expert_weights[0]) + + if is_profile: + if communicator.needs_profile_buffer_reservation: + # 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] = [ + torch.empty_like(w) for w in first_layer_weights + ] + for weight, buffer in zip(expert_weights[0], weights_buffer): + dummy_recv_buffer = [buffer for _ in range(ep_size)] + torch.distributed.barrier() + all_gather( + dummy_recv_buffer, + weight, + group=ep_group, + ) + 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: list[torch.Tensor] = [ - torch.empty_like(w) for w in first_layer_weights - ] - if is_profile: - # Reserve communication buffers via a minimal dummy all_gather on first layer - for weight, buffer in zip(expert_weights[0], weights_buffer): - dummy_recv_buffer = [buffer for _ in range(ep_size)] - torch.distributed.barrier() - all_gather( - dummy_recv_buffer, - weight, - group=ep_group, - ) - return + weights_buffer = [torch.empty_like(w) for w in first_layer_weights] # NOTE(bowen): We need this synchronize to run, but I don't know why. # If you figure out the reason, please let me know -- thank you! diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py index fde99c1e2c2..65c553cfec3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py @@ -15,9 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( PromMetric, PromMetricT, ) -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( - nixlXferTelemetry, -) +from vllm.distributed.nixl_utils import nixlXferTelemetry from vllm.v1.metrics.utils import create_metric_per_engine diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py index 514214347ae..d0b72464a27 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py @@ -3,60 +3,14 @@ """Shared constants, lazy imports and helpers for the NIXL connector.""" import contextlib -import os -import sys from collections.abc import Iterator from typing import Any import zmq -from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.network_utils import make_zmq_socket -logger = init_logger(__name__) - - -# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used -try: - if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ: - # avoid a memory leak in UCX when using NIXL on some models - # see: https://github.com/vllm-project/vllm/issues/24264 - if "nixl" in sys.modules or "rixl" in sys.modules: - logger.warning( - "NIXL was already imported, we can't reset UCX_RCACHE_MAX_UNRELEASED. " - "Please set it to '1024' manually." - ) - else: - logger.info( - "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare " - "memory leak in UCX when using NIXL." - ) - os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" - - if not current_platform.is_rocm(): - from nixl._api import nixl_agent as NixlWrapper - from nixl._bindings import nixlXferTelemetry - else: - from rixl._api import nixl_agent as NixlWrapper - from rixl._bindings import nixlXferTelemetry - - logger.info("NIXL is available") -except ImportError: - logger.warning("NIXL is not available") - NixlWrapper = None - nixlXferTelemetry = None - - -try: - if not current_platform.is_rocm(): - from nixl._api import nixl_agent_config - else: - from rixl._api import nixl_agent_config -except ImportError: - nixl_agent_config = None - logger.warning("NIXL agent config is not available") - # Supported platforms and types of kv transfer buffer. # {device: tuple of supported kv buffer types} _NIXL_SUPPORTED_DEVICE = { diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py index 45aa33033e7..8157df63d1b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py @@ -45,8 +45,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import ( ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import ( _NIXL_SUPPORTED_DEVICE, - NixlWrapper, - nixl_agent_config, zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( @@ -54,6 +52,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import compute_mamba_phys_ratio, derive_mamba_conv_split, ) +from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config from vllm.distributed.parallel_state import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, diff --git a/vllm/distributed/nixl_utils.py b/vllm/distributed/nixl_utils.py new file mode 100644 index 00000000000..b2b433339be --- /dev/null +++ b/vllm/distributed/nixl_utils.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import sys + +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ: + if "nixl" in sys.modules or "rixl" in sys.modules: + logger.warning_once( + "NIXL was already imported, we can't reset " + "UCX_RCACHE_MAX_UNRELEASED. " + "Please set it to '1024' manually." + ) + else: + logger.info_once( + "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare " + "memory leak in UCX when using NIXL." + ) + os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024" + +try: + if current_platform.is_cuda(): + from nixl._api import nixl_agent as NixlWrapper + else: + from rixl._api import nixl_agent as NixlWrapper + + logger.info_once("NIXL is available") +except ImportError: + logger.warning_once("NIXL is not available") + NixlWrapper = None # type: ignore[assignment, misc] + +try: + if current_platform.is_cuda(): + from nixl._api import nixl_agent_config + else: + from rixl._api import nixl_agent_config +except ImportError: + nixl_agent_config = None # type: ignore[assignment] + logger.warning_once("NIXL agent config is not available") + +try: + if current_platform.is_cuda(): + from nixl._bindings import nixlXferTelemetry + else: + from rixl._bindings import nixlXferTelemetry +except ImportError: + nixlXferTelemetry = None # type: ignore[assignment, misc] + +__all__ = ["NixlWrapper", "nixl_agent_config", "nixlXferTelemetry"] diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1e05d68078e..443c0aae242 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4821,6 +4821,23 @@ class GPUModelRunner( ) self.model.set_aux_hidden_state_layers(aux_layers) + + if ( + is_mixture_of_experts(self.model) + and self.parallel_config.enable_eplb + and not load_dummy_weights + ): + logger.info_once( + "EPLB is enabled for model %s.", + self.model_config.model, + ) + assert self.eplb_state is not None + self.eplb_state.add_model( + self.model, + self.model_config, + ) + eplb_models += 1 + time_after_load = time.perf_counter() self.model_memory_usage = m.consumed_memory except torch.cuda.OutOfMemoryError as e: @@ -4860,15 +4877,10 @@ class GPUModelRunner( is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb and not load_dummy_weights + and self.eplb_state is not None + and self.eplb_state.is_async ): - logger.info_once("EPLB is enabled for model %s.", self.model_config.model) - assert self.eplb_state is not None - self.eplb_state.add_model( - self.model, - self.model_config, - ) - if self.eplb_state.is_async: - self.eplb_state.start_async_loop() + self.eplb_state.start_async_loop() if ( self.vllm_config.compilation_config.mode