[Bugfix][PCP] Align MoE routing padding with gathered tokens

[wangyicong52]

Purpose

The non-all-to-all PCP path gathers hidden states and router logits across PCP ranks, but the router still reads the rank-local padding mask. With N local rows and P PCP ranks, routing processes P*N rows while the mask still has N entries. CUDA top-k checks that the mask length matches the token count and can raise an is_padding size mismatch.

This PR aligns the MoE padding mask with the PCP-gathered token rows and restores the original rank-local mask in a finally block. The dispatch/combine interfaces and PCP1/all-to-all paths remain unchanged. Regression tests cover PCP1/2/4, the all-to-all bypass, per-forward caching, and exception cleanup.

A search of open PRs covering PCP padding, PCP MoE, and the relevant padding-cache key found no equivalent fix. Nearby PRs #53639, #60049, and #56258 address different paths.

Implementation

• Non-all-to-all PCP routing receives a padding mask with the same rank-major layout and row count as the gathered hidden states.
• The original local padding mask is restored on normal and exceptional exits.
• The gathered mask is cached per forward context and reused across MoE calls.

Validation

Environment

• Model: official mistralai/Mistral-Small-4-119B-2603 native FP8 checkpoint.
• Image: vllm/vllm-openai:nightly-81198e97ba7eee2a22540caaa756b7fdddcb4d93.
• Hardware: two nodes, each with 4x NVIDIA H20 GPUs.

vLLM startup arguments

# Set MASTER_ADDR to the reachable IP address of the master node.
vllm serve /models/Mistral-Small-4-119B-2603 \
  --host 0.0.0.0 \
  --port 19180 \
  --served-model-name mistral-small4-pcp-ab \
  --config-format mistral \
  --load-format mistral \
  --tokenizer-mode mistral \
  --dtype bfloat16 \
  --tensor-parallel-size 2 \
  --prefill-context-parallel-size 2 \
  --pipeline-parallel-size 1 \
  --data-parallel-size 1 \
  --decode-context-parallel-size 1 \
  --distributed-executor-backend mp \
  --master-addr "${MASTER_ADDR}" \
  --master-port 29511 \
  --language-model-only \
  --enforce-eager \
  --attention-backend TRITON_MLA \
  --moe-backend triton \
  --max-model-len 4096 \
  --max-num-batched-tokens 2048 \
  --max-num-seqs 4 \
  --block-size 64 \
  --kv-cache-dtype auto \
  --gpu-memory-utilization 0.80 \
  --no-enable-prefix-caching \
  --no-async-scheduling

Result

Without this fix (baseline):
With the unmodified image and model, all seven native weight shards loaded (29.34 GiB per rank). The unpatched run then failed during vLLM's standard startup memory profiling, before the API server became ready:

RuntimeError: .../topk_softmax_kernels.cu:785, is_padding size mismatch, expected: 4096

With this PR:
Applying the _sequence_parallel_context change from this PR at startup allowed the TP2/PCP2 deployment to start and generate successfully. It passed all five chat smoke requests, including a 2,235-token input and concurrent requests of different lengths.

Using the official tests/evals/gsm8k/gsm8k_eval.py scoring logic, we evaluated all 1,319 test examples with 5-shot prompts, temperature=0, seed=42, max_tokens=256, and concurrency=4. Results:

Run	GSM8K	Invalid answers	Runtime result
Baseline + TP4/PCP1	1191/1319 (90.2957%)	0	Full evaluation completed; 0 API errors and 0 Pod restarts
Baseline + TP2/PCP2	Could not run	—	Failed during startup profiling with the padding-size mismatch above
This PR applied + TP2/PCP2	1197/1319 (90.7506%)	0	Full evaluation completed; 0 API errors and 0 Pod restarts

AI assistance was used to extract the change, prepare regression tests and run the validation described above. The human submitter retains responsibility for line-by-line review and final sign-off.

[rebecca-bell-2001]

This is a Bytebell Pull request verifier bot.
Here is our verification report for this PR

Full report here

Findings — 9 in 2 files

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

major · bug · lines 661–665

Guard omits do_naive_dispatch_combine, so when naive EP dispatch rewrites the token axis before the PCP gather, the mask is only P×local and can still mismatch.

Evidence:

• change vllm/model_executor/layers/fused_moe/runner/moe_runner.py 661-665 — The new mask-gather condition tests only pcp_size>1 and not use_all2all_kernels; do_naive_dispatch_combine is never consulted.
• base vllm/model_executor/layers/fused_moe/runner/moe_runner.py 811-825 — _maybe_dispatch first runs get_ep_group().dispatch_router_logits and only then all_gathers hidden_states/router_logits over the PCP group, so routing rows are P×(EP-dispatched rows).
• base vllm/model_executor/layers/fused_moe/runner/moe_runner.py 796-800 — do_naive_dispatch_combine is a separate gate (dp_size>1 or is_sequence_parallel, and not supports_internal_mk) that can be true at the same time as the PCP condition.

minor · bug · lines 667–673

Cached gather is never invalidated: additional_kwargs lives for the whole context, so a later MoE call whose local mask changed reuses stale rows.

Evidence:

