[P/D][Bugfix] Fix PD async KV load lookahead handling for MTP spec decode (#46694)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Chauncey committed
Jul 9, 2026 at 18:01 UTC
ae6170f874a7c66512bd24aa714b3a02f8a61838
2 files changed
+32
-7
vllm/v1/core/kv_cache_manager.py
+24
@@ -8,10 +8,14 @@ from typing import Literal, overload
8
9
from vllm.distributed.kv_events import BlockStored, KVCacheEvent
10
from vllm.logger import init_logger
11
+from vllm.utils.math_utils import cdiv
12
from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
13
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
14
from vllm.v1.core.kv_cache_utils import KVCacheBlock
15
from vllm.v1.kv_cache_interface import (
16
+ AttentionSpec,
17
+ CrossAttentionSpec,
18
+ EncoderOnlyAttentionSpec,
19
KVCacheConfig,
20
get_kv_cache_spec_kind,
21
get_kv_cache_spec_sliding_window,
@@ -596,6 +600,26 @@ class KVCacheManager:
600
"""Get the block ids of a request."""
601
return self.get_blocks(request_id).get_block_ids()
602
603
+ def get_block_ids_for_computed_tokens(
604
+ self,
605
+ request_id: str,
606
+ num_computed_tokens: int,
607
+ ) -> tuple[list[int], ...]:
608
+ """Get block ids covering the request's computed tokens."""
609
+ block_ids = self.get_block_ids(request_id)
610
+ clipped_block_ids: list[list[int]] = []
611
+ for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids):
612
+ spec = group.kv_cache_spec
613
+ if not isinstance(spec, AttentionSpec) or isinstance(
614
+ spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec)
615
+ ):
616
+ clipped_block_ids.append(ids)
617
+ continue
618
+
619
+ num_valid_blocks = cdiv(num_computed_tokens, spec.block_size)
620
+ clipped_block_ids.append(ids[:num_valid_blocks])
621
+ return tuple(clipped_block_ids)
622
+
623
def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
624
"""Cache the blocks for the request, if enabled.
625
vllm/v1/core/sched/scheduler.py
+8
-7
@@ -874,12 +874,10 @@ class Scheduler(SchedulerInterface):
874
if num_new_tokens == 0:
875
break
876
877
- # Handles an edge case when P/D Disaggregation
878
- # is used with Spec Decoding where an
879
- # extra block gets allocated which
880
- # creates a mismatch between the number
881
- # of local and remote blocks.
882
- limit_lookahead_tokens = load_kv_async and self.use_eagle
877
+ # During async KV load, no forward pass is run yet.
878
+ # Allocate speculative lookahead slots later to avoid
879
+ # mismatching local and remote block counts.
880
+ limit_lookahead_tokens = load_kv_async and self.num_lookahead_tokens > 0
881
effective_lookahead_tokens = (
882
0 if limit_lookahead_tokens else self.num_lookahead_tokens
883
)
@@ -2379,7 +2377,10 @@ class Scheduler(SchedulerInterface):
2377
num_prompt_tokens=request.num_prompt_tokens,
2378
)
2379
2382
- block_ids = self.kv_cache_manager.get_block_ids(request.request_id)
2380
+ block_ids = self.kv_cache_manager.get_block_ids_for_computed_tokens(
2381
+ request_id=request.request_id,
2382
+ num_computed_tokens=request.num_computed_tokens,
2383
+ )
2384
2385
if not isinstance(self.connector, SupportsHMA):
2386
# NOTE(Kuntai): We should deprecate this code path after we enforce