fix(kv): publish exact recurrent sources for external stores

tests/v1/core/test_kv_connector_block_state.py

@@ -0,0 +1,106 @@
+# 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": 8192},
+        {"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": 3584},
+        {"req": [(1, 777, 4096)]},
+        retention_interval=4096,
+    )
+
+    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=([], [])))
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        _build_kv_connector_block_state(
+            _hybrid_config(),
+            manager,
+            {"req": 0},
+            None,
+            retention_interval=4100,
+        )

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

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
 
@@ -42,6 +42,7 @@
 from vllm.v1.core.sched.output import (
     CachedRequestData,
     GrammarOutput,
+    KVConnectorBlockState,
     NewRequestData,
     ScheduledEncoderInputStats,
     SchedulerOutput,
@@ -73,6 +74,73 @@
 logger = init_logger(__name__)
 
 
+def _build_kv_connector_block_state(
+    kv_cache_config: KVCacheConfig,
+    kv_cache_manager: KVCacheManager,
+    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.
+
+    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.
+
+    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)
+        for req_id in computed_token_extents
+    }
+    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 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(
+        block_ids=current_block_ids,
+        boundary_state_offloads=boundary_state_offloads,
+    )
+
+
 class Scheduler(SchedulerInterface):
     def __init__(
         self,
@@ -1283,6 +1351,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 +1360,18 @@ 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,
+                {
+                    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,
+            )
 
         kv_cache_block_copies, cow_retained_blocks = (
             self.kv_cache_manager.take_kv_cache_block_copies()
@@ -1343,6 +1424,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(),
         )