From 980c16c8e4c66dce0bc4d355e4ec7fa2ae7b5141 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Fri, 11 Sep 2026 01:07:53 -0400 Subject: [PATCH] [PCP][Spec Decode] Adds PCP support for single-module MTP and replicated DSpark. (#56107) Signed-off-by: Summer Yang Signed-off-by: Lucas Wilkinson Co-authored-by: Summer Yang Co-authored-by: OpenAI Codex Co-authored-by: OpenAI Codex --- .../test_eagle_draft_attn_metadata.py | 1 + .../test_gpu_autoregressive_speculator.py | 1 + tests/v1/worker/test_gpu_pcp_manager.py | 11 +- .../layers/sparse_attn_indexer.py | 5 +- vllm/v1/worker/cp_utils.py | 10 +- vllm/v1/worker/gpu/model_runner.py | 16 +- vllm/v1/worker/gpu/pcp_manager.py | 137 +++++++++++------- .../spec_decode/autoregressive/speculator.py | 34 ++++- .../gpu/spec_decode/dflash/speculator.py | 12 ++ .../worker/gpu/spec_decode/mtp/speculator.py | 3 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 8 +- 11 files changed, 159 insertions(+), 79 deletions(-) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index b3006081e5..95e263f9de 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py @@ -48,6 +48,7 @@ def _make_fake_speculator( ) return SimpleNamespace( arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"), + draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool), block_tables=fake_block_tables, input_buffers=fake_input_buffers, attn_groups=[], diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index 6d8b8cbdbc..e9c7a3b290 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -94,6 +94,7 @@ def _make_speculator( speculator = object.__new__(_TestSpeculator) speculator.supports_mm_inputs = False + speculator.pcp_manager = None speculator.vllm_config = None speculator.input_buffers = SimpleNamespace( input_ids=torch.arange(4), diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index eea668aafb..2d8941f6ff 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -11,7 +11,6 @@ from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.pcp_manager import PCPManager -from vllm.v1.worker.gpu.states import RequestState def _copy_to_cpu(value, out=None, device=None): @@ -125,6 +124,7 @@ def _make_global_decode_batch( base = InputBatch.make_dummy(num_reqs, num_tokens, buffers) buffers.seq_lens[:num_reqs] = torch.from_numpy(seq_lens_np).to(device) + buffers.positions[:num_reqs] = torch.tensor(num_computed_tokens, device=device) query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) buffers.query_start_loc[: num_reqs + 1] = torch.from_numpy(query_start_loc_np).to( device @@ -166,19 +166,10 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): the field afterwards from the PCP-owned buffers. """ device = torch.device("cuda:0") - req_states = RequestState( - max_num_reqs=4, - max_model_len=64, - max_num_batched_tokens=8, - num_speculative_steps=1, - vocab_size=8, - device=device, - ) manager = PCPManager( pcp_world_size=2, pcp_rank=0, device=device, - req_states=req_states, max_num_reqs=4, max_num_tokens=8, dcp_world_size=2, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 7d656f06e8..bb00a9cfe1 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -401,9 +401,10 @@ def sparse_attn_indexer( # out-of-bounds reads in the kernel. # Keep PCP padding so every rank contributes the same all-gather shape. num_tokens = slot_mapping.shape[0] - if use_pcp: - num_tokens //= get_pcp_group().world_size if k is not None: + # MTP draft decodes are replicated and use an unexpanded slot mapping. + if use_pcp and num_tokens > k.shape[0]: + num_tokens //= get_pcp_group().world_size k = k[:num_tokens] if not skip_k_cache_insert: diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 7e208da393..79651759aa 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -18,16 +18,20 @@ else: logger = init_logger(__name__) -def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: +def check_attention_cp_compatibility( + vllm_config: VllmConfig, + target_layer_names: set[str] | None = None, +) -> None: pcp_size = vllm_config.parallel_config.prefill_context_parallel_size dcp_size = vllm_config.parallel_config.decode_context_parallel_size interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size if pcp_size * dcp_size > 1: layer_type = cast(type[Any], AttentionLayerBase) layers = get_layers_from_vllm_config(vllm_config, layer_type) - for layer in layers.values(): + for layer_name, layer in layers.items(): + check_pcp = target_layer_names is None or layer_name in target_layer_names get_attn_backend = getattr(layer, "get_attn_backend", None) - if pcp_size > 1 and get_attn_backend is not None: + if pcp_size > 1 and check_pcp and get_attn_backend is not None: backend = get_attn_backend() assert backend.supports_pcp(), ( "PCP requires attention backend support, " diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 533fcbe1e7..6a3d2e34e9 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -656,7 +656,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.vllm_config, self.device, self.supports_mm_inputs, - self.req_states, self.block_tables, cls=self.pcp_manager_cls, ) @@ -668,6 +667,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.kv_cache_config, self.max_num_reqs, ) + if self.speculator is not None: + self.speculator.pcp_manager = self.pcp_manager initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config, @@ -697,7 +698,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.fast_prefill = FastPrefillHelper( self.cudagraph_manager, self.max_num_tokens ) - check_attention_cp_compatibility(self.vllm_config) + check_attention_cp_compatibility(self.vllm_config, target_attn_layer_names) if isinstance(self.speculator, DraftModelSpeculator): # HACK(woosuk) self.speculator.set_attn( @@ -1964,9 +1965,16 @@ class GPUModelRunner(LoRAModelRunnerMixin): return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output) # Last rank: sample tokens + draft_hidden_states = hidden_states + assert draft_hidden_states is not None hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( self.pcp_manager, hidden_states, input_batch ) + if self.pcp_manager is not None and aux_hidden_states is not None: + aux_hidden_states = [ + self.pcp_manager.restore_hidden_states(states) + for states in aux_hidden_states + ] sampler_output, num_sampled, num_rejected = self.sample( hidden_states, input_batch, grammar_output @@ -2043,10 +2051,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The # target returns a persistent buffer sized at max_num_batched_tokens; # slice to the active token count that propose() expects. - spec_hidden_states = hidden_states + spec_hidden_states = draft_hidden_states if hasattr(self.model, "get_mtp_target_hidden_states"): pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)] with use_workspace_lane(self._draft_workspace_lane): draft_tokens = self.speculator.propose( input_batch, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index abec427b37..6b97254621 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -6,7 +6,7 @@ from dataclasses import dataclass, replace import numpy as np import torch -from vllm.config import CUDAGraphMode, VllmConfig +from vllm.config import VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -15,10 +15,7 @@ from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.gpu.input_batch import ( InputBatch, InputBuffers, - combine_sampled_and_draft_tokens, - prepare_pos_seq_lens, ) -from vllm.v1.worker.gpu.states import RequestState logger = init_logger(__name__) @@ -48,7 +45,6 @@ class PCPManager: pcp_world_size: int, pcp_rank: int, device: torch.device, - req_states: RequestState | None = None, max_num_reqs: int | None = None, max_num_tokens: int | None = None, block_tables: BlockTables | None = None, @@ -64,7 +60,9 @@ class PCPManager: self.cp_interleave = cp_interleave self._global_batch: InputBatch | None = None - self._req_states = req_states + 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 @@ -77,11 +75,6 @@ class PCPManager: if max_num_local_reqs is not None and max_num_tokens is not None else None ) - self._local_req_idx = ( - torch.arange(max_num_local_reqs, dtype=torch.int32, device=device) - if max_num_local_reqs is not None - else None - ) 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: @@ -144,19 +137,23 @@ class PCPManager: raise NotImplementedError("MRV2 PCP does not support MM inputs yet.") if vllm_config.lora_config is not None: raise NotImplementedError("MRV2 PCP does not support LoRA yet.") - if vllm_config.speculative_config is not None: - raise NotImplementedError( - "MRV2 PCP does not support speculative decoding yet." - ) - is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk") - if ( - is_sparse_mla - and vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - ): - raise NotImplementedError( - "MRV2 sparse MLA PCP does not support CUDA graphs yet. " - "Set -cc.cudagraph_mode=NONE." - ) + speculative_config = vllm_config.speculative_config + if speculative_config is not None: + if speculative_config.use_dspark(): + dcp_size = parallel_config.decode_context_parallel_size + if dcp_size not in (1, pcp_size): + raise NotImplementedError( + "MRV2 PCP DSpark requires DCP=1 or DCP=PCP; got " + f"DCP={dcp_size}, PCP={pcp_size}." + ) + elif ( + speculative_config.method != "mtp" + or speculative_config.use_multi_module_mtp() + ): + raise NotImplementedError( + "MRV2 PCP only supports DSpark or single-module MTP " + "speculative decoding." + ) if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(): raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") @@ -358,12 +355,8 @@ class PCPManager: input_batch: InputBatch, padded_num_tokens: int | None = None, ) -> InputBatch: - assert self._req_states is not None assert self._input_buffers is not None - req_states = self._req_states input_buffers = self._input_buffers - if input_batch.num_draft_tokens > 0: - raise NotImplementedError("MRV2 PCP does not support spec decode yet.") global_batch = input_batch self._global_batch = global_batch @@ -458,12 +451,21 @@ class PCPManager: local_gather_idx = self._padded_gather_idx[ rank_token_start : rank_token_start + num_local_tokens_padded ] + self._local_gather_idx = local_gather_idx torch.index_select( global_batch.input_ids, 0, local_gather_idx, out=input_buffers.input_ids[:num_local_tokens_padded], ) + # Keep the GPU request-state cursor materialized by prepare_inputs(). + # The CPU cursor can lag after speculative rejection. + torch.index_select( + global_batch.positions, + 0, + local_gather_idx, + out=input_buffers.positions[:num_local_tokens_padded], + ) local_query_start_loc_np = np.empty( input_buffers.max_num_reqs + 1, dtype=np.int32 @@ -478,17 +480,16 @@ class PCPManager: local_to_global_req_idx = async_copy_to_gpu( local_to_global_req_idx_np, device=self.device ) - local_start_pos = async_copy_to_gpu(local_start_pos_np, device=self.device) - - assert self._local_req_idx is not None - prepare_pos_seq_lens( - self._local_req_idx[:num_local_reqs], - local_query_start_loc, - local_start_pos, - input_buffers.positions, - input_buffers.seq_lens[:num_local_reqs], - ) seq_lens = input_buffers.seq_lens[:num_local_reqs] + if num_local_tokens > 0: + local_end_positions = torch.index_select( + input_buffers.positions, + 0, + local_query_start_loc[1:] - 1, + ) + seq_lens.copy_(local_end_positions + 1) + else: + seq_lens.zero_() is_padding = input_buffers.is_padding[:num_local_tokens_padded] is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) @@ -511,18 +512,9 @@ class PCPManager: cu_num_logits = torch.zeros( num_local_reqs + 1, device=self.device, dtype=torch.int32 ) - logits_indices = combine_sampled_and_draft_tokens( - input_buffers.input_ids, - local_to_global_req_idx, - req_states.last_sampled_tokens, - local_query_start_loc, - seq_lens, - req_states.prefill_len.gpu, - req_states.draft_tokens, - cu_num_logits, - total_num_logits, - 1, - ) + # Local logits are never sampled. The complete hidden-state tensor is + # restored first and sampled with the untouched global InputBatch. + logits_indices = local_query_start_loc[1:] - 1 local_prefill_len_np = global_batch.prefill_len_np[ local_to_global_batch_req_idx_np @@ -536,7 +528,7 @@ class PCPManager: seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32) seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens - return replace( + self._local_batch = replace( input_batch, req_ids=local_req_ids, num_reqs=num_local_reqs, @@ -570,6 +562,7 @@ class PCPManager: cu_num_logits_np=cu_num_logits_np, prompt_lens=None, ) + return self._local_batch def prepare_attn( self, input_batch: InputBatch @@ -640,6 +633,46 @@ class PCPManager: gathered = get_pcp_group().all_gather(hidden_states, dim=0) 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 + + def prepare_draft_prefill( + self, + input_batch: InputBatch, + input_ids: torch.Tensor, + ) -> None: + if input_batch is not self._global_batch or self._local_batch is None: + return + local_batch = self._local_batch + assert self._local_gather_idx is not None + num_local_tokens = self._local_gather_idx.shape[0] + torch.index_select( + input_ids, + 0, + self._local_gather_idx, + out=local_batch.input_ids[:num_local_tokens], + ) + self.draft_prefill_batch = local_batch + + def restore_draft_prefill( + self, + last_hidden_states: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.draft_prefill_batch is None: + return last_hidden_states, hidden_states + local_last_hidden_states = last_hidden_states + last_hidden_states = self.restore_hidden_states(local_last_hidden_states) + hidden_states = ( + last_hidden_states + if local_last_hidden_states is hidden_states + else self.restore_hidden_states(hidden_states) + ) + self.draft_prefill_batch = None + return last_hidden_states, hidden_states + def restore_for_sampling( self, hidden_states: torch.Tensor, @@ -686,7 +719,6 @@ def maybe_build_pcp_manager( vllm_config: VllmConfig, device: torch.device, supports_mm_inputs: bool, - req_states: RequestState, block_tables: BlockTables, cls: type[PCPManager] = PCPManager, ) -> PCPManager | None: @@ -705,7 +737,6 @@ def maybe_build_pcp_manager( pcp_world_size=pcp_size, pcp_rank=pcp_rank, device=device, - req_states=req_states, max_num_reqs=vllm_config.scheduler_config.max_num_seqs, max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens, block_tables=block_tables, diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index bba4e2b4bd..cab250ff4e 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -250,8 +250,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): ) else: hidden_states = last_hidden_states - self.hidden_states[:num_tokens_padded].copy_(hidden_states) - self._copy_request_inputs( num_reqs, input_batch.idx_mapping, @@ -272,6 +270,17 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): self.max_num_reqs, ) + if self.pcp_manager is not None: + self.pcp_manager.prepare_draft_prefill( + input_batch, + self.input_buffers.input_ids[:num_tokens_padded], + ) + prefill = self.pcp_manager.draft_prefill_batch + if prefill is not None: + num_tokens = prefill.num_tokens + num_tokens_padded = prefill.num_tokens_after_padding + self.hidden_states[:num_tokens_padded].copy_(hidden_states) + # When all requests are decoding (no true prefills), each has # num_speculative_steps + 1 tokens, enabling FULL graph replay. uniform_token_count = get_uniform_decode_token_count( @@ -324,6 +333,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): # Early exit. return self.draft_tokens[:num_reqs, :1] + if self.pcp_manager is not None and not dummy_run: + self.block_tables.gather_block_tables( + input_batch.idx_mapping, num_reqs_padded=num_reqs + ) + # Prepare the inputs for the decode steps. prepare_decode_inputs( self.draft_tokens[:num_reqs, 0], @@ -387,6 +401,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + 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] batch_descriptor = BatchDescriptor(num_tokens=num_tokens) with set_forward_context( attn_metadata, @@ -396,6 +415,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, + is_padding=is_padding, ): inputs_embeds = None if self.supports_mm_inputs: @@ -406,15 +426,15 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): is_mm_embed.shape[0] if is_mm_embed is not None else num_tokens ) self.inputs_embeds[:num_input_tokens] = self.model.embed_input_ids( - self.input_buffers.input_ids[:num_input_tokens], + input_buffers.input_ids[:num_input_tokens], multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed, ) inputs_embeds = self.inputs_embeds[:num_tokens] model_inputs = dict( - input_ids=self.input_buffers.input_ids[:num_tokens], - positions=self.input_buffers.positions[:num_tokens], + input_ids=input_buffers.input_ids[:num_tokens], + positions=input_buffers.positions[:num_tokens], hidden_states=self.hidden_states[:num_tokens], inputs_embeds=inputs_embeds, ) @@ -463,6 +483,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): cudagraph_runtime_mode=cudagraph_runtime_mode, mm_inputs=mm_inputs, ) + if self.pcp_manager is not None: + last_hidden_states, hidden_states = self.pcp_manager.restore_draft_prefill( + last_hidden_states, hidden_states + ) sample_hidden_states = last_hidden_states[last_token_indices] self.draft_tokens[:num_reqs, 0] = self.sample_draft( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index a3e8df7924..845781fab2 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -38,6 +38,13 @@ class DFlashSpeculator(DraftModelSpeculator): _speculator_name = "DFlash" # For logging, so we can share methods with subclasses def __init__(self, vllm_config: VllmConfig, device: torch.device): + parallel_config = vllm_config.parallel_config + if parallel_config.prefill_context_parallel_size > 1: + vllm_config = copy.copy(vllm_config) + vllm_config.parallel_config = replace( + parallel_config, + prefill_context_parallel_size=1, + ) super().__init__(vllm_config, device) self.hidden_states = torch.zeros( @@ -388,6 +395,11 @@ class DFlashSpeculator(DraftModelSpeculator): ) return self.draft_tokens[:num_reqs] + if self.pcp_manager is not None and not dummy_run: + self.block_tables.gather_block_tables( + input_batch.idx_mapping, num_reqs_padded=num_reqs + ) + # The query slot mapping is written into the shared BlockTables slot_mappings. # That buffer's address is what the captured CUDA graph reads from at replay. assert self.draft_kv_cache_group_id >= 0 diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index 4889e9a750..d94702612d 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -28,7 +28,8 @@ class MTPSpeculator(AutoRegressiveSpeculator): # toggles skip_topk so step 0 computes MTP's own indices and # steps 1+ reuse them. self.share_mtp_topk_indices = ( - getattr(draft_hf_config, "index_share_for_mtp_iteration", False) + self.vllm_config.parallel_config.prefill_context_parallel_size == 1 + and getattr(draft_hf_config, "index_share_for_mtp_iteration", False) and hasattr(draft_model.model, "set_skip_topk") and hasattr(draft_model.model, "compact_topk_indices") ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 3c510e8885..2681677c76 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -3,7 +3,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping from dataclasses import replace -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -29,6 +29,9 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.utils import AttentionGroup +if TYPE_CHECKING: + from vllm.v1.worker.gpu.pcp_manager import PCPManager + logger = init_logger(__name__) @@ -149,6 +152,7 @@ class DraftModelSpeculator(BaseSpeculator): self.arange = torch.arange( self.max_num_reqs + 1, dtype=torch.int32, device="cpu" ) + self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool) self.draft_logits: torch.Tensor | None = None if self.speculative_config.draft_sample_method == "probabilistic": @@ -166,6 +170,7 @@ class DraftModelSpeculator(BaseSpeculator): ) self.supports_mm_inputs = False + self.pcp_manager: PCPManager | None = None @abstractmethod def load_draft_model( @@ -326,6 +331,7 @@ class DraftModelSpeculator(BaseSpeculator): kv_cache_config=self.kv_cache_config, causal=causal, seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound, + is_prefilling=self.draft_is_prefilling[:num_reqs], ) return attn_metadata