@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4f423bd5bc

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 1, 23:45 UTC

[EPLB] Nixl communicator optimization. Zero-copy transfers (#41633)

Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

Ilya Markov committed Jun 4, 2026 at 05:40 UTC 4f423bd5bcf38cb2fc09db5dc9e0bc8a6bfb1638

12 files changed +287 -219

tests/distributed/test_eplb_execute.py

+21 -9

@@ -277,12 +277,15 @@ def assert_verification_synced(local_ok: bool, msg: str) -> None:

277 assert bool(ok_tensor.item()), msg

278

279

280 -def create_eplb_communicator_or_raise(*, group_coordinator, backend, expert_weights):

280 +def create_eplb_communicator_or_raise(

281 + *, group_coordinator, backend, expert_weights, expert_buffer

282 +):

283 try:

284 return create_eplb_communicator(

285 group_coordinator=group_coordinator,

286 backend=backend,

287 expert_weights=expert_weights,

288 + expert_buffer=expert_buffer,

289 )

290 except Exception as exc:

291 raise RuntimeError(

@@ -355,7 +358,8 @@ def _test_async_transfer_layer_without_mtp_worker(

358 communicator = create_eplb_communicator_or_raise(

359 group_coordinator=ep_group_coordinator,

360 backend=eplb_communicator,

358 - expert_weights=expert_weights[0],

361 + expert_weights=expert_weights,

362 + expert_buffer=expert_buffer,

363 )

364 communicator.set_stream(cuda_stream)

365

@@ -368,6 +372,7 @@ def _test_async_transfer_layer_without_mtp_worker(

372 ep_group=ep_group,

373 communicator=communicator,

374 cuda_stream=cuda_stream,

375 + layer_idx=layer_idx,

376 )

377 cuda_stream.synchronize()

378 move_from_buffer(

@@ -460,10 +465,12 @@ def _test_rearrange_expert_weights_with_redundancy(

465 num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices

466 )

467

468 + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

469 communicator = create_eplb_communicator_or_raise(

470 group_coordinator=ep_group_coordinator,

471 backend=eplb_communicator,

466 - expert_weights=expert_weights[0],

472 + expert_weights=expert_weights,

473 + expert_buffer=expert_buffer,

474 )

475

476 # Execute weight rearrangement

@@ -471,9 +478,9 @@ def _test_rearrange_expert_weights_with_redundancy(

478 old_indices,

479 new_indices,

480 expert_weights,

481 + expert_buffer,

482 ep_group,

475 - is_profile=False,

476 - communicator=communicator,

483 + communicator,

484 )

485

486 # Verify the rearrangement result

@@ -593,10 +600,12 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

600 layer_copy.append(weight.clone())

601 original_weights.append(layer_copy)

602

603 + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

604 communicator = create_eplb_communicator_or_raise(

605 group_coordinator=ep_group_coordinator,

606 backend="torch_nccl",

599 - expert_weights=expert_weights[0],

607 + expert_weights=expert_weights,

608 + expert_buffer=expert_buffer,

609 )

610

611 # Execute rearrangement (should be no change)

@@ -604,9 +613,9 @@ def _test_rearrange_expert_weights_no_change(env, world_size) -> None:

613 indices,

614 indices, # Same indices

615 expert_weights,

616 + expert_buffer,

617 ep_group,

618 communicator,

609 - is_profile=False,

619 )

620

621 # Verify that the weights have not changed

@@ -726,10 +735,12 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

735 layer_copy.append(weight.clone())

736 original_weights.append(layer_copy)

737

738 + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

739 communicator = create_eplb_communicator_or_raise(

740 group_coordinator=ep_group_coordinator,

741 backend="torch_nccl",

732 - expert_weights=expert_weights[0],

742 + expert_weights=expert_weights,

743 + expert_buffer=expert_buffer,

744 )

745

746 # Execute profile mode rearrangement

@@ -737,9 +748,10 @@ def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:

748 old_indices,

749 new_indices,

750 expert_weights,

751 + expert_buffer,

752 ep_group,

753 communicator,

742 - is_profile=True, # Profile mode

754 + is_profile=True,

755 )

756

757 # In profile mode, the weights should remain unchanged

tests/distributed/test_eplb_fused_moe_layer.py

+11 -1

@@ -9,9 +9,11 @@ import pytest

9 import torch

11 from vllm.config import VllmConfig, set_current_vllm_config

12 +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator

13 from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace

14 from vllm.distributed.parallel_state import (

15 ensure_model_parallel_initialized,

16 + get_eplb_group,

17 get_tp_group,

18 )

19 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):

215 for lidx in range(test_config.num_layers):

216 shuffled_indices[lidx] = torch.randperm(test_config.num_experts)

217

218 + expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]]

219 + communicator = create_eplb_communicator(

220 + group_coordinator=get_eplb_group(),

221 + backend="torch_nccl",

222 + expert_weights=rank_expert_weights,

223 + expert_buffer=expert_buffer,

224 + )

225 rearrange_expert_weights_inplace(

226 indices,

227 shuffled_indices,

228 rank_expert_weights,

229 + expert_buffer,

230 ep_group,

221 - is_profile=False,

231 + communicator,

232 )

233

234 num_local_experts = test_config.num_local_experts

tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py

+11 -1

@@ -10,11 +10,13 @@ import torch

11 from tests.kernels.moe.utils import make_test_quant_config

12 from vllm.config import VllmConfig, set_current_vllm_config

13 +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator

14 from vllm.distributed.eplb.eplb_state import EplbLayerState

15 from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace

16 from vllm.distributed.parallel_state import (

17 ensure_model_parallel_initialized,

18 get_dp_group,

19 + get_eplb_group,

20 )

21 from vllm.forward_context import set_forward_context

22 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):

173 for lidx in range(test_config.num_layers):

174 shuffled_indices[lidx] = torch.randperm(test_config.num_experts)

175

176 + expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]]

177 + communicator = create_eplb_communicator(

178 + group_coordinator=get_eplb_group(),

179 + backend="torch_nccl",

180 + expert_weights=rank_expert_weights,

181 + expert_buffer=expert_buffer,

182 + )

183 rearrange_expert_weights_inplace(

184 indices,

185 shuffled_indices,

186 rank_expert_weights,

187 + expert_buffer,

188 ep_group,

179 - is_profile=False,

189 + communicator,

190 )

191

192 num_global_experts = test_config.num_experts

tests/kernels/moe/test_moe_layer.py

+4 -1

@@ -1287,10 +1287,12 @@ def _test_body_eplb(

1287

1288 expert_weights = [list(eplb_moe_layer.get_expert_weights())]

1289

1290 + expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]

1291 communicator = create_eplb_communicator(

1292 group_coordinator=get_eplb_group(),

1293 backend=vllm_config.parallel_config.eplb_config.communicator,

1293 - expert_weights=expert_weights[0],

1294 + expert_weights=expert_weights,

1295 + expert_buffer=expert_buffer,

1296 )

1297

1298 # Rearrange expert weights across EP ranks

@@ -1298,6 +1300,7 @@ def _test_body_eplb(

1300 old_global_expert_indices=initial_indices.unsqueeze(0),

1301 new_global_expert_indices=shuffled_indices.unsqueeze(0),

1302 expert_weights=expert_weights,

1303 + expert_buffer=expert_buffer,

1304 ep_group=cpu_group,

1305 communicator=communicator,

1306 )

vllm/distributed/device_communicators/custom_all_reduce.py

+1 -5

@@ -40,11 +40,7 @@ def _can_p2p(rank: int, world_size: int) -> bool:

40 return True

43 -def is_weak_contiguous(inp: torch.Tensor):

44 - return inp.is_contiguous() or (

45 - inp.storage().nbytes() - inp.storage_offset() * inp.element_size()

46 - == inp.numel() * inp.element_size()

47 - )

43 +from vllm.distributed.utils import is_weak_contiguous # noqa: E402

46 class CustomAllreduce:

vllm/distributed/device_communicators/quick_all_reduce.py

+1 -5

@@ -24,11 +24,7 @@ except Exception:

24 quick_ar = False

27 -def is_weak_contiguous(inp: torch.Tensor):

28 - return inp.is_contiguous() or (

29 - inp.storage().nbytes() - inp.storage_offset() * inp.element_size()

30 - == inp.numel() * inp.element_size()

31 - )

27 +from vllm.distributed.utils import is_weak_contiguous # noqa: E402, F401

30 class QuickReduceRegime(Enum):

vllm/distributed/elastic_ep/elastic_execute.py

+5 -1

@@ -470,10 +470,14 @@ class ElasticEPScalingExecutor:

470 module._replace_quant_method(module.quant_method.old_quant_method)

471 prepare_communication_buffer_for_model(self.worker.model_runner.model)

472

473 + eplb_model_state.expert_buffer = [

474 + torch.empty_like(w) for w in model.expert_weights[0]

475 + ]

476 eplb_model_state.communicator = create_eplb_communicator(

477 group_coordinator=get_eplb_group(),

478 backend=parallel_config.eplb_config.communicator,

476 - expert_weights=model.expert_weights[0],

479 + expert_weights=model.expert_weights,

480 + expert_buffer=eplb_model_state.expert_buffer,

481 )

482

483 if (

vllm/distributed/eplb/async_worker.py

@@ -120,6 +120,7 @@ def transfer_run_periodically(

120 ep_group=eplb_group,

121 is_profile=is_profile,

122 cuda_stream=cuda_stream,

123 + layer_idx=layer_idx,

124 )

125

126 # Wait until all writes to expert_buffer have finished before making the

vllm/distributed/eplb/eplb_communicator.py

+199 -186

@@ -30,6 +30,7 @@ from vllm.distributed.parallel_state import (

30 is_local_first_rank,

31 )

32 from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator

33 +from vllm.distributed.utils import is_weak_contiguous

34 from vllm.logger import init_logger

35 from vllm.platforms import current_platform

@@ -63,8 +64,22 @@ class EplbCommunicator(ABC):

64 pass

66 @abstractmethod

66 - def execute(self, old_indices: np.ndarray | None = None) -> None:

67 - pass

67 + def execute(self) -> None:

68 + """Complete all enqueued transfers.

69 +

70 + Some backends perform communication here; others (e.g. NIXL)

71 + issue transfers eagerly in add_recv and only wait here.

72 + On return, all data is available in the destination buffers.

73 + """

74 +

75 + def set_transfer_context( # noqa: B027

76 + self, old_indices: np.ndarray, layer_idx: int

77 + ) -> None:

78 + """Pre-set layer context before add_recv calls.

79 +

80 + Default is a no-op; overridden by backends (e.g. NIXL) that need

81 + layer-level context to issue transfers inside add_recv.

82 + """

84 @property

85 def needs_profile_buffer_reservation(self) -> bool:

@@ -125,7 +140,7 @@ class TorchDistNcclEplbCommunicator(EplbCommunicator):

140 )

