[NPU]fix: restore MoE ALLTOALL expert IDs after NPU memory reuse - #376
Conversation
Signed-off-by: wuxiang <498160096@qq.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the Qwen3-30B-A3B NPU test configuration by reducing the rollout GPUs per engine and enabling the colocate flag. It also introduces a monkey-patching mechanism to restore the ALLTOALL expert-ID template after colocated memory reuse on Ascend A3 devices. The review feedback suggests wrapping third-party imports in a try-except block to handle potential API changes gracefully, and adding defensive checks to prevent AttributeErrors if the expert-ID tensor is uninitialized.
| from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type | ||
|
|
||
| if get_ascend_device_type() != AscendDeviceType.A3: | ||
| return | ||
|
|
||
| from vllm_ascend.ops.fused_moe.token_dispatcher import TokenDispatcherWithAll2AllV | ||
|
|
||
| if getattr(TokenDispatcherWithAll2AllV, "_vime_expert_ids_patched", False): | ||
| return |
There was a problem hiding this comment.
Monkey-patching internal APIs of third-party libraries (like vllm_ascend) is highly prone to breaking when those libraries are updated. If the import paths or attributes change in future versions, this will cause a hard crash during worker startup. Wrapping these imports and checks in a try...except block ensures defensive execution and prevents startup failures.
| from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type | |
| if get_ascend_device_type() != AscendDeviceType.A3: | |
| return | |
| from vllm_ascend.ops.fused_moe.token_dispatcher import TokenDispatcherWithAll2AllV | |
| if getattr(TokenDispatcherWithAll2AllV, "_vime_expert_ids_patched", False): | |
| return | |
| try: | |
| from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type | |
| if get_ascend_device_type() != AscendDeviceType.A3: | |
| return | |
| from vllm_ascend.ops.fused_moe.token_dispatcher import TokenDispatcherWithAll2AllV | |
| except (ImportError, AttributeError): | |
| return | |
| if getattr(TokenDispatcherWithAll2AllV, "_vime_expert_ids_patched", False): | |
| return |
| def _patched_dispatch_preprocess(self, hidden_states, topk_ids): | ||
| generation = TokenDispatcherWithAll2AllV._vime_expert_ids_generation | ||
| if self.num_local_experts > 1 and getattr(self, "_vime_seen_expert_ids_generation", -1) != generation: | ||
| expert_ids = self.expert_ids_per_ep_rank | ||
| self.expert_ids_per_ep_rank = torch.arange( | ||
| self.num_experts, | ||
| device=expert_ids.device, | ||
| dtype=expert_ids.dtype, | ||
| ).remainder(self.num_local_experts) | ||
| self._vime_seen_expert_ids_generation = generation | ||
| return original_dispatch_preprocess(self, hidden_states, topk_ids) |
There was a problem hiding this comment.
If self.expert_ids_per_ep_rank is None or not yet initialized, accessing expert_ids.device or expert_ids.dtype will raise an AttributeError. Adding a defensive check ensures that we only attempt to recreate the tensor if it is already present.
| def _patched_dispatch_preprocess(self, hidden_states, topk_ids): | |
| generation = TokenDispatcherWithAll2AllV._vime_expert_ids_generation | |
| if self.num_local_experts > 1 and getattr(self, "_vime_seen_expert_ids_generation", -1) != generation: | |
| expert_ids = self.expert_ids_per_ep_rank | |
| self.expert_ids_per_ep_rank = torch.arange( | |
| self.num_experts, | |
| device=expert_ids.device, | |
| dtype=expert_ids.dtype, | |
| ).remainder(self.num_local_experts) | |
| self._vime_seen_expert_ids_generation = generation | |
| return original_dispatch_preprocess(self, hidden_states, topk_ids) | |
| def _patched_dispatch_preprocess(self, hidden_states, topk_ids): | |
| generation = TokenDispatcherWithAll2AllV._vime_expert_ids_generation | |
| if self.num_local_experts > 1 and getattr(self, "_vime_seen_expert_ids_generation", -1) != generation: | |
| expert_ids = getattr(self, "expert_ids_per_ep_rank", None) | |
| if expert_ids is not None: | |
| self.expert_ids_per_ep_rank = torch.arange( | |
| self.num_experts, | |
| device=expert_ids.device, | |
| dtype=expert_ids.dtype, | |
| ).remainder(self.num_local_experts) | |
| self._vime_seen_expert_ids_generation = generation | |
| return original_dispatch_preprocess(self, hidden_states, topk_ids) |
Documentation build overview
48 files changed ·
|
miracle0517
changed the title
| NPUWorker._npu_worker_patched = True | ||
|
|
||
| @staticmethod | ||
| def _patch_a3_moe_alltoall_expert_ids() -> None: |
There was a problem hiding this comment.
Could this cause garbled characters in single inference? I think this issue should also be reproducible in single inference as well.
There was a problem hiding this comment.
We noticed that a colleague has already submitted a relevant fix PR in the vllm-ascend repository. vllm-project/vllm-ascend#13352, from a version control perspective, it is reasonable for us to apply a patch in vime to override the older version of vllm-ascend in order to resolve the precision issues. This approach is both justified and manageable.
There was a problem hiding this comment.
Could this cause garbled characters in single inference? I think this issue should also be reproducible in single inference as well.
Not in our current reproduction. Single or low-concurrency requests take the MC2 path and remain normal. The issue only appears when high concurrency triggers the Vime IPC mixed prefill/decode + ALLTOALL path.
We also verified that direct vLLM-Ascend serving can use ALLTOALL normally, so this does not look like a general single-inference or ALLTOALL accuracy issue. See the lightweight comparison in [#374]
4 participants
Summary
Fix an MoE accuracy issue on Ascend A3 when running Qwen3-30B-A3B in colocated mode with expert parallelism and ALLTOALL communication.
Problem
After a colocated weight update and NPU memory reuse, rollout inference could produce incorrect first-token log probabilities and repetitive outputs. The issue did not reproduce in disaggregated mode.
The expert_ids_per_ep_rank tensor used by the ALLTOALL token dispatcher could retain invalid contents after the vLLM sleep/wake and weight-update lifecycle. This caused received tokens to be assigned to the wrong local experts even though the model weights and ALLTOALL payload were correct.
Fix
Test