[NPU] Add extra topk_weights input in deepep ll dispatch (#29480)

2026-07-09 09:23:29 +08:00
parent cf8f1df6e8
commit 0ffed946f2
1 changed files with 3 additions and 0 deletions
@@ -660,6 +660,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
        hidden_states, masked_m, event, hook = self._dispatch_core(
            hidden_states,
            topk_ids,
+            topk_weights,
        )
        return (
            hidden_states,
@@ -706,6 +707,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
        self,
        hidden_states: torch.Tensor,
        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
    ):
        input_global_scale = self.quant_config.get("input_global_scale", None)

@@ -732,6 +734,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
                self.num_max_dispatch_tokens_per_rank,
                self.num_experts,
                use_fp8=self.use_fp8,
+                **(dict(topk_weights=topk_weights) if _is_npu else dict()),
                **(dict(use_nvfp4=True) if self.use_nvfp4 else dict()),
                **(
                    dict(x_global_scale=input_global_scale)