From c1ae634a111dd782e56a143b03f7d89e8c86292a Mon Sep 17 00:00:00 2001 From: derek Date: Tue, 1 Sep 2026 16:53:22 -0400 Subject: [PATCH 1/3] fix(kv): publish exact recurrent store sources --- .../v1/core/test_kv_connector_block_state.py | 91 +++++++++++++++++++ vllm/v1/core/sched/output.py | 20 ++++ vllm/v1/core/sched/scheduler.py | 59 ++++++++++++ 3 files changed, 170 insertions(+) create mode 100644 tests/v1/core/test_kv_connector_block_state.py diff --git a/tests/v1/core/test_kv_connector_block_state.py b/tests/v1/core/test_kv_connector_block_state.py new file mode 100644 index 000000000000..0c1c0fa8110b --- /dev/null +++ b/tests/v1/core/test_kv_connector_block_state.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.v1.core.sched.scheduler import _build_kv_connector_block_state +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) + + +def _hybrid_config() -> KVCacheConfig: + return KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["attention"], + FullAttentionSpec( + block_size=2048, + num_kv_heads=1, + head_size=1, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=512, + shapes=((1, 1),), + dtypes=(torch.float16,), + mamba_cache_mode="align", + ), + ), + ], + ) + + +def test_connector_block_state_has_snapshot_and_retained_boundaries() -> None: + grouped_ids = ( + [201, 202, 203, 204], + [0, 0, 0, 0, 0, 0, 0, 107, 0, 0, 110, 0, 0, 0, 0, 115], + ) + manager = SimpleNamespace(get_block_ids=MagicMock(return_value=grouped_ids)) + + state = _build_kv_connector_block_state( + _hybrid_config(), + manager, + ["req"], + {"req": [(1, 999, 7680)]}, + retention_interval=4096, + ) + + assert state.block_ids == {"req": grouped_ids} + assert state.boundary_state_offloads == { + "req": [(1, 999, 7680), (1, 107, 4096), (1, 115, 8192)] + } + manager.get_block_ids.assert_called_once_with("req") + + +def test_connector_block_state_preserves_explicit_boundary_source() -> None: + grouped_ids = ([201, 202], [0, 0, 0, 0, 0, 0, 0, 107]) + manager = SimpleNamespace(get_block_ids=MagicMock(return_value=grouped_ids)) + + state = _build_kv_connector_block_state( + _hybrid_config(), + manager, + ["req"], + {"req": [(1, 777, 4096)]}, + retention_interval=4096, + ) + + assert state.boundary_state_offloads == {"req": [(1, 777, 4096)]} + + +def test_connector_block_state_rejects_misaligned_retention() -> None: + manager = SimpleNamespace(get_block_ids=MagicMock(return_value=([], []))) + + with pytest.raises(ValueError, match="must be divisible"): + _build_kv_connector_block_state( + _hybrid_config(), + manager, + ["req"], + None, + retention_interval=4100, + ) diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index ad09529e8d8f..9c046df8034f 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -205,6 +205,21 @@ class ScheduledEncoderInputStats: output_tokens: int = 0 +@dataclass +class KVConnectorBlockState: + """Authoritative scheduler-side KV source state for external stores. + + ``block_ids`` is a full per-group snapshot for every request scheduled in + this step. ``boundary_state_offloads`` identifies the exact recurrent + state block at each durable token boundary; connectors must not reconstruct + these sources from allocation deltas because align-mode Mamba tables are + sparse and mutable. + """ + + block_ids: dict[str, tuple[list[int], ...]] + boundary_state_offloads: dict[str, list[tuple[int, int, int]]] + + @dataclass class SchedulerOutput: # list of the requests that are scheduled for the first time. @@ -280,6 +295,11 @@ class SchedulerOutput: # tail (mamba "align" CoW target). None unless partial hash hits are active. partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None = None + # Authoritative source tables and recurrent boundary blocks for external + # KV stores. This is scheduler-only metadata consumed while connector + # metadata is built; workers use the resulting opaque connector metadata. + kv_connector_block_state: KVConnectorBlockState | None = None + # Dynamic speculative decoding: optimal K chosen by scheduler. # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int | None = None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 69c631ef1b89..15337489c290 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -42,6 +42,7 @@ from vllm.v1.core.sched.output import ( CachedRequestData, GrammarOutput, + KVConnectorBlockState, NewRequestData, ScheduledEncoderInputStats, SchedulerOutput, @@ -73,6 +74,55 @@ logger = init_logger(__name__) +def _build_kv_connector_block_state( + kv_cache_config: KVCacheConfig, + kv_cache_manager: KVCacheManager, + request_ids: Iterable[str], + partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None, + retention_interval: int | None, +) -> KVConnectorBlockState: + """Snapshot exact source blocks for connector stores in this step.""" + current_block_ids = { + req_id: kv_cache_manager.get_block_ids(req_id) for req_id in request_ids + } + boundary_state_offloads = { + req_id: list(entries) + for req_id, entries in (partial_tail_offloads or {}).items() + } + if retention_interval is not None and retention_interval > 0: + for group_id, group in enumerate(kv_cache_config.kv_cache_groups): + spec = group.kv_cache_spec + if not isinstance(spec, MambaSpec): + continue + if retention_interval % spec.block_size != 0: + raise ValueError( + "prefix_cache_retention_interval must be divisible by " + f"Mamba block size: {retention_interval=} {spec.block_size=}" + ) + blocks_per_boundary = retention_interval // spec.block_size + for req_id, grouped_ids in current_block_ids.items(): + if group_id >= len(grouped_ids): + continue + entries = boundary_state_offloads.setdefault(req_id, []) + existing = { + (entry_group, boundary_tokens) + for entry_group, _, boundary_tokens in entries + } + for block_index in range( + blocks_per_boundary - 1, + len(grouped_ids[group_id]), + blocks_per_boundary, + ): + block_id = grouped_ids[group_id][block_index] + boundary_tokens = (block_index + 1) * spec.block_size + if block_id > 0 and (group_id, boundary_tokens) not in existing: + entries.append((group_id, block_id, boundary_tokens)) + return KVConnectorBlockState( + block_ids=current_block_ids, + boundary_state_offloads=boundary_state_offloads, + ) + + class Scheduler(SchedulerInterface): def __init__( self, @@ -1283,6 +1333,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # pin); the manager drops stale entries when the request's blocks are # popped for free. pending_partial_tail_offloads = None + kv_connector_block_state = None if ( self.connector is not None and self.vllm_config.kv_transfer_config is not None @@ -1291,6 +1342,13 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: pending_partial_tail_offloads = ( self.kv_cache_manager.take_partial_tail_offloads() or None ) + kv_connector_block_state = _build_kv_connector_block_state( + self.kv_cache_config, + self.kv_cache_manager, + num_scheduled_tokens, + pending_partial_tail_offloads, + self.cache_config.prefix_cache_retention_interval, + ) kv_cache_block_copies, cow_retained_blocks = ( self.kv_cache_manager.take_kv_cache_block_copies() @@ -1343,6 +1401,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: new_block_ids_to_zero=self._get_new_block_ids_to_zero(), kv_cache_block_copies=pending_kv_cache_block_copies, partial_tail_offloads=pending_partial_tail_offloads, + kv_connector_block_state=kv_connector_block_state, num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(), ) From ae8da76136bddbf6a60a0e6ea2ad1d971e103818 Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 2 Sep 2026 23:25:58 -0400 Subject: [PATCH 2/3] fix(kv): exclude uncomputed retention boundaries Co-authored-by: OpenAI Codex --- .../v1/core/test_kv_connector_block_state.py | 21 ++++++++++++-- vllm/v1/core/sched/scheduler.py | 29 +++++++++++++++---- 2 files changed, 42 insertions(+), 8 deletions(-) diff --git a/tests/v1/core/test_kv_connector_block_state.py b/tests/v1/core/test_kv_connector_block_state.py index 0c1c0fa8110b..77b34ce57ef8 100644 --- a/tests/v1/core/test_kv_connector_block_state.py +++ b/tests/v1/core/test_kv_connector_block_state.py @@ -51,7 +51,7 @@ def test_connector_block_state_has_snapshot_and_retained_boundaries() -> None: state = _build_kv_connector_block_state( _hybrid_config(), manager, - ["req"], + {"req": 8192}, {"req": [(1, 999, 7680)]}, retention_interval=4096, ) @@ -70,7 +70,7 @@ def test_connector_block_state_preserves_explicit_boundary_source() -> None: state = _build_kv_connector_block_state( _hybrid_config(), manager, - ["req"], + {"req": 3584}, {"req": [(1, 777, 4096)]}, retention_interval=4096, ) @@ -78,6 +78,21 @@ def test_connector_block_state_preserves_explicit_boundary_source() -> None: assert state.boundary_state_offloads == {"req": [(1, 777, 4096)]} +def test_connector_block_state_excludes_uncomputed_allocated_boundary() -> None: + grouped_ids = ([201, 202], [0, 0, 0, 0, 0, 0, 0, 107]) + manager = SimpleNamespace(get_block_ids=MagicMock(return_value=grouped_ids)) + + state = _build_kv_connector_block_state( + _hybrid_config(), + manager, + {"req": 3584}, + None, + retention_interval=4096, + ) + + assert state.boundary_state_offloads == {"req": []} + + def test_connector_block_state_rejects_misaligned_retention() -> None: manager = SimpleNamespace(get_block_ids=MagicMock(return_value=([], []))) @@ -85,7 +100,7 @@ def test_connector_block_state_rejects_misaligned_retention() -> None: _build_kv_connector_block_state( _hybrid_config(), manager, - ["req"], + {"req": 0}, None, retention_interval=4100, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 15337489c290..2280fe002cd3 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -3,7 +3,7 @@ import itertools import time from collections import defaultdict, deque -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from dataclasses import replace from typing import Any @@ -77,13 +77,25 @@ def _build_kv_connector_block_state( kv_cache_config: KVCacheConfig, kv_cache_manager: KVCacheManager, - request_ids: Iterable[str], + computed_token_extents: Mapping[str, int], partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None, retention_interval: int | None, ) -> KVConnectorBlockState: - """Snapshot exact source blocks for connector stores in this step.""" + """Snapshot exact source blocks for connector stores in this step. + + Args: + kv_cache_config: Physical cache groups owned by the scheduler. + kv_cache_manager: Manager providing the current request block tables. + computed_token_extents: Token extent produced after this scheduled step. + partial_tail_offloads: Explicit copy-on-write recurrent state sources. + retention_interval: Token interval for durable recurrent state boundaries. + + Returns: + Authoritative block tables and recurrent boundary sources. + """ current_block_ids = { - req_id: kv_cache_manager.get_block_ids(req_id) for req_id in request_ids + req_id: kv_cache_manager.get_block_ids(req_id) + for req_id in computed_token_extents } boundary_state_offloads = { req_id: list(entries) @@ -115,6 +127,8 @@ def _build_kv_connector_block_state( ): block_id = grouped_ids[group_id][block_index] boundary_tokens = (block_index + 1) * spec.block_size + if boundary_tokens > computed_token_extents[req_id]: + break if block_id > 0 and (group_id, boundary_tokens) not in existing: entries.append((group_id, block_id, boundary_tokens)) return KVConnectorBlockState( @@ -1345,7 +1359,12 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: kv_connector_block_state = _build_kv_connector_block_state( self.kv_cache_config, self.kv_cache_manager, - num_scheduled_tokens, + { + req_id: ( + self.requests[req_id].num_computed_tokens + num_scheduled_token + ) + for req_id, num_scheduled_token in num_scheduled_tokens.items() + }, pending_partial_tail_offloads, self.cache_config.prefix_cache_retention_interval, ) From 277c984541fecc3f039a06a8cdbeae4008a314f0 Mon Sep 17 00:00:00 2001 From: derek Date: Wed, 2 Sep 2026 23:35:38 -0400 Subject: [PATCH 3/3] docs(kv): document retention alignment failure Co-authored-by: OpenAI Codex --- vllm/v1/core/sched/scheduler.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2280fe002cd3..30bee6f9572c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -92,6 +92,10 @@ def _build_kv_connector_block_state( Returns: Authoritative block tables and recurrent boundary sources. + + Raises: + ValueError: If the retention interval is not aligned to a recurrent + cache block boundary. """ current_block_ids = { req_id: kv_cache_manager.get_block_ids(req_id)