[v1][kvcache] Honor prefix-cache retention interval for Mamba/linear attention (#45845)
Signed-off-by: Dao Le <daole@inferact.ai> Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Dao007forever committed
Jun 22, 2026 at 19:51 UTC
430a95ae3aecedfc8b568ad0d540477bf2fb4344
3 files changed
+105
-9
tests/v1/core/test_prefix_caching.py
+38
@@ -3905,3 +3905,41 @@ def test_pure_swa_retention_dense_default_caches_all(monkeypatch):
3905
is not None
3906
}
3907
assert cached == set(range(16))
3908
+
3909
+
3910
+def test_mamba_reachable_block_mask_sparsifies_retention():
3911
+ """Mamba state-snapshot retention: with VLLM_PREFIX_CACHE_RETENTION_INTERVAL
3912
+ the manager keeps one cached state per interval-sized segment (plus the
3913
+ latest replay boundary) instead of a snapshot per block, which is what
3914
+ lets a small attention block_size avoid Mamba dominating the KV pool."""
3915
+ from vllm.v1.core.single_type_kv_cache_manager import MambaManager
3916
+
3917
+ block_size = 16
3918
+ spec = MambaSpec(
3919
+ block_size=block_size,
3920
+ shapes=(1, 1),
3921
+ dtypes=(torch.float32,),
3922
+ mamba_cache_mode="align",
3923
+ )
3924
+
3925
+ def retained(retention_interval, num_prompt_tokens=256, end_block=16):
3926
+ m = MambaManager.reachable_block_mask(
3927
+ start_block=0,
3928
+ end_block=end_block,
3929
+ alignment_tokens=block_size,
3930
+ kv_cache_spec=spec,
3931
+ use_eagle=False,
3932
+ retention_interval=retention_interval,
3933
+ num_prompt_tokens=num_prompt_tokens,
3934
+ )
3935
+ return None if m is None else {i for i, v in enumerate(m) if v}
3936
+
3937
+ # Dense default (None) -> no mask, every block cached (unchanged behavior).
3938
+ assert retained(None) is None
3939
+ # interval == block_size -> every block is a boundary -> stays dense.
3940
+ assert retained(block_size) is None
3941
+ # interval 64 = 4 blocks: segment tails at i%4==3 -> {3,7,11,15}; latest
3942
+ # replay boundary 240//16 - 1 = 14. Sparse subset of the 16 blocks.
3943
+ assert retained(64) == {3, 7, 11, 14, 15}
3944
+ # interval 0 -> only the latest replay boundary (block 14).
3945
+ assert retained(0) == {14}
vllm/v1/core/kv_cache_coordinator.py
+8
-7
@@ -22,6 +22,7 @@ from vllm.v1.kv_cache_interface import (
22
FullAttentionSpec,
23
KVCacheConfig,
24
KVCacheSpec,
25
+ MambaSpec,
26
SlidingWindowSpec,
27
)
28
from vllm.v1.request import Request
@@ -35,18 +36,18 @@ def _validate_prefix_cache_retention_interval(
36
if retention_interval is None:
37
return
38
38
- # Retention only sparsifies sliding-window checkpoints for now; every other
39
- # manager (full attention, Mamba, chunked-local) caches densely and
40
- # ignores it to be conservative.
41
- # TODO: Support Mamba/linear attention.
39
+ # Retention sparsifies sliding-window and Mamba (linear-attention)
40
+ # checkpoints; full-attention and chunked-local groups cache densely and
41
+ # ignore it (their hit granularity must stay fine).
42
if not any(
43
- isinstance(g.kv_cache_spec, SlidingWindowSpec)
43
+ isinstance(g.kv_cache_spec, (SlidingWindowSpec, MambaSpec))
44
for g in kv_cache_config.kv_cache_groups
45
):
46
raise ValueError(
47
"VLLM_PREFIX_CACHE_RETENTION_INTERVAL is set but this model has "
48
- "no sliding-window KV cache group, so retention has no effect. "
49
- "Unset it (the feature only applies to sliding-window attention)."
48
+ "no sliding-window or Mamba KV cache group, so retention has no "
49
+ "effect. Unset it (it only applies to sliding-window and Mamba "
50
+ "attention)."
51
)
52
53
if retention_interval < 0 or retention_interval % scheduler_block_size != 0:
vllm/v1/core/single_type_kv_cache_manager.py
+59
-2
@@ -1018,6 +1018,60 @@ class MambaManager(SingleTypeKVCacheManager):
1018
1019
return computed_blocks
1020
1021
+ @classmethod
1022
+ def reachable_block_mask(
1023
+ cls,
1024
+ start_block: int,
1025
+ end_block: int,
1026
+ alignment_tokens: int | None,
1027
+ kv_cache_spec: KVCacheSpec,
1028
+ use_eagle: bool,
1029
+ retention_interval: int | None = None,
1030
+ num_prompt_tokens: int | None = None,
1031
+ ) -> list[bool] | None:
1032
+ """Sparse Mamba state-snapshot retention.
1033
+
1034
+ ``retention_interval``:
1035
+
1036
+ ``None`` -> dense (cache every block; default, unchanged behavior)
1037
+ ``0`` -> keep only the latest replay boundary
1038
+ ``> 0`` -> keep one state per ``retention_interval``-sized segment
1039
+ """
1040
+ if retention_interval is None or alignment_tokens is None:
1041
+ # Dense caching (default) or no alignment constraint imposed.
1042
+ return None
1043
+ assert isinstance(kv_cache_spec, MambaSpec)
1044
+ block_size = kv_cache_spec.block_size
1045
+ mask = [False] * (end_block - start_block)
1046
+
1047
+ # (1) Segment-boundary states. A Mamba hit needs exactly the single
1048
+ # state block ending on the boundary (no window, and draft models have
1049
+ # no mamba layers, so no eagle shift). Block ``i`` ends at token
1050
+ # ``(i + 1) * block_size``.
1051
+ segment_tokens = None if retention_interval == 0 else retention_interval
1052
+ if segment_tokens is not None:
1053
+ per_segment = segment_tokens // block_size
1054
+ if per_segment <= 1:
1055
+ # Interval at/below the block size: every block is a boundary.
1056
+ return None
1057
+ first_boundary = (
1058
+ start_block + per_segment
1059
+ ) // per_segment * per_segment - 1
1060
+ for i in range(first_boundary - start_block, len(mask), per_segment):
1061
+ mask[i] = True
1062
+
1063
+ # (2) Replay boundary. ``get_computed_blocks`` caps hits at
1064
+ # ``num_prompt - 1``, so an exact prompt replay lands on the latest
1065
+ # fine-aligned boundary. Sparse retention would otherwise skip its
1066
+ # state, so keep it explicitly.
1067
+ if num_prompt_tokens is not None:
1068
+ latest = (num_prompt_tokens - 1) // alignment_tokens * alignment_tokens
1069
+ boundary_block = latest // block_size - 1
1070
+ if start_block <= boundary_block < end_block:
1071
+ mask[boundary_block - start_block] = True
1072
+
1073
+ return mask
1074
+
1075
def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
1076
assert isinstance(self.kv_cache_spec, MambaSpec)
1077
@@ -1224,9 +1278,12 @@ class MambaManager(SingleTypeKVCacheManager):
1278
for block in self.req_to_blocks[request.request_id][
1279
num_cached_blocks_before:num_cached_blocks_after
1280
]:
1227
- if block.is_null:
1281
+ # Skip null blocks (align-mode skipped states) and blocks that
1282
+ # were not cached this step — with sparse retention
1283
+ # (reachable_block_mask) the intermediate state snapshots carry
1284
+ # no hash and must not be recorded as cached-this-step.
1285
+ if block.is_null or block.block_hash is None:
1286
continue
1229
- assert block.block_hash is not None
1287
self.cached_blocks_this_step.add(block.block_hash)
1288
1289
def new_step_starts(self) -> None: