Skip to content

[Spec Decode] Update extract_hidden_states to use deferred kv_connector clear #37013

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mgoin merged 6 commits into from
Mar 16, 2026
Merged

[Spec Decode] Update extract_hidden_states to use deferred kv_connector clear #37013

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
159911e
Update extract_hidden_states to deferred kv_connector clear
fynnsu Mar 13, 2026
ab64c43
Patch bug in ExampleHiddenStatesConnector that causes occasional crashes
fynnsu Mar 13, 2026
460e38f
lint
fynnsu Mar 13, 2026
c22855b
fix docstring
fynnsu Mar 13, 2026
308b63c
fix missing attribute
fynnsu Mar 13, 2026
30eb4cf
fix shape mismatch
fynnsu Mar 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 21 additions & 33 deletions tests/v1/spec_decode/test_extract_hidden_states.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,29 +252,22 @@ def test_propose():
]

# Sampled token IDs from target model
sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)

# Mock scheduler output
mock_scheduler_output = mock.MagicMock()
sampled_token_ids = torch.tensor(
[42, 60], dtype=torch.int32, device=device
).unsqueeze(-1)

# Call propose
with mock.patch(
"vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
) as mock_has_kv:
mock_has_kv.return_value = False

draft_tokens, kv_connector_output = proposer.propose(
sampled_token_ids=sampled_token_ids,
target_hidden_states=target_hidden_states,
common_attn_metadata=common_attn_metadata,
scheduler_output=mock_scheduler_output,
slot_mappings=None,
)
draft_tokens = proposer.propose(
sampled_token_ids=sampled_token_ids,
target_hidden_states=target_hidden_states,
common_attn_metadata=common_attn_metadata,
slot_mappings=None,
)

# Verify draft tokens match sampled tokens
# Shape should be [batch_size, 1] for num_speculative_tokens=1
assert draft_tokens.shape == (batch_size, 1)
assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
assert torch.equal(draft_tokens, sampled_token_ids)

# Verify the model was called
model_mock.assert_called_once()
Expand Down Expand Up @@ -326,21 +319,16 @@ def test_propose_different_layer_counts(num_hidden_layers):
for _ in range(num_hidden_layers)
]

sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
mock_scheduler_output = mock.MagicMock()

with mock.patch(
"vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
) as mock_has_kv:
mock_has_kv.return_value = False

draft_tokens, _ = proposer.propose(
sampled_token_ids=sampled_token_ids,
target_hidden_states=target_hidden_states,
common_attn_metadata=common_attn_metadata,
scheduler_output=mock_scheduler_output,
slot_mappings=None,
)
sampled_token_ids = torch.tensor(
[42, 60], dtype=torch.int32, device=device
).unsqueeze(-1)

draft_tokens = proposer.propose(
sampled_token_ids=sampled_token_ids,
target_hidden_states=target_hidden_states,
common_attn_metadata=common_attn_metadata,
slot_mappings=None,
)