141 )

142

128 - def execute(self, old_indices: np.ndarray | None = None) -> None:

143 + def execute(self) -> None:

144 if not self._p2p_ops:

145 return

146 try:

@@ -168,7 +183,7 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):

183 for tensor in tensors:

184 self._ops.append(("recv", tensor, src_rank))

185

171 - def execute(self, old_indices: np.ndarray | None = None) -> None:

186 + def execute(self) -> None:

187 if not self._ops:

188 return

189

@@ -229,29 +244,47 @@ class NixlEplbCommunicator(EplbCommunicator):

244 def __init__(

245 self,

246 cpu_group: ProcessGroup,

232 - expert_weights: Sequence[torch.Tensor],

233 - cuda_stream: torch.cuda.Stream | None = None,

247 + all_expert_weights: Sequence[Sequence[torch.Tensor]],

248 + expert_buffer: Sequence[torch.Tensor],

249 ) -> None:

235 - assert expert_weights, "NixlEplbCommunicator requires non-empty expert_weights."

250 + assert all_expert_weights, (

251 + "NixlEplbCommunicator requires non-empty all_expert_weights."

252 + )

253 + assert expert_buffer, "NixlEplbCommunicator requires non-empty expert_buffer."

254 nixl_wrapper_cls = nixl_utils.NixlWrapper

255 if nixl_wrapper_cls is None:

256 raise RuntimeError("NIXL/ RIXL is unavailable.")

257 +

258 self._cpu_group = cpu_group

240 - self._cuda_stream = cuda_stream

259 self._world_size = cpu_group.size()

260 self._rank = cpu_group.rank()

243 - # expert_id -> weight tensors to pack into the send buffer.

244 - self._expert_send_map: dict[int, list[torch.Tensor]] = {}

245 - # src_rank -> expert_id -> weight tensors to unpack after transfer.

246 - self._recv_map: dict[int, dict[int, list[torch.Tensor]]] = {}

247 - self._num_local_experts: int = expert_weights[0].shape[0]

248 - self._device = expert_weights[0].device

249 - for tensor in expert_weights:

