[Bugfix][PCP] Align MoE routing padding with gathered tokens (#60946)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: wangyicong <wangyicong@bytedance.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -132,6 +132,7 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings():
     pcp_block_tables = SimpleNamespace(
         num_kv_cache_groups=1,
         input_block_tables=(torch.zeros(num_reqs * 2, 1, dtype=torch.int32),),
+        cp_interleave=1,
     )
     pcp_manager = PCPManager(
         pcp_world_size=pcp_world_size,

tests/v1/worker/test_gpu_pcp_manager.py

@@ -7,7 +7,9 @@ import numpy as np
 import pytest
 import torch
 
-from vllm.config import CUDAGraphMode, ParallelConfig
+from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens
 from vllm.v1.attention.ops import pcp as attention_pcp
 from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
@@ -48,6 +50,7 @@ def _make_capture_manager(block_table: torch.Tensor):
         input_block_tables=(block_table,),
         num_kv_cache_groups=1,
         block_sizes=(2,),
+        cp_interleave=1,
     )
     manager = PCPManager(
         pcp_world_size=2,
@@ -240,9 +243,10 @@ def test_dummy_draft_does_not_reuse_previous_graph_batch():
 
     manager.prepare_draft_prefill(dummy_batch, dummy_batch.input_ids)
 
-    assert (
-        manager.get_draft_input_buffers(manager.input_buffers) is manager.input_buffers
+    draft_inputs, _ = manager.get_draft_input_buffers(
+        manager.input_buffers, dummy_batch.num_tokens_after_padding
     )
+    assert draft_inputs is manager.input_buffers
 
 
 def test_capture_uses_pcp_persistent_inputs():
@@ -256,7 +260,7 @@ def test_capture_uses_pcp_persistent_inputs():
     )
     dummy_batch.input_ids.copy_(torch.arange(4))
     dummy_batch.positions.fill_(7)
-    input_batch = manager.prepare_inputs_to_capture(dummy_batch)
+    input_batch = manager.prepare_inputs_for_capture_or_profile(dummy_batch)
 
     assert input_batch is not dummy_batch
     assert input_batch.req_ids == dummy_batch.req_ids
@@ -266,6 +270,17 @@ def test_capture_uses_pcp_persistent_inputs():
         assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr()
 
 
+def test_block_table_rebind_updates_interleave():
+    manager, block_tables = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
+    final_block_tables = SimpleNamespace(**vars(block_tables))
+    final_block_tables.cp_interleave = 64
+
+    manager.set_block_tables(final_block_tables)
+
+    assert manager._block_tables is final_block_tables
+    assert manager.cp_interleave == 64
+
+
 def test_dummy_context_updates_pcp_local_block_tables():
     global_block_table = torch.full((4, 4), -1, dtype=torch.int32)
     manager, block_tables = _make_capture_manager(global_block_table)
@@ -275,7 +290,7 @@ def test_dummy_context_updates_pcp_local_block_tables():
         input_buffers=manager.input_buffers,
         max_query_len=1,
     )
-    input_batch = manager.prepare_inputs_to_capture(input_batch)
+    input_batch = manager.prepare_inputs_for_capture_or_profile(input_batch)
     local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs)
 
     set_dummy_context(
@@ -736,6 +751,49 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
     assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
+def test_moe_padding_matches_pcp_all_gather(monkeypatch):
+    """The PCP batch carries the all-gather of every rank's local padding mask,
+    the rows the MoE routes, even when ranks own different numbers of tokens."""
+    monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1")
+    device = torch.device("cuda:0")
+    local_masks, batch_masks = [], []
+    for pcp_rank in range(2):
+        manager = PCPManager(
+            pcp_world_size=2,
+            pcp_rank=pcp_rank,
+            device=device,
+            shard_decode_requests=True,
+            max_num_reqs=4,
+            max_num_tokens=8,
+            moe_gathers_tokens=True,
+        )
+        global_batch = _make_global_decode_batch(
+            [16, 24, 32], InputBuffers(4, 8, device), device
+        )
+        local_batch = manager.partition_batch(
+            global_batch,
+            BatchExecutionDescriptor(
+                cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=4, num_reqs=None
+            ),
+        )
+        num_tokens = local_batch.num_tokens_after_padding
+        local_masks.append(manager.input_buffers.is_padding[:num_tokens].cpu())
+        batch_masks.append(local_batch.is_padding.clone())
+
+    # Sharded decode puts two requests on rank 0 and one on rank 1.
+    assert [mask.logical_not().sum().item() for mask in local_masks] == [2, 1]
+    for batch_mask in batch_masks:
+        assert torch.equal(batch_mask.cpu(), torch.cat(local_masks))
+        num_tokens = batch_mask.numel()
+        hidden_states = torch.randn(num_tokens, 4, device=device)
+        router_logits = torch.randn(num_tokens, 8, device=device)
+        with set_forward_context(None, VllmConfig(), is_padding=batch_mask):
+            _, topk_ids, _ = fused_topk(hidden_states, router_logits, 2, False)
+        assert (topk_ids[~batch_mask] >= 0).all()
+        assert (topk_ids[batch_mask] == -1).all()
+
+
 @pytest.mark.parametrize(
     ("pcp_world_size", "dcp_world_size", "expected"),
     [(1, 1, False), (2, 1, True), (2, 2, False)],

tests/v1/worker/test_gpu_target_dependent_ar_speculator.py

@@ -22,6 +22,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.kv_cache_interface import MLAAttentionSpec
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu.pcp_manager import PCPManager
 from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
 from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import (
     MultiModuleMTPSpeculator,
@@ -397,6 +398,52 @@ def test_run_model_reuses_tensor_return_for_mtp(monkeypatch):
     assert actual_feedback_hidden is hidden
 
 
+@pytest.mark.parametrize("partitioned", [False, True])
+def test_run_model_uses_complete_pcp_padding(monkeypatch, partitioned):
+    hidden = torch.full((4, 3), 1.0)
+    speculator = _make_speculator(monkeypatch, hidden)
+    gathered_is_padding = torch.tensor(
+        [False, False, False, True, False, False, True, True]
+    )
+    manager = PCPManager(
+        pcp_world_size=2,
+        pcp_rank=0,
+        device=torch.device("cpu"),
+        shard_decode_requests=True,
+        max_num_tokens=4,
+        moe_gathers_tokens=True,
+    )
+    if partitioned:
+        manager.draft_prefill_batch = SimpleNamespace(
+            input_ids=torch.arange(4),
+            positions=torch.arange(4),
+            is_padding=gathered_is_padding,
+        )
+    else:
+        local_is_padding = torch.tensor([False, False, False, True])
+        speculator.input_buffers.is_padding = local_is_padding
+        gathered_is_padding = torch.cat([local_is_padding, local_is_padding])
+    seen = {}
+    monkeypatch.setattr(
+        spec_module,
+        "set_forward_context",
+        lambda *args, **kwargs: (seen.update(kwargs), nullcontext())[1],
+    )
+    speculator.pcp_manager = manager
+
+    speculator._run_model(
+        4,
+        attn_metadata=None,
+        slot_mappings=None,
+        num_tokens_across_dp=None,
+        cudagraph_runtime_mode=CUDAGraphMode.NONE,
+    )
+
+    torch.testing.assert_close(seen["is_padding"], gathered_is_padding)
+    if partitioned:
+        assert seen["is_padding"] is gathered_is_padding
+
+
 @pytest.mark.parametrize(
     (
         "method_name",