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

Add comments

Browse Source 
Signed-off-by: ilmarkov <markovilya197@gmail.com>
This commit is contained in:
ilmarkov 2025-12-11 12:56:54 +00:00
parent 208b51bcd8
commit a5ecdc18c0
1 changed files with 17 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -263,31 +263,36 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
has_any = matches.any(axis=1) has_any = matches.any(axis=1)
if np.any(has_any): if np.any(has_any):
first_idx = np.argmax(matches, axis=1) first_idx = np.argmax(matches, axis=1)
rows = np.nonzero(has_any)[0] layer_indices = np.nonzero(has_any)[0]
cols = first_idx[rows] matched_new_positions = first_idx[layer_indices]
post_phy2log_np[rows, start + pos] = new_seg[rows, cols] post_phy2log_np[layer_indices, start + pos] = new_seg[
post_phyrank_np[rows, start + pos] = new_rnk[rows, cols] layer_indices, matched_new_positions
used_new_indices[rows, cols] = True ]
preserved_positions[rows, pos] = True post_phyrank_np[layer_indices, start + pos] = new_rnk[
layer_indices, matched_new_positions
]
used_new_indices[layer_indices, matched_new_positions] = True
preserved_positions[layer_indices, pos] = True
# Second pass: fill remaining slots with remaining new experts # Second pass: fill remaining slots with remaining new experts
remaining_mask = ~used_new_indices # [L, S] remaining_mask = ~used_new_indices # [L, S]
fill_mask = ~preserved_positions # [L, S] fill_mask = ~preserved_positions # [L, S]
if remaining_mask.any() and fill_mask.any(): if remaining_mask.any() and fill_mask.any():
idx_base = np.broadcast_to( idx_base = np.tile(np.arange(slots_per_gpu), (num_layers, 1))
np.arange(slots_per_gpu), (num_layers, slots_per_gpu) # Sentinel value for unavailable positions.
)
large = slots_per_gpu + 1 large = slots_per_gpu + 1
# Priorities: keep original index for available spots, set sentinel
# for unavailable; lower is earlier.
remaining_priority = np.where(remaining_mask, idx_base, large) remaining_priority = np.where(remaining_mask, idx_base, large)
fill_priority = np.where(fill_mask, idx_base, large) fill_priority = np.where(fill_mask, idx_base, large)
# Sort to get per-row ordered indices of True positions # Sort to get ordered indices of available src/dst positions per layer.
remaining_indices = np.argsort(remaining_priority, axis=1) remaining_indices = np.argsort(remaining_priority, axis=1)
fill_indices = np.argsort(fill_priority, axis=1) fill_indices = np.argsort(fill_priority, axis=1)
# How many to fill per row # Fill count per layer (cannot exceed either side).
remaining_counts = remaining_mask.sum(axis=1) remaining_counts = remaining_mask.sum(axis=1)
fill_counts = fill_mask.sum(axis=1) fill_counts = fill_mask.sum(axis=1)
take_counts = np.minimum(remaining_counts, fill_counts) take_counts = np.minimum(remaining_counts, fill_counts)
# Assign per row # Assign remaining new experts to remaining slots per layer.
for layer_idx in range(num_layers): for layer_idx in range(num_layers):
k = int(take_counts[layer_idx]) k = int(take_counts[layer_idx])
if k <= 0: if k <= 0: