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] Consolidate is_unchanged/is_received_locally into TransferMetadata (#37341)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Apr 20, 2026 at 14:12 UTC
3173441b0f73324771c5638e8c123a2cef784198
4 files changed
+36
-64
tests/distributed/test_eplb_execute.py
+2
-4
@@ -361,7 +361,7 @@ def _test_async_transfer_layer_without_mtp_worker(
361
communicator.set_stream(cuda_stream)
362
363
for layer_idx in range(num_layers):
364
- is_unchanged, is_received_locally, recv_metadata = asyncio.run(
364
+ transfer_metadata = asyncio.run(
365
transfer_layer(
366
old_layer_indices=old_indices_cpu[layer_idx],
367
new_layer_indices=new_indices_cpu[layer_idx],
@@ -376,9 +376,7 @@ def _test_async_transfer_layer_without_mtp_worker(
376
move_from_buffer(
377
expert_weights=expert_weights[layer_idx],
378
expert_weights_buffers=expert_buffer,
379
- is_unchanged=is_unchanged,
380
- is_received_locally=is_received_locally,
381
- recv_metadata=recv_metadata,
379
+ transfer_metadata=transfer_metadata,
380
new_indices=new_indices_cpu[layer_idx].numpy(),
381
ep_rank=ep_rank,
382
)
vllm/distributed/eplb/async_worker.py
+2
-8
@@ -118,11 +118,7 @@ async def transfer_run_periodically(
118
# model_state.expert_buffer, which will be consumed by the main thread in
119
# move_to_workspace
120
while model_state.rebalanced and layer_idx < num_layers:
121
- (
122
- is_unchanged,
123
- is_received_locally,
124
- recv_metadata,
125
- ) = await transfer_layer(
121
+ transfer_metadata = await transfer_layer(
122
old_layer_indices=physical_to_logical_map_cpu[layer_idx],
123
new_layer_indices=new_physical_to_logical_map[layer_idx],
124
expert_weights=model_state.model.expert_weights[layer_idx],
@@ -145,9 +141,7 @@ async def transfer_run_periodically(
141
model_state.pending_result = AsyncEplbLayerResult(
142
layer_idx=layer_idx,
143
new_physical_to_logical_map=new_physical_to_logical_map[layer_idx],
148
- is_unchanged=is_unchanged,
149
- is_received_locally=is_received_locally,
150
- recv_metadata=recv_metadata,
144
+ transfer_metadata=transfer_metadata,
145
consumed_event=consumed_event,
146
)
147
vllm/distributed/eplb/eplb_state.py
+1
-3
@@ -1147,9 +1147,7 @@ def _move_to_workspace(
1147
move_from_buffer(
1148
expert_weights=model_state.model.expert_weights[result.layer_idx],
1149
expert_weights_buffers=model_state.expert_buffer,
1150
- is_unchanged=result.is_unchanged,
1151
- is_received_locally=result.is_received_locally,
1152
- recv_metadata=result.recv_metadata,
1150
+ transfer_metadata=result.transfer_metadata,
1151
new_indices=result.new_physical_to_logical_map.numpy(),
1152
ep_rank=ep_rank,
1153
)
vllm/distributed/eplb/rebalance_execute.py
+31
-49
@@ -21,9 +21,13 @@ logger = init_logger(__name__)
21
22
23
@dataclass
24
-class RecvMetadata:
25
- """Metadata describing remote receives during EPLB rebalancing."""
24
+class TransferMetadata:
25
+ """Metadata describing a completed EPLB buffer transfer."""
26
27
+ is_unchanged: np.ndarray
28
+ """Mask of (num_local_experts,) indicating experts unchanged after rebalance."""
29
+ is_received_locally: np.ndarray
30
+ """Mask of (num_local_experts,) indicating experts received from local data."""
31
recv_primary_mask: np.ndarray
32
"""Mask of (num_local_experts,) indicating primary experts received."""
33
recv_count: int
@@ -34,10 +38,6 @@ class RecvMetadata:
38
"""Target expert indices (num_local_experts,) in local tensors to send."""
39
40
37
-# Type alias for the result of move_to_buffer or transfer_layer
38
-MoveToBufferResult = tuple[np.ndarray, np.ndarray, RecvMetadata]
39
-
40
-
41
@dataclass
42
class AsyncEplbLayerResult:
43
"""
@@ -51,11 +51,7 @@ class AsyncEplbLayerResult:
51
New physical→logical mapping for layers_idx, on CPU.
52
Shape: (num_physical_experts)
53
"""
54
- is_unchanged: np.ndarray
55
- """Per-physical-expert flag: weight was not moved during transfer."""
56
- is_received_locally: np.ndarray
57
- """Per-physical-expert flag: weight was received on this rank."""
58
- recv_metadata: RecvMetadata
54
+ transfer_metadata: TransferMetadata
55
"""Metadata describing what was received during transfer_layer."""
56
consumed_event: CpuGpuEvent
57
"""
@@ -182,7 +178,7 @@ def move_to_buffer(
178
cuda_stream: torch.cuda.Stream | None,
179
ep_rank: int,
180
communicator: EplbCommunicator,
185
-) -> MoveToBufferResult:
181
+) -> TransferMetadata:
182
"""
183
Rearranges expert weights during EPLB rebalancing.
184
@@ -199,11 +195,7 @@ def move_to_buffer(
195
communicator: EplbCommunicator instance for P2P communication.
196
197
Returns:
202
- is_unchanged (np.ndarray): (num_local_experts,), True where an expert row
203
- is unchanged after rebalance.
204
- is_received_locally (np.ndarray): (num_local_experts,), True where a row
205
- can be updated from local data.
206
- RecvMetadata: Metadata needed for completing remote weight transfers.
198
+ TransferMetadata: Metadata needed for completing remote weight transfers.
199
"""
200
assert old_indices.shape == new_indices.shape
201
recv_primary_mask = np.zeros((num_local_experts,), dtype=np.bool_)
@@ -339,24 +331,20 @@ def move_to_buffer(
331
# 4. Execute the P2P operations. The real communication happens here.
332
communicator.execute()
333
# wait for the communication to finish
342
- return (
343
- is_unchanged,
344
- is_received_locally,
345
- RecvMetadata(
346
- recv_primary_mask=recv_primary_mask,
347
- recv_count=recv_count,
348
- recv_expert_ids=recv_expert_ids,
349
- recv_dst_rows=recv_dst_rows,
350
- ),
334
+ return TransferMetadata(
335
+ is_unchanged=is_unchanged,
336
+ is_received_locally=is_received_locally,
337
+ recv_primary_mask=recv_primary_mask,
338
+ recv_count=recv_count,
339
+ recv_expert_ids=recv_expert_ids,
340
+ recv_dst_rows=recv_dst_rows,
341
)
342
343
344
def move_from_buffer(
345
expert_weights: Sequence[torch.Tensor],
346
expert_weights_buffers: list[torch.Tensor],
357
- is_unchanged: np.ndarray,
358
- is_received_locally: np.ndarray,
359
- recv_metadata: RecvMetadata,
347
+ transfer_metadata: TransferMetadata,
348
new_indices: np.ndarray,
349
ep_rank: int,
350
) -> None:
@@ -368,17 +356,17 @@ def move_from_buffer(
356
expert_weights: List of the actual MoE layer weights used in the execution.
357
expert_weights_buffers: Intermediate buffers containing the experts weights
358
after the transfer is completed.
371
- is_unchanged: (num_local_experts,), True where an expert row is unchanged.
372
- is_received_locally: (num_local_experts,), True where a row is updated locally.
373
- recv_metadata: RecvMetadata containing remote receive metadata.
359
+ transfer_metadata: TransferMetadata containing transfer metadata.
360
new_indices: (num_experts_total,) mapping from local rows to desired
361
(possibly global) expert id, after rebalance.
362
ep_rank: Rank of the process in the expert parallel group.
363
"""
378
- recv_primary_mask = recv_metadata.recv_primary_mask
379
- recv_count = recv_metadata.recv_count
380
- recv_expert_ids = recv_metadata.recv_expert_ids
381
- recv_dst_rows = recv_metadata.recv_dst_rows
364
+ is_unchanged = transfer_metadata.is_unchanged
365
+ is_received_locally = transfer_metadata.is_received_locally
366
+ recv_primary_mask = transfer_metadata.recv_primary_mask
367
+ recv_count = transfer_metadata.recv_count
368
+ recv_expert_ids = transfer_metadata.recv_expert_ids
369
+ recv_dst_rows = transfer_metadata.recv_dst_rows
370
num_local_experts = is_unchanged.shape[0]
371
372
# Mask for rows to copy back from buffers:
@@ -440,7 +428,7 @@ async def transfer_layer(
428
is_profile: bool = False,
429
cuda_stream: torch.cuda.Stream | None = None,
430
rank_mapping: dict[int, int] | None = None,
443
-) -> MoveToBufferResult:
431
+) -> TransferMetadata:
432
"""
433
Rearranges the expert weights in place according to the new expert indices.
434
@@ -463,11 +451,8 @@ async def transfer_layer(
451
rank_mapping: Optional rank mapping for elastic expert parallelism.
452
453
Returns:
466
- is_unchanged (np.ndarray): (num_local_experts,), True where expert
467
- is left unchanged.
468
- is_received_locally (np.ndarray): (num_local_experts,), True where expert
469
- can be received locally.
470
- RecvMetadata: Metadata needed for completing remote weight transfers.
454
+ TransferMetadata: Metadata needed for completing remote weight transfers,
455
+ including is_unchanged and is_received_locally masks.
456
"""
457
ep_size = ep_group.size()
458
if rank_mapping is not None:
@@ -502,7 +487,7 @@ async def transfer_layer(
487
old_layer_indices_np = old_layer_indices.cpu().numpy()
488
new_layer_indices_np = new_layer_indices.cpu().numpy()
489
505
- is_unchanged, is_received_locally, recv_metadata = move_to_buffer(
490
+ return move_to_buffer(
491
num_local_experts=num_local_physical_experts,
492
old_indices=old_layer_indices_np,
493
new_indices=new_layer_indices_np,
@@ -512,7 +497,6 @@ async def transfer_layer(
497
ep_rank=ep_group.rank(),
498
communicator=communicator,
499
)
515
- return is_unchanged, is_received_locally, recv_metadata
500
501
502
def rearrange_expert_weights_inplace(
@@ -605,7 +589,7 @@ def rearrange_expert_weights_inplace(
589
new_global_expert_indices_cpu = new_global_expert_indices.cpu().numpy()
590
591
for layer_idx in range(num_moe_layers):
608
- is_unchanged, is_received_locally, recv_metadata = move_to_buffer(
592
+ transfer_metadata = move_to_buffer(
593
num_local_experts=num_local_physical_experts,
594
old_indices=old_global_expert_indices_cpu[layer_idx],
595
new_indices=new_global_expert_indices_cpu[layer_idx],
@@ -619,9 +603,7 @@ def rearrange_expert_weights_inplace(
603
move_from_buffer(
604
expert_weights=expert_weights[layer_idx],
605
expert_weights_buffers=weights_buffer,
622
- is_unchanged=is_unchanged,
623
- is_received_locally=is_received_locally,
624
- recv_metadata=recv_metadata,
606
+ transfer_metadata=transfer_metadata,
607
new_indices=new_global_expert_indices_cpu[layer_idx],
608
ep_rank=ep_rank,
609
)
@@ -715,4 +697,4 @@ def _map_new_expert_indices_with_rank_mapping(
697
return mapped_expert_indices
698
699
718
-__all__ = ["transfer_layer", "move_from_buffer", "RecvMetadata"]
700
+__all__ = ["transfer_layer", "move_from_buffer", "TransferMetadata"]