250 - assert tensor.device == self._device, (

251 - "All local EPLB tensors are expected to be on the same device: "

252 - f"expected={self._device}, got={tensor.device}"

261 +

262 + self._all_expert_weights = all_expert_weights

263 + self._expert_buffer = expert_buffer

264 + self._num_local_experts: int = all_expert_weights[0][0].shape[0]

265 + self._device = all_expert_weights[0][0].device

266 +

267 + for layer_tensors in all_expert_weights:

268 + for tensor in layer_tensors:

269 + assert is_weak_contiguous(tensor), (

270 + "Expert weight tensors must be contiguous in memory"

271 + )

272 + assert tensor.device == self._device, (

273 + "All local EPLB tensors are expected to be on the same "

274 + f"device: expected={self._device}, got={tensor.device}"

275 + )

276 + for tensor in expert_buffer:

277 + assert is_weak_contiguous(tensor), (

278 + "expert_buffer tensors must be contiguous in memory"

279 )

280

281 + # (local_dlist, remote_dlist, xfer_handle) for in-flight READs;

282 + # accumulated by add_recv, drained by execute.

283 + self._xfer_entries: list[tuple[int, int, int]] = []

284 + # Per-rank expert_id -> physical row; set by set_transfer_context.

285 + self._expert_to_src_row: list[dict[int, int]] | None = None

286 + self._layer_idx: int | None = None

287 +

288 nixl_agent_config = nixl_utils.nixl_agent_config

289 config = (

290 nixl_agent_config(capture_telemetry=False)

@@ -260,15 +293,16 @@ class NixlEplbCommunicator(EplbCommunicator):

293 )

294 self._nixl_wrapper = nixl_wrapper_cls(self._make_agent_name(), config)

295 self._nixl_memory_type = "VRAM"

263 - self._registered_desc: object | None = None

296 + # NIXL registration handles; deregistered in __del__.

297 + self._registered_descs: list[object] = []

298 self._remote_agents: dict[int, str] = {}

265 - self._remote_send_meta: dict[int, tuple[int, int]] = {}

266 - self._send_buffer: torch.Tensor = torch.empty(0)

267 - self._recv_buffer: torch.Tensor = torch.empty(0)

268 - self._expert_bytes: int = 0

299 + # peer -> (layer, tensor) -> (base_ptr, bytes_per_expert, dev_id).

300 + self._remote_send_meta: dict[

301 + int, dict[tuple[int, int], tuple[int, int, int]]

302 + ] = {}

303

304 self._cuda_device_id = int(self._device.index or 0)

271 - self._init_step("buffers", self._init_registered_buffers, expert_weights)

305 + self._init_step("buffers", self._init_registered_buffers)

306 self._init_step("agents", self._init_remote_agents)

307 self._init_step("send meta", self._exchange_remote_send_meta)

308 self._log_initialized()

@@ -291,19 +325,34 @@ class NixlEplbCommunicator(EplbCommunicator):

325 uid = uuid.uuid4().hex[:8]

326 return f"eplb-{self._rank}{pp_suffix}-{uid}"

327

328 + def set_stream(self, cuda_stream: torch.cuda.Stream | None) -> None:

329 + pass

330 +

331 def add_send(

332 self,

333 tensors: list[torch.Tensor],

334 dst_rank: int,

335 expert_id: int,

336 ) -> None:

300 - assert dst_rank != self._rank, (

301 - "EPLB communicator should not enqueue same-rank sends: "

302 - f"rank={self._rank}, dst_rank={dst_rank}"

337 + # No-op: NIXL READ is receiver-initiated. The sender's expert

338 + # weights are pre-registered and always readable in-place.

339 + pass

340 +

341 + def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:

342 + # Pre-compute expert_id -> src_row mapping for every rank so that

343 + # add_recv can immediately issue NIXL READs.

344 + assert not self._xfer_entries, (

345 + f"set_transfer_context() called with {len(self._xfer_entries)} "

346 + f"pending transfers from layer {self._layer_idx}; "

347 + f"execute() was not called after previous add_recv() calls"

348 )

304 - # An expert sent to multiple peers is packed only once; skip duplicates.

305 - if expert_id not in self._expert_send_map:

306 - self._expert_send_map[expert_id] = tensors

349 + self._layer_idx = layer_idx

350 + n = self._num_local_experts

351 + rank_experts = old_indices[: self._world_size * n].reshape(self._world_size, n)

352 + self._expert_to_src_row = [

353 + {int(eid): i for i, eid in enumerate(row) if eid != -1}

354 + for row in rank_experts

355 + ]

356

357 def add_recv(

358 self,

@@ -311,13 +360,44 @@ class NixlEplbCommunicator(EplbCommunicator):

360 src_rank: int,

361 expert_id: int,

362 ) -> None:

314 - assert src_rank != self._rank, (

315 - "EPLB communicator should not enqueue same-rank recvs: "

316 - f"rank={self._rank}, src_rank={src_rank}"

363 + # Build NIXL descriptors and issue the RDMA READ immediately,

364 + # overlapping the transfer with the remaining Python loop in

365 + # move_to_buffer.

366 + assert self._expert_to_src_row is not None and self._layer_idx is not None, (

367 + "set_transfer_context() must be called before add_recv()"

368 + )

369 + src_row = self._expert_to_src_row[src_rank][expert_id]

370 + layer_idx = self._layer_idx

371 +

372 + local_descs: list[tuple[int, int, int]] = []

373 + remote_descs: list[tuple[int, int, int]] = []

374 + for t_idx, t in enumerate(tensors):

375 + send_base, send_stride, remote_dev = self._remote_send_meta[src_rank][

376 + (layer_idx, t_idx)

377 + ]

378 + assert t.nbytes == send_stride, (

379 + f"tensor {t_idx} size {t.nbytes} != remote stride {send_stride}"

380 + )

381 + local_descs.append(

382 + (

383 + t.data_ptr(),

384 + t.nbytes,

385 + self._cuda_device_id,

386 + )

387 + )

388 + remote_descs.append(

389 + (

390 + send_base + src_row * send_stride,

391 + send_stride,

392 + remote_dev,

393 + )

394 + )

395 +

396 + local_h, remote_h, xfer_h = self._create_peer_xfer(

397 + src_rank, local_descs, remote_descs

398 )

318 - recv_experts = self._recv_map.setdefault(src_rank, {})

319 - if expert_id not in recv_experts:

320 - recv_experts[expert_id] = tensors

399 + self._nixl_wrapper.transfer(xfer_h)

400 + self._xfer_entries.append((local_h, remote_h, xfer_h))

401

402 def _init_remote_agents(self) -> None:

403 local_metadata = self._nixl_wrapper.get_agent_metadata()

@@ -334,73 +414,60 @@ class NixlEplbCommunicator(EplbCommunicator):

414 peer_metadata

415 )

416

337 - def _init_registered_buffers(self, expert_weights: Sequence[torch.Tensor]) -> None:

338 - total_bytes = max(sum(t.nbytes for t in expert_weights), 1)

339 - assert total_bytes % self._num_local_experts == 0, (

340 - f"Number of bytes in moe layer {total_bytes} is not divisible "

341 - f"by number of local experts {self._num_local_experts}"

342 - )

343 - self._expert_bytes = total_bytes // self._num_local_experts

417 + def _init_registered_buffers(self) -> None:

418 + all_tensors: list[torch.Tensor] = []

419 + for layer_tensors in self._all_expert_weights:

420 + all_tensors.extend(layer_tensors)

421 + all_tensors.extend(self._expert_buffer)

422

345 - self._send_buffer = torch.empty(

346 - total_bytes, device=self._device, dtype=torch.uint8

347 - )

348 - self._recv_buffer = torch.empty(

349 - total_bytes, device=self._device, dtype=torch.uint8

350 - )

351 -

352 - descs = self._nixl_wrapper.get_reg_descs([self._send_buffer, self._recv_buffer])

423 + descs = self._nixl_wrapper.get_reg_descs(all_tensors)

424 self._nixl_wrapper.register_memory(descs)

354 - self._registered_desc = descs

425 + self._registered_descs.append(descs)

426

427 def _exchange_remote_send_meta(self) -> None:

357 - """Exchange send-buffer metadata so each rank can build dynamic

358 - descriptors at execute time."""

359 - local_meta: tuple[int, int] = (

360 - self._send_buffer.data_ptr(),

361 - self._cuda_device_id,

362 - )

363 - gathered_meta: list[tuple[int, int] | None] = [None] * self._world_size

428 + """Exchange per-layer per-tensor metadata so receivers can compute

429 + remote RDMA addresses at transfer time."""

430 + local_meta: dict[tuple[int, int], tuple[int, int, int]] = {}

431 + for layer_idx, layer_tensors in enumerate(self._all_expert_weights):

432 + for t_idx, t in enumerate(layer_tensors):

433 + nbytes_per_expert = t.nbytes // self._num_local_experts

434 + local_meta[(layer_idx, t_idx)] = (

435 + t.data_ptr(),

436 + nbytes_per_expert,

437 + self._cuda_device_id,

438 + )

439 +

440 + # Per-rank map: (layer_idx, tensor_idx) -> (base_ptr, bytes_per_expert, dev_id).

441 + # add_recv uses base_ptr + src_row * bytes_per_expert to compute

442 + # the remote RDMA address for each expert.

443 + gathered_meta: list[dict[tuple[int, int], tuple[int, int, int]] | None] = [

444 + None

445 + ] * self._world_size

446 torch.distributed.all_gather_object(

447 gathered_meta, local_meta, group=self._cpu_group

448 )

449

450 + local_keys = set(local_meta.keys())

451 for peer in self._remote_agents:

452 peer_meta = gathered_meta[peer]

453 assert peer_meta is not None

454 + peer_keys = set(peer_meta.keys())

455 + if peer_keys != local_keys:

456 + raise RuntimeError(

457 + f"NIXL EPLB metadata key mismatch with rank {peer}: "

458 + f"local={sorted(local_keys)}, peer={sorted(peer_keys)}"

459 + )

460 + for key in local_keys:

461 + _, local_stride, _ = local_meta[key]

462 + _, peer_stride, _ = peer_meta[key]

463 + if local_stride != peer_stride:

464 + raise RuntimeError(

465 + f"NIXL EPLB nbytes_per_expert mismatch for {key} "

466 + f"with rank {peer}: "

467 + f"local={local_stride}, peer={peer_stride}"

468 + )

469 self._remote_send_meta[peer] = peer_meta

470

373 - @staticmethod

374 - def _pack_send_buffer(

375 - in_tensors: list[torch.Tensor],

376 - send_buffer: torch.Tensor,

377 - byte_offset: int,

378 - ) -> None:

379 - for tensor in in_tensors:

380 - raw = tensor.reshape(-1).view(torch.uint8)

381 - if raw.numel() == 0:

382 - continue

383 - send_buffer[byte_offset : byte_offset + raw.numel()].copy_(

384 - raw, non_blocking=True

385 - )

386 - byte_offset += raw.numel()

387 -

388 - @staticmethod

389 - def _unpack_recv_buffer(

390 - recv_buffer: torch.Tensor,

391 - out_tensors: list[torch.Tensor],

392 - byte_offset: int,

393 - ) -> None:

394 - for tensor in out_tensors:

395 - num_bytes = tensor.numel() * tensor.element_size()

396 - if num_bytes == 0:

397 - continue

398 - tensor.reshape(-1).view(torch.uint8).copy_(

399 - recv_buffer[byte_offset : byte_offset + num_bytes],

400 - non_blocking=True,

401 - )

402 - byte_offset += num_bytes

403 -

471 def _wait_for_all_transfers(self, handles: list[int]) -> None:

472 pending = set(handles)

473 while pending:

@@ -456,110 +523,52 @@ class NixlEplbCommunicator(EplbCommunicator):

523 )

524 return (local_handle, remote_handle, xfer_handle)

525

459 - def execute(self, old_indices: np.ndarray | None = None) -> None:

460 - assert old_indices is not None, (

461 - "NixlEplbCommunicator.execute requires old_indices"

526 + def execute(self) -> None:

527 + assert self._layer_idx is not None or not self._xfer_entries, (

528 + "set_transfer_context() must be called before execute() "

529 + "if any add_recv() calls were made"

530 )

463 -

464 - xfer_entries: list[tuple[int, int, int]] = []

531 try:

466 - n = self._num_local_experts

467 - rank_experts = old_indices[: self._world_size * n].reshape(

468 - self._world_size, n

469 - )

470 - # Build expert_id -> send slot mapping per rank.

471 - expert_to_send_slot: list[dict[int, int]] = [

472 - {int(eid): i for i, eid in enumerate(row) if eid != -1}

473 - for row in rank_experts

474 - ]

532 + self._wait_for_all_transfers([x[2] for x in self._xfer_entries])

533

476 - # Phase 1: pack each expert at its slot offset in the send buffer.

477 - with torch.cuda.stream(self._cuda_stream):

478 - for expert_id, tensors in self._expert_send_map.items():

479 - slot = expert_to_send_slot[self._rank][expert_id]

480 - byte_offset = slot * self._expert_bytes

481 - self._pack_send_buffer(tensors, self._send_buffer, byte_offset)

482 -

483 - # Ensure all packed data is visible in device memory before pulls.

484 - if self._cuda_stream is not None:

485 - self._cuda_stream.synchronize()

486 - else:

487 - torch.cuda.current_stream().synchronize()

488 - # READ is receiver-initiated; synchronize all ranks before transfer.

489 - # We use monitored_barrier so a rank that crashes or exits early

490 - # produces a diagnostic timeout instead of a silent hang.

534 + # Post-READ barrier.

535 + # Correctness fence for zero-copy: prevents overwrite-while-

536 + # remote-read race.

537 torch.distributed.monitored_barrier(

538 group=self._cpu_group,

539 timeout=timedelta(minutes=5),

540 )

495 -

496 - # Phase 2: issue one batched READ per peer.

497 - recv_offsets: dict[tuple[int, int], int] = {}

498 - recv_offset = 0

499 - recv_base = self._recv_buffer.data_ptr()

500 - for src in range(self._world_size):

501 - if src == self._rank:

502 - continue

503 - recv_experts = self._recv_map.get(src)

504 - if not recv_experts:

505 - continue

506 - expert_ids = list(recv_experts.keys())

507 - remote_base, remote_dev = self._remote_send_meta[src]

508 - local_descs: list[tuple[int, int, int]] = []

509 - remote_descs: list[tuple[int, int, int]] = []

510 - for expert_id in expert_ids:

511 - slot = expert_to_send_slot[src][expert_id]

512 - remote_off = slot * self._expert_bytes

513 - recv_offsets[(src, expert_id)] = recv_offset

514 - local_descs.append(

515 - (

516 - recv_base + recv_offset,

517 - self._expert_bytes,

518 - self._cuda_device_id,

519 - )

520 - )

521 - remote_descs.append(

522 - (remote_base + remote_off, self._expert_bytes, remote_dev)

523 - )

524 - recv_offset += self._expert_bytes

525 - assert recv_offset <= self._recv_buffer.nbytes

526 - local_h, remote_h, xfer_h = self._create_peer_xfer(

527 - src, local_descs, remote_descs

528 - )

529 - self._nixl_wrapper.transfer(xfer_h)

530 - xfer_entries.append((local_h, remote_h, xfer_h))

531 -

532 - # Phase 3: wait for all in-flight transfers, then unpack.

533 - self._wait_for_all_transfers([x[2] for x in xfer_entries])

534 -

535 - with torch.cuda.stream(self._cuda_stream):

536 - for (src, expert_id), offset in recv_offsets.items():

537 - self._unpack_recv_buffer(

538 - self._recv_buffer,

539 - self._recv_map[src][expert_id],

540 - offset,

541 - )

541 finally:

543 - for local_h, remote_h, xfer_h in xfer_entries:

542 + for local_h, remote_h, xfer_h in self._xfer_entries:

543 with contextlib.suppress(Exception):

544 self._nixl_wrapper.release_xfer_handle(xfer_h)

545 with contextlib.suppress(Exception):

546 self._nixl_wrapper.release_dlist_handle(local_h)

547 with contextlib.suppress(Exception):

548 self._nixl_wrapper.release_dlist_handle(remote_h)

550 - self._expert_send_map.clear()

551 - self._recv_map.clear()

549 + self._xfer_entries.clear()

550 + self._expert_to_src_row = None

551 + self._layer_idx = None

552

553 def __del__(self) -> None:

554 - try:

555 - if self._registered_desc is not None:

556 - self._nixl_wrapper.deregister_memory(self._registered_desc)

557 - self._registered_desc = None

554 + with contextlib.suppress(Exception):

555 + for local_h, remote_h, xfer_h in self._xfer_entries:

556 + with contextlib.suppress(Exception):

557 + self._nixl_wrapper.release_xfer_handle(xfer_h)

558 + with contextlib.suppress(Exception):

559 + self._nixl_wrapper.release_dlist_handle(local_h)

560 + with contextlib.suppress(Exception):

561 + self._nixl_wrapper.release_dlist_handle(remote_h)

562 + with contextlib.suppress(Exception):

563 + for descs in self._registered_descs:

564 + with contextlib.suppress(Exception):

565 + self._nixl_wrapper.deregister_memory(descs)

566 + self._registered_descs.clear()

567 + with contextlib.suppress(Exception):

568 for agent_name in self._remote_agents.values():

559 - self._nixl_wrapper.remove_remote_agent(agent_name)

569 + with contextlib.suppress(Exception):

570 + self._nixl_wrapper.remove_remote_agent(agent_name)

571 self._remote_agents.clear()

561 - except Exception as e:

562 - logger.warning("Error during NixlEplbCommunicator cleanup: %s", e)

572

573

574 class PyNcclEplbCommunicator(EplbCommunicator):

@@ -600,7 +609,7 @@ class PyNcclEplbCommunicator(EplbCommunicator):

609 for tensor in tensors:

610 self._pynccl_comm.recv(tensor, src_rank, stream=self._cuda_stream)

611

603 - def execute(self, old_indices: np.ndarray | None = None) -> None:

612 + def execute(self) -> None:

613 if self._group_started:

614 self._pynccl_comm.group_end()

615 self._group_started = False

@@ -609,7 +618,8 @@ class PyNcclEplbCommunicator(EplbCommunicator):

618 def create_eplb_communicator(

619 group_coordinator: GroupCoordinator,

620 backend: str | None,

612 - expert_weights: Sequence[torch.Tensor],

621 + expert_weights: Sequence[Sequence[torch.Tensor]],

622 + expert_buffer: Sequence[torch.Tensor],

623 ) -> EplbCommunicator:

624 """Create an EPLB communicator for the given backend.

625

@@ -624,16 +634,18 @@ def create_eplb_communicator(

634 ``"pynccl"`` in that case. When tensors reside on CPU,

635 ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU

636 process group.

627 - expert_weights: Expert weight tensors from *one* MoE layer.

628 - NixlEplbCommunicator pre-allocates send/recv buffers sized

629 - to this layer, so all other MoE layers must have the same

630 - tensor count, shapes, and dtypes.

637 + expert_weights: Expert weight tensors for *all* MoE layers.

638 + Shape ``(num_layers)(num_tensors_per_layer)``.

639 + NixlEplbCommunicator registers all layers with NIXL for

640 + zero-copy RDMA reads.

641 + expert_buffer: Pre-allocated receive buffer tensors (one per

642 + weight tensor in a single layer).

643 """

632 - # Keep a safe default for callers that have not resolved communicator yet.

644 if backend is None:

645 backend = "torch_nccl"

646

636 - tensor_device_type = expert_weights[0].device.type if expert_weights else "cpu"

647 + first_layer = expert_weights[0] if expert_weights else []

648 + tensor_device_type = first_layer[0].device.type if first_layer else "cpu"

649 torch_group = (

650 group_coordinator.cpu_group

651 if tensor_device_type == "cpu"

@@ -649,7 +661,7 @@ def create_eplb_communicator(

661 unsupported_dtypes = sorted(

662 {

663 tensor.dtype

652 - for tensor in expert_weights

664 + for tensor in first_layer

665 if not ncclDataTypeEnum.supports_torch_dtype(tensor.dtype)

666 },

667 key=str,

@@ -704,7 +716,8 @@ def create_eplb_communicator(

716 try:

717 return NixlEplbCommunicator(

718 cpu_group=group_coordinator.cpu_group,

707 - expert_weights=expert_weights,

719 + all_expert_weights=expert_weights,

720 + expert_buffer=expert_buffer,

721 )

722 except Exception as exc:

723 raise RuntimeError(

vllm/distributed/eplb/eplb_state.py

+3 -1

@@ -450,7 +450,8 @@ class EplbState:

450 communicator = create_eplb_communicator(

451 group_coordinator=get_eplb_group(),

452 backend=self.parallel_config.eplb_config.communicator,

453 - expert_weights=model.expert_weights[0],

453 + expert_weights=model.expert_weights,

454 + expert_buffer=expert_buffer,

455 )

456

457 model_state = EplbModelState(

@@ -766,6 +767,7 @@ class EplbState:

767 eplb_model_state.physical_to_logical_map,

768 new_physical_to_logical_map,

769 eplb_model_state.model.expert_weights,

770 + eplb_model_state.expert_buffer,

771 ep_group,

772 eplb_model_state.communicator,

773 is_profile,

vllm/distributed/eplb/rebalance_execute.py

+16 -9

@@ -178,6 +178,7 @@ def move_to_buffer(

178 cuda_stream: torch.cuda.Stream | None,

179 ep_rank: int,

180 communicator: EplbCommunicator,

181 + layer_idx: int = 0,

182 ) -> TransferMetadata:

183 """

184 Rearranges expert weights during EPLB rebalancing.

@@ -193,6 +194,7 @@ def move_to_buffer(

194 cuda_stream: CUDA stream for async copies (can be None for sync mode).

195 ep_rank: Rank of this process in expert parallel group.

196 communicator: EplbCommunicator instance for P2P communication.

197 + layer_idx: Index of the MoE layer being transferred.

198

199 Returns:

200 TransferMetadata: Metadata needed for completing remote weight transfers.

@@ -265,6 +267,8 @@ def move_to_buffer(

267 for w, b in zip(expert_weights, expert_weights_buffers):

268 b[dst].copy_(w[src_local], non_blocking=True)

269

270 + communicator.set_transfer_context(old_indices, layer_idx)

271 +

272 # 2. Post sends

273 if send_count > 0:

274 experts = send_expert_ids[:send_count]

@@ -331,9 +335,8 @@ def move_to_buffer(

335 expert_id=int(expert),

336 )

337

334 - # 4. Execute the P2P operations. The real communication happens here.

335 - communicator.execute(old_indices=old_indices)

336 - # wait for the communication to finish

338 + # 4. Execute transfers and wait for completion.

339 + communicator.execute()

340 return TransferMetadata(

341 is_unchanged=is_unchanged,

342 is_received_locally=is_received_locally,

@@ -431,6 +434,7 @@ def transfer_layer(

434 is_profile: bool = False,

435 cuda_stream: torch.cuda.Stream | None = None,

436 rank_mapping: dict[int, int] | None = None,

437 + layer_idx: int = 0,

438 ) -> TransferMetadata:

439 """

440 Rearranges the expert weights in place according to the new expert indices.

@@ -452,6 +456,7 @@ def transfer_layer(

456 communications to reserve enough memory for the buffers.

457 cuda_stream: CUDA stream for async copies (can be None for sync mode).

458 rank_mapping: Optional rank mapping for elastic expert parallelism.

459 + layer_idx: Index of the MoE layer being transferred.

460

461 Returns:

462 TransferMetadata: Metadata needed for completing remote weight transfers,

@@ -499,6 +504,7 @@ def transfer_layer(

504 cuda_stream=cuda_stream,

505 ep_rank=ep_group.rank(),

506 communicator=communicator,

507 + layer_idx=layer_idx,

508 )

509

510

@@ -506,6 +512,7 @@ def rearrange_expert_weights_inplace(

512 old_global_expert_indices: torch.Tensor,

513 new_global_expert_indices: torch.Tensor,

514 expert_weights: Sequence[Sequence[torch.Tensor]],

515 + expert_buffer: Sequence[torch.Tensor],

516 ep_group: ProcessGroup,

517 communicator: EplbCommunicator,

518 is_profile: bool = False,

@@ -524,6 +531,8 @@ def rearrange_expert_weights_inplace(

531 of tensors of shape (num_local_physical_experts, hidden_size_i).

532 For example, a linear layer may have up and down projection,

533 so weight_count = 2. Each weight's hidden size can be different.

534 + expert_buffer: Pre-allocated receive buffer tensors (one per

535 + weight tensor in a single layer).

536 ep_group: The device process group for expert parallelism.

537 communicator: EplbCommunicator instance for P2P communication.

538 is_profile (bool): If `True`, do not perform any actual weight copy.

@@ -566,10 +575,10 @@ def rearrange_expert_weights_inplace(

575 # Reserve NCCL communication buffers via a dummy all_gather.

576 # Backends that pre-allocate their own transfer buffers

577 # skip this to avoid the extra memory spike during profiling.

569 - weights_buffer: list[torch.Tensor] = [

578 + profile_buffer: list[torch.Tensor] = [

579 torch.empty_like(w) for w in first_layer_weights

580 ]

572 - for weight, buffer in zip(expert_weights[0], weights_buffer):

581 + for weight, buffer in zip(expert_weights[0], profile_buffer):

582 dummy_recv_buffer = [buffer for _ in range(ep_size)]

583 torch.distributed.barrier()

584 all_gather(

@@ -579,10 +588,7 @@ def rearrange_expert_weights_inplace(

588 )

589 return

590

582 - # Buffers to hold the expert weights during the exchange.

583 - # NOTE: Currently we assume the same weights across different layers

584 - # have the same shape.

585 - weights_buffer = [torch.empty_like(w) for w in first_layer_weights]

591 + weights_buffer = list(expert_buffer)

592

593 old_global_expert_indices_cpu = old_global_expert_indices.cpu().numpy()

594 new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()

@@ -597,6 +603,7 @@ def rearrange_expert_weights_inplace(

603 cuda_stream=None,

604 ep_rank=ep_rank,

605 communicator=communicator,

606 + layer_idx=layer_idx,

607 )

608

609 move_from_buffer(

vllm/distributed/utils.py

+14

@@ -64,6 +64,20 @@ def divide(numerator, denominator):

64 return numerator // denominator

67 +def is_weak_contiguous(inp: torch.Tensor) -> bool:

68 + """Check that *inp* occupies a single contiguous block of memory.

69 +

70 + Unlike ``torch.Tensor.is_contiguous()``, this also accepts tensors

71 + whose strides are not strictly C-contiguous (e.g. column-major) as

72 + long as the underlying storage from the tensor's offset onward is

73 + exactly ``numel * element_size`` bytes.

74 + """

75 + return inp.is_contiguous() or (

76 + inp.storage().nbytes() - inp.storage_offset() * inp.element_size()

77 + == inp.numel() * inp.element_size()

78 + )

79 +

80 +

81 def split_tensor_along_last_dim(

82 tensor: torch.Tensor,

83 num_partitions: int,