@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f272dfdce1

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[KV Connector] Mooncake store: prefix-cache retention interval for sparse attention (#44774)

Yifan Qiao committed Jun 10, 2026 at 21:36 UTC f272dfdce1217e56ff859c9ed4b7353e684e2001

4 files changed +92 -37

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

+51 -1

@@ -15,7 +15,7 @@ from vllm.v1.kv_cache_interface import (

15 )

18 -def _make_coord(groups, hash_block_size, use_eagle=False):

18 +def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=None):

19 """Construct a coordinator using the natural LCM of group block sizes as

20 the scheduler block size — mirrors ``resolve_kv_cache_block_sizes`` for

21 the test fixtures."""

@@ -26,6 +26,7 @@ def _make_coord(groups, hash_block_size, use_eagle=False):

26 scheduler_block_size=scheduler_block_size,

27 hash_block_size=hash_block_size,

28 use_eagle=use_eagle,

29 + retention_interval=retention_interval,

30 )

@@ -302,6 +303,55 @@ def test_store_mask_fast_path_single_attention_group():

303 assert masks == ([True] * 4, [True] * 4)

304

305

306 +# ----- store_mask with retention_interval (DSV4 sparse SWA checkpointing) -----

307 +

308 +

309 +def _retention_groups():

310 + """Hybrid full-attn(block=32) + SWA(block=8, sw=8); lcm=32. The SWA group

311 + densely keeps one tail block per 32-token boundary."""

312 + full = _full(32)

313 + swa = _swa(block_size=8, sliding_window=8)

314 + return [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)]

315 +

316 +

317 +def test_store_mask_dense_default_matches_every_lcm_boundary():

318 + """retention_interval=None (default) keeps the SWA tail at every lcm

319 + boundary: tokens 32/64/96/128 -> chunks 3/7/11/15."""

320 + coord = _make_coord(_retention_groups(), hash_block_size=8)

321 + masks = coord.store_mask(128)

322 + assert masks[0] == [True, True, True, True]

323 + assert masks[1] == [i % 4 == 3 for i in range(16)]

324 +

325 +

326 +def test_store_mask_retention_interval_sparsifies_swa_tails():

327 + """retention_interval=64 keeps an SWA tail once per 64-token segment

328 + (chunks 7 and 15) instead of every 32 tokens, dropping the mid-segment

329 + boundaries at 32 and 96."""

330 + coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=64)

331 + masks = coord.store_mask(128)

332 + assert masks[0] == [True, True, True, True] # full attn unaffected

333 + assert masks[1] == [i in (7, 15) for i in range(16)]

334 +

335 +

336 +def test_store_mask_retention_interval_zero_keeps_only_replay_boundary():

337 + """retention_interval=0 drops all segment tails; only the latest replay

338 + boundary (capped at num_prompt-1, aligned down to lcm) is retained."""

339 + coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=0)

340 + # No replay info -> nothing reachable for the SWA group.

341 + assert coord.store_mask(128)[1] == [False] * 16

342 + # num_prompt=100 -> latest hit boundary = (100-1)//32*32 = 96 -> chunk 11.

343 + masks = coord.store_mask(128, num_prompt_tokens=100)

344 + assert masks[1] == [i == 11 for i in range(16)]

345 +

346 +

347 +def test_store_mask_retention_interval_keeps_segment_and_replay_tails():

348 + """Sparse segment tails (interval=64 -> chunks 7,15) plus the replay

349 + boundary tail (num_prompt=100 -> chunk 11) coexist."""

350 + coord = _make_coord(_retention_groups(), hash_block_size=8, retention_interval=64)

351 + masks = coord.store_mask(128, num_prompt_tokens=100)

352 + assert masks[1] == [i in (7, 11, 15) for i in range(16)]

353 +

354 +

355 # ----- Eagle / MTP interaction with load_mask -----

356

357

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

+35 -35

@@ -22,9 +22,6 @@ from vllm.v1.kv_cache_interface import (

22 )

23 from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

25 -# Dummy placeholder hash for store_mask's template computation.

26 -_DUMMY_BLOCK_HASH = BlockHash(b"\x00" * 32)

27 -

26 class ExternalCachedBlockPool:

27 """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set."""

@@ -62,6 +59,7 @@ class MooncakeStoreCoordinator:

59 scheduler_block_size: int,

60 hash_block_size: int,

61 use_eagle: bool = False,

62 + retention_interval: int | None = None,

63 ) -> None:

64 assert all(

65 g.kv_cache_spec.block_size % hash_block_size == 0 for g in kv_cache_groups

@@ -78,6 +76,13 @@ class MooncakeStoreCoordinator:

76 self.hash_block_size = hash_block_size

77 self.lcm_block_size = scheduler_block_size

78 self.use_eagle = use_eagle

79 + # Mirror vLLM core's KVCacheCoordinator.retention_interval.

80 + self.retention_interval = retention_interval

81 + self.eagle_group_ids = {

82 + i for i, g in enumerate(kv_cache_groups) if g.is_eagle_group

83 + }

84 + if use_eagle and not self.eagle_group_ids:

85 + self.eagle_group_ids = set(range(len(kv_cache_groups)))

86 self._verify_and_split_kv_cache_groups()

88 def _verify_and_split_kv_cache_groups(self) -> None:

@@ -163,44 +168,39 @@ class MooncakeStoreCoordinator:

168 )

