From 44e25d50538abc470e3586a90885f3fcab59fb75 Mon Sep 17 00:00:00 2001 From: wuruiyun <1251540982@qq.com> Date: Tue, 13 Oct 2026 09:38:38 +0800 Subject: [PATCH 1/2] [Bugfix][Worker] Fix routing replay capturing TP-local expert ids when DP=1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #15451. The routed-experts (R3 / routing replay) capture patch assumed that `topk_ids` covers the whole batch whenever `ctx.dp_metadata is None`: if ctx.dp_metadata is None: # single dp token_num_per_dp = topk_ids.shape[0] That is wrong. `ctx.dp_metadata` describes data parallelism only, but the All2All and MC2 strategies also split the token dimension inside the TP group: `PrepareAndFinalizeWithAll2All.prepare` pads to the TP size, and `PrepareAndFinalizeWithMC2.prepare` pads to `padded_num_tokens = ceil(max_tokens_across_dp / tp_size) * tp_size`; both then `tensor_split` the result across the TP ranks. The capture hook runs after `prepare()` (`AscendRoutedExperts.forward_impl` -> `prepare()` -> `BaseRouter._select_experts` -> `capture_fn(topk_ids)`), so with TP > 1 it only ever sees this rank's shard. Capturing that shard as the complete batch wrote just `device_buffer[:batch / tp_size]`, and every remaining row kept its previous contents. On a freshly started worker those rows stay zero; after warmup they hold another batch's expert ids. R3 then replays them. Measured on Qwen3-30B-A3B, TP=2 + EP, DP=1, 6-token prefill (shard length 3), before the fix: routed_experts shape=(13, 48, 8) rows(all-zero) = [0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0] Gather the per-rank shards back before writing the buffer: - a new `_TP_SHARDED_COMM_TYPES` set names the strategies that split (`ALLTOALL`, `MC2`, `FUSED_MC2`). `ALLGATHER` is deliberately excluded: it gathers the batch across the EP/DP group before routing, so its `topk_ids` already spans the whole batch. - when `tp_size > 1` and the active comm type is in that set, the DP=1 branch all-gathers the shards into the padded batch first. The number of rows to rebuild is read from the padding target `prepare()` used, not inferred from this rank's shard: the two only agree for an even split, and All2All splits a batch with `num_tokens >= tp_size` unevenly whenever it is not a multiple of the TP size. `_padded_batch_rows` therefore returns - `max(num_tokens, tp_size)` for All2All, which pads to the TP size and only when the batch is smaller than it, and - `max(num_tokens, padded_num_tokens)` for MC2/FUSED_MC2, where `num_tokens` is the pre-pad batch size `prepare()` recorded and `padded_num_tokens` is `_EXTRA_CTX.padded_num_tokens`. Under sequence parallelism the wrapper has already sharded the sequence to a multiple of the TP size, so `padded_num_tokens` is the whole answer there. The helper takes no rank-local input at all, so its result is rank-invariant by construction, and it raises instead of falling back to `shard_rows * tp_size` when the metadata is missing -- guessing there would silently bring the rank-dependent length back. The gather decision itself is made only from rank-invariant values (`tp_size` and the comm type) for the same reason: the previous `total_rows > shard_rows` guard could in principle let some ranks gather while others skipped. Measured for All2All with `num_tokens=241, tp=2` (uneven split): both ranks derive 241, where a `shard_rows * tp_size` derivation gave 242 on rank 0 and 240 on rank 1. Even splits (MC2 with any batch, All2All with a multiple of the TP size, sequence parallel) are unchanged, and the two derivations produce byte-identical captures on the All2All 240- and 241-token prefills. `_gather_tp_sharded_topk_ids` also passes `topk_ids.contiguous()` to the collective. Two alternatives from the issue are deliberately not used: - `n != full_num_tokens` is false for a single-token batch, where MC2 still pads to `tp_size` and splits, so the gather would be skipped and the second rank's row would stay stale. - `prepare_finalize.num_tokens` is the local sequence-parallel shard when `replace_allreduce` is set, not the full batch. Writing the whole gathered tensor, instead of trimming it to a per-rank token count, avoids a second token-count estimate: the padding `prepare()` appends is always a suffix, and the consumer only reads `device_buffer[:num_tokens]` for the step. Note (pre-existing, not changed here): `All2All.prepare` computes `pad_size = tp_size - num_tokens`, which pads only when the batch is smaller than the TP group. A batch with `num_tokens >= tp_size` and `num_tokens % tp_size != 0` is therefore split into unequal chunks. Deriving the total above keeps this capture consistent with `All2All.finalize`, which uses the same pre-split shape, but if the intent of that line was "pad to a multiple of the TP size" it should read `(-num_tokens) % tp_size`. Worth a separate look. Tests: `tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py` gains `test_moe_routing_replay_covers_every_token`, which drives both TP-sharded strategies from one engine and asserts the row count plus that no row is empty. The strategy is picked by `num_tokens <= mc2_tokens_capacity` and, in eager mode without CUDA-graph sizes, that capacity is `max_num_reqs * decode_query_len` (measured: 8 with `max_num_seqs=8`). The test therefore pins `max_num_seqs=8` and uses two prompts: a 6-token one that routes through MC2 and a 49-token one that routes through ALLTOALL (both confirmed with `MoE comm method selected: ... method=MoECommType.鈥). 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 could no longer tell the two apart. Measured on Qwen3-30B-A3B: on the unpatched code both prompts fail, on `rows [3, 4, 5] of 13` for the MC2 prompt and on `rows [25..48] of 56` for the ALLTOALL one; with this patch both pass. The new case is restricted to Qwen3-30B-A3B: what it pins is model independent, and the case above already loads Qwen3.5-35B-A3B, so loading it a second time would only double the file's runtime. No unit tests are added; the existing unit suite is unchanged and still passes. Signed-off-by: wuruiyun <1251540982@qq.com> --- .../consistency/test_moe_routing_replay.py | 66 ++++++++++++++ .../worker/patch_routed_experts_capture.py | 87 ++++++++++++++++++- 2 files changed, 151 insertions(+), 2 deletions(-) diff --git a/tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py b/tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py index a1744271ad47..9460d1fe3357 100644 --- a/tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py +++ b/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" + ) diff --git a/vllm_ascend/patch/worker/patch_routed_experts_capture.py b/vllm_ascend/patch/worker/patch_routed_experts_capture.py index 0b2958bb40a0..1253a8d2a983 100644 --- a/vllm_ascend/patch/worker/patch_routed_experts_capture.py +++ b/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()) From 917fbac4a39d18fbc78e688a8c06ea7c7687ffc8 Mon Sep 17 00:00:00 2001 From: wuruiyun <1251540982@qq.com> Date: Thu, 15 Oct 2026 19:32:46 +0800 Subject: [PATCH 2/2] [CI][Test] Reuse the pre-existing R3 e2e case for both TP-sharded strategies Signed-off-by: wuruiyun <1251540982@qq.com> --- .../consistency/test_moe_routing_replay.py | 85 +++++-------------- 1 file changed, 23 insertions(+), 62 deletions(-) diff --git a/tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py b/tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py index 9460d1fe3357..e3ba5dd3b890 100644 --- a/tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py +++ b/tests/e2e/pull_request/two_card/rlhf/consistency/test_moe_routing_replay.py @@ -23,94 +23,55 @@ "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"}) def test_moe_routing_replay(model): - """Routed experts output is valid and non-empty for MoE models (MRV2).""" - prompts = ["Hello, please introduce yourself."] - with VllmRunner( - model, - tensor_parallel_size=2, - enable_expert_parallel=True, - cudagraph_capture_sizes=[1, 2, 4, 8], - distributed_executor_backend="mp", - enable_return_routed_experts=True, - async_scheduling=False, - ) as vllm_model: - sampling_params = SamplingParams( - max_tokens=5, - temperature=0.8, - top_p=0.95, - output_kind=RequestOutputKind.FINAL_ONLY, - ) - inputs = vllm_model.get_inputs(prompts=prompts) - outputs = vllm_model.model.generate(prompts=inputs, sampling_params=sampling_params) - assert outputs[0].finished - assert len(outputs[0].outputs[0].text) > 0 - 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. + """Every token row carries its own routing decision (regression for #15451). - 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. + With DP == 1 and TP > 1 the TP-sharded strategies still pad and split the + token dimension inside the TP group, so ``select_experts`` only routes this + rank's shard; a capture that treats that shard as the whole batch leaves the + rows owned by the other TP ranks unwritten. - 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``. + Both strategies are covered from this one engine. ``set_moe_comm_method`` + picks MC2 for a batch of at most ``mc2_tokens_capacity`` tokens and ALLTOALL + above it, and ``set_mc2_tokens_capacity`` derives that capacity from + ``max_cudagraph_capture_size`` when ``cudagraph_capture_sizes`` is set -- 8 + tokens here, split over the two TP ranks. The prompts are submitted one at a + time so each prefill is a step of its own: the short one prefills inside the + size-8 graph (still MC2), the long one goes through ALLTOALL. """ - 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 + "Hello, please introduce yourself. " * 8, # 49 tokens -> ALLTOALL ] with VllmRunner( model, tensor_parallel_size=2, enable_expert_parallel=True, - enforce_eager=True, - max_num_seqs=8, + cudagraph_capture_sizes=[1, 2, 4, 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, + max_tokens=5, + temperature=0.8, + top_p=0.95, 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 outputs[0].finished + completion = outputs[0].outputs[0] + assert len(completion.text) > 0 + routed_experts = completion.routed_experts assert routed_experts is not None + assert routed_experts.size > 0 # 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 + expected_rows = len(outputs[0].prompt_token_ids) + len(completion.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, (