Make DeepEP combine recv do not overlap (#11535)

2025-10-14 09:40:42 +08:00
parent 384733639a
commit cb8ed2c09a
1 changed files with 3 additions and 0 deletions
--- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py
@@ -616,6 +616,9 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
        return hidden_states, event, hook, overlap_args

    def combine_b(self, hidden_states, event, hook, overlap_args):
+        if overlap_args is not None:
+            overlap_args.stream.wait_stream(self.device_module.current_stream())
+
        hook() if self.return_recv_hook else event.current_stream_wait()

        if overlap_args is not None: