[Bugfix][Worker] Fix routing replay capturing TP-local expert ids when DP=1

tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py

@@ -23,6 +23,12 @@
     "Qwen/Qwen3-30B-A3B",
 ]
 
+# The TP-sharded capture paths pinned by the case below are model independent, so it
+# runs on a single model: `test_moe_routing_replay` already loads the other one, and
+# loading a 30B+ model a second time would double this file's runtime for no extra
+# coverage.
+TP_SHARDED_MODEL = "Qwen/Qwen3-30B-A3B"
+
 
 @pytest.mark.parametrize("model", MODELS)
 @patch.dict(os.environ, {"OMP_NUM_THREADS": "1", "VLLM_USE_V2_MODEL_RUNNER": "1"})
@@ -51,3 +57,63 @@ def test_moe_routing_replay(model):
         routed_experts = outputs[0].outputs[0].routed_experts
         assert routed_experts is not None
         assert routed_experts.size > 0
+
+
+@pytest.mark.parametrize("model", MODELS)
+@patch.dict(os.environ, {"OMP_NUM_THREADS": "1", "VLLM_USE_V2_MODEL_RUNNER": "1"})
+def test_moe_routing_replay_covers_every_token(model):
+    """Every token row carries its own routing decision, under MC2 and ALLTOALL.
+
+    Regression test for #15451. Both strategies pad and split the token dimension
+    inside the TP group even with DP == 1, so ``select_experts`` only routes this
+    rank's shard. A capture that treats that shard as the complete batch writes
+    only the first ``batch / tp_size`` rows and leaves the rows owned by the other
+    TP ranks unwritten, which routing replay then serves as zeros.
+
+    The two strategies are picked by ``max_num_tokens <= mc2_tokens_capacity``
+    (``set_moe_comm_method``), and in eager mode without CUDA-graph sizes that
+    capacity is ``max_num_reqs * decode_query_len``, i.e. ``max_num_seqs`` here.
+    Pinning ``max_num_seqs`` to 8 therefore gives a deterministic threshold: the
+    short prompt routes through MC2 and the long one through ALLTOALL.
+
+    Eager mode is deliberate: graph warm-up would pre-fill the capture buffer, so
+    an unwritten row would hold stale expert ids instead of zeros and the
+    assertion below could no longer tell the two apart.
+
+    Runs on one model only, see ``TP_SHARDED_MODEL``.
+    """
+    if model != TP_SHARDED_MODEL:
+        pytest.skip(f"TP-sharded capture is already covered on {TP_SHARDED_MODEL}")
+    prompts = [
+        "Hello, please introduce yourself.",  # 6 tokens -> MC2
+        "Hello, please introduce yourself. " * 8,  # > 8 tokens -> ALLTOALL
+    ]
+    with VllmRunner(
+        model,
+        tensor_parallel_size=2,
+        enable_expert_parallel=True,
+        enforce_eager=True,
+        max_num_seqs=8,
+        distributed_executor_backend="mp",
+        enable_return_routed_experts=True,
+        async_scheduling=False,
+    ) as vllm_model:
+        sampling_params = SamplingParams(
+            max_tokens=8,
+            temperature=0.0,
+            output_kind=RequestOutputKind.FINAL_ONLY,
+        )
+        for prompt in prompts:
+            inputs = vllm_model.get_inputs(prompts=[prompt])
+            outputs = vllm_model.model.generate(prompts=inputs, sampling_params=sampling_params)
+            output = outputs[0].outputs[0]
+            routed_experts = output.routed_experts
+            assert routed_experts is not None
+            # The last sampled token has not been forwarded yet, so it has no routing.
+            expected_rows = len(outputs[0].prompt_token_ids) + len(output.token_ids) - 1
+            assert routed_experts.shape[0] == expected_rows
+            empty_rows = [row for row in range(expected_rows) if not routed_experts[row].any()]
+            assert not empty_rows, (
+                f"rows {empty_rows} of {expected_rows} carry no routing data for a "
+                f"{len(outputs[0].prompt_token_ids)}-token prompt"
+            )

vllm_ascend/patch/worker/patch_routed_experts_capture.py

@@ -31,6 +31,72 @@
 
 logger = logging.getLogger(__name__)
 
