[Suggestion] Emit the PCP-gathered MoE padding mask from partition_batch

tests/v1/worker/test_gpu_pcp_manager.py

@@ -256,7 +256,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
@@ -275,7 +275,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 +736,41 @@ 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():
+    """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."""
+    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.cpu())
+
+    # 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, torch.cat(local_masks))
+
+
 @pytest.mark.parametrize(
     ("pcp_world_size", "dcp_world_size", "expected"),
     [(1, 1, False), (2, 1, True), (2, 2, False)],

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -730,6 +730,9 @@ def ubatch_forward_fn(cg_mode: CUDAGraphMode) -> None:
 
             # Capture with dummy rows marked as padding.
             input_buffers.is_padding.fill_(True)
+            is_padding = input_buffers.is_padding[:num_tokens]
+            if pcp_manager is not None:
+                is_padding = pcp_manager.moe_is_padding(is_padding)
 
             def forward_fn(cg_mode: CUDAGraphMode) -> None:
                 batch_descriptor = None
@@ -747,7 +750,7 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None:
                     num_tokens_across_dp=num_tokens_across_dp,
                     slot_mapping=slot_mappings,
                     batch_descriptor=batch_descriptor,
-                    is_padding=input_buffers.is_padding[:num_tokens],
+                    is_padding=is_padding,
                 ):
                     if cg_mode == CUDAGraphMode.PIECEWISE:
                         # PIECEWISE graph (compiled PW or breakable, chosen inside
@@ -803,7 +806,7 @@ def prepare_inputs_to_capture(
         num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
     )
     if pcp_manager is not None:
-        input_batch = pcp_manager.prepare_inputs_to_capture(input_batch)
+        input_batch = pcp_manager.prepare_inputs_for_capture_or_profile(input_batch)
 
     block_table_provider = pcp_manager or block_tables
     input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs)

vllm/v1/worker/gpu/model_runner.py

@@ -307,7 +307,13 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         # Draft tokens propagation - for spec-dec + struct outputs.
         self.draft_tokens_handler = DraftTokensHandler(self.device)
 
-        self.pcp_manager: pcp.PCPManager | None = None
+        # Built before memory profiling; block tables are attached at KV init.
+        self.pcp_manager = pcp.maybe_build_pcp_manager(
+            self.vllm_config,
+            self.device,
+            self.supports_mm_inputs,
+            cls=self.pcp_manager_cls,
+        )
 
         # Pooling models.
         self.is_pooling_model = self.model_config.runner_type == "pooling"
@@ -695,13 +701,8 @@ def initialize_kv_cache(
             cp_rank=self.dcp_rank,
             cp_interleave=self.cp_interleave,
         )
