sglang 0.4.4.post3__py3-none-any.whl → 0.4.4.post4__py3-none-any.whl

sglang/srt/layers/moe/ep_moe/token_dispatcher.py

@@ -1,3 +1,5 @@
+from sglang.srt.utils import DeepEPMode
+
 try:
     from deep_ep import Buffer
 
@@ -21,7 +23,7 @@ _buffer_normal = None
 _buffer_low_latency = None
 
 
-def get_buffer_normal(group: dist.ProcessGroup, hidden_bytes: int):
+def _get_buffer_normal(group: dist.ProcessGroup, hidden_bytes: int):
     """
     Copy from DeepEP example usage in model inference prefilling.
     https://github.com/deepseek-ai/DeepEP?tab=readme-ov-file#example-use-in-model-training-or-inference-prefilling
@@ -51,7 +53,7 @@ def get_buffer_normal(group: dist.ProcessGroup, hidden_bytes: int):
     return _buffer_normal
 
 
-def get_buffer_low_latency(
+def _get_buffer_low_latency(
     group: dist.ProcessGroup,
     num_max_dispatch_tokens_per_rank: int,
     hidden: int,
@@ -76,151 +78,103 @@ def get_buffer_low_latency(
         assert num_experts % group.size() == 0
         _buffer_low_latency = Buffer(
             group,
-            0,
-            num_rdma_bytes,
+            num_rdma_bytes=num_rdma_bytes,
             low_latency_mode=True,
             num_qps_per_rank=num_experts // group.size(),
         )
     return _buffer_low_latency
 
 
-class DeepEPDispatcher:
-    """
-    Copy from Megatron-Core token_dispatcher MoEFlexTokenDispatcher
-    https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/moe/token_dispatcher.py
-    """
-
+class _DeepEPDispatcherImplBase:
     def __init__(
         self,
         group: torch.distributed.ProcessGroup,
         router_topk: int,
-        permute_fusion: bool = False,
-        capacity_factor: float = None,
-        num_experts: int = None,
-        num_local_experts: int = None,
-        hidden_size: int = None,
-        params_dtype: torch.dtype = None,
-        async_finish: bool = False,
+        permute_fusion: bool,
+        num_experts: int,
+        num_local_experts: int,
+        hidden_size: int,
+        params_dtype: torch.dtype,
     ):
+        if not use_deepep:
+            raise ImportError(
+                "DeepEP is not installed. Please install DeepEP package from "
+                "https://github.com/deepseek-ai/deepep."
+            )
+
         self.group = group
         self.router_topk = router_topk
-        self.capacity_factor = capacity_factor
         self.permute_fusion = permute_fusion
         self.num_experts = num_experts
         self.num_local_experts = num_local_experts
         self.hidden_size = hidden_size
-        self.recv_expert_count = None
         self.params_dtype = params_dtype
         self.params_bytes = 2
-        # Metadata
-        self.token_indices = None
-        self.token_probs = None
-        # Handle used for combine operation
+
         self.handle = None
-        self.async_finish = async_finish
 
-        # `num_max_dispatch_tokens_per_rank` (the actual batch size in the decoding engine) should be less than 256
-        # https://github.com/deepseek-ai/DeepEP?tab=readme-ov-file#example-use-in-inference-decoding
-        self.num_max_dispatch_tokens_per_rank = 128
+    def dispatch_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        num_experts: int,
+        num_max_dispatch_tokens_per_rank: int,
+    ):
+        raise NotImplementedError
 
-        if not use_deepep:
-            raise ImportError(
-                "DeepEP is not installed. Please install DeepEP package from "
-                "https://github.com/deepseek-ai/deepep."
-            )
-        self.buffer_normal = get_buffer_normal(
-            self.group, self.hidden_size * self.params_bytes
-        )
-        self.buffer_low_latency = None
-        # Todo: enable low latency dispatch
-        """
-        self.buffer_low_latency = get_buffer_low_latency(
-            self.group,
-            self.num_max_dispatch_tokens_per_rank,
-            self.hidden_size * self.params_bytes,
-            self.num_experts,
-        )
-        """
+    def dispatch_b(self, *args, **kwargs):
+        raise NotImplementedError
 
-    def deepep_permute(
+    def combine_a(
         self,
-        hidden_states,
-        fp8_dtype=None,
-        use_fp8_w8a8=False,
-        use_block_quant=False,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
     ):
-        reorder_topk_ids, src2dst, seg_indptr = deepep_run_moe_deep_preprocess(
-            self.topk_idx, self.num_experts
-        )
-        num_total_tokens = reorder_topk_ids.numel()
-        gateup_input = torch.empty(
-            (int(num_total_tokens), hidden_states.shape[1]),
-            device=hidden_states.device,
-            dtype=(
-                fp8_dtype
-                if (use_fp8_w8a8 and not use_block_quant)
-                else hidden_states.dtype
-            ),
-        )
-        # PreReorder
-        deepep_permute_triton_kernel[(hidden_states.shape[0],)](
-            hidden_states,
-            gateup_input,
-            src2dst,
-            self.topk_idx,
-            None,
-            self.router_topk,
-            hidden_states.shape[1],
-            BLOCK_SIZE=512,
+        raise NotImplementedError
+
+    def combine_b(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
+    def __init__(self, async_finish: bool, **kwargs):
+        super().__init__(**kwargs)
+
+        self.buffer_normal = _get_buffer_normal(
+            self.group, self.hidden_size * self.params_bytes
         )
-        self.src2dst = src2dst
-        return reorder_topk_ids, seg_indptr, gateup_input
+        self.async_finish = async_finish
+        self.src2dst = None
 
-    def dispatch(
+    def dispatch_a(
         self,
         hidden_states: torch.Tensor,
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
         num_experts: int,
-        forward_mode: ForwardMode,
-        num_max_dispatch_tokens_per_rank: int = 128,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        num_max_dispatch_tokens_per_rank: int,
+    ):
         topk_idx = topk_idx.to(torch.int64)
-        # Todo: enable low latency dispatch
-        if True:  # not forward_mode.is_decode():
-            (
-                hidden_states,
-                topk_idx,
-                topk_weights,
-                num_recv_tokens_per_expert_list,
-                handle,
-                event,
-            ) = self.dispatch_normal(hidden_states, topk_idx, topk_weights, num_experts)
-            self.tokens_per_expert = torch.tensor(
-                num_recv_tokens_per_expert_list,
-                device=hidden_states.device,
-                dtype=torch.int64,
-            )
-        else:
-            hidden_states, recv_expert_count, handle, event, hook = (
-                self.dispatch_low_latency(
-                    hidden_states,
-                    topk_idx,
-                    num_max_dispatch_tokens_per_rank,
-                    num_experts,
-                )
-            )
-            self.recv_expert_count = recv_expert_count
-
-        if self.async_finish:
-            event.current_stream_wait()
+        previous_event = Buffer.capture() if self.async_finish else None
+        return hidden_states, topk_idx, topk_weights, num_experts, previous_event
 
-        self.handle = handle
-        self.topk_idx = topk_idx
-        self.topk_weights = topk_weights
+    def dispatch_b(
+        self, hidden_states, topk_idx, topk_weights, num_experts, previous_event
+    ):
+        (
+            hidden_states,
+            topk_idx,
+            topk_weights,
+            event,
+        ) = self._dispatch_core(
+            hidden_states, topk_idx, topk_weights, num_experts, previous_event
+        )
+        event.current_stream_wait() if self.async_finish else ()
         if hidden_states.shape[0] > 0:
-            reorder_topk_ids, seg_indptr, hidden_states = self.deepep_permute(
-                hidden_states, fp8_dtype=hidden_states.dtype
+            reorder_topk_ids, seg_indptr, hidden_states = self._deepep_permute(
+                hidden_states, topk_idx, fp8_dtype=hidden_states.dtype
             )
         else:
             reorder_topk_ids = torch.empty(
@@ -229,17 +183,27 @@ class DeepEPDispatcher:
             seg_indptr = torch.zeros(
                 (num_experts + 1,), device=hidden_states.device, dtype=torch.int64
             )
-        return hidden_states, reorder_topk_ids, seg_indptr
 
-    def dispatch_normal(
+        masked_m = expected_m = None
+
+        return (
+            hidden_states,
+            topk_idx,
+            topk_weights,
+            reorder_topk_ids,
+            seg_indptr,
+            masked_m,
+            expected_m,
+        )
+
+    def _dispatch_core(
         self,
         x: torch.Tensor,
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
         num_experts: int,
+        previous_event,
     ):
-        previous_event = Buffer.capture() if self.async_finish else None
-
         (
             num_tokens_per_rank,
             num_tokens_per_rdma_rank,
@@ -254,12 +218,15 @@ class DeepEPDispatcher:
             allocate_on_comm_stream=previous_event is not None,
         )
 
+        # FIXME: `handle` should be transmitted with tokens from dispatch to combine.
+        # However, doing this would incur an unknown synchronization error, but keeping
+        # `handle` as a member variable works.
         (
             recv_x,
             recv_topk_idx,
             recv_topk_weights,
-            num_recv_tokens_per_expert_list,
-            handle,
+            _,  # num_recv_tokens_per_expert_list
+            self.handle,
             event,
         ) = self.buffer_normal.dispatch(
             x,
@@ -278,29 +245,191 @@ class DeepEPDispatcher:
             recv_x,
             recv_topk_idx,
             recv_topk_weights,
-            num_recv_tokens_per_expert_list,
-            handle,
             event,
         )
 
-    def dispatch_low_latency(
+    def _deepep_permute(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        fp8_dtype: Optional[torch.dtype] = None,
+        use_fp8_w8a8: bool = False,
+        use_block_quant: bool = False,
+    ):
+        """
+        Copy from Megatron-Core token_dispatcher MoEFlexTokenDispatcher
+        https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/moe/token_dispatcher.py
+        """
+
+        reorder_topk_ids, self.src2dst, seg_indptr = deepep_run_moe_deep_preprocess(
+            topk_idx, self.num_experts
+        )
+        num_total_tokens = reorder_topk_ids.numel()
+        gateup_input = torch.empty(
+            (int(num_total_tokens), hidden_states.shape[1]),
+            device=hidden_states.device,
+            dtype=(
+                fp8_dtype
+                if (use_fp8_w8a8 and not use_block_quant)
+                else hidden_states.dtype
+            ),
+        )
+        # PreReorder
+        deepep_permute_triton_kernel[(hidden_states.shape[0],)](
+            hidden_states,
+            gateup_input,
+            self.src2dst,
+            topk_idx,
+            None,
+            self.router_topk,
+            hidden_states.shape[1],
+            BLOCK_SIZE=512,
+        )
+        return reorder_topk_ids, seg_indptr, gateup_input
+
+    def combine_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+    ):
+        if hidden_states.shape[0] > 0:
+            num_tokens = self.src2dst.shape[0] // self.router_topk
+            output = torch.empty(
+                (num_tokens, hidden_states.shape[1]),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+            deepep_post_reorder_triton_kernel[(num_tokens,)](
+                hidden_states,
+                output,
+                self.src2dst,
+                topk_idx,
+                topk_weights,
+                self.router_topk,
+                hidden_states.shape[1],
+                BLOCK_SIZE=512,
+            )
+        else:
+            output = torch.zeros(
+                (0, hidden_states.shape[1]),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype,
+            )
+        previous_event = Buffer.capture() if self.async_finish else None
+        return output, previous_event
+
+    def combine_b(self, output, previous_event):
+        hidden_states, event = self._combine_core(output, previous_event)
+        event.current_stream_wait() if self.async_finish else ()
+        self.handle = None
+        self.src2dst = None
+        return hidden_states
+
+    def _combine_core(self, x: torch.Tensor, previous_event):
+        combined_x, _, event = self.buffer_normal.combine(
+            x,
+            self.handle,
+            async_finish=self.async_finish,
+            previous_event=previous_event,
+            allocate_on_comm_stream=previous_event is not None,
+        )
+        return combined_x, event
+
+
+class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
+    def __init__(self, return_recv_hook: bool, **kwargs):
+        super().__init__(**kwargs)
+
+        """
+        num_max_dispatch_tokens_per_rank: the actual batch size in the decoding engine should be less than 256
+        https://github.com/deepseek-ai/DeepEP?tab=readme-ov-file#example-use-in-inference-decoding
+        """
+        # TODO(ch-wan): allow users to set this value
+        self.num_max_dispatch_tokens_per_rank = 128
+        self.buffer_low_latency = _get_buffer_low_latency(
+            self.group,
+            self.num_max_dispatch_tokens_per_rank,
+            self.hidden_size,
+            self.num_experts,
+        )
+        self.return_recv_hook = return_recv_hook
+
+    def dispatch_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        num_experts: int,
+        num_max_dispatch_tokens_per_rank: int,
+    ):
+        topk_idx = topk_idx.to(torch.int64)
+        expected_m = (
+            hidden_states.shape[0]
+            * self.buffer_low_latency.group_size
+            * topk_idx.shape[1]
+            + num_experts
+        ) // num_experts
+        hidden_states, masked_m, event, hook = self._dispatch_core(
+            hidden_states,
+            topk_idx,
+            num_max_dispatch_tokens_per_rank,
+            num_experts,
+            use_fp8=True,
+        )
+        return (
+            hidden_states,
+            topk_idx,
+            topk_weights,
+            masked_m,
+            expected_m,
+            event,
+            hook,
+        )
+
+    def dispatch_b(
+        self,
+        hidden_states,
+        topk_idx,
+        topk_weights,
+        masked_m,
+        expected_m,
+        event,
+        hook,
+    ):
+        hook() if self.return_recv_hook else event.current_stream_wait()
+
+        reorder_topk_ids = seg_indptr = None
+
+        return (
+            hidden_states,
+            topk_idx,
+            topk_weights,
+            reorder_topk_ids,
+            seg_indptr,
+            masked_m,
+            expected_m,
+        )
+
+    def _dispatch_core(
         self,
         hidden_states: torch.Tensor,
         topk_idx: torch.Tensor,
         num_max_dispatch_tokens_per_rank: int,
         num_experts: int,
+        use_fp8: bool = False,
     ):
         """
