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, 11:15 UTC
[EPLB] Optimize memory overhead in Nixl communicator (#40013)
Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Ilya Markov committed
Apr 30, 2026 at 16:46 UTC
a7fb00851030b6991987c559b885dfd8eb039d15
2 files changed
+212
-204
vllm/distributed/eplb/eplb_communicator.py
+204
-199
@@ -11,6 +11,7 @@ from abc import ABC, abstractmethod
11
from collections.abc import Sequence
12
from datetime import timedelta
13
14
+import numpy as np
15
import torch
16
from torch.distributed import (
17
P2POp,
@@ -47,15 +48,25 @@ class EplbCommunicator(ABC):
48
"""Abstract EPLB communicator for expert weight transfers."""
49
50
@abstractmethod
50
- def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
51
+ def add_send(
52
+ self,
53
+ tensors: list[torch.Tensor],
54
+ dst_rank: int,
55
+ expert_id: int,
56
+ ) -> None:
57
pass
58
59
@abstractmethod
54
- def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
60
+ def add_recv(
61
+ self,
62
+ tensors: list[torch.Tensor],
63
+ src_rank: int,
64
+ expert_id: int,
65
+ ) -> None:
66
pass
67
68
@abstractmethod
58
- def execute(self) -> None:
69
+ def execute(self, old_indices: np.ndarray | None = None) -> None:
70
pass
71
72
@property
@@ -85,27 +96,39 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator):
96
self._p2p_ops: list[P2POp] = []
97
self._log_initialized()
98
88
- def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
89
- self._p2p_ops.append(
90
- P2POp(
91
- torch.distributed.isend,
92
- tensor,
93
- dst_rank,
94
- self._ep_group,
99
+ def add_send(
100
+ self,
101
+ tensors: list[torch.Tensor],
102
+ dst_rank: int,
103
+ expert_id: int, # unused by this backend
104
+ ) -> None:
105
+ for tensor in tensors:
106
+ self._p2p_ops.append(
107
+ P2POp(
108
+ torch.distributed.isend,
109
+ tensor,
110
+ dst_rank,
111
+ self._ep_group,
112
+ )
113
)
96
- )
114
98
- def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
99
- self._p2p_ops.append(
100
- P2POp(
101
- torch.distributed.irecv,
102
- tensor,
103
- src_rank,
104
- self._ep_group,
115
+ def add_recv(
116
+ self,
117
+ tensors: list[torch.Tensor],
118
+ src_rank: int,
119
+ expert_id: int, # unused by this backend
120
+ ) -> None:
121
+ for tensor in tensors:
122
+ self._p2p_ops.append(
123
+ P2POp(
124
+ torch.distributed.irecv,
125
+ tensor,
126
+ src_rank,
127
+ self._ep_group,
128
+ )
129
)
106
- )
130
108
- def execute(self) -> None:
131
+ def execute(self, old_indices: np.ndarray | None = None) -> None:
132
if not self._p2p_ops:
133
return
134
try:
@@ -130,13 +153,25 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):
153
self._ops: list[tuple[str, torch.Tensor, int]] = []
154
self._log_initialized()
155
133
- def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
134
- self._ops.append(("send", tensor, dst_rank))
156
+ def add_send(
157
+ self,
158
+ tensors: list[torch.Tensor],
159
+ dst_rank: int,
160
+ expert_id: int, # unused by this backend
161
+ ) -> None:
162
+ for tensor in tensors:
163
+ self._ops.append(("send", tensor, dst_rank))
164
136
- def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
137
- self._ops.append(("recv", tensor, src_rank))
165
+ def add_recv(
166
+ self,
167
+ tensors: list[torch.Tensor],
168
+ src_rank: int,
169
+ expert_id: int, # unused by this backend
170
+ ) -> None:
171
+ for tensor in tensors:
172
+ self._ops.append(("recv", tensor, src_rank))
173
139
- def execute(self) -> None:
174
+ def execute(self, old_indices: np.ndarray | None = None) -> None:
175
if not self._ops:
176
return
177
@@ -207,17 +242,17 @@ class NixlEplbCommunicator(EplbCommunicator):
242
self._cuda_stream = cuda_stream
243
self._world_size = cpu_group.size()
244
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] = []
245
+ # expert_id -> weight tensors to pack into the send buffer.
246
+ self._expert_send_map: dict[int, list[torch.Tensor]] = {}
247
+ # src_rank -> expert_id -> weight tensors to unpack after transfer.
248
+ self._recv_map: dict[int, dict[int, list[torch.Tensor]]] = {}
249
+ self._num_local_experts: int = expert_weights[0].shape[0]
250
self._device = expert_weights[0].device
251
for tensor in expert_weights:
252
assert tensor.device == self._device, (
253
"All local EPLB tensors are expected to be on the same device: "
254
f"expected={self._device}, got={tensor.device}"
255
)
219
- if tensor.dtype not in self._dtypes:
220
- self._dtypes.append(tensor.dtype)
256
257
config = (
258
nixl_agent_config(capture_telemetry=False)
@@ -228,13 +263,12 @@ class NixlEplbCommunicator(EplbCommunicator):
263
self._nixl_memory_type = "VRAM"
264
self._registered_desc: object | None = None
265
self._remote_agents: dict[int, str] = {}
231
- self._remote_send_meta: dict[int, tuple[int, int, int]] = {}
266
+ self._remote_send_meta: dict[int, tuple[int, int]] = {}
267
self._send_buffer: torch.Tensor = torch.empty(0)
268
self._recv_buffer: torch.Tensor = torch.empty(0)
234
- self._peer_partition_bytes: int = 0
235
- self._dtype_max_bytes: dict[torch.dtype, int] = {}
269
+ self._expert_bytes: int = 0
270
+
271
self._cuda_device_id = int(self._device.index or 0)
237
- self._xfer_cache: dict[tuple[int, int, int], tuple[int, int, int]] = {}
272
self._init_step("buffers", self._init_registered_buffers, expert_weights)
273
self._init_step("agents", self._init_remote_agents)
274
self._init_step("send meta", self._exchange_remote_send_meta)
@@ -258,34 +292,33 @@ class NixlEplbCommunicator(EplbCommunicator):
292
uid = uuid.uuid4().hex[:8]
293
return f"eplb-{self._rank}{pp_suffix}-{uid}"
294
261
- def _get_peer_buckets(
295
+ def add_send(
296
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:
297
+ tensors: list[torch.Tensor],
298
+ dst_rank: int,
299
+ expert_id: int,
300
+ ) -> None:
301
assert dst_rank != self._rank, (
302
"EPLB communicator should not enqueue same-rank sends: "
303
f"rank={self._rank}, dst_rank={dst_rank}"
304
)
277
- self._get_peer_buckets(self._send_tensors, tensor.dtype)[dst_rank].append(
278
- tensor
279
- )
305
+ # An expert sent to multiple peers is packed only once; skip duplicates.
306
+ if expert_id not in self._expert_send_map:
307
+ self._expert_send_map[expert_id] = tensors
308
281
- def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
309
+ def add_recv(
310
+ self,
311
+ tensors: list[torch.Tensor],
312
+ src_rank: int,
313
+ expert_id: int,
314
+ ) -> None:
315
assert src_rank != self._rank, (
316
"EPLB communicator should not enqueue same-rank recvs: "
317
f"rank={self._rank}, src_rank={src_rank}"
318
)
286
- self._get_peer_buckets(self._recv_tensors, tensor.dtype)[src_rank].append(
287
- tensor
288
- )
319
+ recv_experts = self._recv_map.setdefault(src_rank, {})
320
+ if expert_id not in recv_experts:
321
+ recv_experts[expert_id] = tensors
322
323
def _init_remote_agents(self) -> None:
324
local_metadata = self._nixl_wrapper.get_agent_metadata()
@@ -303,30 +336,18 @@ class NixlEplbCommunicator(EplbCommunicator):
336
)
337
338
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
339
+ total_bytes = max(sum(t.nbytes for t in expert_weights), 1)
340
+ assert total_bytes % self._num_local_experts == 0, (
341
+ f"Number of bytes in moe layer {total_bytes} is not divisible "
342
+ f"by number of local experts {self._num_local_experts}"
343
+ )
344
+ self._expert_bytes = total_bytes // self._num_local_experts
345
346
self._send_buffer = torch.empty(
326
- send_total_bytes, device=self._device, dtype=torch.uint8
347
+ total_bytes, device=self._device, dtype=torch.uint8
348
)
349
self._recv_buffer = torch.empty(
329
- self._peer_partition_bytes, device=self._device, dtype=torch.uint8
350
+ total_bytes, device=self._device, dtype=torch.uint8
351
)
352
353
descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer])
@@ -336,12 +357,11 @@ class NixlEplbCommunicator(EplbCommunicator):
357
def _exchange_remote_send_meta(self) -> None:
358
"""Exchange send-buffer metadata so each rank can build dynamic
359
descriptors at execute time."""
339
- local_meta: tuple[int, int, int] = (
360
+ local_meta: tuple[int, int] = (
361
self._send_buffer.data_ptr(),
341
- self._peer_partition_bytes,
362
self._cuda_device_id,
363
)
344
- gathered_meta: list[tuple[int, int, int] | None] = [None] * self._world_size
364
+ gathered_meta: list[tuple[int, int] | None] = [None] * self._world_size
365
torch.distributed.all_gather_object(
366
gathered_meta, local_meta, group=self._cpu_group
367
)
@@ -353,14 +373,11 @@ class NixlEplbCommunicator(EplbCommunicator):
373
374
@staticmethod
375
def _pack_send_buffer(
356
- peer_tensors: list[torch.Tensor],
376
+ in_tensors: list[torch.Tensor],
377
send_buffer: torch.Tensor,
378
byte_offset: int,
359
- ) -> int:
360
- """
361
- Returns the byte offset after the last written byte.
362
- """
363
- for tensor in peer_tensors:
379
+ ) -> None:
380
+ for tensor in in_tensors:
381
raw = tensor.reshape(-1).view(torch.uint8)
382
if raw.numel() == 0:
383
continue
@@ -368,18 +385,14 @@ class NixlEplbCommunicator(EplbCommunicator):
385
raw, non_blocking=True
386
)
387
byte_offset += raw.numel()
371
- return byte_offset
388
389
@staticmethod
390
def _unpack_recv_buffer(
391
recv_buffer: torch.Tensor,
376
- peer_tensors: list[torch.Tensor],
392
+ out_tensors: list[torch.Tensor],
393
byte_offset: int,
378
- ) -> int:
379
- """
380
- Returns the byte offset after the last read byte.
381
- """
382
- for tensor in peer_tensors:
394
+ ) -> None:
395
+ for tensor in out_tensors:
396
num_bytes = tensor.numel() * tensor.element_size()
397
if num_bytes == 0:
398
continue
@@ -388,19 +401,6 @@ class NixlEplbCommunicator(EplbCommunicator):
401
non_blocking=True,
402
)
403
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)
@@ -418,82 +418,68 @@ class NixlEplbCommunicator(EplbCommunicator):
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]
421
+ def _create_peer_xfer(
422
+ self,
423
+ src: int,
424
+ local_descs: list[tuple[int, int, int]],
425
+ remote_descs: list[tuple[int, int, int]],
426
+ ) -> tuple[int, int, int]:
427
+ """Create a batched xfer for multiple descriptors from one peer.
428
+
429
+ Each element in *local_descs* / *remote_descs* is an
430
+ ``(address, size, device_id)`` tuple.
431
428
- recv_base = self._recv_buffer.data_ptr()
432
+ Returns ``(local_dlist, remote_dlist, xfer_handle)``.
433
+ """
434
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,
435
+ local_descs, self._nixl_memory_type
436
)
437
local_handle = self._nixl_wrapper.prep_xfer_dlist(
438
"NIXL_INIT_AGENT",
439
local_desc,
440
)
441
444
- remote_base, remote_part_bytes, remote_dev = self._remote_send_meta[src]
445
- agent_name = self._remote_agents[src]
442
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,
443
+ remote_descs, self._nixl_memory_type
444
)
445
remote_handle = self._nixl_wrapper.prep_xfer_dlist(
457
- agent_name,
446
+ self._remote_agents[src],
447
remote_desc,
448
)
449
450
+ indices = list(range(len(local_descs)))
451
xfer_handle = self._nixl_wrapper.make_prepped_xfer(
452
"READ",
453
local_handle,
464
- [0],
454
+ indices,
455
remote_handle,
466
- [0],
456
+ indices,
457
+ )
458
+ return (local_handle, remote_handle, xfer_handle)
459
+
460
+ def execute(self, old_indices: np.ndarray | None = None) -> None:
461
+ assert old_indices is not None, (
462
+ "NixlEplbCommunicator.execute requires old_indices"
463
)
468
- self._xfer_cache[key] = (local_handle, remote_handle, xfer_handle)
469
- return xfer_handle
464
471
- def execute(self) -> None:
472
- xfer_handles: list[int] = []
465
+ xfer_entries: list[tuple[int, int, int]] = []
466
try:
474
- # Phase 1: pack send buffers.
467
+ n = self._num_local_experts
468
+ rank_experts = old_indices[: self._world_size * n].reshape(
469
+ self._world_size, n
470
+ )
471
+ # Build expert_id -> send slot mapping per rank.
472
+ expert_to_send_slot: list[dict[int, int]] = [
473
+ {int(eid): i for i, eid in enumerate(row) if eid != -1}
474
+ for row in rank_experts
475
+ ]
476
+
477
+ # Phase 1: pack each expert at its slot offset in the send buffer.
478
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
- )
479
+ for expert_id, tensors in self._expert_send_map.items():
480
+ slot = expert_to_send_slot[self._rank][expert_id]
481
+ byte_offset = slot * self._expert_bytes
482
+ self._pack_send_buffer(tensors, self._send_buffer, byte_offset)
483
484
# Ensure all packed data is visible in device memory before pulls.
485
if self._cuda_stream is not None:
@@ -508,58 +494,65 @@ class NixlEplbCommunicator(EplbCommunicator):
494
timeout=timedelta(minutes=5),
495
)
496
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] = {}
497
+ # Phase 2: issue one batched READ per peer.
498
+ recv_offsets: dict[tuple[int, int], int] = {}
499
recv_offset = 0
500
+ recv_base = self._recv_buffer.data_ptr()
501
for src in range(self._world_size):
502
if src == self._rank:
503
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:
504
+ recv_experts = self._recv_map.get(src)
505
+ if not recv_experts:
506
continue
529
-
530
- recv_offsets[src] = recv_offset
531
- xfer_handle = self._get_or_create_xfer(
532
- src, actual_total_bytes, recv_offset
507
+ expert_ids = list(recv_experts.keys())
508
+ remote_base, remote_dev = self._remote_send_meta[src]
509
+ local_descs: list[tuple[int, int, int]] = []
510
+ remote_descs: list[tuple[int, int, int]] = []
511
+ for expert_id in expert_ids:
512
+ slot = expert_to_send_slot[src][expert_id]
513
+ remote_off = slot * self._expert_bytes
514
+ recv_offsets[(src, expert_id)] = recv_offset
515
+ local_descs.append(
516
+ (
517
+ recv_base + recv_offset,
518
+ self._expert_bytes,
519
+ self._cuda_device_id,
520
+ )
521
+ )
522
+ remote_descs.append(
523
+ (remote_base + remote_off, self._expert_bytes, remote_dev)
524
+ )
525
+ recv_offset += self._expert_bytes
526
+ assert recv_offset <= self._recv_buffer.nbytes
527
+ local_h, remote_h, xfer_h = self._create_peer_xfer(
528
+ src, local_descs, remote_descs
529
)
534
- self._nixl_wrapper.transfer(xfer_handle)
535
- xfer_handles.append(xfer_handle)
536
- recv_offset += actual_total_bytes
530
+ self._nixl_wrapper.transfer(xfer_h)
531
+ xfer_entries.append((local_h, remote_h, xfer_h))
532
538
- # Phase 3: single wait for all in-flight transfers, then unpack.
539
- self._wait_for_all_transfers(xfer_handles)
533
+ # Phase 3: wait for all in-flight transfers, then unpack.
534
+ self._wait_for_all_transfers([x[2] for x in xfer_entries])
535
536
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
537
+ for (src, expert_id), offset in recv_offsets.items():
538
+ self._unpack_recv_buffer(
539
+ self._recv_buffer,
540
+ self._recv_map[src][expert_id],
541
+ offset,
542
+ )
543
finally:
557
- self._send_tensors.clear()
558
- self._recv_tensors.clear()
544
+ for local_h, remote_h, xfer_h in xfer_entries:
545
+ with contextlib.suppress(Exception):
546
+ self._nixl_wrapper.release_xfer_handle(xfer_h)
547
+ with contextlib.suppress(Exception):
548
+ self._nixl_wrapper.release_dlist_handle(local_h)
549
+ with contextlib.suppress(Exception):
550
+ self._nixl_wrapper.release_dlist_handle(remote_h)
551
+ self._expert_send_map.clear()
552
+ self._recv_map.clear()
553
554
def __del__(self) -> None:
555
try:
562
- self._release_all_cached_handles()
556
if self._registered_desc is not None:
557
self._nixl_wrapper.deregister_memory(self._registered_desc)
558
self._registered_desc = None
@@ -588,15 +581,27 @@ class PyNcclEplbCommunicator(EplbCommunicator):
581
self._pynccl_comm.group_start()
582
self._group_started = True
583
591
- def add_send(self, tensor: torch.Tensor, dst_rank: int) -> None:
584
+ def add_send(
585
+ self,
586
+ tensors: list[torch.Tensor],
587
+ dst_rank: int,
588
+ expert_id: int, # unused by this backend
589
+ ) -> None:
590
self._ensure_group_started()
593
- self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)
591
+ for tensor in tensors:
592
+ self._pynccl_comm.send(tensor, dst_rank, stream=self._cuda_stream)
593
595
- def add_recv(self, tensor: torch.Tensor, src_rank: int) -> None:
594
+ def add_recv(
595
+ self,
596
+ tensors: list[torch.Tensor],
597
+ src_rank: int,
598
+ expert_id: int, # unused by this backend
599
+ ) -> None:
600
self._ensure_group_started()
597
- self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)
601
+ for tensor in tensors:
602
+ self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)
603
599
- def execute(self) -> None:
604
+ def execute(self, old_indices: np.ndarray | None = None) -> None:
605
if self._group_started:
606
self._pynccl_comm.group_end()
607
self._group_started = False
vllm/distributed/eplb/rebalance_execute.py
+8
-5
@@ -294,9 +294,9 @@ def move_to_buffer(
294
recver_pos = remainder_start + sender_pos
295
if recver_pos < len(ranks_to_recv):
296
recv_ranks.append(ranks_to_recv[recver_pos])
297
+ expert_tensors = [w[src] for w in expert_weights]
298
for dst in recv_ranks:
298
- for w in expert_weights:
299
- communicator.add_send(w[src], dst)
299
+ communicator.add_send(expert_tensors, dst, expert_id=int(expert))
300
301
# 3. Post recvs
302
if recv_count > 0:
@@ -325,11 +325,14 @@ def move_to_buffer(
325
src = ranks_to_send[recver_pos // num_dst_per_sender]
326
else:
327
src = ranks_to_send[recver_pos - remainder_start]
328
- for b in expert_weights_buffers:
329
- communicator.add_recv(b[dst], src)
328
+ communicator.add_recv(
329
+ [b[dst] for b in expert_weights_buffers],
330
+ src,
331
+ expert_id=int(expert),
332
+ )
333
334
# 4. Execute the P2P operations. The real communication happens here.
332
- communicator.execute()
335
+ communicator.execute(old_indices=old_indices)
336
# wait for the communication to finish
337
return TransferMetadata(
338
is_unchanged=is_unchanged,