perf: overlap shared expert with routed branch on ep decode

Author: HongminTan

rtp_llm/models_py/model_desc/generic_moe.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 from torch import nn
@@ -33,6 +33,24 @@
 
 logger = logging.getLogger(__name__)
 
+# One shared-expert side stream per device, reused by every MoE layer. It has to
+# exist before the first captured forward, because a captured fork records an edge
+# to an already-running stream, and reusing one avoids per-forward churn.
+_SHARED_EXPERT_STREAMS: Dict[int, torch.cuda.Stream] = {}
+
+
+def _ensure_shared_expert_stream(
+    device: torch.device,
+) -> Optional[torch.cuda.Stream]:
+    if device.type != "cuda" or not torch.cuda.is_available():
+        return None
+    index = device.index if device.index is not None else torch.cuda.current_device()
+    stream = _SHARED_EXPERT_STREAMS.get(index)
+    if stream is None:
+        stream = torch.cuda.Stream(device=torch.device("cuda", index))
+        _SHARED_EXPERT_STREAMS[index] = stream
+    return stream
+
 
 class GenericMoeLayer(nn.Module):
     """Generic MoE layer supporting both Qwen3 and internal model."""
@@ -131,18 +149,17 @@ def __init__(
         #     use_ep_unified_allreduce   routed partial, scattered   \
         #                                shared partial              /  add -> all_reduce
         #
-        #     these two differ only in how routed becomes a full-width partial:
-        #     pure TP leaves it that way, EP scatter-adds this rank's token slice
+        #     pure TP leaves routed a full-width partial; EP has to scatter-add
+        #     this rank's token slice to get one.
         #
         #   two collectives
         #     use_ep_shared_allreduce    routed complete via all_gather  \
         #                                shared partial -> all_reduce    /  add
         #
-        #     no flag, ffn_tp_size > 1   routed complete                 \
-        #                                shared complete                 /  add
-        #
-        #   no collective
-        #     no flag, ffn_tp_size == 1  routed and shared already complete -> add
+        #   no flag: each branch completes on its own, then a local add. The
+        #   router reduces or gathers according to its own tp size, which follows
+        #   the attention view and is independent of ffn_tp_size; DenseMLP
+        #   all-reduces the shared output unless ffn_tp_size == 1.
         #
         #   counts are Group.TP only; the EP all_to_all is not included
         self.use_ep_unified_allreduce = (
@@ -165,6 +182,14 @@ def __init__(
             and self.ffn_tp_size == router_tp_size
             and router.supports_skip_tp_allreduce
         )
+        # The shared branch has no data dependency on the routed branch, so it
+        # can run on a side stream alongside the whole routed chain.
+        self.shared_expert_stream: Optional[torch.cuda.Stream] = None
+        self.shared_expert_ready: Optional[torch.cuda.Event] = None
+        if self.use_ep_unified_allreduce:
+            self.shared_expert_stream = _ensure_shared_expert_stream(self.w1.device)
+            if self.shared_expert_stream is not None:
+                self.shared_expert_ready = torch.cuda.Event()
         if logger.isEnabledFor(logging.DEBUG):
             logger.debug(
                 "GenericMoE unified TP all-reduce %s, EP unified all-reduce %s "
@@ -204,6 +229,41 @@ def _gate_shared_expert_output(
             return torch.sigmoid(gate_output) * shared_expert_output
         return shared_expert_output
 
+    def _shared_expert_partial(
+        self, hidden_states: torch.Tensor
+    ) -> Tuple[torch.Tensor, Optional[torch.cuda.Event]]:
+        """Produce the gated TP-partial shared output for use_ep_unified_allreduce.
+
+        Returns the buffer plus, when the work was issued on the side stream,
+        the event the router must wait on before accumulating into it.
+        """
+        assert self.shared_expert is not None
+        if self.shared_expert_stream is None:
+            return (
+                self._gate_shared_expert_output(
+                    hidden_states,
+                    self.shared_expert(hidden_states, skip_allreduce=True),
+                ),
+                None,
+            )
+
+        stream = self.shared_expert_stream
+        current = torch.cuda.current_stream(hidden_states.device)
+        stream.wait_stream(current)
+        # record_stream stops the allocator from handing these blocks to a later
+        # allocation once their references drop, while the side stream still
+        # reads them.
+        hidden_states.record_stream(stream)
+        with torch.cuda.stream(stream):
+            shared_partial = self._gate_shared_expert_output(
+                hidden_states,
+                self.shared_expert(hidden_states, skip_allreduce=True),
+            )
+        shared_partial.record_stream(current)
+        assert self.shared_expert_ready is not None
+        self.shared_expert_ready.record(stream)
+        return shared_partial, self.shared_expert_ready
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, _ = hidden_states.shape
         router_logits = self.gate(hidden_states)
@@ -250,17 +310,29 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
 
         if self.use_ep_unified_allreduce:
             assert self.shared_expert is not None
-            row_scatter_target = self._gate_shared_expert_output(
-                hidden_states,
-                self.shared_expert(hidden_states, skip_allreduce=True),
-            )
-            experts_output = self.fused_moe(
-                hidden_states=hidden_states,
-                topk_weights=topk_weights,
-                topk_ids=topk_ids,
-                activation="SiGLU",
-                row_scatter_target=row_scatter_target,
+            # Issued before fused_moe and on a side stream, so the shared branch
+            # overlaps the routed chain instead of delaying it.
+            row_scatter_target, shared_ready = self._shared_expert_partial(
+                hidden_states
             )
+            try:
+                experts_output = self.fused_moe(
+                    hidden_states=hidden_states,
+                    topk_weights=topk_weights,
+                    topk_ids=topk_ids,
+                    activation="SiGLU",
+                    row_scatter_target=row_scatter_target,
+                    row_scatter_ready=shared_ready,
+                )
+            except Exception:
+                # The router joins on the success path. When the routed chain
+                # raises instead, the side stream is still writing the buffer,
+                # so anything the caller does after unwinding would race with it.
+                if shared_ready is not None:
+                    torch.cuda.current_stream(hidden_states.device).wait_event(
+                        shared_ready
+                    )
+                raise
             return all_reduce(experts_output, group=Group.TP, inplace=True)
 
         # In pure-TP mode both the routed experts and the shared expert produce
@@ -292,9 +364,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
                 )
                 experts_output = all_reduce(experts_output, group=Group.TP)
             elif self.use_ep_shared_allreduce:
-                # Skipping DenseMLP's reduce to redo it here saves no collective;
-                # it just puts the gate ahead of the reduction, which is
-                # equivalent for a rank-consistent gate.
+                # The router already reassembled the routed output with an
+                # all_gather, so it is complete and cannot join a merged
+                # reduction; only the shared branch still needs one.
                 shared_expert_output = self._gate_shared_expert_output(
                     hidden_states, shared_expert_output
                 )

rtp_llm/models_py/model_desc/test/generic_moe_allreduce_test.py

@@ -1,17 +1,31 @@
-"""CPU contract tests for GenericMoeLayer's unified TP all-reduce path."""
+"""Contract tests for GenericMoeLayer's unified TP all-reduce paths.
 
+Mostly CPU tensors; the side-stream shared expert needs a real CUDA graph.
+"""
+
+import inspect
+from functools import partial
 from types import SimpleNamespace
-from unittest import TestCase, main
+from unittest import TestCase, main, skipUnless
 from unittest.mock import MagicMock, Mock, patch
 
 import torch
 
 from rtp_llm.models_py.distributed.collective_torch import Group
-from rtp_llm.models_py.model_desc.generic_moe import GenericMoeLayer
+from rtp_llm.models_py.model_desc.generic_moe import (
+    _SHARED_EXPERT_STREAMS,
+    GenericMoeLayer,
+    _ensure_shared_expert_stream,
+)
 from rtp_llm.models_py.modules.factory.fused_moe.defs.fused_moe import FusedMoe
 from rtp_llm.models_py.modules.hybrid.dense_mlp import DenseMLP
 from rtp_llm.utils.model_weight import W
 
+# This file also runs as generic_moe_allreduce_test_rocm. Only the CUDA
+# low-latency router advertises supports_row_scatter_finalize, so the side
+# stream never runs on ROCm.
+_CUDA_GRAPH_READY = torch.cuda.is_available() and torch.version.hip is None
+
 
 def _make_layer(
     *,
@@ -288,6 +302,149 @@ def test_row_scatter_reduces_both_branches_once(self, mock_all_reduce):
         torch.testing.assert_close(result, expected_input * 2)
         self.assertTrue(layer.shared_expert.call_args.kwargs["skip_allreduce"])
         self.assertNotIn("skip_tp_allreduce", fused_moe.call_args.kwargs)
+        # CPU weights mean no side stream, so the router needs no join.
+        self.assertIsNone(fused_moe.call_args.kwargs["row_scatter_ready"])
+
+    @patch("rtp_llm.models_py.model_desc.generic_moe.all_reduce")
+    def test_a_failed_routed_chain_joins_the_shared_branch(self, _):
+        # The router does the join on the success path. When the routed chain
+        # raises, forward has to do it, or the side stream stays outstanding.
+        layer = _make_layer(ep_size=2, supports_row_scatter_finalize=True)
+        hidden_states, _, _, _, fused_moe = _configure_forward(layer)
+        ready = object()
+        layer._shared_expert_partial = Mock(
+            return_value=(torch.zeros_like(hidden_states), ready)
+        )
+        fused_moe.side_effect = RuntimeError("dispatch failed")
+        waited = []
+        stream = SimpleNamespace(wait_event=waited.append)
+
+        with patch("torch.cuda.current_stream", return_value=stream):
+            with self.assertRaisesRegex(RuntimeError, "dispatch failed"):
+                layer(hidden_states)
+
+        self.assertEqual(waited, [ready])
+
+    @patch("rtp_llm.models_py.model_desc.generic_moe.all_reduce")
+    def test_layer_only_passes_keywords_fused_moe_accepts(self, _):
+        # The fused_moe mock accepts any keyword, so a layer-only keyword passes
+        # every other test here and only fails once a real FusedMoe is called.
+        accepted = set(inspect.signature(FusedMoe.forward).parameters)
+        cases = (
+            ("ep_unified", dict(ep_size=2, supports_row_scatter_finalize=True)),
+            ("ep_shared", dict(ep_size=2, supports_row_scatter_finalize=False)),
+            ("pure_tp", dict(ep_size=1)),
+            ("ffn_tp_one", dict(ffn_tp_size=1)),
+        )
+        for name, kwargs in cases:
+            with self.subTest(name=name):
+                layer = _make_layer(**kwargs)
+                hidden_states, _, _, _, fused_moe = _configure_forward(layer)
+
+                layer(hidden_states)
+
+                unexpected = set(fused_moe.call_args.kwargs) - accepted
+                self.assertFalse(
+                    unexpected, f"not FusedMoe.forward parameters: {unexpected}"
+                )
+
+
+@skipUnless(_CUDA_GRAPH_READY, "needs a CUDA device")
+class SharedExpertSideStreamCaptureTest(TestCase):
+    """Covers _shared_expert_partial's side-stream branch under real capture.
+
+    CPU weights leave shared_expert_stream unset, so only a CUDA device reaches
+    this branch, and an unclosed fork only shows up at capture end.
+    """
+
+    SHARED_FACTOR = 3.0
+
+    def setUp(self):
+        _SHARED_EXPERT_STREAMS.clear()
+        self.device = torch.device("cuda", torch.cuda.current_device())
+
+    def _stub_layer(self):
+        """Stand-in for ``self``: this path only needs these four attributes."""
+        stream = _ensure_shared_expert_stream(self.device)
+        self.assertIsNotNone(stream)
+        stub = SimpleNamespace(
+            # A row-parallel shared expert leaves a TP-partial sum; scaling
+            # stands in for that work without needing real weights.
+            shared_expert=lambda x, skip_allreduce: x * self.SHARED_FACTOR,
+            shared_expert_gate=None,
+            shared_expert_stream=stream,
+            shared_expert_ready=torch.cuda.Event(),
+        )
+        stub._gate_shared_expert_output = partial(
+            GenericMoeLayer._gate_shared_expert_output, stub
+        )
+        return stub
+
+    def _fork_and_join(self, stub, hidden):
+        shared_partial, ready = GenericMoeLayer._shared_expert_partial(stub, hidden)
+        self.assertIs(ready, stub.shared_expert_ready)
+        # The join the router performs in _finalize_row_scatter.
+        torch.cuda.current_stream().wait_event(ready)
+        return shared_partial
+
+    def test_one_stream_per_device_is_reused(self):
+        # Every MoE layer shares one side stream per device.
+        self.assertIs(
+            _ensure_shared_expert_stream(self.device),
+            _ensure_shared_expert_stream(self.device),
+        )
+
+    def test_capture_closes_the_fork_and_replay_recomputes(self):
+        stub = self._stub_layer()
+        hidden = torch.ones(8, 16, device=self.device)
+        out = torch.empty_like(hidden)
+
+        # torch.cuda.graph requires the work to have run once outside capture.
+        out.copy_(self._fork_and_join(stub, hidden))
+        torch.cuda.synchronize()
+
+        # Capture fails here if the side stream is still outstanding.
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            out.copy_(self._fork_and_join(stub, hidden))
+
+        # A new input value distinguishes a real replay from leftover contents.
+        hidden.fill_(2.0)
+        graph.replay()
+        torch.cuda.synchronize()
+        torch.testing.assert_close(out, torch.full_like(out, 2.0 * self.SHARED_FACTOR))
+
+    def test_graphs_per_batch_size_share_the_stream_and_event(self):
+        # Decode captures one graph per batch size into a single memory pool
+        # (cuda_graph_runner.cc passes one shared_graph_pool_ to every capture),
+        # and every graph reuses the device's side stream and the layer's event.
+        stub = self._stub_layer()
+        pool = torch.cuda.graph_pool_handle()
+        captured = []
+        forked_on = set()
+        for num_tokens in (4, 8):
+            hidden = torch.ones(num_tokens, 16, device=self.device)
+            out = torch.empty_like(hidden)
+            out.copy_(self._fork_and_join(stub, hidden))
+            torch.cuda.synchronize()
+
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=pool):
+                out.copy_(self._fork_and_join(stub, hidden))
+            forked_on.add((id(stub.shared_expert_stream), id(stub.shared_expert_ready)))
+            captured.append((graph, hidden, out))
+
+        self.assertEqual(len(forked_on), 1, "captures used different stream/event")
+
+        for index, (graph, hidden, _) in enumerate(captured):
+            hidden.fill_(index + 2.0)
+            graph.replay()
+        torch.cuda.synchronize()
+
+        for index, (_, _, out) in enumerate(captured):
+            torch.testing.assert_close(
+                out, torch.full_like(out, (index + 2.0) * self.SHARED_FACTOR)
+            )
 
 
 if __name__ == "__main__":