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

Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 15:15 UTC

[Core][DSV4] Skip caching SWA blocks that can never serve a prefix-cache hit (#42258)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>

Yifan Qiao committed May 15, 2026 at 00:59 UTC 4b364f810e12a689dd1862a81a163ccb6df07839

4 files changed +232 -9

tests/v1/core/test_prefix_caching.py

+139

@@ -2528,6 +2528,145 @@ def test_different_block_size():

2528 assert num_computed_tokens == 4 * 16

2529

2530

2531 +def test_hybrid_cache_blocks_swa_tail_window_only():

2532 + """Within each lcm-aligned segment, SWA's ``find_longest_cache_hit`` only

2533 + returns the trailing ``ceil((sliding_window - 1) / block_size)`` blocks

2534 + (its right-to-left scan stops once a contiguous match is found). Blocks

2535 + earlier in the segment can never serve a hit, so

2536 + ``HybridKVCacheCoordinator.cache_blocks`` should skip them rather than

2537 + polluting the prefix-cache hash map."""

2538 + block_size = 8

2539 + # Full attn block_size=32, SWA block_size=8, sw=8 -> lcm=32.

2540 + # tail = ceil(7/8) = 1; per_segment = 32/8 = 4.

2541 + # Per-segment template = [F, F, F, T]; only the last SWA block in each

2542 + # 32-token segment ends up in the prefix-cache hash map.

2543 + kv_cache_config = KVCacheConfig(

2544 + num_blocks=100,

2545 + kv_cache_tensors=[],

2546 + kv_cache_groups=[

2547 + KVCacheGroupSpec(

2548 + ["layer1"],

2549 + FullAttentionSpec(

2550 + block_size=4 * block_size,

2551 + num_kv_heads=1,

2552 + head_size=1,

2553 + dtype=torch.float16,

2554 + ),

2555 + ),

2556 + KVCacheGroupSpec(

2557 + ["layer2"],

2558 + SlidingWindowSpec(

2559 + block_size=block_size,

2560 + num_kv_heads=1,

2561 + head_size=1,

2562 + dtype=torch.float32,

2563 + sliding_window=block_size,

2564 + ),

2565 + ),

2566 + ],

2567 + )

2568 + manager = KVCacheManager(

2569 + kv_cache_config=kv_cache_config,

2570 + max_model_len=8192,

2571 + enable_caching=True,

2572 + hash_block_size=block_size,

2573 + )

2574 +

2575 + # 8 hash-blocks of 8 tokens (64 tokens, two lcm-aligned segments).

2576 + token_ids = [i for i in range(8) for _ in range(block_size)]

2577 + req = make_request("0", token_ids, block_size, sha256)

2578 + computed_blocks, _ = manager.get_computed_blocks(req)

2579 + blocks = manager.allocate_slots(

2580 + req,

2581 + 8 * block_size,

2582 + len(computed_blocks.blocks[0]) * block_size,

2583 + computed_blocks,

2584 + )

2585 + assert blocks is not None

2586 + assert len(req.block_hashes) == 8

2587 +

2588 + pool = manager.block_pool

2589 + # SWA group_id=1: only hash 3 and hash 7 (the last block of each

2590 + # 32-token segment) should be cached. Hashes 0,1,2,4,5,6 cannot serve

2591 + # a hit at any lcm-aligned length, so they must NOT be cached.

2592 + expected_cached = {3, 7}

2593 + for i in range(8):

2594 + cached = pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1])

2595 + if i in expected_cached:

2596 + assert cached is not None, f"SWA hash {i} should be cached"

2597 + else:

2598 + assert cached is None, (

2599 + f"SWA hash {i} cannot serve any lcm-aligned hit; should not be cached"

2600 + )

2601 +

2602 +

2603 +def test_hybrid_cache_blocks_clamped_to_lcm():