• change vllm/model_executor/layers/fused_moe/runner/moe_runner.py 667-673 — The gathered mask is stored under "pcp_moe_is_padding" and only ever read back; the finally at 686-687 restores ctx.is_padding but leaves the cached tensor in place.
• base vllm/forward_context.py 197-197 — additional_kwargs is a plain per-ForwardContext dict with default_factory=dict; nothing prunes keys from it.
• base vllm/forward_context.py 238-258 — create_forward_context is the only writer of the field and just forwards the caller's dict, so a reused context keeps the previous pass's cached gather.
• base vllm/models/deepseek_v41/nvidia/model.py 951-955 — Model code rewrites forward_context.is_padding per call via sp_padding_mask, so the field the cache keyed on is mutable while the cache is not.

minor · performance · lines 668–672

Gather runs whenever a mask is set, yet every mask consumer gates on VLLM_MOE_SKIP_PADDING, so with that flag off it adds a per-pass PCP collective nobody reads.

Evidence:

• change vllm/model_executor/layers/fused_moe/runner/moe_runner.py 661-672 — Condition is is_padding is not None and PCP non-all2all; envs.VLLM_MOE_SKIP_PADDING is not part of the guard.
• base vllm/model_executor/layers/fused_moe/router/fused_topk_router.py 19-23 — The only reader of ctx.is_padding in the router returns None unless envs.VLLM_MOE_SKIP_PADDING is set.
• base vllm/models/kimi_k3/nvidia/model.py 459-463 — The model-side consumer is guarded the same way, so with the flag off the gathered tensor has no reader.

tests/kernels/moe/test_routing.py

minor · tests · line 80

Every case supplies a tensor mask, so the new local_is_padding is not None guard is never exercised; deleting it leaves this test green.

Evidence:

• change tests/kernels/moe/test_routing.py 53-56,80 — parametrization varies only pcp_size/use_all2all and the mask is always torch.tensor([False, True])
• change vllm/model_executor/layers/fused_moe/runner/moe_runner.py 661-664 — the guard's false branch skips the gather and the restores; no case in this test reaches it
• change tests/kernels/moe/test_routing.py 98-104 — the assertions pin the gathered mask and the one-gather count only for a mask that is present

minor · tests · line 82

dp_metadata is left None, so the sp_local_sizes branch of the wrapped context — entered in real PCP forwards — is never combined with the gathered mask.

Evidence:

• change tests/kernels/moe/test_routing.py 82-99 — ForwardContext({}, {}, {}, is_padding=padding) sets no dp_metadata, so only the nullcontext branch is entered
• base vllm/forward_context.py 299-328 — set_forward_context builds dp_metadata whenever dp_size>1 or use_sequence_parallel_moe, the configurations this fix targets
• base vllm/forward_context.py 109-122 — sp_local_sizes is the CM the new with-block wraps; it rewrites local_sizes, the SP token chunking

minor · tests · line 83

The test re-declares the gather predicate and hard-codes do_naive_dispatch_combine=False, so it cannot catch drift between the mask-gather condition and the real dispatch condition.

Evidence:

• change tests/kernels/moe/test_routing.py 62-67,83-86 — runner is a SimpleNamespace with do_naive_dispatch_combine=False written in, and dispatched restates the predicate instead of reading the runner
• base vllm/model_executor/layers/fused_moe/runner/moe_runner.py 796-800 — do_naive_dispatch_combine is a computed property (dp_size>1 or is_sequence_parallel) and not supports_internal_mk, not the constant the test substitutes
• base vllm/model_executor/layers/fused_moe/runner/moe_runner.py 811-825 — the dispatch predicate under test lives here; the test's copy of it is what the assertions are compared against

minor · tests · lines 94–104

New test only asserts ctx.is_padding and the gather count; nothing exercises the router/top-k path the fix targets, so the reported mismatch stays untested.

Evidence:

• change tests/kernels/moe/test_routing.py 88-104 — Two loop iterations enter the context and assert only mask contents numel, restore identity and padding_gathers; no select_experts/topk call.
• change vllm/model_executor/layers/fused_moe/runner/moe_runner.py 666-673 — The behaviour under test is the cached gather installed as ctx.is_padding, consumed in production by the routers.

minor · tests · line 98

Mask order is pinned only against the mock's own concatenation; a real gather that is not rank-major in the hidden-state order still passes.

Evidence:

• change tests/kernels/moe/test_routing.py 70-75,84-86 — the fake inverts each simulated peer's mask and the expected list is built from that same inversion
• base vllm/model_executor/layers/fused_moe/runner/moe_runner.py 820-825 — the hidden-state gather this mask must line up with is dim=0 rank-major; the test never derives the expectation from it

minor · tests · line 100

This assertion cannot fail: the mock gathers hidden+4 with the same rank offsets as hidden, so logits always equals states+4 by construction.

Evidence:

• change tests/kernels/moe/test_routing.py 70-75 — the fake all_gather adds rank*100 to every tensor it is handed, identically for both tensors
• change tests/kernels/moe/test_routing.py 95-100 — logits is gather(hidden+4) and states is gather(hidden), so logits == states+4 is an identity, not a check

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[wangyicong52]

@LucasWilkinson Could you have a look when you have time? Thanks!

[LucasWilkinson]

Apologies for the delay; it's a holiday here in Canada, I will try to look tomorrow or the day after 👍

[wangyicong52]

Apologies for the delay; it's a holiday here in Canada, I will try to look tomorrow or the day after 👍

No worries at all—thanks for letting me know! Just enjoy the Thanksgiving holiday!