@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 50dd4cb427
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 22:30 UTC

[EPLB] Add nixl-based eplb communicator (#36276)

Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com>

Ilya Markov committed Apr 20, 2026 at 12:24 UTC 50dd4cb42726777635acda9ed4f5440ae4a2e281
10 files changed +556 -80
docs/serving/expert_parallel_deployment.md
+1
@@ -153,6 +153,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
153 | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
154 | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` |
155 | `policy` | The policy type for expert parallel load balancing | `"default"` |
156 +| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` |
157
158 For example:
159
tests/distributed/test_eplb_execute.py
+14 -3
@@ -9,7 +9,10 @@ import torch
9 import torch.distributed
10
11 from vllm.config import VllmConfig, set_current_vllm_config
12 -from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
12 +from vllm.distributed.eplb.eplb_communicator import (
13 + create_eplb_communicator,
14 + has_nixl,
15 +)
16 from vllm.distributed.eplb.rebalance_execute import (
17 move_from_buffer,
18 rearrange_expert_weights_inplace,
@@ -527,7 +530,9 @@ def _test_rearrange_expert_weights_with_redundancy(
530 (4, 8, 8, 16),
531 ],
532 )
530 -@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])
533 +@pytest.mark.parametrize(
534 + "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
535 +)
536 def test_rearrange_expert_weights_with_redundancy(
537 world_size,
538 num_layers,
@@ -537,6 +542,8 @@ def test_rearrange_expert_weights_with_redundancy(
542 ):
543 """Test the functionality of rearranging expert weights with redundancy."""
544
545 + if eplb_communicator == "nixl" and not has_nixl():
546 + pytest.skip("NIXL is not available")
547 if torch.accelerator.device_count() < world_size:
548 pytest.skip(f"Need at least {world_size} GPUs to run the test")
549 distributed_run(
@@ -633,7 +640,9 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
640 (2, 2, 2, 3),
641 ],
642 )
636 -@pytest.mark.parametrize("eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl"])
643 +@pytest.mark.parametrize(
644 + "eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
645 +)
646 def test_async_transfer_layer_without_mtp(
647 world_size: int,
648 num_layers: int,
@@ -643,6 +652,8 @@ def test_async_transfer_layer_without_mtp(
652 ):
653 """Exercise async EPLB transfer path without MTP/spec decode."""
654
655 + if eplb_communicator == "nixl" and not has_nixl():
656 + pytest.skip("NIXL is not available")
657 if torch.accelerator.device_count() < world_size:
658 pytest.skip(f"Need at least {world_size} GPUs to run the test")
659
vllm/config/parallel.py
+2 -1
@@ -36,7 +36,7 @@ DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]
36 DataParallelBackend = Literal["ray", "mp"]
37 EPLBPolicyOption = Literal["default"]
38 DCPCommBackend = Literal["ag_rs", "a2a"]
39 -EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "pynccl"]
39 +EPLBCommunicatorBackend = Literal["torch_nccl", "torch_gloo", "nixl", "pynccl"]
40 All2AllBackend = Literal[
41 "naive",
42 "pplx",
@@ -90,6 +90,7 @@ class EPLBConfig:
90 Backend for EPLB expert weight communication:
91 - "torch_nccl": Use torch.distributed on the device process group
92 - "torch_gloo": Use torch.distributed gloo with CPU staging
93 + - "nixl": Use NIXL/ RIXL with staged send/recv buffers
94 - "pynccl": Use PyNccl send/recv
95 - None: Auto-select backend ("torch_gloo" for async, "torch_nccl" for sync)
96 """
vllm/distributed/eplb/eplb_communicator.py
+443 -3
@@ -4,8 +4,12 @@
4 EPLB communicator implementations and factory.
5 """
6
7 +import contextlib
8 +import time
9 +import uuid
10 from abc import ABC, abstractmethod
11 from collections.abc import Sequence
12 +from datetime import timedelta
13
14 import torch
15 from torch.distributed import (
@@ -18,13 +22,27 @@ from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
22 from vllm.distributed.device_communicators.pynccl_wrapper import (
23 ncclDataTypeEnum,
24 )
21 -from vllm.distributed.parallel_state import GroupCoordinator, is_local_first_rank
25 +from vllm.distributed.nixl_utils import (
26 + NixlWrapper,
27 + nixl_agent_config,
28 +)
29 +from vllm.distributed.parallel_state import (
30 + GroupCoordinator,
31 + get_pp_group,
32 + is_local_first_rank,
33 +)
34 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator
35 from vllm.logger import init_logger
36 +from vllm.platforms import current_platform
37
38 logger = init_logger(__name__)
39
40
41 +def has_nixl() -> bool:
42 + """Whether the optional NIXL / RIXL package is available."""
43 + return NixlWrapper is not None
44 +
45 +
46 class EplbCommunicator(ABC):
47 """Abstract EPLB communicator for expert weight transfers."""
48
@@ -40,6 +58,12 @@ class EplbCommunicator(ABC):
58 def execute(self) -> None:
59 pass
60
61 + @property
62 + def needs_profile_buffer_reservation(self) -> bool:
63 + """Whether the profile path must run a dummy collective operation to reserve
64 + communication buffers."""
65 + return True
66 +
67 def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:
68 self._cuda_stream = cuda_stream
69
@@ -167,6 +191,385 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):
191 dst_tensor.copy_(cpu_tensor, non_blocking=True)
192
193
194 +class NixlEplbCommunicator(EplbCommunicator):
195 + """EPLB communicator backed by NIXL READ transfers."""
196 +
197 + def __init__(
198 + self,
199 + cpu_group: ProcessGroup,
200 + expert_weights: Sequence[torch.Tensor],
201 + cuda_stream: torch.cuda.Stream | None = None,
202 + ) -> None:
203 + assert expert_weights, "NixlEplbCommunicator requires non-empty expert_weights."
204 + if NixlWrapper is None:
205 + raise RuntimeError("NIXL/ RIXL is unavailable.")
206 + self._cpu_group = cpu_group
207 + self._cuda_stream = cuda_stream
208 + self._world_size = cpu_group.size()
209 + self._rank = cpu_group.rank()
210 + self._send_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {}
211 + self._recv_tensors: dict[torch.dtype, list[list[torch.Tensor]]] = {}
212 + self._dtypes: list[torch.dtype] = []
213 + self._device = expert_weights[0].device
214 + for tensor in expert_weights:
215 + assert tensor.device == self._device, (
216 + "All local EPLB tensors are expected to be on the same device: "
217 + f"expected={self._device}, got={tensor.device}"
218 + )
219 + if tensor.dtype not in self._dtypes:
220 + self._dtypes.append(tensor.dtype)
221 +
222 + config = (
223 + nixl_agent_config(capture_telemetry=False)
224 + if nixl_agent_config is not None
225 + else None
226 + )
227 + self._nixl_wrapper = NixlWrapper(self._make_agent_name(), config)
228 + self._nixl_memory_type = "VRAM"
229 + self._registered_desc: object | None = None
230 + self._remote_agents: dict[int, str] = {}
231 + self._remote_send_meta: dict[int, tuple[int, int, int]] = {}
232 + self._send_buffer: torch.Tensor = torch.empty(0)
233 + self._recv_buffer: torch.Tensor = torch.empty(0)
234 + self._peer_partition_bytes: int = 0
235 + self._dtype_max_bytes: dict[torch.dtype, int] = {}
236 + self._cuda_device_id = int(self._device.index or 0)
237 + self._xfer_cache: dict[tuple[int, int, int], tuple[int, int, int]] = {}
238 + self._init_step("buffers", self._init_registered_buffers, expert_weights)
239 + self._init_step("agents", self._init_remote_agents)
240 + self._init_step("send meta", self._exchange_remote_send_meta)
241 + self._log_initialized()
242 +
243 + @property
244 + def needs_profile_buffer_reservation(self) -> bool:
245 + return False
246 +
247 + @staticmethod
248 + def _init_step(name: str, fn: object, *args: object, **kwargs: object) -> None:
249 + try:
250 + fn(*args, **kwargs) # type: ignore[operator]
251 + except Exception as exc:
252 + raise RuntimeError(f"NIXL EPLB init failed: {name}") from exc
253 +
254 + def _make_agent_name(self) -> str:
255 + """Build a deployment-unique nixl agent name."""
256 + pp_size = get_pp_group().world_size
257 + pp_suffix = f"-pp{get_pp_group().rank_in_group}" if pp_size > 1 else ""
258 + uid = uuid.uuid4().hex[:8]
259 + return f"eplb-{self._rank}{pp_suffix}-{uid}"
260 +
261 + def _get_peer_buckets(
262 + self,
263 + bucket_map: dict[torch.dtype, list[list[torch.Tensor]]],
264 + dtype: torch.dtype,
265 + ) -> list[list[torch.Tensor]]:
266 + peer_buckets = bucket_map.get(dtype)
267 + if peer_buckets is None:
268 + peer_buckets = [[] for _ in range(self._world_size)]
269 + bucket_map[dtype] = peer_buckets
270 + return peer_buckets
271 +
272 + def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
273 + assert dst_rank != self._rank, (
274 + "EPLB communicator should not enqueue same-rank sends: "
275 + f"rank={self._rank}, dst_rank={dst_rank}"
276 + )
277 + self._get_peer_buckets(self._send_tensors, tensor.dtype)[dst_rank].append(
278 + tensor
279 + )
280 +
281 + def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
282 + assert src_rank != self._rank, (
283 + "EPLB communicator should not enqueue same-rank recvs: "
284 + f"rank={self._rank}, src_rank={src_rank}"
285 + )
286 + self._get_peer_buckets(self._recv_tensors, tensor.dtype)[src_rank].append(
287 + tensor
288 + )
289 +
290 + def _init_remote_agents(self) -> None:
291 + local_metadata = self._nixl_wrapper.get_agent_metadata()
292 + gathered_metadata: list[bytes | None] = [None] * self._world_size
293 + torch.distributed.all_gather_object(
294 + gathered_metadata, local_metadata, group=self._cpu_group
295 + )
296 + for peer in range(self._world_size):
297 + if peer == self._rank:
298 + continue
299 + peer_metadata = gathered_metadata[peer]
300 + assert peer_metadata is not None
301 + self._remote_agents[peer] = self._nixl_wrapper.add_remote_agent(
302 + peer_metadata
303 + )
304 +
305 + def _init_registered_buffers(self, expert_weights: Sequence[torch.Tensor]) -> None:
306 + total_max_bytes = 0
307 + for dtype in self._dtypes:
308 + max_numel = max(
309 + sum(t.numel() for t in expert_weights if t.dtype == dtype), 1
310 + )
311 + max_bytes = max_numel * dtype.itemsize
312 + self._dtype_max_bytes[dtype] = max_bytes
313 + total_max_bytes += max_bytes
314 +
315 + self._peer_partition_bytes = total_max_bytes
316 +
317 + # The send buffer needs world_size partitions because remote peers
318 + # READ from fixed offsets (rank * partition_bytes).
319 + # This allocates world_size * partition_bytes
320 + # which can cause OOM on large models.
321 + # TODO(ilmarkov): shrink to const * partition_bytes and execute
322 + # communication in multiple steps dealing with the worst case.
323 + send_total_bytes = self._peer_partition_bytes * self._world_size
324 +
325 + self._send_buffer = torch.empty(
326 + send_total_bytes, device=self._device, dtype=torch.uint8
327 + )
328 + self._recv_buffer = torch.empty(
329 + self._peer_partition_bytes, device=self._device, dtype=torch.uint8
330 + )
331 +
332 + descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer])
333 + self._nixl_wrapper.register_memory(descs)
334 + self._registered_desc = descs
335 +
336 + def _exchange_remote_send_meta(self) -> None:
337 + """Exchange send-buffer metadata so each rank can build dynamic
338 + descriptors at execute time."""
339 + local_meta: tuple[int, int, int] = (
340 + self._send_buffer.data_ptr(),
341 + self._peer_partition_bytes,
342 + self._cuda_device_id,
343 + )
344 + gathered_meta: list[tuple[int, int, int] | None] = [None] * self._world_size
345 + torch.distributed.all_gather_object(
346 + gathered_meta, local_meta, group=self._cpu_group
347 + )
348 +
349 + for peer in self._remote_agents:
350 + peer_meta = gathered_meta[peer]
351 + assert peer_meta is not None
352 + self._remote_send_meta[peer] = peer_meta
353 +
354 + @staticmethod
355 + def _pack_send_buffer(
356 + peer_tensors: list[torch.Tensor],
357 + send_buffer: torch.Tensor,
358 + byte_offset: int,
359 + ) -> int:
360 + """
361 + Returns the byte offset after the last written byte.
362 + """
363 + for tensor in peer_tensors:
364 + raw = tensor.reshape(-1).view(torch.uint8)
365 + if raw.numel() == 0:
366 + continue
367 + send_buffer[byte_offset : byte_offset + raw.numel()].copy_(
368 + raw, non_blocking=True
369 + )
370 + byte_offset += raw.numel()
371 + return byte_offset
372 +
373 + @staticmethod
374 + def _unpack_recv_buffer(
375 + recv_buffer: torch.Tensor,
376 + peer_tensors: list[torch.Tensor],
377 + byte_offset: int,
378 + ) -> int:
379 + """
380 + Returns the byte offset after the last read byte.
381 + """
382 + for tensor in peer_tensors:
383 + num_bytes = tensor.numel() * tensor.element_size()
384 + if num_bytes == 0:
385 + continue
386 + tensor.reshape(-1).view(torch.uint8).copy_(
387 + recv_buffer[byte_offset : byte_offset + num_bytes],
388 + non_blocking=True,
389 + )
390 + byte_offset += num_bytes
391 + return byte_offset
392 +
393 + def _release_all_cached_handles(self) -> None:
394 + """Best-effort release of every cached dlist and xfer handle."""
395 + for local_dlist, remote_dlist, xfer in self._xfer_cache.values():
396 + for release_fn, handle in (
397 + (self._nixl_wrapper.release_xfer_handle, xfer),
398 + (self._nixl_wrapper.release_dlist_handle, local_dlist),
399 + (self._nixl_wrapper.release_dlist_handle, remote_dlist),
400 + ):
401 + with contextlib.suppress(Exception):
402 + release_fn(handle)
403 + self._xfer_cache.clear()
404 +
405 + def _wait_for_all_transfers(self, handles: list[int]) -> None:
406 + pending = set(handles)
407 + while pending:
408 + completed: list[int] = []
409 + for handle in pending:
410 + state = self._nixl_wrapper.check_xfer_state(handle)
411 + if state == "DONE":
412 + completed.append(handle)
413 + continue
414 + if state != "PROC":
415 + raise RuntimeError(f"NIXL transfer failed with state={state}")
416 + for handle in completed:
417 + pending.remove(handle)
418 + if pending:
419 + time.sleep(0.0005)
420 +
421 + def _get_or_create_xfer(self, src: int, total_bytes: int, recv_offset: int) -> int:
422 + """Return a cached xfer handle or create and cache a new one."""
423 + key = (src, total_bytes, recv_offset)
424 + cached = self._xfer_cache.get(key)
425 + if cached is not None:
426 + return cached[2]
427 +
428 + recv_base = self._recv_buffer.data_ptr()
429 + local_desc = self._nixl_wrapper.get_xfer_descs(
430 + [
431 + (
432 + recv_base + recv_offset,
433 + total_bytes,
434 + self._cuda_device_id,
435 + )
436 + ],
437 + self._nixl_memory_type,
438 + )
439 + local_handle = self._nixl_wrapper.prep_xfer_dlist(
440 + "NIXL_INIT_AGENT",
441 + local_desc,
442 + )
443 +
444 + remote_base, remote_part_bytes, remote_dev = self._remote_send_meta[src]
445 + agent_name = self._remote_agents[src]
446 + remote_desc = self._nixl_wrapper.get_xfer_descs(
447 + [
448 + (
449 + remote_base + self._rank * remote_part_bytes,
450 + total_bytes,
451 + remote_dev,
452 + )
453 + ],
454 + self._nixl_memory_type,
455 + )
456 + remote_handle = self._nixl_wrapper.prep_xfer_dlist(
457 + agent_name,
458 + remote_desc,
459 + )
460 +
461 + xfer_handle = self._nixl_wrapper.make_prepped_xfer(
462 + "READ",
463 + local_handle,
464 + [0],
465 + remote_handle,
466 + [0],
467 + )
468 + self._xfer_cache[key] = (local_handle, remote_handle, xfer_handle)
469 + return xfer_handle
470 +
471 + def execute(self) -> None:
472 + xfer_handles: list[int] = []
473 + try:
474 + # Phase 1: pack send buffers.
475 + with torch.cuda.stream(self._cuda_stream):
476 + for dst in range(self._world_size):
477 + byte_offset = dst * self._peer_partition_bytes
478 + for dtype in self._dtypes:
479 + peer_tensors = self._send_tensors.get(
480 + dtype, [[] for _ in range(self._world_size)]
481 + )[dst]
482 + actual_bytes = sum(
483 + t.numel() * t.element_size() for t in peer_tensors
484 + )
485 + if actual_bytes > self._dtype_max_bytes[dtype]:
486 + raise RuntimeError(
487 + "NIXL EPLB send overflow for dtype "
488 + f"{dtype}: peer={dst}, "
489 + f"required={actual_bytes}, "
490 + f"capacity={self._dtype_max_bytes[dtype]}"
491 + )
492 + byte_offset = self._pack_send_buffer(
493 + peer_tensors,
494 + self._send_buffer,
495 + byte_offset,
496 + )
497 +
498 + # Ensure all packed data is visible in device memory before pulls.
499 + if self._cuda_stream is not None:
500 + self._cuda_stream.synchronize()
501 + else:
502 + torch.cuda.current_stream().synchronize()
503 + # READ is receiver-initiated; synchronize all ranks before transfer.
504 + # We use monitored_barrier so a rank that crashes or exits early
505 + # produces a diagnostic timeout instead of a silent hang.
506 + torch.distributed.monitored_barrier(
507 + group=self._cpu_group,
508 + timeout=timedelta(minutes=5),
509 + )
510 +
511 + # Phase 2: look up or create descriptors and issue all READs.
512 + # Data from all peers is packed sequentially into the single
513 + # partition-sized recv buffer at running offsets.
514 + recv_offsets: dict[int, int] = {}
515 + recv_offset = 0
516 + for src in range(self._world_size):
517 + if src == self._rank:
518 + continue
519 + actual_total_bytes = 0
520 + for dtype in self._dtypes:
521 + peer_tensors = self._recv_tensors.get(
522 + dtype, [[] for _ in range(self._world_size)]
523 + )[src]
524 + actual_total_bytes += sum(
525 + t.numel() * t.element_size() for t in peer_tensors
526 + )
527 + if actual_total_bytes == 0:
528 + continue
529 +
530 + recv_offsets[src] = recv_offset
531 + xfer_handle = self._get_or_create_xfer(
532 + src, actual_total_bytes, recv_offset
533 + )
534 + self._nixl_wrapper.transfer(xfer_handle)
535 + xfer_handles.append(xfer_handle)
536 + recv_offset += actual_total_bytes
537 +
538 + # Phase 3: single wait for all in-flight transfers, then unpack.
539 + self._wait_for_all_transfers(xfer_handles)
540 +
541 + with torch.cuda.stream(self._cuda_stream):
542 + for src, offset in recv_offsets.items():
543 + byte_offset = offset
544 + for dtype in self._dtypes:
545 + peer_tensors = self._recv_tensors.get(
546 + dtype, [[] for _ in range(self._world_size)]
547 + )[src]
548 + byte_offset = self._unpack_recv_buffer(
549 + self._recv_buffer,
550 + peer_tensors,
551 + byte_offset,
552 + )
553 + except Exception:
554 + self._release_all_cached_handles()
555 + raise
556 + finally:
557 + self._send_tensors.clear()
558 + self._recv_tensors.clear()
559 +
560 + def __del__(self) -> None:
561 + try:
562 + self._release_all_cached_handles()
563 + if self._registered_desc is not None:
564 + self._nixl_wrapper.deregister_memory(self._registered_desc)
565 + self._registered_desc = None
566 + for agent_name in self._remote_agents.values():
567 + self._nixl_wrapper.remove_remote_agent(agent_name)
568 + self._remote_agents.clear()
569 + except Exception as e:
570 + logger.warning("Error during NixlEplbCommunicator cleanup: %s", e)
571 +
572 +
573 class PyNcclEplbCommunicator(EplbCommunicator):
574 """EPLB communicator backed by PyNcclCommunicator using ncclSend/ncclRecv."""
575
@@ -204,6 +607,24 @@ def create_eplb_communicator(
607 backend: str | None,
608 expert_weights: Sequence[torch.Tensor],
609 ) -> EplbCommunicator:
610 + """Create an EPLB communicator for the given backend.
611 +
612 + Args:
613 + group_coordinator: Process-group coordinator that provides the
614 + device and CPU communication groups.
615 + backend: Communicator backend name (``"torch_nccl"``,
616 + ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``).
617 + Falls back to ``"torch_nccl"`` when *None*.
618 + Stateless (elastic EP) groups only support ``"torch_nccl"``
619 + and ``"pynccl"``; ``"torch_nccl"`` is silently promoted to
620 + ``"pynccl"`` in that case. When tensors reside on CPU,
621 + ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU
622 + process group.
623 + expert_weights: Expert weight tensors from *one* MoE layer.
624 + NixlEplbCommunicator pre-allocates send/recv buffers sized
625 + to this layer, so all other MoE layers must have the same
626 + tensor count, shapes, and dtypes.
627 + """
628 # Keep a safe default for callers that have not resolved communicator yet.
629 if backend is None:
630 backend = "torch_nccl"
@@ -256,7 +677,7 @@ def create_eplb_communicator(
677 if backend not in ("torch_nccl", "pynccl"):
678 raise ValueError(
679 f"Elastic EP requires 'torch_nccl' or 'pynccl' EPLB communicator "
259 - f"(got '{backend}'). torch_gloo is not supported with stateless groups."
680 + f"(got '{backend}')."
681 )
682 if backend == "torch_nccl":
683 logger.warning(
@@ -266,7 +687,26 @@ def create_eplb_communicator(
687 backend = "pynccl"
688 return _create_pynccl()
689
269 - if backend == "torch_gloo":
690 + if backend == "nixl":
691 + if not has_nixl():
692 + raise RuntimeError(
693 + "EPLB communicator 'nixl' requested but NIXL is unavailable."
694 + )
695 + if not (current_platform.is_cuda_alike() and tensor_device_type != "cpu"):
696 + raise RuntimeError(
697 + "EPLB communicator 'nixl' supports only cuda-like devices "
698 + f"(got {tensor_device_type})."
699 + )
700 + try:
701 + return NixlEplbCommunicator(
702 + cpu_group=group_coordinator.cpu_group,
703 + expert_weights=expert_weights,
704 + )
705 + except Exception as exc:
706 + raise RuntimeError(
707 + f"Failed to initialize NixlEplbCommunicator ({exc})."
708 + ) from exc
709 + elif backend == "torch_gloo":
710 return TorchDistGlooStagedEplbCommunicator(
711 cpu_group=group_coordinator.cpu_group,
712 )
vllm/distributed/eplb/rebalance_execute.py
+20 -14
@@ -541,23 +541,29 @@ def rearrange_expert_weights_inplace(
541 assert num_physical_experts == ep_size * num_local_physical_experts
542
543 first_layer_weights = list(expert_weights[0])
544 +
545 + if is_profile:
546 + if communicator.needs_profile_buffer_reservation:
547 + # Reserve NCCL communication buffers via a dummy all_gather.
548 + # Backends that pre-allocate their own transfer buffers
549 + # skip this to avoid the extra memory spike during profiling.
550 + weights_buffer: list[torch.Tensor] = [
551 + torch.empty_like(w) for w in first_layer_weights
552 + ]
553 + for weight, buffer in zip(expert_weights[0], weights_buffer):
554 + dummy_recv_buffer = [buffer for _ in range(ep_size)]
555 + torch.distributed.barrier()
556 + all_gather(
557 + dummy_recv_buffer,
558 + weight,
559 + group=ep_group,
560 + )
561 + return
562 +
563 # Buffers to hold the expert weights during the exchange.
564 # NOTE: Currently we assume the same weights across different layers
565 # have the same shape.
547 - weights_buffer: list[torch.Tensor] = [
548 - torch.empty_like(w) for w in first_layer_weights
549 - ]
550 - if is_profile:
551 - # Reserve communication buffers via a minimal dummy all_gather on first layer
552 - for weight, buffer in zip(expert_weights[0], weights_buffer):
553 - dummy_recv_buffer = [buffer for _ in range(ep_size)]
554 - torch.distributed.barrier()
555 - all_gather(
556 - dummy_recv_buffer,
557 - weight,
558 - group=ep_group,
559 - )
560 - return
566 + weights_buffer = [torch.empty_like(w) for w in first_layer_weights]
567
568 # NOTE(bowen): We need this synchronize to run, but I don't know why.
569 # If you figure out the reason, please let me know -- thank you!
vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py
+1 -3
@@ -15,9 +15,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
15 PromMetric,
16 PromMetricT,
17 )
18 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
19 - nixlXferTelemetry,
20 -)
18 +from vllm.distributed.nixl_utils import nixlXferTelemetry
19 from vllm.v1.metrics.utils import create_metric_per_engine
20
21
vllm/distributed/kv_transfer/kv_connector/v1/nixl/utils.py
-46
@@ -3,60 +3,14 @@
3 """Shared constants, lazy imports and helpers for the NIXL connector."""
4
5 import contextlib
6 -import os
7 -import sys
6 from collections.abc import Iterator
7 from typing import Any
8
9 import zmq
10
13 -from vllm.logger import init_logger
11 from vllm.platforms import current_platform
12 from vllm.utils.network_utils import make_zmq_socket
13
17 -logger = init_logger(__name__)
18 -
19 -
20 -# Lazy import nixl_wrapper to avoid loading nixl_bindings if nixl is not used
21 -try:
22 - if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ:
23 - # avoid a memory leak in UCX when using NIXL on some models
24 - # see: https://github.com/vllm-project/vllm/issues/24264
25 - if "nixl" in sys.modules or "rixl" in sys.modules:
26 - logger.warning(
27 - "NIXL was already imported, we can't reset UCX_RCACHE_MAX_UNRELEASED. "
28 - "Please set it to '1024' manually."
29 - )
30 - else:
31 - logger.info(
32 - "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare "
33 - "memory leak in UCX when using NIXL."
34 - )
35 - os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
36 -
37 - if not current_platform.is_rocm():
38 - from nixl._api import nixl_agent as NixlWrapper
39 - from nixl._bindings import nixlXferTelemetry
40 - else:
41 - from rixl._api import nixl_agent as NixlWrapper
42 - from rixl._bindings import nixlXferTelemetry
43 -
44 - logger.info("NIXL is available")
45 -except ImportError:
46 - logger.warning("NIXL is not available")
47 - NixlWrapper = None
48 - nixlXferTelemetry = None
49 -
50 -
51 -try:
52 - if not current_platform.is_rocm():
53 - from nixl._api import nixl_agent_config
54 - else:
55 - from rixl._api import nixl_agent_config
56 -except ImportError:
57 - nixl_agent_config = None
58 - logger.warning("NIXL agent config is not available")
59 -
14 # Supported platforms and types of kv transfer buffer.
15 # {device: tuple of supported kv buffer types}
16 _NIXL_SUPPORTED_DEVICE = {
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+1 -2
@@ -45,8 +45,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
45 )
46 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (
47 _NIXL_SUPPORTED_DEVICE,
48 - NixlWrapper,
49 - nixl_agent_config,
48 zmq_ctx,
49 )
50 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
52 compute_mamba_phys_ratio,
53 derive_mamba_conv_split,
54 )
55 +from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
56 from vllm.distributed.parallel_state import (
57 get_tensor_model_parallel_rank,
58 get_tensor_model_parallel_world_size,
vllm/distributed/nixl_utils.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import os
5 +import sys
6 +
7 +from vllm.logger import init_logger
8 +from vllm.platforms import current_platform
9 +
10 +logger = init_logger(__name__)
11 +
12 +if "UCX_RCACHE_MAX_UNRELEASED" not in os.environ:
13 + if "nixl" in sys.modules or "rixl" in sys.modules:
14 + logger.warning_once(
15 + "NIXL was already imported, we can't reset "
16 + "UCX_RCACHE_MAX_UNRELEASED. "
17 + "Please set it to '1024' manually."
18 + )
19 + else:
20 + logger.info_once(
21 + "Setting UCX_RCACHE_MAX_UNRELEASED to '1024' to avoid a rare "
22 + "memory leak in UCX when using NIXL."
23 + )
24 + os.environ["UCX_RCACHE_MAX_UNRELEASED"] = "1024"
25 +
26 +try:
27 + if current_platform.is_cuda():
28 + from nixl._api import nixl_agent as NixlWrapper
29 + else:
30 + from rixl._api import nixl_agent as NixlWrapper
31 +
32 + logger.info_once("NIXL is available")
33 +except ImportError:
34 + logger.warning_once("NIXL is not available")
35 + NixlWrapper = None # type: ignore[assignment, misc]
36 +
37 +try:
38 + if current_platform.is_cuda():
39 + from nixl._api import nixl_agent_config
40 + else:
41 + from rixl._api import nixl_agent_config
42 +except ImportError:
43 + nixl_agent_config = None # type: ignore[assignment]
44 + logger.warning_once("NIXL agent config is not available")
45 +
46 +try:
47 + if current_platform.is_cuda():
48 + from nixl._bindings import nixlXferTelemetry
49 + else:
50 + from rixl._bindings import nixlXferTelemetry
51 +except ImportError:
52 + nixlXferTelemetry = None # type: ignore[assignment, misc]
53 +
54 +__all__ = ["NixlWrapper", "nixl_agent_config", "nixlXferTelemetry"]
vllm/v1/worker/gpu_model_runner.py
+20 -8
@@ -4821,6 +4821,23 @@ class GPUModelRunner(
4821 )
4822
4823 self.model.set_aux_hidden_state_layers(aux_layers)
4824 +
4825 + if (
4826 + is_mixture_of_experts(self.model)
4827 + and self.parallel_config.enable_eplb
4828 + and not load_dummy_weights
4829 + ):
4830 + logger.info_once(
4831 + "EPLB is enabled for model %s.",
4832 + self.model_config.model,
4833 + )
4834 + assert self.eplb_state is not None
4835 + self.eplb_state.add_model(
4836 + self.model,
4837 + self.model_config,
4838 + )
4839 + eplb_models += 1
4840 +
4841 time_after_load = time.perf_counter()
4842 self.model_memory_usage = m.consumed_memory
4843 except torch.cuda.OutOfMemoryError as e:
@@ -4860,15 +4877,10 @@ class GPUModelRunner(
4877 is_mixture_of_experts(self.model)
4878 and self.parallel_config.enable_eplb
4879 and not load_dummy_weights
4880 + and self.eplb_state is not None
4881 + and self.eplb_state.is_async
4882 ):
4864 - logger.info_once("EPLB is enabled for model %s.", self.model_config.model)
4865 - assert self.eplb_state is not None
4866 - self.eplb_state.add_model(
4867 - self.model,
4868 - self.model_config,
4869 - )
4870 - if self.eplb_state.is_async:
4871 - self.eplb_state.start_async_loop()
4883 + self.eplb_state.start_async_loop()
4884
4885 if (
4886 self.vllm_config.compilation_config.mode