2604 + """HybridKVCacheCoordinator.cache_blocks() clamps to lcm_block_size.

2605 + Chunks past the last lcm-aligned boundary can never participate in a

2606 + cache hit (find_longest_cache_hit always returns lcm-aligned hits), so

2607 + caching them only pollutes the prefix-cache hash map and keeps blocks

2608 + on the LRU list that could otherwise return to the free pool."""

2609 + block_size = 16

2610 + # Full attn block_size=32, SWA block_size=16 -> lcm=32.

2611 + kv_cache_config = KVCacheConfig(

2612 + num_blocks=100,

2613 + kv_cache_tensors=[],

2614 + kv_cache_groups=[

2615 + KVCacheGroupSpec(

2616 + ["layer1"],

2617 + FullAttentionSpec(

2618 + block_size=block_size * 2,

2619 + num_kv_heads=1,

2620 + head_size=1,

2621 + dtype=torch.float16,

2622 + ),

2623 + ),

2624 + KVCacheGroupSpec(

2625 + ["layer2"],

2626 + SlidingWindowSpec(

2627 + block_size=block_size,

2628 + num_kv_heads=1,

2629 + head_size=1,

2630 + dtype=torch.float32,

2631 + sliding_window=2 * block_size,

2632 + ),

2633 + ),

2634 + ],

2635 + )

2636 + manager = KVCacheManager(

2637 + kv_cache_config=kv_cache_config,

2638 + max_model_len=8192,

2639 + enable_caching=True,

2640 + hash_block_size=block_size,

2641 + )

2642 +

2643 + # 7 hash-blocks of 16 tokens (112 tokens). With lcm=32 the clamp truncates

2644 + # to 96 tokens — SWA caches 6 hashes, full-attn caches 3.

2645 + token_ids = [i for i in range(7) for _ in range(block_size)]

2646 + req = make_request("0", token_ids, block_size, sha256)

2647 + computed_blocks, _ = manager.get_computed_blocks(req)

2648 + blocks = manager.allocate_slots(

2649 + req,

2650 + 7 * block_size,

2651 + len(computed_blocks.blocks[0]) * block_size,

2652 + computed_blocks,

2653 + )

2654 + assert blocks is not None

2655 + assert len(req.block_hashes) == 7

2656 +

2657 + pool = manager.block_pool

2658 + # SWA group_id=1: hashes 0..5 cached (6 blocks * 16 tokens = 96), hash 6

2659 + # spans tokens [96, 112) past the lcm boundary and must NOT be cached.

2660 + for i in range(6):

2661 + assert (

2662 + pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1])

2663 + is not None

2664 + ), f"SWA hash {i} should be cached"

2665 + assert pool.get_cached_block(req.block_hashes[6], kv_cache_group_ids=[1]) is None, (

2666 + "SWA hash 6 spans tokens past the lcm boundary; should not be cached"

2667 + )

2668 +

2669 +

2670 def test_block_lookup_cache_single_block_per_key():

2671 cache = BlockHashToBlockMap()

2672 key0 = BlockHashWithGroupId(b"hash0")

vllm/v1/core/block_pool.py

+16 -5

@@ -216,6 +216,7 @@ class BlockPool:

216 num_full_blocks: int,

217 block_size: int,

218 kv_cache_group_id: int,

219 + block_mask: list[bool] | None = None,

220 ) -> None:

221 """Cache a list of full blocks for prefix caching.

222 This function takes a list of blocks that will have their block hash

@@ -233,11 +234,19 @@ class BlockPool:

234 be cached after this function.

235 block_size: Number of tokens in each block.

236 kv_cache_group_id: The id of the KV cache group.

237 + block_mask: Optional mask aligned with

238 + ``blocks[num_cached_blocks:num_full_blocks]``. When provided,

239 + blocks where the mask is False are skipped (treated like null

240 + blocks). Used by groups whose ``find_longest_cache_hit`` only

241 + consults a subset of blocks (e.g. SWA tail-window), so blocks

242 + that can never serve a hit stay out of the prefix-cache hash

243 + map.

244 """

