diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 89061c6bc1..4b62915ede 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -8,10 +8,14 @@ from typing import Literal, overload from vllm.distributed.kv_events import BlockStored, KVCacheEvent from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import KVCacheBlock from vllm.v1.kv_cache_interface import ( + AttentionSpec, + CrossAttentionSpec, + EncoderOnlyAttentionSpec, KVCacheConfig, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, @@ -596,6 +600,26 @@ class KVCacheManager: """Get the block ids of a request.""" return self.get_blocks(request_id).get_block_ids() + def get_block_ids_for_computed_tokens( + self, + request_id: str, + num_computed_tokens: int, + ) -> tuple[list[int], ...]: + """Get block ids covering the request's computed tokens.""" + block_ids = self.get_block_ids(request_id) + clipped_block_ids: list[list[int]] = [] + for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids): + spec = group.kv_cache_spec + if not isinstance(spec, AttentionSpec) or isinstance( + spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec) + ): + clipped_block_ids.append(ids) + continue + + num_valid_blocks = cdiv(num_computed_tokens, spec.block_size) + clipped_block_ids.append(ids[:num_valid_blocks]) + return tuple(clipped_block_ids) + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: """Cache the blocks for the request, if enabled. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 483cc2f543..9507140887 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -874,12 +874,10 @@ class Scheduler(SchedulerInterface): if num_new_tokens == 0: break - # Handles an edge case when P/D Disaggregation - # is used with Spec Decoding where an - # extra block gets allocated which - # creates a mismatch between the number - # of local and remote blocks. - limit_lookahead_tokens = load_kv_async and self.use_eagle + # During async KV load, no forward pass is run yet. + # Allocate speculative lookahead slots later to avoid + # mismatching local and remote block counts. + limit_lookahead_tokens = load_kv_async and self.num_lookahead_tokens > 0 effective_lookahead_tokens = ( 0 if limit_lookahead_tokens else self.num_lookahead_tokens ) @@ -2379,7 +2377,10 @@ class Scheduler(SchedulerInterface): num_prompt_tokens=request.num_prompt_tokens, ) - block_ids = self.kv_cache_manager.get_block_ids(request.request_id) + block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens( + request_id=request.request_id, + num_computed_tokens=request.num_computed_tokens, + ) if not isinstance(self.connector, SupportsHMA): # NOTE(Kuntai): We should deprecate this code path after we enforce