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
[BugFix] Register expert_map as named buffer for wake_up and sleep (#25458)
Signed-off-by: wuxibin <wuxibin@bytedance.com> Signed-off-by: youkaichao <youkaichao@gmail.com> Co-authored-by: youkaichao <youkaichao@gmail.com>
Joel committed
Sep 23, 2025 at 21:49 UTC
61d1b35561c07589331fedf5eba58eaaedcd4423
1 file changed
+7
-2
vllm/model_executor/layers/fused_moe/layer.py
+7
-2
@@ -972,12 +972,15 @@ class FusedMoE(CustomOp):
972
"experts. Falling back to linear expert placement.")
973
expert_placement_strategy = "linear"
974
975
- self.local_num_experts, self.expert_map = determine_expert_map(
975
+ self.expert_map: Optional[torch.Tensor]
976
+ local_num_experts, expert_map = determine_expert_map(
977
ep_size=self.ep_size,
978
ep_rank=self.ep_rank,
979
global_num_experts=self.global_num_experts,
980
expert_placement_strategy=expert_placement_strategy,
981
)
982
+ self.local_num_experts = local_num_experts
983
+ self.register_buffer("expert_map", expert_map)
984
logger.info_once(
985
"[EP Rank %s/%s] Expert parallelism is enabled. Expert "
986
"placement strategy: %s. Local/global"
@@ -1154,10 +1157,12 @@ class FusedMoE(CustomOp):
1157
# ep_size and ep_rank should already be updated
1158
assert self.expert_map is not None
1159
with self.expert_map.device:
1157
- self.local_num_experts, self.expert_map = determine_expert_map(
1160
+ local_num_experts, expert_map = determine_expert_map(
1161
ep_size=self.ep_size,
1162
ep_rank=self.ep_rank,
1163
global_num_experts=self.global_num_experts)
1164
+ self.local_num_experts = local_num_experts
1165
+ self.register_buffer("expert_map", expert_map)
1166
1167
def _load_per_tensor_weight_scale(self, shard_id: str,
1168
param: torch.nn.Parameter,