169 return masks

170

166 - def store_mask(self, aligned_token_len: int) -> tuple[list[bool], ...]:

171 + def store_mask(

172 + self,

173 + aligned_token_len: int,

174 + num_prompt_tokens: int | None = None,

175 + ) -> tuple[list[bool], ...]:

176 """Per-group store masks: ``mask[g][i]`` is True iff chunk ``i`` of

168 - group ``g`` would be populated by some future cache hit at length

169 - ``L = N * lcm_block_size <= aligned_token_len``.

177 + group ``g`` should be written to the store so a future cache hit can

178 + consume it.

179 +

180 + Reuses the engine's ``SingleTypeKVCacheManager.reachable_block_mask``

181 + so the store retains exactly the blocks the local prefix cache would.

182 """

183 assert aligned_token_len % self.lcm_block_size == 0, (

184 f"aligned_token_len ({aligned_token_len}) must be a multiple of "

185 f"lcm_block_size ({self.lcm_block_size})"

186 )

175 - if aligned_token_len == 0:

176 - return tuple([] for _ in self.kv_cache_groups)

177 -

178 - num_chunks_per_group = [

179 - aligned_token_len // g.kv_cache_spec.block_size

180 - for g in self.kv_cache_groups

181 - ]

182 -

183 - # Fast path: single group or full attn groups or uniform block_sizes

184 - if all(

185 - isinstance(spec, FullAttentionSpec)

186 - or spec.block_size == self.lcm_block_size

187 - for spec, _, _ in self.attention_groups

188 - ):

189 - return tuple([True] * n for n in num_chunks_per_group)

190 -

191 - n_segments = aligned_token_len // self.lcm_block_size

192 - dummy_hashes: list[BlockHash] = [_DUMMY_BLOCK_HASH] * (

193 - self.lcm_block_size // self.hash_block_size

194 - )

195 - template_masks, _ = self.find_longest_cache_hit(

196 - dummy_hashes,

197 - max_length=self.lcm_block_size,

198 - cached_block_pool=ExternalCachedBlockPool(),

199 - )

200 - return tuple(

201 - list(template_masks[g]) * n_segments

202 - for g in range(len(self.kv_cache_groups))

203 - )

187 + masks: list[list[bool]] = []

188 + for g_idx, g in enumerate(self.kv_cache_groups):

189 + spec = _unwrap_spec(g.kv_cache_spec)

190 + num_chunks = aligned_token_len // spec.block_size

191 + manager_cls = KVCacheSpecRegistry.get_manager_class(spec)

192 + assert manager_cls is not None

193 + mask = manager_cls.reachable_block_mask(

194 + start_block=0,

195 + end_block=num_chunks,

196 + alignment_tokens=self.lcm_block_size,

197 + kv_cache_spec=spec,

198 + use_eagle=g_idx in self.eagle_group_ids,

199 + retention_interval=self.retention_interval,

200 + num_prompt_tokens=num_prompt_tokens,

201 + )

202 + masks.append([True] * num_chunks if mask is None else mask)

203 + return tuple(masks)

204

205 def block_hashes_for_spec(

206 self, block_hashes: list[BlockHash], spec: KVCacheSpec

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

@@ -213,6 +213,7 @@ class ReqMeta:

213 current_event: torch.cuda.Event | None = None

214

215 token_ids: list[int] | None = None

216 + num_prompt_tokens: int | None = None

217

218 @staticmethod

219 def from_request_tracker(

@@ -272,6 +273,7 @@ class ReqMeta:

273 block_hashes=block_hashes,

274 is_last_chunk=is_last_chunk,

275 token_ids=token_ids,

276 + num_prompt_tokens=tracker.prefill_end_tokens,

277 )

278

279

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+4 -1

@@ -535,7 +535,9 @@ class KVCacheStoreSendingThread(KVTransferThread):

535

536 # Within each lcm region only per-spec relevant chunks are loaded

537 # (e.g., SWA or linear attn), so mask out irrelevant chunks

538 - store_masks = self.coord.store_mask(token_len)

538 + store_masks = self.coord.store_mask(

539 + token_len, num_prompt_tokens=req_meta.num_prompt_tokens

540 + )

541 starts: list[int] = []

542 ends: list[int] = []

543 keys: list[str] = []

@@ -1091,6 +1093,7 @@ class MooncakeStoreWorker:

1093 scheduler_block_size=self.block_size,

1094 hash_block_size=self.hash_block_size,

1095 use_eagle=use_eagle,

1096 + retention_interval=envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL,

1097 )

1098 # One ChunkedTokenDatabase per group; addresses populated in

1099 # register_kv_caches once the kv-cache layout is known.