xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-09-06 05:16:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Remove expensive checks

Browse Source 
Signed-off-by: ilmarkov <markovilya197@gmail.com>
This commit is contained in:
ilmarkov 2025-12-15 18:08:15 +00:00
parent dcf4783967
commit 699ea97108
1 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/distributed/eplb/policy/default.py
View File

@ -59,15 +59,16 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
items_row = pack_items[layer_idx] items_row = pack_items[layer_idx]
for group in indices[layer_idx]: for group in indices[layer_idx]:
# Select the lightest pack that still has capacity. # Pick the lightest pack; full packs are masked out by inf.
available = items_row < groups_per_pack pack = int(np.argmin(weights_row))
assert np.any(available)
pack = int(np.argmin(np.where(available, weights_row, np.inf)))
pack_index[layer_idx, group] = pack pack_index[layer_idx, group] = pack
rank_in_pack[layer_idx, group] = items_row[pack] rank_in_pack[layer_idx, group] = items_row[pack]
weights_row[pack] += weight[layer_idx, group] weights_row[pack] += weight[layer_idx, group]
items_row[pack] += 1 items_row[pack] += 1
if items_row[pack] == groups_per_pack:
# Mark as unavailable for future selections.
weights_row[pack] = np.inf
return pack_index, rank_in_pack return pack_index, rank_in_pack