-        self.pcp_manager = pcp.maybe_build_pcp_manager(
-            self.vllm_config,
-            self.device,
-            self.supports_mm_inputs,
-            self.block_tables,
-            cls=self.pcp_manager_cls,
-        )
+        if self.pcp_manager is not None:
+            self.pcp_manager.set_block_tables(self.block_tables)
         self.ubatch_runner = maybe_build_ubatch_runner(
             self.vllm_config,
             self.device,
@@ -1840,7 +1841,9 @@ def execute_model(
                 # All-zero input_ids route every token to the same experts.
                 input_batch.input_ids.random_(0, self.vocab_size)
             if self.pcp_manager is not None:
-                input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)
+                input_batch = self.pcp_manager.prepare_inputs_for_capture_or_profile(
+                    input_batch
+                )
             if skip_attn_for_dummy_run:
                 assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
                     "Attention metadata must be prepared for dummy runs when using "

vllm/v1/worker/gpu/pcp_manager.py

@@ -56,6 +56,7 @@ def __init__(
         dcp_world_size: int = 1,
         dcp_rank: int = 0,
         cp_interleave: int = 1,
+        moe_gathers_tokens: bool = False,
     ) -> None:
         self.pcp_world_size = pcp_world_size
         self.pcp_rank = pcp_rank
@@ -69,18 +70,33 @@ def __init__(
         self._local_batch: InputBatch | None = None
         self._local_gather_idx: torch.Tensor | None = None
         self.draft_prefill_batch: InputBatch | None = None
-        self._block_tables = block_tables
         self._hidden_restore_idx: torch.Tensor | None = None
         self._padded_gather_idx: torch.Tensor | None = None
         self._gathered_kv_write_mask: torch.Tensor | None = None
         self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device)
+        # Padding mask over the PCP-gathered rows the MoE routes (rank-major).
+        self._moe_is_padding = (
+            torch.ones(max_num_tokens * pcp_world_size, dtype=torch.bool, device=device)
+            if moe_gathers_tokens and max_num_tokens is not None
+            else None
+        )
 
         max_num_local_reqs = 2 * max_num_reqs if max_num_reqs is not None else None
         self._input_buffers = (
             InputBuffers(max_num_local_reqs, max_num_tokens, device)
             if max_num_local_reqs is not None and max_num_tokens is not None
             else None
         )
+        self._max_num_local_reqs = max_num_local_reqs
+        self._max_num_tokens = max_num_tokens
+        self.set_block_tables(block_tables)
+
+    def set_block_tables(self, block_tables: BlockTables | None) -> None:
+        """Attach the KV block tables, which only exist after KV cache init."""
+        self._block_tables = block_tables
+        max_num_local_reqs = self._max_num_local_reqs
+        max_num_tokens = self._max_num_tokens
+        device, pcp_world_size = self.device, self.pcp_world_size
         self._local_block_tables: tuple[torch.Tensor, ...] | None
         self._local_block_table_ptrs: torch.Tensor | None
         if block_tables is not None and max_num_local_reqs is not None:
@@ -593,6 +609,12 @@ def partition_batch(
             input_buffers.positions[:num_local_tokens_padded].masked_fill_(
                 is_padding, 0
             )
+        if self._moe_is_padding is not None:
+            rows = np.arange(num_local_tokens_padded)
+            gathered = rows >= np.array(per_rank_num_tokens)[:, None]
+            is_padding = async_tensor_h2d(
+                gathered.ravel(), out=self._moe_is_padding[: gathered.size]
+            )
 
         total_num_logits = num_local_reqs if num_local_tokens > 0 else 0
         if total_num_logits > 0:
@@ -683,7 +705,9 @@ def partition_batch(
         )
         return self._local_batch
 
-    def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch:
+    def prepare_inputs_for_capture_or_profile(
+        self, input_batch: InputBatch
+    ) -> InputBatch:
         """Stage a capture or dummy batch in persistent PCP input buffers."""
         input_buffers = self.input_buffers
         num_reqs = input_batch.num_reqs_after_padding
@@ -692,8 +716,8 @@ def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch:
             input_batch,
             input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids),
             positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions),
-            is_padding=input_buffers.is_padding[:num_tokens].copy_(
-                input_batch.is_padding
+            is_padding=self.moe_is_padding(
+                input_buffers.is_padding[:num_tokens].copy_(input_batch.is_padding)
             ),
             query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_(
                 input_batch.query_start_loc
@@ -702,6 +726,15 @@ def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch:
         )
         return input_batch
 
+    def moe_is_padding(self, is_padding: torch.Tensor) -> torch.Tensor:
+        """Expand a capture/dummy padding mask, identical on every PCP rank, to
+        the rows the MoE routes."""
+        if self._moe_is_padding is None:
+            return is_padding
+        out = self._moe_is_padding[: self.pcp_world_size * is_padding.shape[0]]
+        out.view(self.pcp_world_size, -1).copy_(is_padding)
+        return out
+
     def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
         assert self._local_block_tables is not None
         return tuple(
@@ -831,7 +864,7 @@ def maybe_build_pcp_manager(
     vllm_config: VllmConfig,
     device: torch.device,
     supports_mm_inputs: bool,
-    block_tables: BlockTables,
+    block_tables: BlockTables | None = None,
     cls: type[PCPManager] = PCPManager,
 ) -> PCPManager | None:
     parallel_config = vllm_config.parallel_config
@@ -856,4 +889,6 @@ def maybe_build_pcp_manager(
         dcp_world_size=dcp_size,
         dcp_rank=dcp_rank,
         cp_interleave=parallel_config.cp_kv_cache_interleave_size,
+        # Without EP all2all kernels, MoERunner all-gathers PCP tokens to route.
+        moe_gathers_tokens=not parallel_config.enable_expert_parallel,
     )