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

7
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(
self,
model: MixtureOfExperts, model: MixtureOfExperts,
is_profile: bool = False, is_profile: bool = False,
execute_shuffle: bool = True, execute_shuffle: bool = True,
global_expert_load: Optional[torch.Tensor] = None, global_expert_load: Optional[torch.Tensor] = None,
rank_mapping: Optional[dict[int, int]] = None) -> 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]: