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

Futher optimize rearrange

Browse Source 
Signed-off-by: ilmarkov <markovilya197@gmail.com>
This commit is contained in:
ilmarkov 2025-12-11 11:35:10 +00:00
parent f28720db88
commit ab0ca861da
1 changed files with 3 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/distributed/eplb/rebalance_execute.py
View File

@ -233,10 +233,7 @@ def move_to_buffer(
recv_dst_rows[:recv_count] = dst_rows recv_dst_rows[:recv_count] = dst_rows
recv_primary_mask[dst_rows] = True recv_primary_mask[dst_rows] = True
eligible_local_buffer_mask = np.logical_and( eligible_local_buffer_mask = np.logical_and(~is_unchanged, is_received_locally)
np.logical_and(~is_unchanged, is_received_locally),
new_local_expert_ids != -1,
)
# 1. Local moves into tmp buffers # 1. Local moves into tmp buffers
if bool(eligible_local_buffer_mask.any()) and send_count > 0: if bool(eligible_local_buffer_mask.any()) and send_count > 0:
@ -396,7 +393,7 @@ def move_from_buffer(
dest_indices = np.nonzero(dest_mask_np)[0].tolist() dest_indices = np.nonzero(dest_mask_np)[0].tolist()
for dst in dest_indices: for dst in dest_indices:
for w, b in zip(expert_weights, expert_weights_buffers): for w, b in zip(expert_weights, expert_weights_buffers):
w[dst].copy_(b[dst]) w[dst].copy_(b[dst], non_blocking=True)
# Duplicate remote received rows to non-primary duplicate dsts # Duplicate remote received rows to non-primary duplicate dsts
if recv_count == 0: if recv_count == 0:
@ -433,7 +430,7 @@ def move_from_buffer(
for dst, src in zip(matched_dst_rows.tolist(), matched_src_rows.tolist()): for dst, src in zip(matched_dst_rows.tolist(), matched_src_rows.tolist()):
for w in expert_weights: for w in expert_weights:
w[dst].copy_(w[src]) w[dst].copy_(w[src], non_blocking=True)
async def transfer_layer( async def transfer_layer(