A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Misc] enhance type hint for rearrange return value (#23519)
Signed-off-by: Andy Xie <andy.xning@gmail.com>
Ning Xie committed
Aug 30, 2025 at 21:43 UTC
68a349114f2a90f17645a8172496b0803d677f29
1 file changed
+10
-7
vllm/distributed/eplb/eplb_state.py
+10
-7
@@ -409,12 +409,14 @@ class EplbState:
409
self.expert_rearrangement_step = 0
410
self.rearrange(model)
411
412
- def rearrange(self,
413
- model: MixtureOfExperts,
414
- is_profile: bool = False,
415
- execute_shuffle: bool = True,
416
- global_expert_load: Optional[torch.Tensor] = None,
417
- rank_mapping: Optional[dict[int, int]] = None) -> None:
412
+ def rearrange(
413
+ self,
414
+ model: MixtureOfExperts,
415
+ is_profile: bool = False,
416
+ execute_shuffle: bool = True,
417
+ global_expert_load: Optional[torch.Tensor] = None,
418
+ rank_mapping: Optional[dict[int,
419
+ int]] = None) -> Optional[torch.Tensor]:
420
"""
421
Rearrange the experts according to the current load.
422
"""
@@ -548,6 +550,7 @@ class EplbState:
550
" (profile) " if is_profile else " ",
551
time_end - time_start,
552
)
553
+ return None
554
555
@staticmethod
556
def recv_state() -> tuple[torch.Tensor, torch.Tensor]:
@@ -613,4 +616,4 @@ def _node_count_with_rank_mapping(
616
if is_same_node and node_assignment[other_rank] == 0:
617
node_assignment[other_rank] = next_node_id
618
616
- return next_node_id
\ No newline at end of file
619
+ return next_node_id