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

Fork 0

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

[Mooncake] Skip KV lookup for non-reachable SWA blocks (#45444)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>

Wei Zhao committed Jun 18, 2026 at 05:23 UTC 5fd3b276f8fa34b70d3c83314700f626c66f9a22

4 files changed +120 -18

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

+10 -10

@@ -200,7 +200,7 @@ def test_store_mask_full_attention_all_true():

200 groups = [KVCacheGroupSpec(["L0"], _full(16))]

201 coord = _make_coord(groups, hash_block_size=16)

202 masks = coord.store_mask(64)

203 - assert masks == ([True, True, True, True],)

203 + assert masks == (None,)

204

205

206 def test_store_mask_zero_aligned_returns_empty_per_group():

@@ -210,7 +210,7 @@ def test_store_mask_zero_aligned_returns_empty_per_group():

210 ]

211 coord = _make_coord(groups, hash_block_size=16)

212 masks = coord.store_mask(0)

213 - assert masks == ([], [])

213 + assert masks == (None, None)

214

215

216 def test_store_mask_swa_only_window_around_each_lcm_boundary():

@@ -224,7 +224,7 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary():

224 coord = _make_coord(groups, hash_block_size=8)

225 masks = coord.store_mask(64)

226 # Full-attn: 2 chunks * 32 tokens.

227 - assert masks[0] == [True, True]

227 + assert masks[0] is None

228 # SWA: 8 chunks * 8 tokens. Only chunks ending at 32 and 64 are stored.

229 assert masks[1] == [False, False, False, True, False, False, False, True]

230

@@ -237,7 +237,7 @@ def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm():

237 groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)]

238 coord = _make_coord(groups, hash_block_size=8)

239 masks = coord.store_mask(64)

240 - assert masks[0] == [True, True]

240 + assert masks[0] is None

241 # Boundary at 32: blocks ending in [16, 32) — chunks 2 and 3.

242 # Boundary at 64: chunks 6 and 7. Others stay False.

243 assert masks[1] == [False, False, True, True, False, False, True, True]

@@ -265,12 +265,12 @@ def test_store_mask_dsv4_5_groups_full_mla_plus_4_swa():

265 masks = coord.store_mask(512)

266

267 # Full-MLA: 2 chunks of 256, both stored.

268 - assert masks[0] == [True, True]

268 + assert masks[0] is None

269 # SWA(64, sw=128): tail = ceil(127/64) = 2; C = 256/64 = 4.

270 # Per-segment template = [F,F,T,T]; tiled twice.

271 assert masks[1] == [False, False, True, True] * 2

272 # SWA(64, sw=512): tail = 8 >= C = 4 → entire segment True.

273 - assert masks[2] == [True] * 8

273 + assert masks[2] is None

274 # SWA(4, sw=16): tail = ceil(15/4) = 4; C = 256/4 = 64.

275 # Last 4 of each 64-chunk segment True.

276 assert masks[3] == ([False] * 60 + [True] * 4) * 2

@@ -289,7 +289,7 @@ def test_store_mask_fast_path_all_block_sizes_equal_lcm():

289 assert coord.lcm_block_size == 64

290 masks = coord.store_mask(256)

291 # Every block in every group is True — no sub-lcm filtering possible.

292 - assert masks == ([True] * 4, [True] * 4)

292 + assert masks == (None, None)

293

294

295 def test_store_mask_fast_path_single_attention_group():

@@ -300,7 +300,7 @@ def test_store_mask_fast_path_single_attention_group():

300 coord = _make_coord(groups, hash_block_size=16)

301 assert len(coord.attention_groups) == 1

302 masks = coord.store_mask(64)

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

303 + assert masks == (None, None)

304

305

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

@@ -319,7 +319,7 @@ def test_store_mask_dense_default_matches_every_lcm_boundary():

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]

322 + assert masks[0] is None

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

324

325

@@ -329,7 +329,7 @@ def test_store_mask_retention_interval_sparsifies_swa_tails():

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

332 + assert masks[0] is None # full attn unaffected

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

334

335

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+61

@@ -1233,6 +1233,67 @@ def test_lookup_swa_single_group_returns_full_when_tail_window_present():

1233 assert worker.lookup(64, [b"h0", b"h1", b"h2", b"h3"]) == 64

1234

1235

1236 +def test_lookup_checks_all_potential_swa_hit_boundaries():

1237 + """Lookup should skip SWA chunks that can never validate a hit, but still

1238 + check earlier aligned boundaries when sparse retention stores only the

1239 + current request's replay boundary.

1240 + """

1241 + from vllm.v1.kv_cache_interface import (

1242 + FullAttentionSpec,

1243 + KVCacheGroupSpec,

1244 + SlidingWindowSpec,

1245 + )

1246 +

1247 + worker = _make_bare_worker(block_size=8)

1248 + full = FullAttentionSpec(block_size=32, num_kv_heads=8, head_size=64, dtype=None)

1249 + swa = SlidingWindowSpec(

1250 + block_size=8, num_kv_heads=8, head_size=64, dtype=None, sliding_window=8

1251 + )

1252 + worker._kv_cache_groups = [

1253 + KVCacheGroupSpec(["full"], full),

1254 + KVCacheGroupSpec(["swa"], swa),

1255 + ]

1256 + worker.token_dbs = [

1257 + ChunkedTokenDatabase(

1258 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0),

1259 + block_size=32,

1260 + hash_block_size=8,

1261 + ),

