Skip to content

Commit 474fa73

Browse files
weijinqian0weijinqian_v1
andauthored
[bugfix] Fix moe bug: allgather error. (#3279)
It will crash when deepseek model executed in A2. - vLLM version: v0.11.0rc3 - vLLM main: vllm-project/vllm@releases/v0.11.0 --------- Signed-off-by: weijinqian_v1 <[email protected]> Co-authored-by: weijinqian_v1 <[email protected]>
1 parent b8c58d6 commit 474fa73

File tree

2 files changed

+2
-1
lines changed

2 files changed

+2
-1
lines changed

tests/ut/ops/test_token_dispatcher.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,7 @@ def test_token_dispatch_with_quant(self):
269269
def test_token_combine_with_expert_map(self):
270270
self.dispatcher.expert_map = torch.tensor([0, 1, 2, 3])
271271
self.dispatcher.sorted_token_indices = torch.tensor([0, 1, 1, 1, 1, 1])
272+
self.dispatcher.expanded_row_idx = torch.tensor([0, 1, 1, 1, 1, 1])
272273
self.dispatcher.sorted_weights = torch.tensor(
273274
[0.5, 0.5, 0.5, 0.5, 0.5, 0.5])
274275
self.dispatcher.original_shape = (3, 128)

vllm_ascend/ops/moe/token_dispatcher.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -383,7 +383,7 @@ def token_combine(self,
383383
assert self.original_shape is not None
384384
final_hidden_states = torch_npu.npu_moe_token_unpermute(
385385
permuted_tokens=hidden_states,
386-
sorted_indices=self.expanded_row_idx,
386+
sorted_indices=torch.abs(self.expanded_row_idx),
387387
probs=self.topk_weights)
388388
if len(self.original_shape) == 3:
389389
final_hidden_states = final_hidden_states.view(self.original_shape)

0 commit comments

Comments
 (0)