[Bugfix][MoE] Restore TP-sharded routed-expert capture for DP=1 (#61388)

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>
Signed-off-by: Ao Shen <281812371+aoshen02@users.noreply.github.com>
Signed-off-by: Ao Shen <aoshen@inferact.ai>
Co-authored-by: Ao Shen <281812371+aoshen02@users.noreply.github.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Ao Shen <aoshen@inferact.ai>

tests/model_executor/test_routed_experts_capture.py

@@ -7,9 +7,16 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 
-from vllm.config import VllmConfig
+from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.config.compilation import CompilationMode
 from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.forward_context import (
+    BatchDescriptor,
+    DPMetadata,
+    ForwardContext,
+    get_forward_context,
+    set_forward_context,
+)
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import CPUExpertsInt4
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -36,6 +43,12 @@ pytestmark = pytest.mark.cpu_test
 _REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"
 
 
+def _forward_context(**kwargs):
+    return ForwardContext(
+        no_compile_layers={}, attn_metadata={}, slot_mapping={}, **kwargs
+    )
+
+
 def _capturer_with_buffer(
     *,
     max_tokens: int = 8,
@@ -319,13 +332,120 @@ def test_routed_experts_capturer_single_dp_no_metadata():
     """dp_metadata is None: capture writes the full topk_ids rows."""
     capturer = _capturer_with_buffer(dp_rank=0)
     topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
-    ctx = SimpleNamespace(dp_metadata=None)
+    ctx = _forward_context()
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
         capturer.capture(layer_id=0, topk_ids=topk)
     assert torch.equal(capturer.device_buffer[:3, 0, :], topk)
     assert capturer.device_buffer[3, 0, 0].item() == -1
 
 
+@pytest.mark.parametrize("tp_size, num_tokens", [(2, 6), (2, 3), (2, 1), (4, 2)])
+@pytest.mark.parametrize(
+    "cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE]
+)
+def test_single_dp_ep_capture_restores_tp_shards(tp_size, num_tokens, cg_mode):
+    """Capture all current rows after model-level SP, including buffer reuse."""
+    capturer = _capturer_with_buffer(tp_size=tp_size)
+    config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            data_parallel_size=1,
+            tensor_parallel_size=tp_size,
+            enable_expert_parallel=True,
+            use_sequence_parallel_moe=False,
+            is_moe_model=True,
+        ),
+        compilation_config=SimpleNamespace(
+            fast_moe_cold_start=False, static_forward_context={}
+        ),
+    )
+    tp_group = Mock()
+    shard_n = (num_tokens + tp_size - 1) // tp_size
+    with (
+        set_forward_context(
+            None,
+            config,
+            num_tokens=num_tokens,
+            cudagraph_runtime_mode=cg_mode,
+        ),
+        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
+    ):
+        assert get_forward_context().dp_metadata is None
+        for offset in (20, 1):
+            expected = torch.arange(num_tokens * 2, dtype=torch.int32).reshape(-1, 2)
+            expected += offset
+            padded = torch.full((shard_n * tp_size, 2), 99, dtype=torch.int32)
+            padded[:num_tokens] = expected
+            tp_group.all_gather.return_value = padded
+            capturer.capture(0, padded[:shard_n])
+            assert torch.equal(
+                capturer.snapshot_routing_data(num_tokens)[:, 0].int(), expected
+            )
+            assert (capturer.device_buffer[num_tokens:, 0] == -1).all()
+    if num_tokens == 1:
+        tp_group.all_gather.assert_not_called()
+    else:
+        assert tp_group.all_gather.call_count == 2
+
+
+@pytest.mark.parametrize("num_tokens, num_rows", [(6, 6), (3, 4), (0, 0)])
+def test_single_dp_tp_capture_does_not_regather_full_batch(num_tokens, num_rows):
+    capturer = _capturer_with_buffer(tp_size=2)
+    topk = torch.arange(num_rows * 2, dtype=torch.int32).reshape(-1, 2)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+    ):
+        capturer.capture(0, topk)
+    get_tp_group.assert_not_called()
+    assert torch.equal(capturer.device_buffer[:num_tokens, 0], topk[:num_tokens])
+    assert (capturer.device_buffer[num_tokens:, 0] == -1).all()
+
+
+@pytest.mark.parametrize("num_tokens, num_rows", [(6, 2), (0, 1)])
+def test_single_dp_tp_capture_rejects_unexpected_batch_without_writing(
+    num_tokens, num_rows
+):
+    capturer = _capturer_with_buffer(tp_size=2)
+    topk = torch.ones(num_rows, 2, dtype=torch.int32)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),
+    ):
+        capturer.capture(0, topk)
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize("tp_rank", [0, 1])
+def test_capture_uses_pre_sp_size_after_model_shards_padding_mask(tp_rank):
+    """Model-level SP shortens the mask; capture still needs all three rows."""
+    from vllm.models.common.ops import sequence_parallel as sp
+
+    capturer = _capturer_with_buffer(tp_size=2)
+    expected = torch.arange(8, dtype=torch.int32).reshape(4, 2)
+    ctx = _forward_context(
+        is_padding=torch.zeros(3, dtype=torch.bool),
+        batch_descriptor=BatchDescriptor(num_tokens=3),
+    )
+    with (
+        patch.object(sp, "get_tensor_model_parallel_world_size", return_value=2),
+        patch.object(sp, "get_tensor_model_parallel_rank", return_value=tp_rank),
+    ):
+        ctx.is_padding = sp.sp_padding_mask(ctx.is_padding, torch.empty(3, 1))
+    assert len(ctx.is_padding) == 2
+    tp_group = Mock()
+    tp_group.all_gather.return_value = expected
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
+    ):
+        capturer.capture(0, expected[2 * tp_rank : 2 * (tp_rank + 1)])
+    tp_group.all_gather.assert_called_once()
+    torch.testing.assert_close(capturer.device_buffer[:3, 0], expected[:3])
+    assert (capturer.device_buffer[3:, 0] == -1).all()
+
+
 @pytest.mark.parametrize(
     ("num_experts", "dtype_name", "torch_dtype"),
     [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)],
