[li-lizhe]

Motivation

DlinferFusedMoEImpl.__init__ builds the expert-id map of the current rank with a hardcoded CUDA
device index:

if self.ep_size > 1:
    self.expert_ids_per_ep_rank = torch.tensor(
        [i % (self.num_experts // self.ep_size) for i in range(num_experts)],
        dtype=torch.int32,
        device=torch.cuda.current_device(),
    )

This backend (lmdeploy/pytorch/backends/dlinfer/) is the one that serves the non-CUDA
accelerators (Ascend NPU, MACA, Cambricon), where torch.cuda.current_device() raises
AssertionError: Torch not compiled with CUDA enabled. As a result, building any fused-MoE layer
with ep_size > 1 aborts during construction, before the first forward.

Modification

Build the map lazily, on the device of the first input tensor, through a small helper:

• self.expert_ids_per_ep_rank is no longer created in __init__; the new
  _get_expert_ids_per_ep_rank(device) builds it on the given device and caches it (it is rebuilt
  only when the device changes).
• DlinferFusedMoEImpl.forward and DlinferFusedMoEW8A8Impl.forward attach the map from
  hidden_states.device instead of the value frozen at construction time.

CUDA behaviour is unchanged: the same tensor with the same contents ends up on the same device, it
is just created from the input device instead of from a CUDA device index.

BC-breaking (Optional)

None.

Verification

Ascend NPU (Ascend 910, torch 2.15.0.dev20260917+cpu + torch_npu 2.15.0.dev20260917+gitec69335,
torch.cuda.is_available() == False), with the class body of this file executed verbatim (the
container has no dlinfer / lmdeploy install, so the dlinfer op wrappers and the base spec classes
were stubbed; only the file under change is real):

Before:

torch.cuda.current_device() RAISED: AssertionError Torch not compiled with CUDA enabled
DlinferFusedMoEImpl(ep_size=2) CTOR CRASHED: AssertionError Torch not compiled with CUDA enabled

After:

DlinferFusedMoEImpl(ep_size=2) CTOR OK
forward -> fused_moe
moe_metadata.expert_ids_per_ep_rank = ([0, 1, 2, 3], 'npu:0', torch.int32)
w8a8 forward -> fused_moe_w8a8
moe_metadata.expert_ids_per_ep_rank (w8a8) = ([0, 1, 2, 3], 'npu:0', torch.int32)
ep_size=1 ids = None | forward -> fused_moe

Not tested: an end-to-end MoE model with EP on Ascend (no dlinfer wheel available on that
machine), and the CUDA / MACA / Cambricon paths — the change is device-agnostic and keeps the
previous CUDA device (hidden_states.device is the current CUDA device there).

[li-lizhe]

Gentle ping on this one — it's been quiet for 7 days since the last activity, so I wanted to check whether there's anything else you'd like me to adjust or add before this can move forward. Happy to make changes if so. Thanks!