A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 17:00 UTC
[Spec Decode] Update extract_hidden_states to use deferred kv_connector clear (#37013)
Fynn Schmitt-Ulms committed
Mar 16, 2026 at 09:53 UTC
04bf5a35fa2692aa75e0442791849dd976014ce8
4 files changed
+35
-70
tests/v1/spec_decode/test_extract_hidden_states.py
+21
-33
@@ -252,29 +252,22 @@ def test_propose():
252
]
253
254
# Sampled token IDs from target model
255
- sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
256
-
257
- # Mock scheduler output
258
- mock_scheduler_output = mock.MagicMock()
255
+ sampled_token_ids = torch.tensor(
256
+ [42, 60], dtype=torch.int32, device=device
257
+ ).unsqueeze(-1)
258
259
# Call propose
261
- with mock.patch(
262
- "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
263
- ) as mock_has_kv:
264
- mock_has_kv.return_value = False
265
-
266
- draft_tokens, kv_connector_output = proposer.propose(
267
- sampled_token_ids=sampled_token_ids,
268
- target_hidden_states=target_hidden_states,
269
- common_attn_metadata=common_attn_metadata,
270
- scheduler_output=mock_scheduler_output,
271
- slot_mappings=None,
272
- )
260
+ draft_tokens = proposer.propose(
261
+ sampled_token_ids=sampled_token_ids,
262
+ target_hidden_states=target_hidden_states,
263
+ common_attn_metadata=common_attn_metadata,
264
+ slot_mappings=None,
265
+ )
266
267
# Verify draft tokens match sampled tokens
268
# Shape should be [batch_size, 1] for num_speculative_tokens=1
269
assert draft_tokens.shape == (batch_size, 1)
277
- assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
270
+ assert torch.equal(draft_tokens, sampled_token_ids)
271
272
# Verify the model was called
273
model_mock.assert_called_once()
@@ -326,21 +319,16 @@ def test_propose_different_layer_counts(num_hidden_layers):
319
for _ in range(num_hidden_layers)
320
]
321
329
- sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
330
- mock_scheduler_output = mock.MagicMock()
331
-
332
- with mock.patch(
333
- "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
334
- ) as mock_has_kv:
335
- mock_has_kv.return_value = False
336
-
337
- draft_tokens, _ = proposer.propose(
338
- sampled_token_ids=sampled_token_ids,
339
- target_hidden_states=target_hidden_states,
340
- common_attn_metadata=common_attn_metadata,
341
- scheduler_output=mock_scheduler_output,
342
- slot_mappings=None,
343
- )
322
+ sampled_token_ids = torch.tensor(
323
+ [42, 60], dtype=torch.int32, device=device
324
+ ).unsqueeze(-1)
325
+
326
+ draft_tokens = proposer.propose(
327
+ sampled_token_ids=sampled_token_ids,
328
+ target_hidden_states=target_hidden_states,
329
+ common_attn_metadata=common_attn_metadata,
330
+ slot_mappings=None,
331
+ )
332
333
assert draft_tokens.shape == (batch_size, 1)
346
- assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
334
+ assert torch.equal(draft_tokens, sampled_token_ids)
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+3
-1
@@ -286,7 +286,9 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1):
286
cached_req = self._active_requests[req_id]
287
req_block_ids = self._req_blocks[req_id]
288
289
- assert new_block_ids is not None
289
+ if new_block_ids is None:
290
+ continue
291
+
292
block_ids = new_block_ids[0]
293
294
req_block_ids.extend(block_ids)
vllm/v1/spec_decode/extract_hidden_states.py
+10
-24
@@ -3,26 +3,21 @@
3
4
from __future__ import annotations
5
6
-from contextlib import nullcontext
6
from typing import TYPE_CHECKING
7
8
import torch
9
import torch.nn as nn
10
11
from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config
13
-from vllm.distributed.kv_transfer import has_kv_transfer_group
12
from vllm.forward_context import set_forward_context
13
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
14
from vllm.model_executor.model_loader import get_model
15
from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata
16
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
19
-from vllm.v1.outputs import KVConnectorOutput
17
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
18
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
22
-from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
19
20
if TYPE_CHECKING:
25
- from vllm.v1.core.sched.output import SchedulerOutput
21
from vllm.v1.kv_cache_interface import KVCacheConfig
22
23
PADDING_SLOT_ID = -1
@@ -79,11 +74,10 @@ class ExtractHiddenStatesProposer:
74
sampled_token_ids: torch.Tensor,
75
target_hidden_states: list[torch.Tensor],
76
common_attn_metadata: CommonAttentionMetadata,
82
- scheduler_output: SchedulerOutput,
77
slot_mappings: dict[str, torch.Tensor]
78
| list[dict[str, torch.Tensor]]
79
| None = None,
86
- ) -> tuple[torch.Tensor, KVConnectorOutput | None]:
80
+ ) -> torch.Tensor:
81
"""Propose draft tokens by calling the ExtractHiddenStatesModel model.
82
83
The ExtractHiddenStatesModel caches the hidden states in the KV cache
@@ -99,7 +93,6 @@ class ExtractHiddenStatesProposer:
93
target_hidden_states: List of hidden state tensors from target model
94
(one per aux hidden state layer)
95
common_attn_metadata: Attention metadata
102
- scheduler_output: Scheduler output for KV connector
96
slot_mappings: Slot mappings for KV cache (unused, provided for
97
interface compatibility)
98
@@ -136,22 +129,15 @@ class ExtractHiddenStatesProposer:
129
if num_tokens_across_dp is not None:
130
num_tokens_across_dp[self.dp_rank] = num_input_tokens
131
139
- with (
140
- set_forward_context(
141
- per_layer_attn_metadata,
142
- self.vllm_config,
143
- num_tokens=num_input_tokens,
144
- num_tokens_across_dp=num_tokens_across_dp,
145
- cudagraph_runtime_mode=cudagraph_runtime_mode,
146
- slot_mapping=self._get_slot_mapping(
147
- num_input_tokens, common_attn_metadata.slot_mapping
148
- ),
132
+ with set_forward_context(
133
+ per_layer_attn_metadata,
134
+ self.vllm_config,
135
+ num_tokens=num_input_tokens,
136
+ num_tokens_across_dp=num_tokens_across_dp,
137
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
138
+ slot_mapping=self._get_slot_mapping(
139
+ num_input_tokens, common_attn_metadata.slot_mapping
140
),
150
- (
151
- KVConnectorModelRunnerMixin._get_kv_connector_output(scheduler_output)
152
- if has_kv_transfer_group()
153
- else nullcontext()
154
- ) as kv_connector_output,
141
):
142
self.model(
143
hidden_states=self.hidden_states[:num_input_tokens],
@@ -159,7 +145,7 @@ class ExtractHiddenStatesProposer:
145
146
# Return the sampled tokens as "draft" tokens
147
# Shape: [batch_size, 1] to match num_speculative_tokens=1
162
- return sampled_token_ids.unsqueeze(-1), kv_connector_output
148
+ return sampled_token_ids
149
150
def _get_slot_mapping(
151
self,
vllm/v1/worker/gpu_model_runner.py
+1
-12
@@ -4328,23 +4328,12 @@ class GPUModelRunner(
4328
)
4329
target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states]
4330
4331
- draft_token_ids, drafter_kv_connector_output = self.drafter.propose(
4331
+ draft_token_ids = self.drafter.propose(
4332
sampled_token_ids=sampled_token_ids,
4333
target_hidden_states=target_hidden_states,
4334
common_attn_metadata=common_attn_metadata,
4335
- scheduler_output=scheduler_output,
4335
slot_mappings=slot_mappings,
4336
)
4338
- # Combine KVConnectorOutputs or select the non-empty one
4339
- if self.kv_connector_output and drafter_kv_connector_output:
4340
- self.kv_connector_output = KVConnectorOutput.merge(
4341
- self.kv_connector_output, drafter_kv_connector_output
4342
- )
4343
- else:
4344
- self.kv_connector_output = (
4345
- self.kv_connector_output or drafter_kv_connector_output
4346
- )
4347
-
4337
next_token_ids, valid_sampled_tokens_count = (
4338
self.drafter.prepare_next_token_ids_padded(
4339
common_attn_metadata,