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

Fork 0

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

[Bugfix] Two-phase KV allocation for cross-group prefix cache hits (supersedes #33775) (#44409)

Signed-off-by: Saddss <2872669061@qq.com>

Saddss committed Jun 15, 2026 at 22:39 UTC 588db1836245bb40589d45b3e87048a53a13cfe5

4 files changed +255 -34

tests/v1/core/test_prefix_caching.py

+175 -1

@@ -22,7 +22,7 @@ from vllm.multimodal.inputs import (

22 from vllm.sampling_params import SamplingParams

23 from vllm.utils.hashing import sha256, sha256_cbor

24 from vllm.v1.core.block_pool import BlockHashToBlockMap, BlockPool

25 -from vllm.v1.core.kv_cache_manager import KVCacheManager, Request

25 +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager, Request

26 from vllm.v1.core.kv_cache_utils import (

27 BlockHash,

28 BlockHashWithGroupId,

@@ -3519,6 +3519,180 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized():

3519 assert manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is None

3520

3521

3522 +def test_cache_hit_local_and_external():

3523 + # Regression test for #33775: when a request hits the local prefix cache

3524 + # in one KV cache group and needs external (connector) blocks in another,

3525 + # the external allocation of an earlier group must not evict the local

3526 + # cache-hit blocks of a later group. Otherwise the same physical block can

3527 + # be handed out twice, producing duplicate block IDs / ref_cnt corruption.

3528 + block_size = 16

3529 + kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 100)

3530 + del kv_cache_config.kv_cache_groups[2:]

3531 + req_id = "test"

3532 + manager = make_kv_cache_manager(

3533 + kv_cache_config,

3534 + max_model_len=8192,

3535 + enable_caching=True,

3536 + hash_block_size=block_size,

3537 + use_eagle=True,

3538 + )

3539 +

3540 + top_blocks = []

3541 + head = manager.block_pool.free_block_queue.fake_free_list_head

3542 + for _ in range(10):

3543 + top_blocks.append(head.next_free_block)

3544 + head = head.next_free_block

3545 + cache_hit = KVCacheBlocks((top_blocks[:5], top_blocks[5:]))

3546 +

3547 + manager.allocate_slots(

3548 + make_request(req_id, [0] * (8 * block_size), block_size, sha256),

3549 + 16,

3550 + 5 * block_size,

3551 + cache_hit,

3552 + 0,

3553 + 2 * block_size,

3554 + )

3555 +

3556 + req_blocks = manager.get_blocks(req_id)

3557 + req_block_ids = req_blocks.get_block_ids()

3558 + all_block_ids = req_block_ids[0] + req_block_ids[1]

3559 + assert len(set(all_block_ids)) == len(all_block_ids), "Block IDs are not unique"

3560 +

3561 +

3562 +def _take_free_blocks(manager: KVCacheManager, num_blocks: int) -> list[KVCacheBlock]:

3563 + """Grab the first ``num_blocks`` blocks at the head of the free queue

3564 + without removing them. These ref_cnt==0 blocks stand in for evictable

3565 + cache-hit blocks left behind by a previous (e.g. preempted) request, and

3566 + sitting at the head guarantees a later group's external ``get_new_blocks``

3567 + would contend for them on unpatched code (issue #33775)."""

3568 + blocks: list[KVCacheBlock] = []

3569 + head = manager.block_pool.free_block_queue.fake_free_list_head

3570 + for _ in range(num_blocks):

3571 + head = head.next_free_block

3572 + blocks.append(head)

3573 + return blocks

3574 +

3575 +

3576 +def _assert_no_double_allocation(manager: KVCacheManager, req_id: str) -> None:

3577 + """No physical block may be handed out twice across groups, and every

3578 + block referenced by the request must have a live ref_cnt."""

3579 + block_ids = manager.get_blocks(req_id).get_block_ids()

3580 + flat = [block_id for group in block_ids for block_id in group]

3581 + assert len(set(flat)) == len(flat), "Block IDs are not unique across groups"

3582 + null_id = manager.block_pool.null_block.block_id

3583 + for block_id in flat:

3584 + if block_id == null_id:

3585 + continue

3586 + assert manager.block_pool.blocks[block_id].ref_cnt >= 1, (

3587 + f"block {block_id} referenced by the request has ref_cnt 0"

3588 + )

3589 +

3590 +

3591 +def _two_phase_block_size(manager: KVCacheManager) -> int:

3592 + return manager.kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size

3593 +

3594 +

3595 +def _cross_group_cache_hit(

3596 + manager: KVCacheManager,

3597 + req_id: str,

3598 + num_groups: int,

3599 + local_blocks_per_group: int = 5,

3600 + num_external_blocks: int = 2,

3601 + num_new_blocks: int = 1,

3602 +) -> Request:

3603 + """Allocate ``req_id`` with a per-group local prefix hit plus external

3604 + (connector) computed tokens, driving the coordinator's two-phase path.

3605 + Returns the allocated request so callers can free it (e.g. to preempt)."""

3606 + block_size = _two_phase_block_size(manager)

3607 + hit_blocks = _take_free_blocks(manager, num_groups * local_blocks_per_group)

3608 + cache_hit = KVCacheBlocks(

3609 + tuple(

3610 + hit_blocks[i * local_blocks_per_group : (i + 1) * local_blocks_per_group]

3611 + for i in range(num_groups)

3612 + )

3613 + )

3614 + prompt_blocks = local_blocks_per_group + num_external_blocks + num_new_blocks

3615 + request = make_request(

3616 + req_id, [0] * (prompt_blocks * block_size), block_size, sha256

3617 + )

3618 + manager.allocate_slots(

3619 + request,

3620 + num_new_blocks * block_size,

3621 + local_blocks_per_group * block_size,

3622 + cache_hit,

3623 + 0,

3624 + num_external_blocks * block_size,

3625 + )

3626 + return request

3627 +

3628 +

3629 +def _make_two_phase_manager(num_groups: int) -> KVCacheManager:

3630 + assert num_groups in (2, 3)

3631 + block_size = 16

3632 + kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 100)

3633 + del kv_cache_config.kv_cache_groups[num_groups:]

3634 + return make_kv_cache_manager(

3635 + kv_cache_config,

3636 + max_model_len=8192,

3637 + enable_caching=True,

3638 + hash_block_size=block_size,

3639 + use_eagle=True,

3640 + )

3641 +

3642 +

3643 +def test_cache_hit_local_and_external_three_groups():

3644 + # Scenario 1 (issue #33775): SWA + full attention with *three* KV cache

3645 + # groups (1 full + 2 sliding-window). A local prefix hit in some groups

3646 + # combined with external (connector) blocks in others must not let one

3647 + # group's external `get_new_blocks` evict another group's not-yet-touched

3648 + # cache-hit blocks, which would hand the same physical block out twice.

3649 + manager = _make_two_phase_manager(num_groups=3)

3650 + _cross_group_cache_hit(manager, "test", num_groups=3)

3651 + _assert_no_double_allocation(manager, "test")

3652 +

3653 +

3654 +def test_cache_hit_local_and_external_three_groups_preempt_and_reallocate():

3655 + # Scenario 2: the same 3-group hybrid config, but the request is preempted

3656 + # (freed) and then reallocated. After the free, the coordinator must treat

3657 + # the request as new again so external blocks are re-allocated, and the

3658 + # two-phase ordering must still prevent cross-group double allocation when

3659 + # reallocating against the now-evictable cache-hit blocks.

3660 + manager = _make_two_phase_manager(num_groups=3)

3661 +

3662 + request = _cross_group_cache_hit(manager, "test", num_groups=3)

3663 + _assert_no_double_allocation(manager, "test")

3664 +

3665 + # Preempt: free the request; its blocks return to the pool (full ones stay

3666 + # cached/evictable) and the coordinator forgets it.

3667 + manager.free(request)

3668 + assert manager.get_blocks("test").get_block_ids() == ([], [], [])

3669 +

3670 + # Reallocate the same request id against fresh cache-hit blocks taken from

3671 + # the current free-queue head, mirroring a preempted request being

3672 + # scheduled again. Because the request is no longer known, the coordinator

3673 + # re-arms `is_new_request` and re-runs external allocation, which must still

3674 + # not double-allocate across groups.

3675 + _cross_group_cache_hit(manager, "test", num_groups=3)

3676 + _assert_no_double_allocation(manager, "test")

3677 + assert manager.get_blocks("test").get_block_ids() != ([], [], [])

3678 +

3679 +

3680 +def test_cache_hit_local_and_external_two_groups_preempt_and_reallocate():

3681 + # Scenario 3: the minimal 2-group hybrid config (1 full + 1 sliding-window)

3682 + # exercised through the same preempt -> reallocate cycle as scenario 2.

3683 + manager = _make_two_phase_manager(num_groups=2)

3684 +

3685 + request = _cross_group_cache_hit(manager, "test", num_groups=2)

3686 + _assert_no_double_allocation(manager, "test")

3687 +

3688 + manager.free(request)

3689 + assert manager.get_blocks("test").get_block_ids() == ([], [])

3690 +

3691 + _cross_group_cache_hit(manager, "test", num_groups=2)

3692 + _assert_no_double_allocation(manager, "test")

3693 + assert manager.get_blocks("test").get_block_ids() != ([], [])

3694 +

3695 +

3696 def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(monkeypatch):

3697 """Default path (no retention): freeing an SWA request must place its

3698 uncached scratch blocks at the front of the free queue (recycled first)

tests/v1/core/test_single_type_kv_cache_manager.py

+1 -1

@@ -390,7 +390,7 @@ def test_evictable_cached_blocks_not_double_allocated():

390 # should only allocate the truly new block.

391 assert num_blocks_to_allocate == 2

392

393 - manager.allocate_new_computed_blocks(

393 + manager.add_local_computed_blocks(

394 request_id,

395 [evictable_block],

396 num_local_computed_tokens=block_size,

vllm/v1/core/kv_cache_coordinator.py

+21 -1

@@ -201,13 +201,33 @@ class KVCacheCoordinator(ABC):

201 num_local_computed_tokens: The number of local computed tokens.

202 num_external_computed_tokens: The number of external computed tokens.

203 """

204 + # A running request is already tracked in num_cached_block and won't

205 + # have new prefix-cache hits, so this is a no-op for it.

206 + if any(

207 + request_id in manager.num_cached_block

208 + for manager in self.single_type_managers

209 + ):

210 + assert all(len(blocks) == 0 for blocks in new_computed_blocks)

211 + return

212 +

213 + # Two-phase allocation (issue #33775): first touch every group's local

214 + # cache-hit blocks, then allocate external blocks for every group. This

215 + # ensures an earlier group's external `get_new_blocks` cannot evict a

216 + # later group's not-yet-touched cache-hit blocks.

217 for i, manager in enumerate(self.single_type_managers):

205 - manager.allocate_new_computed_blocks(

218 + manager.add_local_computed_blocks(

219 request_id,

220 new_computed_blocks[i],

221 num_local_computed_tokens,

222 num_external_computed_tokens,

223 )

224 + if num_external_computed_tokens > 0:

225 + for manager in self.single_type_managers:

226 + manager.allocate_external_computed_blocks(

227 + request_id,

228 + num_local_computed_tokens,

229 + num_external_computed_tokens,

230 + )

231

232 def allocate_new_blocks(

233 self,

vllm/v1/core/single_type_kv_cache_manager.py

+58 -31

@@ -179,7 +179,7 @@ class SingleTypeKVCacheManager(ABC):

179 )

180 return num_new_blocks + num_evictable_blocks

181

182 - def allocate_new_computed_blocks(

182 + def add_local_computed_blocks(

183 self,

184 request_id: str,

185 new_computed_blocks: Sequence[KVCacheBlock],

@@ -187,12 +187,11 @@ class SingleTypeKVCacheManager(ABC):

187 num_external_computed_tokens: int,

188 ) -> None:

189 """

190 - Add the new computed blocks to the request. This involves three steps:

191 - 1. Touch the computed blocks to make sure they won't be evicted.

192 - 1.5. (Optional) For sliding window, skip blocks are padded with null blocks.

190 + Add the locally cached (prefix-hit) blocks to the request:

191 + 1. Touch the computed blocks (paired with adding them to `req_blocks`)

192 + so their ref_cnt exactly tracks the referencing requests.

193 + 1.5. (Optional) For sliding window, skipped blocks are padded with nulls.

194 2. Add the remaining computed blocks.

194 - 3. (Optional) For KV connectors, allocate new blocks for external computed

195 - tokens (if any).

195

196 Args:

197 request_id: The request ID.

@@ -201,14 +200,8 @@ class SingleTypeKVCacheManager(ABC):

200 num_local_computed_tokens: The number of local computed tokens.

201 num_external_computed_tokens: The number of external computed tokens.

202 """

204 -

205 - if request_id in self.num_cached_block:

206 - # Fast-path: a running request won't have any new prefix-cache hits.

207 - # It should not have any new computed blocks.

208 - assert len(new_computed_blocks) == 0

209 - return

210 -

211 - # A new request.

203 + # The coordinator only calls this for first-time allocations (running

204 + # requests are short-circuited there), so the request has no blocks yet.

205 req_blocks = self.req_to_blocks[request_id]

206 assert len(req_blocks) == 0

207 num_total_computed_tokens = (

@@ -220,11 +213,6 @@ class SingleTypeKVCacheManager(ABC):

213 # It is possible that all new computed blocks are skipped when

214 # num_skipped_blocks > len(new_computed_blocks).

215 new_computed_blocks = new_computed_blocks[num_skipped_blocks:]

223 - # Some external computed tokens may be skipped too.

224 - num_external_computed_tokens = min(

225 - num_total_computed_tokens - num_skipped_tokens,

226 - num_external_computed_tokens,

227 - )

216

217 # Touch the computed blocks to make sure they won't be evicted.

218 if self.enable_caching:

@@ -243,18 +231,48 @@ class SingleTypeKVCacheManager(ABC):

231 # have a block_hash set.

232 self.num_cached_block[request_id] = len(req_blocks)

233

246 - if num_external_computed_tokens > 0:

247 - # Allocate new blocks for external computed tokens.

248 - allocated_blocks = self.block_pool.get_new_blocks(

249 - cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)

234 + def allocate_external_computed_blocks(

235 + self,

236 + request_id: str,

237 + num_local_computed_tokens: int,

238 + num_external_computed_tokens: int,

239 + ) -> None:

240 + """

241 + Allocate new blocks for external (KV-connector) computed tokens.

242 +

243 + Must run only after every group's local blocks have been touched via

244 + `add_local_computed_blocks`, so this group's `get_new_blocks` cannot

245 + evict another group's cache-hit blocks (issue #33775).

246 +

247 + Args:

248 + request_id: The request ID.

249 + num_local_computed_tokens: The number of local computed tokens.

250 + num_external_computed_tokens: The number of external computed tokens.

251 + """

252 + num_total_computed_tokens = (

253 + num_local_computed_tokens + num_external_computed_tokens

254 + )

255 + num_skipped_tokens = self.get_num_skipped_tokens(num_total_computed_tokens)

256 + if num_skipped_tokens > 0:

257 + # Some external computed tokens may be skipped too.

258 + num_external_computed_tokens = min(

259 + num_total_computed_tokens - num_skipped_tokens,

260 + num_external_computed_tokens,

261 )

251 - req_blocks.extend(allocated_blocks)

252 - if type(self.kv_cache_spec) in (

253 - FullAttentionSpec,

254 - TQFullAttentionSpec,

255 - MLAAttentionSpec,

256 - ):

257 - self.new_block_ids.extend(b.block_id for b in allocated_blocks)

262 + if num_external_computed_tokens <= 0:

263 + return

264 +

265 + req_blocks = self.req_to_blocks[request_id]

266 + allocated_blocks = self.block_pool.get_new_blocks(

267 + cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)

268 + )

269 + req_blocks.extend(allocated_blocks)

270 + if type(self.kv_cache_spec) in (

271 + FullAttentionSpec,

272 + TQFullAttentionSpec,

273 + MLAAttentionSpec,

274 + ):

275 + self.new_block_ids.extend(b.block_id for b in allocated_blocks)

276

277 def allocate_new_blocks(

278 self, request_id: str, num_tokens: int, num_tokens_main_model: int

@@ -1233,7 +1251,7 @@ class MambaManager(SingleTypeKVCacheManager):

1251 class CrossAttentionManager(SingleTypeKVCacheManager):

1252 """Manager for cross-attention KV cache in encoder-decoder models."""

1253

1236 - def allocate_new_computed_blocks(

1254 + def add_local_computed_blocks(

1255 self,

1256 request_id: str,

1257 new_computed_blocks: Sequence[KVCacheBlock],

@@ -1244,6 +1262,15 @@ class CrossAttentionManager(SingleTypeKVCacheManager):

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

1263 assert len(new_computed_blocks) == 0

1264

1265 + def allocate_external_computed_blocks(

1266 + self,

1267 + request_id: str,

1268 + num_local_computed_tokens: int,

1269 + num_external_computed_tokens: int,

1270 + ) -> None:

1271 + # Cross-attention does not use prefix caching / external KV loads.

1272 + return

1273 +

1274 def cache_blocks(

1275 self,

1276 request: Request,