[bugfix][NPU] Fix startup bug in olmoe 1b 7b (#31782)

2026-07-22 10:51:15 +08:00
parent 394b0dd13a
commit d708969f68
1 changed files with 2 additions and 1 deletions
@@ -95,6 +95,7 @@ class AscendTPDispatcher(BaseDispatcher):
        topk_weights, topk_ids, _ = topk_output
        topk_weights = topk_weights.to(hidden_states.dtype)
        topk_ids = topk_ids.to(torch.int32)
+        top_k = topk_weights.shape[-1]

        (
            permuted_hidden_states,
@@ -105,7 +106,7 @@ class AscendTPDispatcher(BaseDispatcher):
            hidden_states,
            topk_ids,
            self.num_experts,
-            self.top_k,
+            top_k,
        )

        self._dispatch_output = AscendTPDispatchOutput(