-        # For H20, there will be an CUDA error: DeepEP/csrc/kernels/internode_ll.cu:337 'too many blocks in cooperative launch'
-        # Please please make sure to change DeepEP code in internode_ll.cu dispatch / combine first and then reinstall!
+        # For H20, there will be an CUDA error: DeepEP/csrc/kernels/internode_ll.cu:337 'too many blocks in cooperative launch'.
+        # Please make sure to change DeepEP code in internode_ll.cu dispatch / combine as below first and then reinstall.
         # More details refer: https://github.com/deepseek-ai/DeepEP/issues/15#issuecomment-2709715782
-        +
+
         diff --git a/csrc/kernels/internode_ll.cu b/csrc/kernels/internode_ll.cu
-        index f60e933..cddaabf 100644
+        index 76ae2e2..8ecd08f 100644
         --- a/csrc/kernels/internode_ll.cu
         +++ b/csrc/kernels/internode_ll.cu
-        @@ -307,14 +307,14 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
-                    int num_topk, int num_experts, int rank, int num_ranks,
+        @@ -310,8 +310,8 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
+                    int num_topk, int num_experts, int rank, int num_ranks, bool use_fp8,
                     void* workspace, cudaStream_t stream, int phases) {
             constexpr int kNumMaxTopK = 9;
         -    constexpr int kNumWarpsPerGroup = 10;
@@ -308,16 +437,9 @@ class DeepEPDispatcher:
         +    constexpr int kNumWarpsPerGroup = 8;
         +    constexpr int kNumWarpGroups = 4;
             EP_STATIC_ASSERT(kNumMaxTopK + 1 <= kNumWarpGroups * kNumWarpsPerGroup, "Too many top-k selections");
-        +
+
             const auto num_warps = kNumWarpGroups * kNumWarpsPerGroup;
-            const auto num_sms = cell_div(num_experts, kNumWarpGroups);
-            EP_HOST_ASSERT(num_topk <= kNumMaxTopK);
-        -    EP_HOST_ASSERT(cell_div(static_cast<int>(hidden * 2 / sizeof(int4)), 32 * (num_warps - 1)) <= 2);
-        +    // EP_HOST_ASSERT(cell_div(static_cast<int>(hidden * 2 / sizeof(int4)), 32 * (num_warps - 1)) <= 2);
-        +
-            // Workspace checks
-            auto atomic_counter_per_expert = reinterpret_cast<int*>(workspace);
-        @@ -505,8 +505,8 @@ void combine(void* combined_x,
+        @@ -501,8 +501,8 @@ void combine(void* combined_x,
                     int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
                     int num_topk, int num_experts, int rank, int num_ranks,
                     void* workspace, cudaStream_t stream, int phases) {
@@ -326,91 +448,152 @@ class DeepEPDispatcher:
         +    constexpr int kNumWarpsPerGroup = 8;
         +    constexpr int kNumWarpGroups = 4;
             constexpr int kNumMaxTopk = 9;
-        +
+
             const auto num_warps = kNumWarpGroups * kNumWarpsPerGroup;
         """
 
-        recv_hidden_states, recv_expert_count, handle, event, hook = (
+        packed_recv_hidden, packed_recv_count, self.handle, event, hook = (
             self.buffer_low_latency.low_latency_dispatch(
                 hidden_states,
                 topk_idx,
                 num_max_dispatch_tokens_per_rank,
                 num_experts,
-                async_finish=self.async_finish,
-                return_recv_hook=False,  # True for double-batch overlapping, need call hook()
+                use_fp8=use_fp8,
+                async_finish=not self.return_recv_hook,
+                return_recv_hook=self.return_recv_hook,
             )
         )
-        # hook()
-        return recv_hidden_states, recv_expert_count, handle, event, hook
-
-    def combine(
-        self, hidden_states: torch.Tensor, forward_mode: ForwardMode
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
-        # Todo: enable low latency combine
-        if True:  # not forward_mode.is_decode():
-            if hidden_states.shape[0] > 0:
-                num_tokens = self.src2dst.shape[0] // self.router_topk
-                output = torch.empty(
-                    (num_tokens, hidden_states.shape[1]),
-                    device=hidden_states.device,
-                    dtype=hidden_states.dtype,
-                )
-                deepep_post_reorder_triton_kernel[(num_tokens,)](
-                    hidden_states,
-                    output,
-                    self.src2dst,
-                    self.topk_idx,
-                    self.topk_weights,
-                    self.router_topk,
-                    hidden_states.shape[1],
-                    BLOCK_SIZE=512,
-                )
-            else:
-                output = torch.zeros(
-                    (0, hidden_states.shape[1]),
-                    device=hidden_states.device,
-                    dtype=hidden_states.dtype,
-                )
-            hidden_states, event = self.combine_normal(output, self.handle)
-        else:
-            hidden_states, event, hook = self.combine_low_latency(
-                hidden_states, self.topk_idx, self.topk_weights, self.handle
-            )
+        return packed_recv_hidden, packed_recv_count, event, hook
 
-        if self.async_finish:
-            event.current_stream_wait()
+    def combine_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+    ):
+        hidden_states, event, hook = self._combine_core(
+            hidden_states,
+            topk_idx,
+            topk_weights,
+        )
+        return hidden_states, event, hook
 
-        self.handle = None
+    def combine_b(self, hidden_states, event, hook):
+        hook() if self.return_recv_hook else event.current_stream_wait()
         return hidden_states
 
-    def combine_normal(self, x: torch.Tensor, handle: Tuple):
-        previous_event = Buffer.capture() if self.async_finish else None
-
-        combined_x, _, event = self.buffer_normal.combine(
-            x,
-            handle,
-            async_finish=self.async_finish,
-            previous_event=previous_event,
-            allocate_on_comm_stream=previous_event is not None,
-        )
-        return combined_x, event
-
-    def combine_low_latency(
+    def _combine_core(
         self,
         hidden_states: torch.Tensor,
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
-        handle: Tuple,
     ):
-        combined_hidden_states, event_overlap, hook = (
+        combined_hidden_states, event, hook = (
             self.buffer_low_latency.low_latency_combine(
                 hidden_states,
                 topk_idx,
                 topk_weights,
-                handle,
-                async_finish=self.async_finish,
-                return_recv_hook=False,  # True for double-batch overlapping, need call hook()
+                self.handle,
+                async_finish=not self.return_recv_hook,
+                return_recv_hook=self.return_recv_hook,
             )
         )
-        # hook()
-        return combined_hidden_states, event_overlap, hook
+        self.handle = None
+        return combined_hidden_states, event, hook
+
+
+class DeepEPDispatcher:
+    def __init__(
+        self,
+        group: torch.distributed.ProcessGroup,
+        router_topk: int,
+        permute_fusion: bool = False,
+        num_experts: int = None,
+        num_local_experts: int = None,
+        hidden_size: int = None,
+        params_dtype: torch.dtype = None,
+        deepep_mode: DeepEPMode = DeepEPMode.auto,
+        async_finish: bool = False,
+        return_recv_hook: bool = False,
+    ):
+        self.deepep_mode = deepep_mode
+
+        common_kwargs = dict(
+            group=group,
+            router_topk=router_topk,
+            permute_fusion=permute_fusion,
+            num_experts=num_experts,
+            num_local_experts=num_local_experts,
+            hidden_size=hidden_size,
+            params_dtype=params_dtype,
+        )
+
+        if self.deepep_mode.enable_normal():
+            self._normal_dispatcher = _DeepEPDispatcherImplNormal(
+                async_finish=async_finish,
+                **common_kwargs,
+            )
+        if self.deepep_mode.enable_low_latency():
+            self._low_latency_dispatcher = _DeepEPDispatcherImplLowLatency(
+                return_recv_hook=return_recv_hook,
+                **common_kwargs,
+            )
+
+    def dispatch(self, *args, **kwargs) -> Tuple:
+        self.dispatch_a(*args, **kwargs)
+        return self.dispatch_b()
+
+    def dispatch_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        num_experts: int,
+        num_max_dispatch_tokens_per_rank: int = 128,
+        forward_mode: ForwardMode = None,
+    ):
+        inner_state = self._get_impl(forward_mode).dispatch_a(
+            hidden_states=hidden_states,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+            num_experts=num_experts,
+            num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank,
+        )
+        self._dispatch_intermediate_state = forward_mode, inner_state
+
+    def dispatch_b(self):
+        forward_mode, inner_state = self._dispatch_intermediate_state
+        del self._dispatch_intermediate_state
+        return self._get_impl(forward_mode).dispatch_b(*inner_state)
+
+    def combine(self, *args, **kwargs) -> Tuple:
+        self.combine_a(*args, **kwargs)
+        return self.combine_b()
+
+    def combine_a(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        forward_mode: ForwardMode,
+    ):
+        inner_state = self._get_impl(forward_mode).combine_a(
+            hidden_states=hidden_states,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+        )
+        self._combine_intermediate_state = forward_mode, inner_state
+
+    def combine_b(self):
+        forward_mode, inner_state = self._combine_intermediate_state
+        del self._combine_intermediate_state
+        return self._get_impl(forward_mode).combine_b(*inner_state)
+
+    def _get_impl(self, forward_mode: ForwardMode) -> "_DeepEPDispatcherImplBase":
+        resolved_deepep_mode = self.deepep_mode.resolve(forward_mode)
+        if resolved_deepep_mode == DeepEPMode.normal:
+            return self._normal_dispatcher
+        elif resolved_deepep_mode == DeepEPMode.low_latency:
+            return self._low_latency_dispatcher
+        else:
+            raise ValueError(f"Invalid deepep_mode: {self.deepep_mode}")