245 if num_cached_blocks >= num_full_blocks:

246 return

247 new_full_blocks = blocks[num_cached_blocks:num_full_blocks]

248 assert len(request.block_hashes) >= num_full_blocks

249 + assert block_mask is None or len(block_mask) == len(new_full_blocks)

250 if block_size == self.hash_block_size:

251 # Common case.

252 block_hashes: BlockHashList = request.block_hashes

@@ -256,10 +265,10 @@ class BlockPool:

265 [] if self.enable_kv_cache_events else None

266 )

267 for i, blk in enumerate(new_full_blocks):

259 - # Some blocks may be null blocks when enabling sparse attention like

260 - # sliding window attention, or Mamba models with prefix-caching in

261 - # align mode. We skip null blocks here.

262 - if blk.is_null:

268 + # Some blocks may be null or masked out when enabling sparse attention

269 + # like sliding window attention, or Mamba models with prefix-caching

270 + # in align mode. We skip null blocks here.

271 + if blk.is_null or (block_mask is not None and not block_mask[i]):

272 continue

273 assert blk.block_hash is None

274 block_hash = new_block_hashes[i]

@@ -288,12 +297,14 @@ class BlockPool:

297 # Generate extra keys for each block individually.

298 # Each block may have different extra_keys (e.g., different MM

299 # features, or cache_salt only for the first block).

291 - # Skip null blocks to match the length of new_hashes.

300 + # Skip null/masked-out blocks to match the length of new_hashes.

301 extra_keys_list: list[tuple[Any, ...] | None] = []

302 curr_mm_idx = 0

303 for i in range(num_cached_blocks, num_full_blocks):

304 if blocks[i].is_null:

305 continue

306 + if block_mask is not None and not block_mask[i - num_cached_blocks]:

307 + continue

308 block_start = i * block_size

309 block_end = block_start + block_size

310 extra_keys, curr_mm_idx = generate_block_hash_extra_keys(

vllm/v1/core/kv_cache_coordinator.py

+16

@@ -484,6 +484,22 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):

484 if any(gid in self.eagle_group_ids for gid in group_ids)

485 }

486

487 + def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:

488 + # Cache hits in this coordinator are always a multiple of

489 + # ``lcm_block_size`` tokens (see ``find_longest_cache_hit``). Within an

490 + # aligned region, SWA groups only consult a subset of blocks per

491 + # ``lcm_block_size``-segment so the unused blocks also stay out of the

492 + # prefix-cache hash map.

493 + num_computed_tokens = (

494 + num_computed_tokens // self.lcm_block_size * self.lcm_block_size

495 + )

496 + for manager in self.single_type_managers:

497 + manager.cache_blocks(

498 + request,

499 + num_computed_tokens,

500 + alignment_tokens=self.lcm_block_size,

501 + )

502 +