+# Strategies whose ``prepare()`` splits the token dimension across the TP group.
+# ALLGATHER is absent on purpose: it gathers the batch before routing.
+_TP_SHARDED_COMM_TYPES = frozenset(
+    {
+        MoECommType.ALLTOALL,
+        MoECommType.MC2,
+        MoECommType.FUSED_MC2,
+    }
+)
+
+
+def _padded_batch_rows(tp_size: int) -> int:
+    """Row count the MoE dispatch padded the token dimension to before splitting it.
+
+    Taken from the target ``prepare()`` padded to rather than from this rank's
+    shard, which only adds up to the total when the split is even: All2All splits
+    a batch with ``num_tokens >= tp_size`` unevenly whenever it is not a multiple
+    of the TP size. There is deliberately no ``shard_rows * tp_size`` fallback --
+    it would silently bring that rank-dependent length back.
+    """
+    prepare_finalize = getattr(_EXTRA_CTX.moe_comm_method, "prepare_finalize", None)
+    padded_num_tokens = getattr(_EXTRA_CTX, "padded_num_tokens", None)
+    num_tokens = getattr(prepare_finalize, "num_tokens", None)
+
+    if getattr(prepare_finalize, "replace_allreduce", False):
+        # Sequence parallel: the model wrapper already chunked the sequence to a
+        # multiple of the TP size, so this batch-wide target is the padded length.
+        total_rows = padded_num_tokens if isinstance(padded_num_tokens, int) else None
+    elif _EXTRA_CTX.moe_comm_type == MoECommType.ALLTOALL:
+        # All2All pads to the TP size, and only when the batch is smaller than it.
+        total_rows = max(num_tokens, tp_size) if isinstance(num_tokens, int) else None
+    elif isinstance(num_tokens, int) and isinstance(padded_num_tokens, int):
+        # MC2/FUSED_MC2 pad to ``padded_num_tokens``, and only when the batch is
+        # smaller than that target.
+        total_rows = max(num_tokens, padded_num_tokens)
+    else:
+        total_rows = None
+
+    if total_rows is None:
+        raise AssertionError(
+            "RoutedExpertsCapturer: cannot rebuild the TP-sharded routing batch, the MoE dispatch "
+            f"metadata is missing (comm_type={_EXTRA_CTX.moe_comm_type}, num_tokens={num_tokens!r}, "
+            f"padded_num_tokens={padded_num_tokens!r}). Refusing to infer the length from this rank's "
+            "shard, which is only valid when the split is even."
+        )
+    return total_rows
+
+
+def _gather_tp_sharded_topk_ids(topk_ids: torch.Tensor, tp_size: int, total_rows: int) -> torch.Tensor:
+    """Rebuild the pre-split batch from every TP rank's routing shard.
+
+    ``tensor_split`` reproduces the chunks ``prepare()``/``finalize()`` use, so
+    the all-gather restores the original row order.
+    """
+    gathered = torch.empty(
+        (total_rows, topk_ids.shape[1]),
+        dtype=topk_ids.dtype,
+        device=topk_ids.device,
+    )
+    dist.all_gather(
+        list(torch.tensor_split(gathered, tp_size, dim=0)),
+        topk_ids.contiguous(),
+        get_tp_group().device_group,
+    )
+    return gathered
+
 
 def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
     """Capture expert routing decisions for a specific layer.
@@ -61,6 +127,10 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
         few trailing padding rows which are trimmed by the downstream
         ``[:token_num_per_dp]`` slice.
 
+    The TP split also happens with ``DP == 1``: ``ctx.dp_metadata`` does not
+    describe it, but All2All/MC2 still split inside the TP group, so the same
+    all-gather is needed there. ALLGATHER does not split and needs none.
+
     Args:
         layer_id: The layer index.
         topk_ids: Tensor of shape (batch_size, num_routed_experts).
@@ -69,8 +139,21 @@ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
     ctx = get_forward_context()
     if ctx.dp_metadata is None:  # single dp
         start_loc = 0
-        end_loc = topk_ids.shape[0]
-        token_num_per_dp = topk_ids.shape[0]
+        end_loc = token_num_per_dp = topk_ids.shape[0]
+
+        # All2All/MC2 split the token dimension across the TP group even with
+        # DP=1, so ``topk_ids`` only holds this rank's shard. Capturing it as the
+        # whole batch leaves the other ``device_buffer`` rows stale, and routing
+        # replay then serves wrong expert ids for them. The condition is
+        # rank-invariant, so every TP rank gathers in lockstep.
+        if self.tp_size > 1 and _EXTRA_CTX.moe_comm_type in _TP_SHARDED_COMM_TYPES:
+            total_rows = _padded_batch_rows(self.tp_size)
+            topk_ids = _gather_tp_sharded_topk_ids(topk_ids, self.tp_size, total_rows)
+            # Rows past the batch are the padding ``prepare()`` appended, and the
+            # consumer only reads ``device_buffer[:num_tokens]``. Writing the whole
+            # gathered tensor avoids a second token-count estimate; the clamp only
+            # guards an over-padded gather against the buffer size.
+            end_loc = token_num_per_dp = min(topk_ids.shape[0], self.device_buffer.shape[0])
     else:  # multi dp
         num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu
         token_num_per_dp = int(num_tokens_dp[self.dp_rank].item())