From 8fae3fdf3de4162db8e3c18a3a005d571f22b656 Mon Sep 17 00:00:00 2001 From: li-lizhe <147392333@qq.com> Date: Fri, 23 Oct 2026 06:32:07 +0800 Subject: [PATCH] Rebase onto main@aba987ad22 Resolved conflicts with `git merge-file`; change set unchanged. Signed-off-by: li-lizhe <147392333@qq.com> --- lmdeploy/pytorch/backends/dlinfer/moe.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/moe.py b/lmdeploy/pytorch/backends/dlinfer/moe.py index 0cbd34a546..d7a5e1ebbd 100644 --- a/lmdeploy/pytorch/backends/dlinfer/moe.py +++ b/lmdeploy/pytorch/backends/dlinfer/moe.py @@ -52,12 +52,22 @@ def __init__(self, self.ep_group = ep_group self.expert_ids_per_ep_rank = None self.chunked_moe_layout = None - if self.ep_size > 1: + + def _get_expert_ids_per_ep_rank(self, device: torch.device): + """Build the expert-id map of the current rank on the given device. + + The map is created from the input device instead of a CUDA device index, so the non-CUDA accelerators served by + this backend (Ascend NPU, MACA, Cambricon) do not have to be CUDA at construction time. + """ + if self.ep_size <= 1: + return None + if self.expert_ids_per_ep_rank is None or self.expert_ids_per_ep_rank.device != device: self.expert_ids_per_ep_rank = torch.tensor( - [i % (self.num_experts // self.ep_size) for i in range(num_experts)], + [i % (self.num_experts // self.ep_size) for i in range(self.num_experts)], dtype=torch.int32, - device=torch.cuda.current_device(), + device=device, ) + return self.expert_ids_per_ep_rank def ep_expert_list(self, world_size: int, rank: int): """Experts list of current rank.""" @@ -110,7 +120,7 @@ def forward(self, step_context = get_step_ctx_manager().current_context() moe_metadata = getattr(step_context, 'moe_metadata', None) if moe_metadata is not None: - moe_metadata.expert_ids_per_ep_rank = self.expert_ids_per_ep_rank + moe_metadata.expert_ids_per_ep_rank = self._get_expert_ids_per_ep_rank(hidden_states.device) return fused_moe(hidden_states, gate_up_weights, down_weights, topk_weights, topk_ids, self.top_k, self.renormalize, moe_metadata, self.chunked_moe_layout) @@ -169,7 +179,7 @@ def forward(self, moe_metadata = getattr(step_context, 'moe_metadata', None) if moe_metadata is None: raise RuntimeError('Dlinfer W8A8 MoE requires moe_metadata in the current step context.') - moe_metadata.expert_ids_per_ep_rank = self.expert_ids_per_ep_rank + moe_metadata.expert_ids_per_ep_rank = self._get_expert_ids_per_ep_rank(hidden_states.device) return fused_moe_w8a8(hidden_states, gate_up_weights, gate_up_scale, down_weights, down_scale, topk_weights, topk_ids, self.top_k, self.renormalize, moe_metadata)