diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 33a51a89d..2d7192e2e 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -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)