1262 + ChunkedTokenDatabase(

1263 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),

1264 + block_size=8,

1265 + hash_block_size=8,

1266 + ),

1267 + ]

1268 + worker.coord = mooncake_store_worker.MooncakeStoreCoordinator(

1269 + worker._kv_cache_groups,

1270 + scheduler_block_size=32,

1271 + hash_block_size=8,

1272 + retention_interval=0,

1273 + )

1274 + # Candidate order: 3 full-attention chunks, then SWA chunks 3, 7, 11.

1275 + # Only the first full chunk and the SWA chunk ending at token 32 exist, so

1276 + # lookup should recover a 32-token external prefix hit. A sparse

1277 + # prompt-specific store mask for num_prompt_tokens=96 would only check SWA

1278 + # chunk 7 and miss this earlier reusable prefix.

1279 + worker.store.batch_is_exist.return_value = [1, 0, 0, 1, 0, 0]

1280 +

1281 + result = worker.lookup(

1282 + 96,

1283 + [f"h{i}".encode() for i in range(12)],

1284 + )

1285 +

1286 + assert result == 32

1287 + keys = worker.store.batch_is_exist.call_args.args[0]

1288 + assert len(keys) == 6

1289 + swa_keys = [key for key in keys if "@group:1@" in key]

1290 + assert swa_keys == [

1291 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6833",

1292 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6837",

1293 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@683131",

1294 + ]

1295 +

1296 +

1297 # ---------------------------------------------------------------------------

1298 # register_kv_caches tests

1299 # ---------------------------------------------------------------------------

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

+40 -7

@@ -172,19 +172,50 @@ class MooncakeStoreCoordinator:

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

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

178 - consume it.

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

176 + """Per-group store masks.

177 +

178 + ``mask[g][i]`` is True iff chunk ``i`` of group ``g`` should be

179 + written to the store so a future cache hit can consume it. ``None`` is

180 + the all-True sentinel.

181

182 Reuses the engine's ``SingleTypeKVCacheManager.reachable_block_mask``

183 so the store retains exactly the blocks the local prefix cache would.

184 """

185 + return self._reachable_masks(

186 + aligned_token_len,

187 + retention_interval=self.retention_interval,

188 + num_prompt_tokens=num_prompt_tokens,

189 + )

190 +

191 + def lookup_mask(

192 + self,

193 + aligned_token_len: int,

194 + ) -> tuple[list[bool] | None, ...]:

195 + """Per-group lookup masks.

196 +

197 + ``mask[g][i]`` is True iff chunk ``i`` of group ``g`` should be

198 + looked up as an aligned hit boundary. ``None`` is the all-True

199 + sentinel.

200 + """

201 + return self._reachable_masks(

202 + aligned_token_len,

203 + retention_interval=None,

204 + num_prompt_tokens=None,

205 + )

206 +

207 + def _reachable_masks(

208 + self,

209 + aligned_token_len: int,

210 + *,

211 + retention_interval: int | None,

212 + num_prompt_tokens: int | None,

213 + ) -> tuple[list[bool] | None, ...]:

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

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

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

217 )

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

218 + masks: list[list[bool] | None] = []

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

220 spec = _unwrap_spec(g.kv_cache_spec)

221 num_chunks = aligned_token_len // spec.block_size

@@ -196,10 +227,12 @@ class MooncakeStoreCoordinator:

227 alignment_tokens=self.lcm_block_size,

228 kv_cache_spec=spec,

229 use_eagle=g_idx in self.eagle_group_ids,

199 - retention_interval=self.retention_interval,

230 + retention_interval=retention_interval,

231 num_prompt_tokens=num_prompt_tokens,

232 )

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

233 + if mask is not None:

234 + assert len(mask) == num_chunks

235 + masks.append(mask)

236 return tuple(masks)

237

238 def block_hashes_for_spec(

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

+9 -1

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

548 for chunk_idx, (start, end, key) in enumerate(

549 db.process_tokens(token_len, req_meta.block_hashes)

550 ):

551 - if chunk_idx >= len(mask) or not mask[chunk_idx]:

551 + if mask is not None and (

552 + chunk_idx >= len(mask) or not mask[chunk_idx]

553 + ):

554 continue

555 starts.append(start)

556 ends.append(end)

@@ -1375,9 +1377,11 @@ class MooncakeStoreWorker:

1377 # candidate_meta[i] is the (group_id, hash_bytes) for candidate_keys[i].

1378 candidate_keys: list[str] = []

1379 candidate_meta: list[tuple[int, bytes]] = []

1380 + lookup_masks = self.coord.lookup_mask(token_len)

1381 tp_count = min(self.tp_size, self.num_kv_head)

1382 for g_idx, db in enumerate(self.token_dbs):

1383 spec_block_size = db.block_size

1384 + lookup_mask = lookup_masks[g_idx]

1385 group_hashes = self.coord.block_hashes_for_spec(

1386 block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec

1387 )

@@ -1385,6 +1389,10 @@ class MooncakeStoreWorker:

1389 start_idx = chunk_id * spec_block_size

1390 if start_idx >= token_len:

1391 break

1392 + if lookup_mask is not None and (

1393 + chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id]

1394 + ):

1395 + continue

1396 for tp in range(tp_count):

1397 for pp in range(self.pp_size):

1398 md = dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)