[KVConnector] Add retention interval to OffloadingConnector

tests/v1/core/test_prefix_caching.py

@@ -4364,6 +4364,161 @@ def retained(retention, boundary, window, end_block=16):
     assert retained(0, 0, block_size) == {14}
 
 
+def test_swa_reachable_block_mask_with_dcp_scaling():
+    """DCP shards each block's KV across ranks, scaling the effective block size.
+    Verify that dcp_world_size > 1 scales block size in reachability calculations,
+    producing different masks than dcp_world_size=1."""
+    from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
+
+    block_size = 16
+    sliding_window = 32  # need = cdiv(31, 16) = 2
+
+    def get_mask(dcp_world_size):
+        spec = SlidingWindowSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=1,
+            dtype=torch.float32,
+            sliding_window=sliding_window,
+        )
+        m = SlidingWindowManager.reachable_block_mask(
+            start_block=0,
+            end_block=16,
+            alignment_tokens=64,  # alignment constraint
+            kv_cache_spec=spec,
+            use_eagle=False,
+            retention_interval=64,  # sparse: 1 state per 64-token segment
+            reachable_boundaries=(255,),
+            dcp_world_size=dcp_world_size,
+        )
+        return None if m is None else {i for i, v in enumerate(m) if v}
+
+    # With dcp_world_size=1: effective block_size = 16
+    # per_segment = 64 // 16 = 4 blocks per segment
+    # need = 2, so 2 < 4 -> sparse retention with coarser granularity
+    mask_no_dcp = get_mask(dcp_world_size=1)
+
+    # With dcp_world_size=2: effective block_size = 32
+    # per_segment = 64 // 32 = 2 blocks per segment
+    # need = 2, so 2 >= 2 -> different retention granularity
+    mask_with_dcp = get_mask(dcp_world_size=2)
+
+    # DCP scaling changes block_size, altering segment granularity
+    assert mask_no_dcp is not None
+    assert mask_with_dcp is not None
+    # Different DCP world sizes should produce different masks
+    assert mask_no_dcp != mask_with_dcp, (
+        "DCP scaling should change retention granularity"
+    )
+
+
+@pytest.mark.parametrize(
+    "block_size,dcp_world_size,alignment_tokens",
+    [
+        (16, 1, 24),  # 24 % 16 != 0
+        (16, 3, 64),  # effective block_size 48; 64 % 48 != 0
+    ],
+)
+def test_swa_reachable_block_mask_sub_block_alignment_is_dense(
+    block_size, dcp_world_size, alignment_tokens
+):
+    """The mask is block-granular, so a sub-block alignment_tokens (not a
+    multiple of the DCP-scaled block size) cannot be represented exactly. This
+    happens for hybrid offloading (e.g. Gemma), where alignment_tokens is the
+    full-attention chunk size and need not divide the SWA block size. The mask
+    must fall back to dense (None) rather than raise."""
+    from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
+
+    spec = SlidingWindowSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        sliding_window=32,
+    )
+
+    mask = SlidingWindowManager.reachable_block_mask(
+        start_block=0,
+        end_block=16,
+        alignment_tokens=alignment_tokens,
+        kv_cache_spec=spec,
+        use_eagle=False,
+        retention_interval=alignment_tokens,
+        reachable_boundaries=(),
+        dcp_world_size=dcp_world_size,
+    )
+    assert mask is None
+
+
+def test_mamba_reachable_block_mask_ignores_dcp():
+    """Mamba uses TP, not DCP: each rank holds the full recurrent state, so a
+    state block spans kv_cache_spec.block_size tokens regardless of DCP. The
+    mask must not scale by dcp_world_size, so retention granularity is
+    identical for any DCP world size."""
+    from vllm.v1.core.single_type_kv_cache_manager import MambaManager
+
+    spec = MambaSpec(
+        block_size=16,
+        shapes=(1, 1),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+
+    def get_mask(dcp_world_size):
+        m = MambaManager.reachable_block_mask(
+            start_block=0,
+            end_block=16,
+            alignment_tokens=64,
+            kv_cache_spec=spec,
+            use_eagle=False,
+            retention_interval=64,  # sparse: one state per 64-token segment
+            reachable_boundaries=(255,),
+            dcp_world_size=dcp_world_size,
+        )
+        return None if m is None else {i for i, v in enumerate(m) if v}
+
+    # block_size stays 16 regardless of DCP: per_segment = 64 // 16 = 4.
+    mask_no_dcp = get_mask(dcp_world_size=1)
+    mask_with_dcp = get_mask(dcp_world_size=2)
+
+    assert mask_no_dcp is not None
+    assert mask_with_dcp is not None
+    assert mask_no_dcp == mask_with_dcp, "DCP must not change Mamba retention"
+
+
+def test_mamba_reachable_block_mask_large_dcp_stays_sparse():
+    """A large DCP world size must not scale the Mamba block size, so it can
+    never collapse sparse retention into dense caching."""
+    from vllm.v1.core.single_type_kv_cache_manager import MambaManager
+
+    spec = MambaSpec(
+        block_size=8,
+        shapes=(1, 1),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+
+    def get_mask(dcp_world_size):
+        return MambaManager.reachable_block_mask(
+            start_block=0,
+            end_block=16,
+            alignment_tokens=64,
+            kv_cache_spec=spec,
+            use_eagle=False,
+            retention_interval=64,
+            reachable_boundaries=(255,),
+            dcp_world_size=dcp_world_size,
+        )
+
+    # block_size stays 8 regardless of DCP: per_segment = 64 // 8 = 8, sparse.
+    mask_no_dcp = get_mask(dcp_world_size=1)
+    mask_large_dcp = get_mask(dcp_world_size=16)
+
+    assert mask_no_dcp is not None, "Sparse retention expected"
+    assert mask_large_dcp is not None, "Large DCP must not force dense caching"
+    assert mask_no_dcp == mask_large_dcp, "DCP must not change Mamba retention"
+
+
 def test_swa_shared_prefix_reuse_under_zero_retention():
     """SWA cross-request analog: a partial shared prefix's sliding-window tail
     must stay reusable under ``prefix_cache_retention_interval=0``. Without

tests/v1/kv_connector/unit/offloading_connector/test_config.py

@@ -15,7 +15,6 @@
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
     SchedulerOffloadConfig,
-    is_store_reachable_swa_chunk,
 )
 from vllm.platforms import current_platform
 from vllm.v1.kv_cache_interface import (
@@ -31,6 +30,7 @@
     SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
 )
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 
 
 def _make_vllm_config(
@@ -46,6 +46,7 @@ def _make_vllm_config(
     config.cache_config.enable_prefix_caching = True
     config.cache_config.prefix_match_unit = None
     config.cache_config.cache_dtype = torch.float16
+    config.cache_config.prefix_cache_retention_interval = None
     config.model_config.model = "test-model"
     config.model_config.use_mla = False
     # _full_attention_spec's heads at tp=1: the parallelism-agnostic gate
@@ -375,18 +376,18 @@ def test_dcp_scales_attention_but_not_mamba_group_blocks():
         _make_mamba_hybrid_kv_cache_config(),
     )
     mamba_group = scheduler_config.kv_group_configs[1]
-    assert mamba_group.alignment_chunk_count == 2
-    assert [
-        chunk_idx
-        for chunk_idx in range(4)
-        if is_store_reachable_swa_chunk(
-            chunk_idx,
-            4,
-            mamba_group.alignment_chunk_count,
-            mamba_group.sliding_window_size_in_chunks,
-            mamba_group.is_eagle_group,
-        )
-    ] == [1, 3]
+    assert mamba_group.sliding_window_size_in_chunks == 1
+    assert scheduler_config.alignment_tokens is not None
+    manager_cls = KVCacheSpecRegistry.get_manager_class(mamba_group.kv_cache_spec)
+    assert manager_cls is not None
+    block_mask = manager_cls.reachable_block_mask(
+        start_block=0,
+        end_block=4,
+        alignment_tokens=scheduler_config.alignment_tokens,
+        kv_cache_spec=mamba_group.kv_cache_spec,
+        use_eagle=mamba_group.is_eagle_group,
+    )
+    assert block_mask is None or [i for i in range(4) if block_mask[i]] == [1, 3]
 
 
 @pytest.mark.parametrize("dcp_size,expected", [(1, 16), (2, 32)])

tests/v1/kv_connector/unit/offloading_connector/test_events.py

@@ -29,6 +29,7 @@
     KVCacheGroupSpec,
     KVCacheSpecKind,
 )
+from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
 from vllm.v1.kv_offload.base import (
     Locality,
     Medium,
@@ -88,12 +89,22 @@ def _group_config(
         tokens_per_hash = block_size
     tokens_per_chunk = block_size * blocks_per_chunk
     assert tokens_per_chunk % tokens_per_hash == 0
+    kv_spec = FullAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+    )
+    manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec)
+    assert manager_cls is not None
     return GroupOffloadConfig(
         group_idx=group_idx,
         tokens_per_block=block_size,
         tokens_per_chunk=tokens_per_chunk,
         hashes_per_chunk=tokens_per_chunk // tokens_per_hash,
         sliding_window_size_in_chunks=sliding_window_size_in_chunks,
+        kv_cache_spec=kv_spec,
+        manager_cls=manager_cls,
         kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,
     )