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

Keep memory copy blocking for elastic EP case

Browse Source 
Signed-off-by: ilmarkov <markovilya197@gmail.com>
This commit is contained in:
ilmarkov 2025-12-18 15:10:43 +00:00
parent 67ff54997d
commit f7b8992020
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/distributed/eplb/eplb_state.py
View File

@ -954,10 +954,10 @@ class EplbState:
target_device = model_state.physical_to_logical_map.device
new_physical = model_state.new_physical_to_logical_map
# In order to avoid race condition with async eplb worker,
# we need to copy blocking in case of updated EP size.
if model_state.physical_to_logical_map.shape[1] != new_physical.shape[1]:
model_state.physical_to_logical_map = new_physical.to(
target_device, non_blocking=True
)
model_state.physical_to_logical_map = new_physical.to(target_device)
else:
model_state.physical_to_logical_map[layer].copy_(
new_physical[layer].to(target_device, non_blocking=True)