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 10, 06:30 UTC
[Refactor] EPLB rebalance algo to NumPy (#30697)
Signed-off-by: ilmarkov <markovilya197@gmail.com>
Ilya Markov committed
Jan 12, 2026 at 19:13 UTC
1eb61ab34b1415023a35afeb40f1cf85256b5fa5
3 files changed
+128
-130
tests/distributed/test_eplb_algo.py
+16
-15
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import numpy as np
5
import pytest
6
import torch
7
@@ -312,9 +313,9 @@ if __name__ == "__main__":
313
test_basic_rebalance()
314
315
315
-def _make_phy_replicas_idx_from_phy2log(phy2log: torch.Tensor) -> torch.Tensor:
316
- """Create replicas indices mapping from phy2log"""
317
- pr = torch.zeros_like(phy2log)
316
+def _make_phy_replicas_idx_from_phy2log(phy2log: np.ndarray) -> np.ndarray:
317
+ """Create replicas indices mapping from phy2log."""
318
+ pr = np.zeros_like(phy2log, dtype=np.int64)
319
for layer in range(phy2log.shape[0]):
320
seen: dict[int, int] = {}
321
row = phy2log[layer].tolist()
@@ -326,11 +327,11 @@ def _make_phy_replicas_idx_from_phy2log(phy2log: torch.Tensor) -> torch.Tensor:
327
328
329
def _validate_intragpu_rearrangement(
329
- old_global_expert_indices: torch.Tensor,
330
- new_phy2log: torch.Tensor,
331
- new_phy_replicas_idx: torch.Tensor,
332
- post_phy2log: torch.Tensor,
333
- post_phy_replicas_idx: torch.Tensor,
330
+ old_global_expert_indices: np.ndarray,
331
+ new_phy2log: np.ndarray,
332
+ new_phy_replicas_idx: np.ndarray,
333
+ post_phy2log: np.ndarray,
334
+ post_phy_replicas_idx: np.ndarray,
335
num_ranks: int,
336
slots_per_gpu: int,
337
):
@@ -345,7 +346,7 @@ def _validate_intragpu_rearrangement(
346
post_rnk = post_phy_replicas_idx[0, start:end]
347
348
# Pairwise equality for (expert, rank) pairs to ensure nothing is lost
348
- def sorted_pairs(seg: torch.Tensor, rnk: torch.Tensor):
349
+ def sorted_pairs(seg, rnk):
350
pairs = list(zip(seg.tolist(), rnk.tolist()))
351
pairs.sort()
352
return pairs
@@ -386,8 +387,8 @@ def _validate_intragpu_rearrangement(
387
# GPU0 new -> [1,5,0,4]; GPU1 new -> [6,2,7,3]
388
2,
389
4,
389
- torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7]]),
390
- torch.tensor([[1, 5, 0, 4, 6, 2, 7, 3]]),
390
+ np.array([[0, 1, 2, 3, 4, 5, 6, 7]]),
391
+ np.array([[1, 5, 0, 4, 6, 2, 7, 3]]),
392
id="simple",
393
),
394
pytest.param(
@@ -401,8 +402,8 @@ def _validate_intragpu_rearrangement(
402
# GPU1 new -> [6, 2, 3, 2, 1] (expert 2 duplicated)
403
2,
404
5,
404
- torch.tensor([[0, 1, 0, 2, 3, 4, 5, 6, 1, 2]]),
405
- torch.tensor([[0, 5, 4, 0, 1, 6, 2, 3, 2, 1]]),
405
+ np.array([[0, 1, 0, 2, 3, 4, 5, 6, 1, 2]]),
406
+ np.array([[0, 5, 4, 0, 1, 6, 2, 3, 2, 1]]),
407
id="duplicates",
408
),
409
pytest.param(
@@ -418,8 +419,8 @@ def _validate_intragpu_rearrangement(
419
# GPU2 new -> [1, 2, 3, 0]
420
3,
421
4,
421
- torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]),
422
- torch.tensor([[0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 0]]),
422
+ np.array([[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]),
423
+ np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 0]]),
424
id="skewed_expert",
425
),
426
],
tests/distributed/test_eplb_execute.py
+1
-1
@@ -311,7 +311,7 @@ def _test_async_transfer_layer_without_mtp_worker(
311
is_unchanged=is_unchanged,
312
is_received_locally=is_received_locally,
313
recv_metadata=recv_metadata,
314
- new_indices=new_indices_cpu[layer_idx],
314
+ new_indices=new_indices_cpu[layer_idx].numpy(),
315
ep_rank=ep_rank,
316
)
317
vllm/distributed/eplb/policy/default.py
+111
-114
@@ -21,8 +21,8 @@ from .abstract import AbstractEplbPolicy
21
class DefaultEplbPolicy(AbstractEplbPolicy):
22
@classmethod
23
def balanced_packing(
24
- cls, weight: torch.Tensor, num_packs: int
25
- ) -> tuple[torch.Tensor, torch.Tensor]:
24
+ cls, weight: np.ndarray, num_packs: int
25
+ ) -> tuple[np.ndarray, np.ndarray]:
26
"""
27
Pack n weighted objects to m packs, such that each bin contains exactly
28
n/m objects and the weights of all packs are as balanced as possible.
@@ -39,50 +39,43 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
39
assert num_groups % num_packs == 0
40
groups_per_pack = num_groups // num_packs
41
42
- device = weight.device
43
-
42
if groups_per_pack == 1:
45
- pack_index = torch.arange(
46
- weight.size(-1), dtype=torch.int64, device=device
47
- ).expand(weight.shape)
48
- rank_in_pack = torch.zeros_like(weight, dtype=torch.int64, device=device)
43
+ pack_index = np.tile(np.arange(num_groups, dtype=np.int64), (num_layers, 1))
44
+ rank_in_pack = np.zeros_like(pack_index, dtype=np.int64)
45
return pack_index, rank_in_pack
46
51
- weight_np = weight.cpu().numpy()
52
-
47
# Sort and get indices in decending order
54
- indices_np = np.argsort(-weight_np, axis=-1)
48
+ indices = np.argsort(-weight, axis=-1)
49
56
- pack_index_np = np.full((num_layers, num_groups), -1, dtype=np.int64)
57
- rank_in_pack_np = np.full((num_layers, num_groups), -1, dtype=np.int64)
58
-
59
- # Run the packing algorithm
60
- for i in range(num_layers):
61
- pack_weights = [0.0] * num_packs
62
- pack_items = [0] * num_packs
63
-
64
- for group in indices_np[i]:
65
- # Find a pack with capacity that has the lowest weight
66
- pack = min(
67
- (j for j in range(num_packs) if pack_items[j] < groups_per_pack),
68
- key=pack_weights.__getitem__,
69
- )
50
+ pack_index = np.full((num_layers, num_groups), -1, dtype=np.int64)
51
+ rank_in_pack = np.full((num_layers, num_groups), -1, dtype=np.int64)
52
71
- assert pack_items[pack] < groups_per_pack
72
- pack_index_np[i, group] = pack
73
- rank_in_pack_np[i, group] = pack_items[pack]
74
- pack_weights[pack] += weight_np[i, group]
75
- pack_items[pack] += 1
53
+ pack_weights = np.zeros((num_layers, num_packs), dtype=np.float64)
54
+ pack_items = np.zeros((num_layers, num_packs), dtype=np.int64)
55
77
- pack_index = torch.from_numpy(pack_index_np).to(device)
78
- rank_in_pack = torch.from_numpy(rank_in_pack_np).to(device)
56
+ # Run the packing algorithm
57
+ for layer_idx in range(num_layers):
58
+ weights_row = pack_weights[layer_idx]
59
+ items_row = pack_items[layer_idx]
60
+
61
+ for group in indices[layer_idx]:
62
+ # Pick the lightest pack; full packs are masked out by inf.
63
+ pack = int(np.argmin(weights_row))
64
+
65
+ pack_index[layer_idx, group] = pack
66
+ rank_in_pack[layer_idx, group] = items_row[pack]
67
+ weights_row[pack] += weight[layer_idx, group]
68
+ items_row[pack] += 1
69
+ if items_row[pack] == groups_per_pack:
70
+ # Mark as unavailable for future selections.
71
+ weights_row[pack] = np.inf
72
73
return pack_index, rank_in_pack
74
75
@classmethod
76
def replicate_experts(
84
- cls, weight: torch.Tensor, num_phy: int
85
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
77
+ cls, weight: np.ndarray, num_phy: int
78
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
79
"""
80
Replicate `num_log` experts to `num_phy` replicas, such that the maximum
81
load of all replicas is minimized.
@@ -99,13 +92,12 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
92
n, num_log = weight.shape
93
num_redundant = num_phy - num_log
94
assert num_redundant >= 0
102
- device = weight.device
103
- phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1)
104
- replica_idx = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
105
- logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
106
- arangen = torch.arange(n, dtype=torch.int64, device=device)
95
+ phy2log = np.tile(np.arange(num_phy, dtype=np.int64), (n, 1))
96
+ replica_idx = np.zeros((n, num_phy), dtype=np.int64)
97
+ logcnt = np.ones((n, num_log), dtype=np.int64)
98
+ arangen = np.arange(n, dtype=np.int64)
99
for i in range(num_log, num_phy):
108
- redundant_indices = (weight / logcnt).max(dim=-1).indices
100
+ redundant_indices = np.argmax(weight / logcnt, axis=-1)
101
phy2log[:, i] = redundant_indices
102
replica_idx[:, i] = logcnt[arangen, redundant_indices]
103
logcnt[arangen, redundant_indices] += 1
@@ -114,12 +106,12 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
106
@classmethod
107
def rebalance_experts_hierarchical(
108
cls,
117
- weight: torch.Tensor,
109
+ weight: np.ndarray,
110
num_physical_experts: int,
111
num_groups: int,
112
num_nodes: int,
113
num_gpus: int,
122
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
114
+ ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
115
"""
116
Parameters:
117
weight: [num_moe_layers, num_logical_experts]
@@ -146,35 +138,33 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
138
assert num_physical_experts % num_gpus == 0
139
phy_experts_per_gpu = num_physical_experts // num_gpus
140
149
- def inverse(perm: torch.Tensor) -> torch.Tensor:
150
- inv = torch.empty_like(perm)
151
- inv.scatter_(
152
- 1,
153
- perm,
154
- torch.arange(
155
- perm.size(1), dtype=torch.int64, device=perm.device
156
- ).expand(perm.shape),
157
- )
141
+ def inverse(perm: np.ndarray) -> np.ndarray:
142
+ inv = np.empty_like(perm)
143
+ row_idx = np.arange(perm.shape[0])[:, None]
144
+ col_idx = np.arange(perm.shape[1], dtype=np.int64)
145
+ inv[row_idx, perm] = col_idx
146
return inv
147
148
# Step 1: pack groups to nodes
161
- tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
149
+ tokens_per_group = weight.reshape(num_layers, num_groups, group_size).sum(
150
+ axis=-1
151
+ )
152
group_pack_index, group_rank_in_pack = cls.balanced_packing(
153
tokens_per_group, num_nodes
154
)
155
+ # Map each logical expert into a node-local ordering based on packed groups.
156
log2mlog = (
157
(
167
- (group_pack_index * groups_per_node + group_rank_in_pack) * group_size
168
- ).unsqueeze(-1)
169
- + torch.arange(
170
- group_size, dtype=torch.int64, device=group_pack_index.device
158
+ (group_pack_index * groups_per_node + group_rank_in_pack)[..., None]
159
+ * group_size
160
)
172
- ).flatten(-2)
161
+ + np.arange(group_size, dtype=np.int64)
162
+ ).reshape(num_layers, num_logical_experts)
163
mlog2log = inverse(log2mlog)
164
165
# Step 2: construct redundant experts within nodes
176
- # [num_layers * num_nodes, num_logical_experts // num_nodes]
177
- tokens_per_mlog = weight.gather(-1, mlog2log).view(
166
+ # Reorder weights into the node-local layout so replication is done per node.
167
+ tokens_per_mlog = np.take_along_axis(weight, mlog2log, axis=1).reshape(
168
-1, num_logical_experts // num_nodes
169
)
170
phy2mlog, replicas_idx, mlogcnt = cls.replicate_experts(
@@ -182,39 +172,43 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
172
)
173
174
# Step 3: pack physical_experts to GPUs
185
- # [num_layers * num_nodes, num_physical_experts // num_nodes]
186
- tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
175
+ # Effective per-physical load = logical load divided by replica count.
176
+ tokens_per_phy = np.take_along_axis(tokens_per_mlog / mlogcnt, phy2mlog, axis=1)
177
pack_index, rank_in_pack = cls.balanced_packing(
178
tokens_per_phy, num_gpus // num_nodes
179
)
180
phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
181
pphy2phy = inverse(phy2pphy)
182
193
- pphy2mlog = phy2mlog.gather(
194
- -1, pphy2phy
195
- ) # [num_layers * num_nodes, num_log_per_nodes]
183
+ # Reorder node-local logical indices into the post-packing physical order.
184
+ pphy2mlog = np.take_along_axis(phy2mlog, pphy2phy, axis=1)
185
pphy2mlog = (
197
- pphy2mlog.view(num_layers, num_nodes, -1)
198
- + torch.arange(
186
+ pphy2mlog.reshape(num_layers, num_nodes, -1)
187
+ + np.arange(
188
0,
189
num_logical_experts,
190
num_logical_experts // num_nodes,
202
- device=group_pack_index.device,
203
- ).view(1, -1, 1)
204
- ).flatten(-2)
205
- pphy2log = mlog2log.gather(-1, pphy2mlog)
206
- pphy_replicas_idx = replicas_idx.gather(-1, pphy2phy).view(num_layers, -1)
207
- logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
191
+ dtype=np.int64,
192
+ )[None, :, None]
193
+ ).reshape(num_layers, -1)
194
+ # Map node-local logical indices back to global logical expert ids.
195
+ pphy2log = np.take_along_axis(mlog2log, pphy2mlog, axis=1)
196
+ # Reorder replica ranks to the post-packing physical ordering.
197
+ pphy_replicas_idx = np.take_along_axis(replicas_idx, pphy2phy, axis=1).reshape(
198
+ num_layers, -1
199
+ )
200
+ # Convert replica counts back to the original logical ordering.
201
+ logcnt = np.take_along_axis(mlogcnt.reshape(num_layers, -1), log2mlog, axis=1)
202
return pphy2log, pphy_replicas_idx, logcnt
203
204
@classmethod
205
def preserve_intragpu_slots(
206
cls,
213
- phy2log: torch.Tensor,
214
- phy_replicas_idx: torch.Tensor,
207
+ phy2log: np.ndarray,
208
+ phy_replicas_idx: np.ndarray,
209
num_ranks: int,
216
- old_global_expert_indices: torch.Tensor,
217
- ) -> tuple[torch.Tensor, torch.Tensor]:
210
+ old_phy2log: np.ndarray,
211
+ ) -> tuple[np.ndarray, np.ndarray]:
212
"""
213
Reorder the new mapping per GPU so that experts that remain on the same GPU
214
keep their previous slot positions when possible. Incoming experts to that GPU
@@ -222,29 +216,24 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
216
is unchanged and the slots per GPU remain the same between
217
the old and new mappings.
218
"""
225
- device = phy2log.device
219
num_phy_experts = phy2log.shape[1]
220
if num_ranks <= 0 or num_phy_experts % num_ranks != 0:
221
return phy2log, phy_replicas_idx
222
223
# Move to CPU and convert to NumPy for processing
231
- new_phy2log_np = phy2log.cpu().numpy()
232
- replicas_idx_np = phy_replicas_idx.cpu().numpy()
233
- old_phy2log_np = old_global_expert_indices.cpu().numpy()
234
-
224
slots_per_gpu = num_phy_experts // num_ranks
236
- num_layers = new_phy2log_np.shape[0]
225
+ num_layers = phy2log.shape[0]
226
238
- post_phy2log_np = new_phy2log_np.copy()
239
- post_phy_replicas_idx_np = replicas_idx_np.copy()
227
+ post_phy2log = phy2log.copy()
228
+ post_phy_replicas_idx = phy_replicas_idx.copy()
229
230
for gpu_idx in range(num_ranks):
231
start = gpu_idx * slots_per_gpu
232
end = start + slots_per_gpu
233
# Experts across all layers for this GPU
245
- old_local = old_phy2log_np[:, start:end] # [layers, slots]
246
- new_local = new_phy2log_np[:, start:end] # [layers, slots]
247
- new_ridx = replicas_idx_np[:, start:end] # [layers, slots]
234
+ old_local = old_phy2log[:, start:end] # [layers, slots]
235
+ new_local = phy2log[:, start:end] # [layers, slots]
236
+ new_ridx = phy_replicas_idx[:, start:end] # [layers, slots]
237
238
used_new_indices = np.zeros((num_layers, slots_per_gpu), dtype=bool)
239
preserved_positions = np.zeros((num_layers, slots_per_gpu), dtype=bool)
@@ -261,12 +250,12 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
250
first_idx = np.argmax(matches, axis=1)
251
layer_indices = np.nonzero(has_any)[0]
252
matched_new_positions = first_idx[layer_indices]
264
- post_phy2log_np[layer_indices, start + slot_idx] = new_local[
253
+ post_phy2log[layer_indices, start + slot_idx] = new_local[
254
+ layer_indices, matched_new_positions
255
+ ]
256
+ post_phy_replicas_idx[layer_indices, start + slot_idx] = new_ridx[
257
layer_indices, matched_new_positions
258
]
267
- post_phy_replicas_idx_np[layer_indices, start + slot_idx] = (
268
- new_ridx[layer_indices, matched_new_positions]
269
- )
259
used_new_indices[layer_indices, matched_new_positions] = True
260
preserved_positions[layer_indices, slot_idx] = True
261
@@ -295,16 +284,13 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
284
continue
285
src_pos = remaining_indices[layer_idx, :k]
286
dst_pos = fill_indices[layer_idx, :k]
298
- post_phy2log_np[layer_idx, start + dst_pos] = new_local[
287
+ post_phy2log[layer_idx, start + dst_pos] = new_local[
288
layer_idx, src_pos
289
]
301
- post_phy_replicas_idx_np[layer_idx, start + dst_pos] = new_ridx[
290
+ post_phy_replicas_idx[layer_idx, start + dst_pos] = new_ridx[
291
layer_idx, src_pos
292
]
293
305
- # Convert back to torch and move to original device
306
- post_phy2log = torch.from_numpy(post_phy2log_np).to(device)
307
- post_phy_replicas_idx = torch.from_numpy(post_phy_replicas_idx_np).to(device)
294
return post_phy2log, post_phy_replicas_idx
295
296
@classmethod
@@ -340,40 +326,51 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
326
logcnt: [layers, num_logical_experts], number of
327
physical replicas for each logical expert
328
"""
329
+ device = weight.device
330
num_layers, num_logical_experts = weight.shape
344
- weight = weight.float()
331
+ weight_np = weight.float().cpu().numpy()
332
+ old_phy2log_np = (
333
+ old_global_expert_indices.cpu().numpy()
334
+ if old_global_expert_indices is not None
335
+ else None
336
+ )
337
+
338
if num_groups % num_nodes == 0:
339
# use hierarchical load-balance policy
347
- phy2log, phy_replicas_idx, logcnt = cls.rebalance_experts_hierarchical(
348
- weight, num_replicas, num_groups, num_nodes, num_ranks
340
+ phy2log_np, phy_replicas_idx_np, logcnt_np = (
341
+ cls.rebalance_experts_hierarchical(
342
+ weight_np, num_replicas, num_groups, num_nodes, num_ranks
343
+ )
344
)
345
else:
346
# use global load-balance policy
352
- phy2log, phy_replicas_idx, logcnt = cls.rebalance_experts_hierarchical(
353
- weight, num_replicas, 1, 1, num_ranks
347
+ phy2log_np, phy_replicas_idx_np, logcnt_np = (
348
+ cls.rebalance_experts_hierarchical(
349
+ weight_np, num_replicas, 1, 1, num_ranks
350
+ )
351
)
352
+
353
# Optional postprocessing to preserve slots for experts moving
354
# within the same GPU
355
# Only apply when the number of GPUs and slots per GPU remain unchanged.
356
# Helps to avoid unnecessary weight copying when experts move
357
# within the same GPU.
358
if old_global_expert_indices is not None:
361
- phy2log, phy_replicas_idx = cls.preserve_intragpu_slots(
362
- phy2log, phy_replicas_idx, num_ranks, old_global_expert_indices
359
+ phy2log_np, phy_replicas_idx_np = cls.preserve_intragpu_slots(
360
+ phy2log_np, phy_replicas_idx_np, num_ranks, old_phy2log_np
361
)
362
num_redundant_experts = num_replicas - num_logical_experts
363
maxlogcnt = num_redundant_experts + 1
366
- log2phy: torch.Tensor = torch.full(
367
- (num_layers, num_logical_experts, maxlogcnt),
368
- -1,
369
- dtype=torch.int64,
370
- device=logcnt.device,
364
+ log2phy_np = np.full(
365
+ (num_layers, num_logical_experts, maxlogcnt), -1, dtype=np.int64
366
)
372
- log2phy.view(num_layers, -1).scatter_(
373
- -1,
374
- phy2log * maxlogcnt + phy_replicas_idx,
375
- torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(
376
- num_layers, -1
377
- ),
367
+ layer_indices = np.arange(num_layers)[:, None]
368
+ replica_indices = np.tile(
369
+ np.arange(num_replicas, dtype=np.int64), (num_layers, 1)
370
)
371
+ log2phy_np[layer_indices, phy2log_np, phy_replicas_idx_np] = replica_indices
372
+
373
+ phy2log = torch.from_numpy(phy2log_np).to(device)
374
+ log2phy = torch.from_numpy(log2phy_np).to(device)
375
+ logcnt = torch.from_numpy(logcnt_np).to(device)
376
return phy2log, log2phy, logcnt