diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index c2d6e4ef7b..cfe77253e2 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/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, diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index fcf171a89c..1adcea91ee 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/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)], diff --git a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py b/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py index 42678aa9d1..580410de5d 100644 --- a/tests/v1/worker/test_gpu_target_dependent_ar_speculator.py +++ b/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", diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index dfbb0aa82c..0a74b685b4 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -743,6 +743,9 @@ class ModelCudaGraphManager(CudaGraphManager): # 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 @@ -760,7 +763,7 @@ class ModelCudaGraphManager(CudaGraphManager): 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 @@ -837,7 +840,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) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 9cd7299642..5ba90b792e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -305,7 +305,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): # 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" @@ -700,13 +706,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): 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, @@ -1848,7 +1849,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): # 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 " diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index aa245fa267..16f2886358 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -56,6 +56,7 @@ class PCPManager: 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,11 +70,16 @@ class PCPManager: 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 = ( @@ -81,6 +87,18 @@ class PCPManager: 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 + if block_tables is not None: + self.cp_interleave = block_tables.cp_interleave + 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 +611,12 @@ class PCPManager: 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 +707,9 @@ class PCPManager: ) 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 +718,8 @@ class PCPManager: 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 +728,15 @@ class PCPManager: ) 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( @@ -777,9 +812,13 @@ class PCPManager: return gathered[self._hidden_restore_idx] def get_draft_input_buffers( - self, input_buffers: InputBuffers - ) -> InputBatch | InputBuffers: - return self.draft_prefill_batch or input_buffers + self, input_buffers: InputBuffers, num_tokens: int + ) -> tuple[InputBatch | InputBuffers, torch.Tensor]: + """Return draft inputs and their MoE routing padding mask.""" + batch = self.draft_prefill_batch + if batch is not None: + return batch, batch.is_padding + return input_buffers, self.moe_is_padding(input_buffers.is_padding[:num_tokens]) def prepare_draft_prefill( self, input_batch: InputBatch, input_ids: torch.Tensor @@ -831,7 +870,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 +895,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, ) diff --git a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py index fdbedd78f6..d5034a31b4 100644 --- a/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py @@ -376,8 +376,9 @@ class TargetDependentARSpeculator(DraftModelSpeculator): input_buffers: InputBatch | InputBuffers = self.input_buffers is_padding = None if self.pcp_manager is not None: - input_buffers = self.pcp_manager.get_draft_input_buffers(self.input_buffers) - is_padding = input_buffers.is_padding[:num_tokens] + input_buffers, is_padding = self.pcp_manager.get_draft_input_buffers( + self.input_buffers, num_tokens + ) batch_descriptor = BatchDescriptor(num_tokens=num_tokens) with set_forward_context( attn_metadata,