@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 980c16c8e4

[PCP][Spec Decode] Adds PCP support for single-module MTP and replicated DSpark. (#56107)

Signed-off-by: Summer Yang <girasoleyang@gmail.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Lucas Wilkinson committed Sep 11, 2026 at 01:07 UTC 980c16c8e4c66dce0bc4d355e4ec7fa2ae7b5141

11 files changed +159 -79

tests/v1/spec_decode/test_eagle_draft_attn_metadata.py

@@ -48,6 +48,7 @@ def _make_fake_speculator(

48 )

49 return SimpleNamespace(

50 arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),

51 + draft_is_prefilling=torch.zeros(max_num_reqs, dtype=torch.bool),

52 block_tables=fake_block_tables,

53 input_buffers=fake_input_buffers,

54 attn_groups=[],

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -94,6 +94,7 @@ def _make_speculator(

95 speculator = object.__new__(_TestSpeculator)

96 speculator.supports_mm_inputs = False

97 + speculator.pcp_manager = None

98 speculator.vllm_config = None

99 speculator.input_buffers = SimpleNamespace(

100 input_ids=torch.arange(4),

tests/v1/worker/test_gpu_pcp_manager.py

+1 -10

@@ -11,7 +11,6 @@ from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils

11 from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module

12 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers

13 from vllm.v1.worker.gpu.pcp_manager import PCPManager

14 -from vllm.v1.worker.gpu.states import RequestState

16 def _copy_to_cpu(value, out=None, device=None):

@@ -125,6 +124,7 @@ def _make_global_decode_batch(

124

125 base = InputBatch.make_dummy(num_reqs, num_tokens, buffers)

126 buffers.seq_lens[:num_reqs] = torch.from_numpy(seq_lens_np).to(device)

127 + buffers.positions[:num_reqs] = torch.tensor(num_computed_tokens, device=device)

128 query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32)

129 buffers.query_start_loc[: num_reqs + 1] = torch.from_numpy(query_start_loc_np).to(

130 device

@@ -166,19 +166,10 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():

166 the field afterwards from the PCP-owned buffers.

167 """

168 device = torch.device("cuda:0")

169 - req_states = RequestState(

170 - max_num_reqs=4,

171 - max_model_len=64,

172 - max_num_batched_tokens=8,

173 - num_speculative_steps=1,

174 - vocab_size=8,

175 - device=device,

176 - )

169 manager = PCPManager(

170 pcp_world_size=2,

171 pcp_rank=0,

172 device=device,

181 - req_states=req_states,

173 max_num_reqs=4,

174 max_num_tokens=8,

175 dcp_world_size=2,

vllm/model_executor/layers/sparse_attn_indexer.py

+3 -2

@@ -401,9 +401,10 @@ def sparse_attn_indexer(

401 # out-of-bounds reads in the kernel.

402 # Keep PCP padding so every rank contributes the same all-gather shape.

403 num_tokens = slot_mapping.shape[0]

404 - if use_pcp:

405 - num_tokens //= get_pcp_group().world_size

404 if k is not None:

405 + # MTP draft decodes are replicated and use an unexpanded slot mapping.

406 + if use_pcp and num_tokens > k.shape[0]:

407 + num_tokens //= get_pcp_group().world_size

408 k = k[:num_tokens]

409

410 if not skip_k_cache_insert:

vllm/v1/worker/cp_utils.py

+7 -3

@@ -18,16 +18,20 @@ else:

18 logger = init_logger(__name__)

21 -def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:

21 +def check_attention_cp_compatibility(

22 + vllm_config: VllmConfig,

23 + target_layer_names: set[str] | None = None,

24 +) -> None:

25 pcp_size = vllm_config.parallel_config.prefill_context_parallel_size

26 dcp_size = vllm_config.parallel_config.decode_context_parallel_size

27 interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size

28 if pcp_size * dcp_size > 1:

29 layer_type = cast(type[Any], AttentionLayerBase)

30 layers = get_layers_from_vllm_config(vllm_config, layer_type)

28 - for layer in layers.values():

31 + for layer_name, layer in layers.items():

32 + check_pcp = target_layer_names is None or layer_name in target_layer_names

33 get_attn_backend = getattr(layer, "get_attn_backend", None)

30 - if pcp_size > 1 and get_attn_backend is not None:

34 + if pcp_size > 1 and check_pcp and get_attn_backend is not None:

35 backend = get_attn_backend()

36 assert backend.supports_pcp(), (

37 "PCP requires attention backend support, "

vllm/v1/worker/gpu/model_runner.py

+12 -4

@@ -656,7 +656,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

656 self.vllm_config,

657 self.device,

658 self.supports_mm_inputs,

659 - self.req_states,

659 self.block_tables,

660 cls=self.pcp_manager_cls,

661 )

@@ -668,6 +667,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

667 self.kv_cache_config,

668 self.max_num_reqs,

669 )

670 + if self.speculator is not None:

671 + self.speculator.pcp_manager = self.pcp_manager

672 initialize_mamba_ssu_backend(

673 self.vllm_config.mamba_config,

674 self.kv_cache_config,

@@ -697,7 +698,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

698 self.fast_prefill = FastPrefillHelper(

699 self.cudagraph_manager, self.max_num_tokens

700 )

700 - check_attention_cp_compatibility(self.vllm_config)

701 + check_attention_cp_compatibility(self.vllm_config, target_attn_layer_names)

702 if isinstance(self.speculator, DraftModelSpeculator):

703 # HACK(woosuk)

704 self.speculator.set_attn(

@@ -1964,9 +1965,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1965 return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)

1966

1967 # Last rank: sample tokens

1968 + draft_hidden_states = hidden_states

1969 + assert draft_hidden_states is not None

1970 hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(

1971 self.pcp_manager, hidden_states, input_batch

1972 )

1973 + if self.pcp_manager is not None and aux_hidden_states is not None:

1974 + aux_hidden_states = [

1975 + self.pcp_manager.restore_hidden_states(states)

1976 + for states in aux_hidden_states

1977 + ]

1978

1979 sampler_output, num_sampled, num_rejected = self.sample(

1980 hidden_states, input_batch, grammar_output

@@ -2043,10 +2051,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2051 # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The

2052 # target returns a persistent buffer sized at max_num_batched_tokens;

2053 # slice to the active token count that propose() expects.

2046 - spec_hidden_states = hidden_states

2054 + spec_hidden_states = draft_hidden_states

2055 if hasattr(self.model, "get_mtp_target_hidden_states"):

2056 pre_hc_hidden_states = self.model.get_mtp_target_hidden_states()

2049 - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr]

2057 + spec_hidden_states = pre_hc_hidden_states[: draft_hidden_states.size(0)]

2058 with use_workspace_lane(self._draft_workspace_lane):

2059 draft_tokens = self.speculator.propose(

2060 input_batch,

vllm/v1/worker/gpu/pcp_manager.py

+84 -53

@@ -6,7 +6,7 @@ from dataclasses import dataclass, replace

6 import numpy as np

7 import torch

9 -from vllm.config import CUDAGraphMode, VllmConfig

9 +from vllm.config import VllmConfig

10 from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group

11 from vllm.logger import init_logger

12 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

15 from vllm.v1.worker.gpu.input_batch import (

16 InputBatch,

17 InputBuffers,

18 - combine_sampled_and_draft_tokens,

19 - prepare_pos_seq_lens,

18 )

21 -from vllm.v1.worker.gpu.states import RequestState

20 logger = init_logger(__name__)

@@ -48,7 +45,6 @@ class PCPManager:

45 pcp_world_size: int,

46 pcp_rank: int,

47 device: torch.device,

51 - req_states: RequestState | None = None,

48 max_num_reqs: int | None = None,

49 max_num_tokens: int | None = None,

50 block_tables: BlockTables | None = None,

@@ -64,7 +60,9 @@ class PCPManager:

60 self.cp_interleave = cp_interleave

62 self._global_batch: InputBatch | None = None

67 - self._req_states = req_states

63 + self._local_batch: InputBatch | None = None

64 + self._local_gather_idx: torch.Tensor | None = None

65 + self.draft_prefill_batch: InputBatch | None = None

66 self._block_tables = block_tables

67 self._hidden_restore_idx: torch.Tensor | None = None

68 self._padded_gather_idx: torch.Tensor | None = None

@@ -77,11 +75,6 @@ class PCPManager:

75 if max_num_local_reqs is not None and max_num_tokens is not None

76 else None

77 )

80 - self._local_req_idx = (

81 - torch.arange(max_num_local_reqs, dtype=torch.int32, device=device)

82 - if max_num_local_reqs is not None

83 - else None

84 - )

78 self._local_block_tables: tuple[torch.Tensor, ...] | None

79 self._local_block_table_ptrs: torch.Tensor | None

80 if block_tables is not None and max_num_local_reqs is not None:

@@ -144,19 +137,23 @@ class PCPManager:

137 raise NotImplementedError("MRV2 PCP does not support MM inputs yet.")

138 if vllm_config.lora_config is not None:

139 raise NotImplementedError("MRV2 PCP does not support LoRA yet.")

147 - if vllm_config.speculative_config is not None:

148 - raise NotImplementedError(

149 - "MRV2 PCP does not support speculative decoding yet."

150 - )

151 - is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk")

152 - if (

153 - is_sparse_mla

154 - and vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

155 - ):

156 - raise NotImplementedError(

157 - "MRV2 sparse MLA PCP does not support CUDA graphs yet. "

158 - "Set -cc.cudagraph_mode=NONE."

159 - )

140 + speculative_config = vllm_config.speculative_config

141 + if speculative_config is not None:

142 + if speculative_config.use_dspark():

143 + dcp_size = parallel_config.decode_context_parallel_size

144 + if dcp_size not in (1, pcp_size):

145 + raise NotImplementedError(

146 + "MRV2 PCP DSpark requires DCP=1 or DCP=PCP; got "

147 + f"DCP={dcp_size}, PCP={pcp_size}."

148 + )

149 + elif (

150 + speculative_config.method != "mtp"

151 + or speculative_config.use_multi_module_mtp()

152 + ):

153 + raise NotImplementedError(

154 + "MRV2 PCP only supports DSpark or single-module MTP "

155 + "speculative decoding."

156 + )

157 if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs():

158 raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.")

159

@@ -358,12 +355,8 @@ class PCPManager:

355 input_batch: InputBatch,

356 padded_num_tokens: int | None = None,

357 ) -> InputBatch:

361 - assert self._req_states is not None

358 assert self._input_buffers is not None

363 - req_states = self._req_states

359 input_buffers = self._input_buffers

365 - if input_batch.num_draft_tokens > 0:

366 - raise NotImplementedError("MRV2 PCP does not support spec decode yet.")

360

361 global_batch = input_batch

362 self._global_batch = global_batch

@@ -458,12 +451,21 @@ class PCPManager:

451 local_gather_idx = self._padded_gather_idx[

452 rank_token_start : rank_token_start + num_local_tokens_padded

453 ]

454 + self._local_gather_idx = local_gather_idx

455 torch.index_select(

456 global_batch.input_ids,

457 0,

458 local_gather_idx,

459 out=input_buffers.input_ids[:num_local_tokens_padded],

460 )

461 + # Keep the GPU request-state cursor materialized by prepare_inputs().

462 + # The CPU cursor can lag after speculative rejection.

463 + torch.index_select(

464 + global_batch.positions,

465 + 0,

466 + local_gather_idx,

467 + out=input_buffers.positions[:num_local_tokens_padded],

468 + )

469

470 local_query_start_loc_np = np.empty(

471 input_buffers.max_num_reqs + 1, dtype=np.int32

@@ -478,17 +480,16 @@ class PCPManager:

480 local_to_global_req_idx = async_copy_to_gpu(

481 local_to_global_req_idx_np, device=self.device

482 )

481 - local_start_pos = async_copy_to_gpu(local_start_pos_np, device=self.device)

482 -

483 - assert self._local_req_idx is not None

484 - prepare_pos_seq_lens(

485 - self._local_req_idx[:num_local_reqs],

486 - local_query_start_loc,

487 - local_start_pos,

488 - input_buffers.positions,

489 - input_buffers.seq_lens[:num_local_reqs],

490 - )

483 seq_lens = input_buffers.seq_lens[:num_local_reqs]

484 + if num_local_tokens > 0:

485 + local_end_positions = torch.index_select(

486 + input_buffers.positions,

487 + 0,

488 + local_query_start_loc[1:] - 1,

489 + )

490 + seq_lens.copy_(local_end_positions + 1)

491 + else:

492 + seq_lens.zero_()

493 is_padding = input_buffers.is_padding[:num_local_tokens_padded]

494 is_padding[:num_local_tokens].fill_(False)

495 is_padding[num_local_tokens:].fill_(True)

@@ -511,18 +512,9 @@ class PCPManager:

512 cu_num_logits = torch.zeros(

513 num_local_reqs + 1, device=self.device, dtype=torch.int32

514 )

514 - logits_indices = combine_sampled_and_draft_tokens(

515 - input_buffers.input_ids,

516 - local_to_global_req_idx,

517 - req_states.last_sampled_tokens,

518 - local_query_start_loc,

519 - seq_lens,

520 - req_states.prefill_len.gpu,

521 - req_states.draft_tokens,

522 - cu_num_logits,

523 - total_num_logits,

524 - 1,

525 - )

515 + # Local logits are never sampled. The complete hidden-state tensor is

516 + # restored first and sampled with the untouched global InputBatch.

517 + logits_indices = local_query_start_loc[1:] - 1

518

519 local_prefill_len_np = global_batch.prefill_len_np[

520 local_to_global_batch_req_idx_np

@@ -536,7 +528,7 @@ class PCPManager:

528 seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32)

529 seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens

530

539 - return replace(

531 + self._local_batch = replace(

532 input_batch,

533 req_ids=local_req_ids,

534 num_reqs=num_local_reqs,

@@ -570,6 +562,7 @@ class PCPManager:

562 cu_num_logits_np=cu_num_logits_np,

563 prompt_lens=None,

564 )

565 + return self._local_batch

566

567 def prepare_attn(

568 self, input_batch: InputBatch

@@ -640,6 +633,46 @@ class PCPManager:

633 gathered = get_pcp_group().all_gather(hidden_states, dim=0)

634 return gathered[self._hidden_restore_idx]

635

636 + def get_draft_input_buffers(

637 + self, input_buffers: InputBuffers

638 + ) -> InputBatch | InputBuffers:

639 + return self.draft_prefill_batch or input_buffers

640 +

641 + def prepare_draft_prefill(

642 + self,

643 + input_batch: InputBatch,

644 + input_ids: torch.Tensor,

645 + ) -> None:

646 + if input_batch is not self._global_batch or self._local_batch is None:

647 + return

648 + local_batch = self._local_batch

649 + assert self._local_gather_idx is not None

650 + num_local_tokens = self._local_gather_idx.shape[0]

651 + torch.index_select(

652 + input_ids,

653 + 0,

654 + self._local_gather_idx,

655 + out=local_batch.input_ids[:num_local_tokens],

656 + )

657 + self.draft_prefill_batch = local_batch

658 +

659 + def restore_draft_prefill(

660 + self,

661 + last_hidden_states: torch.Tensor,

662 + hidden_states: torch.Tensor,

663 + ) -> tuple[torch.Tensor, torch.Tensor]:

664 + if self.draft_prefill_batch is None:

665 + return last_hidden_states, hidden_states

666 + local_last_hidden_states = last_hidden_states

667 + last_hidden_states = self.restore_hidden_states(local_last_hidden_states)

668 + hidden_states = (

669 + last_hidden_states

670 + if local_last_hidden_states is hidden_states

671 + else self.restore_hidden_states(hidden_states)

672 + )

673 + self.draft_prefill_batch = None

674 + return last_hidden_states, hidden_states

675 +

676 def restore_for_sampling(

677 self,

678 hidden_states: torch.Tensor,

@@ -686,7 +719,6 @@ def maybe_build_pcp_manager(

719 vllm_config: VllmConfig,

720 device: torch.device,

721 supports_mm_inputs: bool,

689 - req_states: RequestState,

722 block_tables: BlockTables,

723 cls: type[PCPManager] = PCPManager,

724 ) -> PCPManager | None:

@@ -705,7 +737,6 @@ def maybe_build_pcp_manager(

737 pcp_world_size=pcp_size,

738 pcp_rank=pcp_rank,

739 device=device,

708 - req_states=req_states,

740 max_num_reqs=vllm_config.scheduler_config.max_num_seqs,

741 max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,

742 block_tables=block_tables,

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

+29 -5

@@ -250,8 +250,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

250 )

251 else:

252 hidden_states = last_hidden_states

253 - self.hidden_states[:num_tokens_padded].copy_(hidden_states)

254 -

253 self._copy_request_inputs(

254 num_reqs,

255 input_batch.idx_mapping,

@@ -272,6 +270,17 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

270 self.max_num_reqs,

271 )

272

273 + if self.pcp_manager is not None:

274 + self.pcp_manager.prepare_draft_prefill(

275 + input_batch,

276 + self.input_buffers.input_ids[:num_tokens_padded],

277 + )

278 + prefill = self.pcp_manager.draft_prefill_batch

279 + if prefill is not None:

280 + num_tokens = prefill.num_tokens

281 + num_tokens_padded = prefill.num_tokens_after_padding

282 + self.hidden_states[:num_tokens_padded].copy_(hidden_states)

283 +

284 # When all requests are decoding (no true prefills), each has

285 # num_speculative_steps + 1 tokens, enabling FULL graph replay.

286 uniform_token_count = get_uniform_decode_token_count(

@@ -324,6 +333,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

333 # Early exit.

334 return self.draft_tokens[:num_reqs, :1]

335

336 + if self.pcp_manager is not None and not dummy_run:

337 + self.block_tables.gather_block_tables(

338 + input_batch.idx_mapping, num_reqs_padded=num_reqs

339 + )

340 +

341 # Prepare the inputs for the decode steps.

342 prepare_decode_inputs(

343 self.draft_tokens[:num_reqs, 0],

@@ -387,6 +401,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

401 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

402 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

403 ) -> tuple[torch.Tensor, torch.Tensor]:

404 + input_buffers: InputBatch | InputBuffers = self.input_buffers

405 + is_padding = None

406 + if self.pcp_manager is not None:

407 + input_buffers = self.pcp_manager.get_draft_input_buffers(self.input_buffers)

408 + is_padding = input_buffers.is_padding[:num_tokens]

409 batch_descriptor = BatchDescriptor(num_tokens=num_tokens)

410 with set_forward_context(

411 attn_metadata,

@@ -396,6 +415,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

415 num_tokens_across_dp=num_tokens_across_dp,

416 slot_mapping=slot_mappings,

417 batch_descriptor=batch_descriptor,

418 + is_padding=is_padding,

419 ):

420 inputs_embeds = None

421 if self.supports_mm_inputs:

@@ -406,15 +426,15 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

426 is_mm_embed.shape[0] if is_mm_embed is not None else num_tokens

427 )

428 self.inputs_embeds[:num_input_tokens] = self.model.embed_input_ids(

409 - self.input_buffers.input_ids[:num_input_tokens],

429 + input_buffers.input_ids[:num_input_tokens],

430 multimodal_embeddings=mm_embeds,

431 is_multimodal=is_mm_embed,

432 )

433 inputs_embeds = self.inputs_embeds[:num_tokens]

434

435 model_inputs = dict(

416 - input_ids=self.input_buffers.input_ids[:num_tokens],

417 - positions=self.input_buffers.positions[:num_tokens],

436 + input_ids=input_buffers.input_ids[:num_tokens],

437 + positions=input_buffers.positions[:num_tokens],

438 hidden_states=self.hidden_states[:num_tokens],

439 inputs_embeds=inputs_embeds,

440 )

@@ -463,6 +483,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

483 cudagraph_runtime_mode=cudagraph_runtime_mode,

484 mm_inputs=mm_inputs,

485 )

486 + if self.pcp_manager is not None:

487 + last_hidden_states, hidden_states = self.pcp_manager.restore_draft_prefill(

488 + last_hidden_states, hidden_states

489 + )

490

491 sample_hidden_states = last_hidden_states[last_token_indices]

492 self.draft_tokens[:num_reqs, 0] = self.sample_draft(

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

+12

@@ -38,6 +38,13 @@ class DFlashSpeculator(DraftModelSpeculator):

38 _speculator_name = "DFlash" # For logging, so we can share methods with subclasses

40 def __init__(self, vllm_config: VllmConfig, device: torch.device):

41 + parallel_config = vllm_config.parallel_config

42 + if parallel_config.prefill_context_parallel_size > 1:

43 + vllm_config = copy.copy(vllm_config)

44 + vllm_config.parallel_config = replace(

45 + parallel_config,

46 + prefill_context_parallel_size=1,

47 + )

48 super().__init__(vllm_config, device)

50 self.hidden_states = torch.zeros(

@@ -388,6 +395,11 @@ class DFlashSpeculator(DraftModelSpeculator):

395 )

396 return self.draft_tokens[:num_reqs]

397

398 + if self.pcp_manager is not None and not dummy_run:

399 + self.block_tables.gather_block_tables(

400 + input_batch.idx_mapping, num_reqs_padded=num_reqs

401 + )

402 +

403 # The query slot mapping is written into the shared BlockTables slot_mappings.

404 # That buffer's address is what the captured CUDA graph reads from at replay.

405 assert self.draft_kv_cache_group_id >= 0

vllm/v1/worker/gpu/spec_decode/mtp/speculator.py

+2 -1

@@ -28,7 +28,8 @@ class MTPSpeculator(AutoRegressiveSpeculator):

28 # toggles skip_topk so step 0 computes MTP's own indices and

29 # steps 1+ reuse them.

30 self.share_mtp_topk_indices = (

31 - getattr(draft_hf_config, "index_share_for_mtp_iteration", False)

31 + self.vllm_config.parallel_config.prefill_context_parallel_size == 1

32 + and getattr(draft_hf_config, "index_share_for_mtp_iteration", False)

33 and hasattr(draft_model.model, "set_skip_topk")

34 and hasattr(draft_model.model, "compact_topk_indices")

35 )

vllm/v1/worker/gpu/spec_decode/speculator.py

+7 -1

@@ -3,7 +3,7 @@

3 from abc import ABC, abstractmethod

4 from collections.abc import Mapping

5 from dataclasses import replace

6 -from typing import Any

6 +from typing import TYPE_CHECKING, Any

8 import numpy as np

9 import torch

@@ -29,6 +29,9 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState

29 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

30 from vllm.v1.worker.utils import AttentionGroup

32 +if TYPE_CHECKING:

33 + from vllm.v1.worker.gpu.pcp_manager import PCPManager

34 +

35 logger = init_logger(__name__)

@@ -149,6 +152,7 @@ class DraftModelSpeculator(BaseSpeculator):

152 self.arange = torch.arange(

153 self.max_num_reqs + 1, dtype=torch.int32, device="cpu"

154 )

155 + self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool)

156

157 self.draft_logits: torch.Tensor | None = None

158 if self.speculative_config.draft_sample_method == "probabilistic":

@@ -166,6 +170,7 @@ class DraftModelSpeculator(BaseSpeculator):

170 )

171

172 self.supports_mm_inputs = False

173 + self.pcp_manager: PCPManager | None = None

174

175 @abstractmethod

176 def load_draft_model(

@@ -326,6 +331,7 @@ class DraftModelSpeculator(BaseSpeculator):

331 kv_cache_config=self.kv_cache_config,

332 causal=causal,

333 seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound,

334 + is_prefilling=self.draft_is_prefilling[:num_reqs],

335 )

336 return attn_metadata

337