assert draft_tokens.shape == (batch_size, 1)
assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
assert torch.equal(draft_tokens, sampled_token_ids)
4 changes: 3 additions & 1 deletion vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,9 @@ def build_connector_meta(
cached_req = self._active_requests[req_id]
req_block_ids = self._req_blocks[req_id]

assert new_block_ids is not None
if new_block_ids is None:
continue
Comment on lines +289 to +290

@tlrmchlsmth tlrmchlsmth Mar 13, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why can new_block_ids be None?

@fynnsu fynnsu Mar 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Having difficulty tracing down the root cause for this. It happens when running a lot of async requests at once and infrequently (i.e. w/ 32 simultaneous requests I hit this after processing 700 requests successfully). It is a rare case on a rare section of the code.

I have verified that it is a real request that triggers it (isn't empty or something), but the request still produces the hidden states file output as expected with the correct token ids and hidden state shapes.

@benchislett benchislett Mar 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@fynnsu input_fits_in_drafter failing due to long sequences is not a concern for this style of mock-drafting, right?

@fynnsu fynnsu Mar 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@fynnsu input_fits_in_drafter failing due to long sequences is not a concern for this style of mock-drafting, right?

I think that is actually a risk with the current setup. I believe it would only be triggered if max_seq_len == max_model_size since num_speculative_tokens=1 but with Eagle3 we just fall back to not drafting, which doesn't work for the extract_hidden_states method.

Claude's suggestion is to just always set this to True for this draft method:

There's a more direct fix to the original problem: since extract_hidden_states doesn't actually speculate (the "draft" tokens are always the sampled tokens, always verify), the
input_fits_in_drafter guard is protecting against a scenario that doesn't apply. The check exists because real drafters (Eagle, DraftModel) would write KV entries beyond
max_model_len. But the extract_hidden_states drafter's KV cache is a separate cache-only layer whose writes are indexed by the target model's slot_mapping — it writes at the same
positions the target already wrote at, so the max_model_len concern doesn't arise.

A targeted fix would be to make input_fits_in_drafter always True for extract_hidden_states:

 input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
      spec_config.uses_extract_hidden_states()  # always fits
      or spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
      <= self.effective_drafter_max_model_len
  )

That logic does make sense to me and we shouldn't actually need the extra position of the "drafted" token, we just can't set num_speculative_tokens=0 without updating a whole bunch of guards. What do you think?


block_ids = new_block_ids[0]

req_block_ids.extend(block_ids)
Expand Down
34 changes: 10 additions & 24 deletions vllm/v1/spec_decode/extract_hidden_states.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,26 +3,21 @@

from __future__ import annotations

from contextlib import nullcontext
from typing import TYPE_CHECKING

import torch
import torch.nn as nn

from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config
from vllm.distributed.kv_transfer import has_kv_transfer_group
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
from vllm.model_executor.model_loader import get_model
from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
from vllm.v1.outputs import KVConnectorOutput
from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin

if TYPE_CHECKING:
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.kv_cache_interface import KVCacheConfig

PADDING_SLOT_ID = -1
Expand Down Expand Up @@ -79,11 +74,10 @@ def propose(
sampled_token_ids: torch.Tensor,
target_hidden_states: list[torch.Tensor],
common_attn_metadata: CommonAttentionMetadata,
scheduler_output: SchedulerOutput,
slot_mappings: dict[str, torch.Tensor]
| list[dict[str, torch.Tensor]]
| None = None,
) -> tuple[torch.Tensor, KVConnectorOutput | None]:
) -> torch.Tensor:
"""Propose draft tokens by calling the ExtractHiddenStatesModel model.

The ExtractHiddenStatesModel caches the hidden states in the KV cache
Expand All @@ -99,7 +93,6 @@ def propose(
target_hidden_states: List of hidden state tensors from target model
(one per aux hidden state layer)
common_attn_metadata: Attention metadata
scheduler_output: Scheduler output for KV connector
slot_mappings: Slot mappings for KV cache (unused, provided for
interface compatibility)

Expand Down Expand Up @@ -136,30 +129,23 @@ def propose(
if num_tokens_across_dp is not None:
num_tokens_across_dp[self.dp_rank] = num_input_tokens

with (
set_forward_context(
per_layer_attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
cudagraph_runtime_mode=cudagraph_runtime_mode,
slot_mapping=self._get_slot_mapping(
num_input_tokens, common_attn_metadata.slot_mapping
),
with set_forward_context(
per_layer_attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
cudagraph_runtime_mode=cudagraph_runtime_mode,
slot_mapping=self._get_slot_mapping(
num_input_tokens, common_attn_metadata.slot_mapping
),
(
KVConnectorModelRunnerMixin._get_kv_connector_output(scheduler_output)
if has_kv_transfer_group()
else nullcontext()
) as kv_connector_output,
):
self.model(
hidden_states=self.hidden_states[:num_input_tokens],
)

# Return the sampled tokens as "draft" tokens
# Shape: [batch_size, 1] to match num_speculative_tokens=1
return sampled_token_ids.unsqueeze(-1), kv_connector_output
return sampled_token_ids

def _get_slot_mapping(
self,
Expand Down
13 changes: 1 addition & 12 deletions vllm/v1/worker/gpu_model_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4337,23 +4337,12 @@ def propose_draft_token_ids(
)
target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states]

draft_token_ids, drafter_kv_connector_output = self.drafter.propose(
draft_token_ids = self.drafter.propose(
sampled_token_ids=sampled_token_ids,
target_hidden_states=target_hidden_states,
common_attn_metadata=common_attn_metadata,
scheduler_output=scheduler_output,
slot_mappings=slot_mappings,
)
# Combine KVConnectorOutputs or select the non-empty one
if self.kv_connector_output and drafter_kv_connector_output:
self.kv_connector_output = KVConnectorOutput.merge(
self.kv_connector_output, drafter_kv_connector_output
)
else:
self.kv_connector_output = (
self.kv_connector_output or drafter_kv_connector_output
)

next_token_ids, valid_sampled_tokens_count = (
self.drafter.prepare_next_token_ids_padded(
common_attn_metadata,
Expand Down
Loading