frozenleaves/vllm-dev
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[Misc] enhance type hint for rearrange return value (#23519)

Browse Source 
Signed-off-by: Andy Xie <andy.xning@gmail.com>
This commit is contained in:
Ning Xie
2025-08-30 21:43:33 +08:00
committed by GitHub
parent e80bca309e
commit 68a349114f
1 changed files with 10 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -409,12 +409,14 @@ class EplbState:
self.expert_rearrangement_step = 0
self.rearrange(model)
def rearrange(self,
model: MixtureOfExperts,
is_profile: bool = False,
execute_shuffle: bool = True,
global_expert_load: Optional[torch.Tensor] = None,
rank_mapping: Optional[dict[int, int]] = None) -> None:
def rearrange(
self,
model: MixtureOfExperts,
is_profile: bool = False,
execute_shuffle: bool = True,
global_expert_load: Optional[torch.Tensor] = None,
rank_mapping: Optional[dict[int,
int]] = None) -> Optional[torch.Tensor]:
"""
Rearrange the experts according to the current load.
"""
@ -548,6 +550,7 @@ class EplbState:
" (profile) " if is_profile else " ",
time_end - time_start,
)
return None
@staticmethod
def recv_state() -> tuple[torch.Tensor, torch.Tensor]:
@ -613,4 +616,4 @@ def _node_count_with_rank_mapping(
if is_same_node and node_assignment[other_rank] == 0:
node_assignment[other_rank] = next_node_id
return next_node_id
return next_node_id