biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 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.expert_rearrangement_step = 0
self.rearrange(model) self.rearrange(model)
def rearrange(self, def rearrange(
model: MixtureOfExperts, self,
is_profile: bool = False, model: MixtureOfExperts,
execute_shuffle: bool = True, is_profile: bool = False,
global_expert_load: Optional[torch.Tensor] = None, execute_shuffle: bool = True,
rank_mapping: Optional[dict[int, int]] = None) -> None: 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. Rearrange the experts according to the current load.
""" """
@@ -548,6 +550,7 @@ class EplbState:
" (profile) " if is_profile else " ", " (profile) " if is_profile else " ",
time_end - time_start, time_end - time_start,
) )
return None
@staticmethod @staticmethod
def recv_state() -> tuple[torch.Tensor, torch.Tensor]: 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: if is_same_node and node_assignment[other_rank] == 0:
node_assignment[other_rank] = next_node_id node_assignment[other_rank] = next_node_id
return next_node_id return next_node_id