@@ -359,7 +479,7 @@ def test_routed_experts_capturer_narrows_snapshot(output_dtype):
     capturer = _capturer_with_buffer(dtype=torch.int32)
     capturer.output_dtype = output_dtype
     topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64)
-    ctx = SimpleNamespace(dp_metadata=None)
+    ctx = _forward_context()
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
         capturer.capture(layer_id=0, topk_ids=topk)
 
@@ -375,8 +495,8 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
     """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     # Concatenated order: rank0 rows then rank1 rows.
     topk = torch.tensor(
@@ -392,8 +512,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
     """N == token_num_per_dp: topk is already local to this DP rank."""
     capturer = _capturer_with_buffer(dp_rank=1)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)
     with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
@@ -404,8 +524,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens():
 def test_routed_experts_capturer_dp_ep_gathered_shards():
     capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
     num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
             num_tokens_across_dp_cpu=num_tokens_dp,
             local_sizes=[2, 2, 1, 1],
         )
@@ -421,13 +541,17 @@ def test_routed_experts_capturer_dp_ep_gathered_shards():
     assert torch.equal(capturer.device_buffer[:2, 0], topk[4:])
 
 
-def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
-    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+@pytest.mark.parametrize("single_dp", [False, True])
+def test_routed_experts_capturer_sp_modular_gathers_tp_shards(single_dp):
+    capturer = _capturer_with_buffer(dp_rank=0 if single_dp else 1, tp_size=2)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=(
+            None if single_dp else DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
+        ),
+        batch_descriptor=BatchDescriptor(num_tokens=3),
     )
-    local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
+    local_shard = torch.tensor([[10, 12], [11, 13]], dtype=torch.int32).t()
     gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32)
     tp_group = Mock()
     tp_group.all_gather.return_value = gathered
@@ -438,16 +562,83 @@ def test_routed_experts_capturer_sp_modular_gathers_tp_shards():
     ):
         capturer.capture(layer_id=0, topk_ids=local_shard)
 
-    tp_group.all_gather.assert_called_once_with(local_shard, dim=0)
+    tp_group.all_gather.assert_called_once()
+    gather_input = tp_group.all_gather.call_args.args[0]
+    assert gather_input.is_contiguous()
+    torch.testing.assert_close(gather_input, local_shard)
+    assert tp_group.all_gather.call_args.kwargs == {"dim": 0}
     assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3])
 
 
+def test_routed_experts_capturer_rejects_incomplete_dp_shard_groups():
+    """Invalid gather metadata must not silently select another rank's rows."""
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
+            num_tokens_across_dp_cpu=torch.tensor([3, 2]),
+            local_sizes=[2, 2, 2],
+        )
+    )
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(ValueError, match="not divisible by DP rank count"),
+    ):
+        capturer.capture(0, torch.ones(6, 2, dtype=torch.int32))
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize("layer_id", [-1, 4])
+def test_routed_experts_capturer_rejects_layer_before_restoring_rows(layer_id):
+    capturer = _capturer_with_buffer(tp_size=2)
+    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=3))
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+        pytest.raises(IndexError, match="exceeds capture buffer"),
+    ):
+        capturer.capture(layer_id, torch.ones(2, 2, dtype=torch.int32))
+    get_tp_group.assert_not_called()
+    assert (capturer.device_buffer == -1).all()
+
+
+@pytest.mark.parametrize(
+    ("counts", "local_sizes", "start", "stop"),
+    [
+        ([3, 0], [2, 2, 0, 0], 3, 3),
+        ([0, 3], None, 0, 3),
+        ([0, 0], [0, 0, 0, 0], 0, 0),
+        ([2, 2], [1, 1, 1, 1], 2, 4),
+    ],
+    ids=["idle-rank", "only-active-rank", "all-idle", "unpadded-gather-tie"],
+)
+def test_routed_experts_capturer_dp_preserves_empty_and_unpadded_spans(
+    counts, local_sizes, start, stop
+):
+    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2)
+    topk = torch.arange(sum(counts) * 2, dtype=torch.int32).reshape(-1, 2)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(
+            num_tokens_across_dp_cpu=torch.tensor(counts), local_sizes=local_sizes
+        )
+    )
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group,
+    ):
+        capturer.capture(0, topk)
+    get_tp_group.assert_not_called()
+    torch.testing.assert_close(
+        capturer.device_buffer[: stop - start, 0], topk[start:stop]
+    )
+    assert (capturer.device_buffer[stop - start :, 0] == -1).all()
+
+
 def test_routed_experts_capturer_dp_unexpected_batch_raises():
     """Mismatch between topk batch dim and DP layout: fail fast."""
     capturer = _capturer_with_buffer(dp_rank=0)
     num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