503 def find_longest_cache_hit(

504 self,

505 block_hashes: list[BlockHash],

vllm/v1/core/single_type_kv_cache_manager.py

+61 -4

@@ -275,7 +275,12 @@ class SingleTypeKVCacheManager(ABC):

275 self.new_block_ids = []

276 return ids

277

278 - def cache_blocks(self, request: Request, num_tokens: int) -> None:

278 + def cache_blocks(

279 + self,

280 + request: Request,

281 + num_tokens: int,

282 + alignment_tokens: int | None = None,

283 + ) -> None:

284 """

285 Cache the blocks for the request.

286

@@ -283,6 +288,12 @@ class SingleTypeKVCacheManager(ABC):

288 request: The request.

289 num_tokens: The total number of tokens that need to be cached

290 (including tokens that are already cached).

291 + alignment_tokens: The cache-hit alignment (in tokens) used by the

292 + coordinator's ``find_longest_cache_hit``. When greater than

293 + this group's ``block_size``, managers whose hit logic only

294 + returns a subset of blocks per alignment-aligned segment

295 + (SWA) skip the rest since they can never participate in a

296 + future cache hit.

297 """

298 num_cached_blocks = self.num_cached_block.get(request.request_id, 0)

299 num_full_blocks = num_tokens // self.block_size

@@ -290,6 +301,13 @@ class SingleTypeKVCacheManager(ABC):

301 if num_cached_blocks >= num_full_blocks:

302 return

303

304 + # Fast path: when the coordinator imposes no alignment constraint

305 + if alignment_tokens is None or alignment_tokens <= self.block_size:

306 + block_mask = None

307 + else:

308 + block_mask = self._cache_block_mask(

309 + num_cached_blocks, num_full_blocks, alignment_tokens

310 + )

311 self.block_pool.cache_full_blocks(

312 request=request,

313 blocks=self.req_to_blocks[request.request_id],

@@ -297,10 +315,26 @@ class SingleTypeKVCacheManager(ABC):

315 num_full_blocks=num_full_blocks,

316 block_size=self.block_size,

317 kv_cache_group_id=self.kv_cache_group_id,

318 + block_mask=block_mask,

319 )

320

321 self.num_cached_block[request.request_id] = num_full_blocks

322

323 + def _cache_block_mask(

324 + self,

325 + num_cached_blocks: int,

326 + num_full_blocks: int,

327 + alignment_tokens: int,

328 + ) -> list[bool] | None:

329 + """Per-block mask for ``cache_full_blocks``. ``None`` means cache

330 + every (non-null) block — the default for full attention.

331 +

332 + Subclasses with sparse hit semantics (SWA) override this to skip

333 + blocks that can never serve a hit at any alignment-aligned prefix

334 + length.

335 + """

336 + return None

337 +

338 def free(self, request_id: str) -> None:

339 """

340 Free the blocks for the request.

@@ -604,6 +638,19 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

638 computed.pop()

639 return computed_blocks

640

641 + def _cache_block_mask(

642 + self, num_cached_blocks: int, num_full_blocks: int, alignment_tokens: int

643 + ) -> list[bool] | None:

644 + assert alignment_tokens > self.block_size

645 + per_segment = alignment_tokens // self.block_size

646 + tail = cdiv(self.sliding_window - 1, self.block_size)

647 + if tail >= per_segment:

648 + return None

649 + skip = per_segment - tail

650 + return [

651 + i % per_segment >= skip for i in range(num_cached_blocks, num_full_blocks)

652 + ]

653 +

654 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:

655 """

656 Get the number of tokens that will be skipped for attention computation.

@@ -1050,9 +1097,14 @@ class MambaManager(SingleTypeKVCacheManager):

1097 """

1098 return num_computed_tokens - 1

1099

1053 - def cache_blocks(self, request: Request, num_tokens: int) -> None:

1100 + def cache_blocks(

1101 + self,

1102 + request: Request,

1103 + num_tokens: int,

1104 + alignment_tokens: int | None = None,

1105 + ) -> None:

1106 num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0)

1055 - super().cache_blocks(request, num_tokens)

1107 + super().cache_blocks(request, num_tokens, alignment_tokens=alignment_tokens)

1108 num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)

1109 if num_cached_blocks_after > num_cached_blocks_before:

1110 for block in self.req_to_blocks[request.request_id][

@@ -1081,7 +1133,12 @@ class CrossAttentionManager(SingleTypeKVCacheManager):

1133 # requests, so `new_computed_blocks` should always be empty.

1134 assert len(new_computed_blocks) == 0

1135

1084 - def cache_blocks(self, request: Request, num_tokens: int) -> None:

1136 + def cache_blocks(

1137 + self,

1138 + request: Request,

1139 + num_tokens: int,

1140 + alignment_tokens: int | None = None,

1141 + ) -> None:

1142 # We do not cache blocks for cross-attention to be shared between

1143 # requests, so this method is not relevant.

1144 raise ValueError("Should not be called as prefix caching is disabled.")