-    ctx = SimpleNamespace(
-        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    ctx = _forward_context(
+        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp)
     )
     # total=5, local=2: n=1 matches neither naive (5) nor modular (2).
     topk = torch.tensor([[1, 2]], dtype=torch.int32)

tests/models/test_deepseek_v41_replay_batch.py

@@ -4,7 +4,7 @@
 
 from dataclasses import replace
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
@@ -16,6 +16,9 @@ from vllm.forward_context import (
     ForwardContext,
     override_forward_context,
 )
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+)
 from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers
 from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
@@ -173,6 +176,31 @@ def test_replay_batch_keeps_each_request_window(state):
     assert context.dp_metadata is None and context.is_padding is None
 
 
+def test_replay_context_restores_tp_routing_rows(state):
+    """Capture uses the compact replay batch's size, not the outer batch's."""
+    batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING)
+    replay, _ = _prepare(state, batch, CUDAGraphMode.NONE)
+    context = replay.forward_context
+    assert context.is_padding is None and context.dp_metadata is None
+    count = len(REPLAY_ROWS)
+    assert context.batch_descriptor.num_tokens == count
+    capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
+    capturer.dp_rank, capturer.tp_size = 0, 2
+    capturer.device_buffer = torch.full((401, 1, 2), -1, device=DEVICE)
+    expected = torch.tensor(REPLAY_ROWS, device=DEVICE)[:, None].expand(-1, 2)
+    gathered = torch.cat((expected, torch.full((1, 2), -1, device=DEVICE)))
+    with (
+        override_forward_context(context),
+        patch(
+            "vllm.model_executor.layers.fused_moe.routed_experts_capturer.get_tp_group"
+        ) as group,
+    ):
+        group.return_value.all_gather.return_value = gathered
+        capturer.capture(0, gathered[: (count + 1) // 2])
+    torch.testing.assert_close(capturer.device_buffer[:count, 0], expected)
+    assert (capturer.device_buffer[count:] == -1).all()
+
+
 def test_replay_batch_keeps_device_decode_boundaries(state):
     """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device
     [1, 3]): their rows stay and the boundaries shift only by trimmed rows."""

vllm/forward_context.py

@@ -327,10 +327,8 @@ def set_forward_context(
             vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp
         )
 
-    # Convenience: if cudagraph is used and num_tokens is given, we can just
-    # create a batch descriptor here if not given (there's no harm since if it
-    # doesn't match in the wrapper it'll fall through).
-    if cudagraph_runtime_mode != CUDAGraphMode.NONE and num_tokens is not None:
+    # Preserve the batch size before model-level sequence parallelism.
+    if num_tokens is not None:
         batch_descriptor = batch_descriptor or BatchDescriptor(num_tokens=num_tokens)
 
     additional_kwargs = current_platform.set_additional_forward_context(

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

@@ -14,7 +14,7 @@ import torch
 
 from vllm.config import VllmConfig
 from vllm.distributed.parallel_state import get_tp_group
-from vllm.forward_context import get_forward_context
+from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
 from vllm.platforms import current_platform
 
@@ -85,103 +85,74 @@ class RoutedExpertsCapturer:
         self.tp_size = vllm_config.parallel_config.tensor_parallel_size
 
     def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
-        """Capture expert routing decisions for a specific layer.
-
-        Under data parallelism, ``topk_ids`` may have four different batch
-        layouts depending on where the DP combine happens and whether
-        Expert Parallelism (EP) or Sequence Parallelism (SP) is active for the
-        MoE layer:
-          - ``n == total`` (naive dispatch): all DP ranks' tokens are
-            concatenated before routing; we slice out this rank's span
-            using the cumulative per-rank counts.
-          - ``n == token_num_per_dp`` (modular-kernel path): DP combine
-            happens inside ``quant_method.apply``; ``select_experts`` only
-            ever sees this rank's tokens, so we take the whole tensor.
-          - ``n == sum(dp_metadata.local_sizes)`` (naive DP+EP dispatch):
-            sequence-parallel shards from every DP rank are gathered through
-            the flattened EP group. The shard sizes include CUDA-graph / SP
-            padding, so we use them to locate this DP rank's unpadded rows.
-          - ``n == ceil(token_num_per_dp / tp_size)`` (SP + modular-kernel
-            path): tokens were split along dim=0 across the TP group by
-            ``_sequence_parallel_context``
-            (``moe_runner_base.py:_sequence_parallel_context``), so each
-            TP rank only sees its shard. We all-gather along dim=0 to
-            reconstruct this DP rank's full routing tensor. SP pads with
-            ceil-div (see ``_compute_sp_num_tokens`` in
-            ``forward_context.py``), so the gathered tensor may contain a
-            few trailing padding rows which are trimmed by the downstream
-            ``[:token_num_per_dp]`` slice.
-
-        Args:
-            layer_id: The layer index.
-            topk_ids: Tensor of shape (batch_size, num_routed_experts).
-
-        """
-        ctx = get_forward_context()
-        if ctx.dp_metadata is None:
-            local_topk_ids = topk_ids
-        else:
-            num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu
-            num_local_tokens = int(num_tokens_dp[self.dp_rank].item())
-            total = int(num_tokens_dp.sum().item())
-            n = topk_ids.shape[0]
-            shard_sizes = getattr(ctx.dp_metadata, "local_sizes", None)
-            gathered_size = sum(shard_sizes) if shard_sizes is not None else -1
-
-            if n == total:
-                end = int(num_tokens_dp[: self.dp_rank + 1].sum().item())
-                local_topk_ids = topk_ids[end - num_local_tokens : end]
-            elif n == num_local_tokens:
-                local_topk_ids = topk_ids
-            elif shard_sizes is not None and n == gathered_size:
-                # Naive DP+EP dispatch gathers the sequence-parallel shards in
-                # flattened DP-then-TP order. ``local_sizes`` is the exact
-                # all-gatherv layout, including padding. Select this DP rank's
-                # contiguous shard group, then trim its trailing padding.
-                num_dp_ranks = len(num_tokens_dp)
-                assert len(shard_sizes) % num_dp_ranks == 0
-                shards_per_dp_rank = len(shard_sizes) // num_dp_ranks
-                first_shard = self.dp_rank * shards_per_dp_rank
-                start = sum(shard_sizes[:first_shard])
-                local_topk_ids = topk_ids[start : start + num_local_tokens]
-            elif (
-                self.tp_size > 1
-                and n == (num_local_tokens + self.tp_size - 1) // self.tp_size
-            ):
-                # SP + modular-kernel path. All-gather across the TP
-                # group along dim=0 to reconstruct the full per-DP-rank
-                # tensor; keep only the first ``token_num_per_dp`` rows
-                # (trailing rows are SP ceil-div padding). The TP group
-                # is always initialized on real rollout workers, and
-                # every rank in the group reaches this branch in
-                # lockstep (bind is per-FusedMoEFactory layer, SP is a global
-                # condition), so a bare all_gather here will not
-                # deadlock -- let it raise if the precondition is
-                # violated rather than skip silently.
-                #
-                # ``topk_ids`` is already whatever the router produced
-                # (typically int32/int64, both supported by NCCL); the
-                # downstream buffer assignment narrows to the capture dtype.
-                local_topk_ids = get_tp_group().all_gather(topk_ids, dim=0)[
-                    :num_local_tokens
-                ]
-            else:
-                sp_expected = (num_local_tokens + self.tp_size - 1) // self.tp_size
-                raise AssertionError(
-                    "RoutedExpertsCapturer: unexpected topk_ids batch "
-                    f"dim {n} (expected {total}, {num_local_tokens}, "
-                    f"{gathered_size}, or {sp_expected} for "
-                    f"dp_rank={self.dp_rank}, tp_size={self.tp_size})"
-                )
-
         if not 0 <= layer_id < self.device_buffer.shape[1]:
             raise IndexError(
                 f"routed-experts layer {layer_id} exceeds capture buffer "
                 f"layer count {self.device_buffer.shape[1]}"
             )
 
+        local_topk_ids = self._restore_local_topk_ids(topk_ids, get_forward_context())
         self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids
 
+    def _restore_local_topk_ids(
+        self, topk_ids: torch.Tensor, ctx: ForwardContext
+    ) -> torch.Tensor:
+        """Restore this DP rank's token rows, trimming trailing SP padding.
+
+        Layout priority: cross-DP batch, local batch, gathered DP/TP shards,
+        then a local TP shard. Gathered rows are DP-major/TP-minor; preserve
+        this priority when layout sizes coincide. PCP is unsupported.
+
+        Without DP metadata, use the pre-SP batch descriptor. Unknown counts
+        pass through; known counts must match a full, sharded or padded batch.
+        A full local batch avoids gathering, including zero/one-token batches
+        where only the publishing TP0 rank needs the real rows.
+        """
+        n = topk_ids.shape[0]
+        dp_metadata = ctx.dp_metadata
+        if dp_metadata is None:
+            if ctx.batch_descriptor is None:
+                return topk_ids
+            num_local_tokens = ctx.batch_descriptor.num_tokens
+        else:
+            num_tokens_dp = dp_metadata.num_tokens_across_dp_cpu
+            num_local_tokens = int(num_tokens_dp[self.dp_rank].item())
+            if n == int(num_tokens_dp.sum().item()):
+                end = int(num_tokens_dp[: self.dp_rank + 1].sum().item())
+                return topk_ids[end - num_local_tokens : end]
+
+        if n == num_local_tokens:
+            return topk_ids
+
+        if dp_metadata is not None:
+            shard_sizes = dp_metadata.local_sizes
+            if shard_sizes is not None and n == sum(shard_sizes):
+                num_dp_ranks = len(num_tokens_dp)
+                if len(shard_sizes) % num_dp_ranks != 0:
+                    raise ValueError(
+                        "RoutedExpertsCapturer: gathered shard count "
+                        f"{len(shard_sizes)} is not divisible by DP rank count "
+                        f"{num_dp_ranks}"
+                    )
+                shards_per_dp_rank = len(shard_sizes) // num_dp_ranks
+                first_shard = self.dp_rank * shards_per_dp_rank
+                start = sum(shard_sizes[:first_shard])
+                return topk_ids[start : start + num_local_tokens]
+
+        shard_n = (num_local_tokens + self.tp_size - 1) // self.tp_size
+        if n == shard_n:
+            return get_tp_group().all_gather(topk_ids.contiguous(), dim=0)[
+                :num_local_tokens
+            ]
+        if dp_metadata is None and n == shard_n * self.tp_size:
+            return topk_ids[:num_local_tokens]
+
+        raise AssertionError(
+            f"RoutedExpertsCapturer: unexpected topk_ids batch dim {n} "
+            f"(local_tokens={num_local_tokens}, "
+            f"dp_rank={self.dp_rank}, tp_size={self.tp_size})"
+        )
+
     def snapshot_routing_data(self, num_tokens: int) -> torch.Tensor:
         """Return a stable snapshot of the current routing data."""
         return self.device_buffer[:num_tokens].to(self.output_dtype)