From 481e481be786c1ca3229e26aa34c15ffd22375af Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Sun, 12 Jul 2026 13:37:51 +0800 Subject: [PATCH] [2/N][Core] support partial prefix cache hit for hybrid model (#46384) Signed-off-by: zjy0516 Signed-off-by: Yifan Qiao Co-authored-by: Yifan Qiao --- .../test_partial_prefix_cache_hits.py | 816 ++++++++++++++++++ tests/v1/core/test_deferred_block_free.py | 62 ++ tests/v1/core/test_kv_cache_utils.py | 60 ++ tests/v1/core/test_prefix_caching.py | 5 +- .../core/test_single_type_kv_cache_manager.py | 14 +- .../unit/test_mooncake_store_coordinator.py | 30 +- .../unit/test_mooncake_store_hma_e2e.py | 1 + vllm/config/cache.py | 20 +- .../v1/mooncake/store/coordinator.py | 26 +- .../kv_connector/v1/mooncake/store/worker.py | 4 +- vllm/engine/arg_utils.py | 5 + vllm/v1/core/block_pool.py | 26 +- vllm/v1/core/kv_cache_coordinator.py | 116 ++- vllm/v1/core/kv_cache_manager.py | 23 +- vllm/v1/core/kv_cache_utils.py | 57 +- vllm/v1/core/sched/output.py | 5 + vllm/v1/core/sched/scheduler.py | 76 +- vllm/v1/core/single_type_kv_cache_manager.py | 422 +++++++-- vllm/v1/worker/gpu/model_runner.py | 11 +- vllm/v1/worker/gpu_model_runner.py | 7 + vllm/v1/worker/utils.py | 44 +- 21 files changed, 1653 insertions(+), 177 deletions(-) create mode 100644 tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py new file mode 100644 index 00000000000..51aefe63aaa --- /dev/null +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -0,0 +1,816 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fine-grained partial prefix-cache hits for hybrid (full attention + mamba +"align") models: scheduler chunk splitting, partial tail registration, CoW +on partial hits, and same-step deferral.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tests.v1.core.test_prefix_caching import make_kv_cache_manager, make_request +from vllm.utils.hashing import sha256 +from vllm.v1.core.kv_cache_utils import ( + KVCacheBlockCopy, + get_block_hash, + get_group_id, + init_none_hash, +) +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) + + +@pytest.fixture(autouse=True) +def _auto_init_hash_fn(): + init_none_hash(sha256) + + +def test_mamba_align_split_partial_tail_schedule(): + """Chunk ends with partial hits on: block-aligned chunks, one extra stop + at the prompt's last hash boundary (registering the partial tail), then + the remaining tokens. block=512, hash=32, prompt=10000, budget=8192: + 0 -> 8192 -> 9728 -> 9984 -> 10000.""" + block_size = 512 + hash_block_size = 32 + mock = SimpleNamespace( + cache_config=SimpleNamespace(block_size=block_size), + use_eagle=False, + hash_block_size=hash_block_size, + mamba_partial_cache_hit=True, + ) + split = Scheduler._mamba_block_aligned_split + + req = make_request("0", [0] * 10000, hash_block_size, sha256) + req.num_computed_tokens = 0 + assert split(self=mock, request=req, num_new_tokens=8192) == 8192 + req.num_computed_tokens = 8192 + # Stop at the last block boundary (9728). + assert split(self=mock, request=req, num_new_tokens=1808) == 1536 + req.num_computed_tokens = 9728 + # Extra stop at the prompt's last hash boundary (9984). + assert split(self=mock, request=req, num_new_tokens=272) == 256 + req.num_computed_tokens = 9984 + # Final 16 tokens run unchanged (no mid-block-resume stop: the next + # block boundary is past the last block boundary). + assert split(self=mock, request=req, num_new_tokens=16) == 16 + + # Partial hits off: no extra stop, the tail runs in one chunk. + mock.mamba_partial_cache_hit = False + req.num_computed_tokens = 9728 + assert split(self=mock, request=req, num_new_tokens=272) == 272 + mock.mamba_partial_cache_hit = True + + # A request resumed mid-block (partial hash hit at 9984): the first chunk + # stops at the next block boundary (10240), later chunk ends re-align. + req2 = make_request("1", [0] * 12000, hash_block_size, sha256) + req2.num_computed_tokens = 9984 + assert split(self=mock, request=req2, num_new_tokens=2016) == 256 + req2.num_computed_tokens = 10240 + assert split(self=mock, request=req2, num_new_tokens=1000) == 512 + + +def test_hybrid_mamba_align_partial_hash_hit(): + hash_block_size = 2 + mamba_block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=20, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert num_computed == 0 + blocks = manager.allocate_slots(req0, 6, num_computed, computed_blocks) + assert blocks is not None + manager.free(req0) + manager.new_step_starts() + + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1] + partial_mamba_block = manager.block_pool.get_cached_block( + partial_mamba_hash, kv_cache_group_ids=[1] + ) + assert partial_mamba_block is not None + assert partial_mamba_block[0].block_hash_num_tokens == 6 + + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req1) + assert num_computed == 6 + assert [len(group) for group in computed_blocks.blocks] == [3, 2] + + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks) + assert new_blocks is not None + mamba_new_block_ids = new_blocks.get_block_ids()[1] + assert len(mamba_new_block_ids) == 1 + assert mamba_new_block_ids[0] != partial_mamba_block[0].block_id + assert manager.get_blocks("1").get_block_ids()[1][1] == mamba_new_block_ids[0] + assert partial_mamba_block[0].block_hash is not None + assert get_block_hash(partial_mamba_block[0].block_hash) == partial_mamba_hash + assert get_group_id(partial_mamba_block[0].block_hash) == 1 + assert partial_mamba_block[0].block_hash_num_tokens == 6 + copies, _ = manager.take_kv_cache_block_copies() + assert ( + KVCacheBlockCopy( + src_block_id=partial_mamba_block[0].block_id, + dst_block_id=mamba_new_block_ids[0], + ) + in copies + ) + assert manager.get_blocks("1").blocks[1][1].block_hash_num_tokens == 8 + + +def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert num_computed == 0 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1] + partial_mamba_block = manager.block_pool.get_cached_block( + partial_mamba_hash, kv_cache_group_ids=[1] + ) + assert partial_mamba_block is not None + partial_mamba_block_id = partial_mamba_block[0].block_id + assert manager.get_blocks("0").get_block_ids()[1][1] == partial_mamba_block_id + + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + new_blocks = manager.allocate_slots(req0, 1) + assert new_blocks is not None + + # Reversed CoW for the owning request: it keeps its own block (the + # worker's block table is append-only), and no new mamba block is handed + # to the worker. The prefix-cache entry is moved to a private copy that + # the queued block copy fills before the next forward. + assert new_blocks.get_block_ids()[1] == [] + assert manager.get_blocks("0").get_block_ids()[1][1] == partial_mamba_block_id + copies, _ = manager.take_kv_cache_block_copies() + cow_copy = next(c for c in copies if c.src_block_id == partial_mamba_block_id) + assert cow_copy.dst_block_id != partial_mamba_block_id + # The source block gave up the hash; the copy target now owns the entry. + assert partial_mamba_block[0].block_hash is None + moved = manager.block_pool.get_cached_block( + partial_mamba_hash, kv_cache_group_ids=[1] + ) + assert moved is not None + assert moved[0].block_id == cow_copy.dst_block_id + assert get_block_hash(moved[0].block_hash) == partial_mamba_hash + assert get_group_id(moved[0].block_hash) == 1 + assert moved[0].block_hash_num_tokens == 6 + + +def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit(): + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert num_computed == 0 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + + partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1] + partial_mamba_block = manager.block_pool.get_cached_block( + partial_mamba_hash, kv_cache_group_ids=[1] + ) + assert partial_mamba_block is not None + partial_mamba_block_id = partial_mamba_block[0].block_id + + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + # The owner moved the prefix-cache entry to a private copy; capture its id. + owner_copies, _ = manager.take_kv_cache_block_copies() + cow_copy = next(c for c in owner_copies if c.src_block_id == partial_mamba_block_id) + moved_block_id = cow_copy.dst_block_id + manager.new_step_starts() + + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req1) + assert num_computed == 6 + # The later request hits the moved (private-copy) entry, not the source. + assert computed_blocks.get_block_ids()[1][1] == moved_block_id + + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks) + assert new_blocks is not None + mamba_new_block_ids = new_blocks.get_block_ids()[1] + assert len(mamba_new_block_ids) == 1 + assert mamba_new_block_ids[0] != moved_block_id + # The hitting request CoWs from the moved entry into its own private block. + copies, _ = manager.take_kv_cache_block_copies() + assert ( + KVCacheBlockCopy( + src_block_id=moved_block_id, + dst_block_id=mamba_new_block_ids[0], + ) + in copies + ) + + +def test_hybrid_mamba_moved_partial_entry_defers_same_step_hit(): + """The owner's move re-arms the same-step guard: the moved entry is + filled by this step's copy, and chained same-step copies read stale + sources, so a request hitting it in the move step must be deferred.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert num_computed == 0 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + manager.new_step_starts() + + # The owning request continues decoding: the partial entry moves to a + # private copy in this step. + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + + # A request hitting the moved entry in the SAME step must be deferred. + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 4, 4], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req1) + assert num_computed == 6 + assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is None + + # Next step the moved entry is consumable. + manager.new_step_starts() + computed_blocks, num_computed = manager.get_computed_blocks(req1) + assert num_computed == 6 + assert manager.allocate_slots(req1, 2, num_computed, computed_blocks) is not None + + +def test_hybrid_full_attention_partial_hash_hit_uses_cow(): + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert num_computed == 0 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + partial_full_hash = req0.block_hashes[6 // hash_block_size - 1] + partial_full_block = manager.block_pool.get_cached_block( + partial_full_hash, kv_cache_group_ids=[0] + ) + assert partial_full_block is not None + + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req1) + assert num_computed == 6 + assert [len(group) for group in computed_blocks.blocks] == [2, 2] + + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks) + assert new_blocks is not None + full_new_block_ids = new_blocks.get_block_ids()[0] + assert len(full_new_block_ids) == 1 + assert full_new_block_ids[0] != partial_full_block[0].block_id + assert partial_full_block[0].block_hash is not None + assert get_block_hash(partial_full_block[0].block_hash) == partial_full_hash + assert get_group_id(partial_full_block[0].block_hash) == 0 + assert partial_full_block[0].block_hash_num_tokens == 6 + copies, retained = manager.take_kv_cache_block_copies() + assert ( + KVCacheBlockCopy( + src_block_id=partial_full_block[0].block_id, + dst_block_id=full_new_block_ids[0], + ) + in copies + ) + assert partial_full_block[0].ref_cnt == 1 + manager.block_pool.free_blocks(retained) + assert partial_full_block[0].ref_cnt == 0 + + +def test_hybrid_partial_hit_cow_target_starts_uncached(): + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert num_computed == 0 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + partial_hash = req0.block_hashes[6 // hash_block_size - 1] + partial_full_block = manager.block_pool.get_cached_block( + partial_hash, kv_cache_group_ids=[0] + ) + partial_mamba_block = manager.block_pool.get_cached_block( + partial_hash, kv_cache_group_ids=[1] + ) + assert partial_full_block is not None + assert partial_mamba_block is not None + + req1 = make_request("1", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req1) + assert num_computed == 6 + + new_blocks = manager.allocate_slots( + req1, + 2, + num_computed, + computed_blocks, + delay_cache_blocks=True, + ) + assert new_blocks is not None + + full_cow_block = manager.get_blocks("1").blocks[0][1] + mamba_cow_block = manager.get_blocks("1").blocks[1][1] + assert full_cow_block.block_id != partial_full_block[0].block_id + assert mamba_cow_block.block_id != partial_mamba_block[0].block_id + assert full_cow_block.block_hash is None + assert full_cow_block.block_hash_num_tokens is None + assert mamba_cow_block.block_hash is None + assert mamba_cow_block.block_hash_num_tokens is None + + assert partial_full_block[0].block_hash is not None + assert get_block_hash(partial_full_block[0].block_hash) == partial_hash + assert get_group_id(partial_full_block[0].block_hash) == 0 + assert partial_full_block[0].block_hash_num_tokens == 6 + assert partial_mamba_block[0].block_hash is not None + assert get_block_hash(partial_mamba_block[0].block_hash) == partial_hash + assert get_group_id(partial_mamba_block[0].block_hash) == 1 + assert partial_mamba_block[0].block_hash_num_tokens == 6 + + +def test_hybrid_partial_hash_truncates_full_attention_hit_length(): + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + pool = manager.block_pool + req = make_request( + "0", + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5], + hash_block_size, + sha256, + ) + + full_blocks = pool.get_new_blocks(3) + pool.cache_full_blocks( + request=req, + blocks=full_blocks, + num_cached_blocks=0, + num_full_blocks=2, + block_size=block_size, + kv_cache_group_id=0, + ) + pool.cache_partial_block( + request=req, + block=full_blocks[2], + num_tokens=10, + kv_cache_group_id=0, + block_size=block_size, + ) + + mamba_block = pool.get_new_blocks(1)[0] + pool.cache_partial_block( + request=req, + block=mamba_block, + num_tokens=6, + kv_cache_group_id=1, + block_size=block_size, + ) + + computed_blocks, num_computed = manager.get_computed_blocks(req) + assert num_computed == 6 + assert [len(group) for group in computed_blocks.blocks] == [2, 2] + + +def test_cow_retained_blocks_returned_for_release(): + """new_step_starts returns the CoW copy retentions instead of freeing + them; the scheduler owns releasing them once the copy has run.""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=24, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + req0 = make_request("0", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + + # The owner's move queues a copy and retains both endpoints. + req0.num_computed_tokens = 6 + req0.append_output_token_ids([3]) + assert manager.allocate_slots(req0, 1) is not None + (cow_copy,), retained = manager.take_kv_cache_block_copies() + assert {b.block_id for b in retained} == { + cow_copy.src_block_id, + cow_copy.dst_block_id, + } + # Not freed yet: the retention refs are still held. + assert all(b.ref_cnt > 0 for b in retained) + manager.block_pool.free_blocks(retained) + + +def test_free_cow_retained_blocks_defers_until_copy_step_processed(): + """Scheduler releases CoW retentions immediately when the copy's step has + been processed (or deferral is off), and defers them otherwise.""" + from collections import deque + + freed: list = [] + blocks = [SimpleNamespace(block_id=7), SimpleNamespace(block_id=9)] + mock = SimpleNamespace( + kv_cache_manager=SimpleNamespace( + block_pool=SimpleNamespace(free_blocks=freed.extend) + ), + deferred_frees=deque(), + defer_block_free=True, + processed_step_seq=2, + ) + free = Scheduler._free_cow_retained_blocks + + # Copy step still in flight: deferred with its fence. + free(mock, list(blocks), fence_seq=3) + assert not freed + assert mock.deferred_frees == deque([(3, blocks[::-1])]) + + # Copy step processed: freed immediately. + mock.processed_step_seq = 3 + free(mock, list(blocks), fence_seq=3) + assert freed == blocks + + # Deferral disabled: freed immediately regardless of the fence. + freed.clear() + mock.deferred_frees.clear() + mock.defer_block_free = False + mock.processed_step_seq = 0 + free(mock, list(blocks), fence_seq=3) + assert freed == blocks + + +def test_full_attention_eagle_drops_one_hash_unit(): + """With fine-grained partial hits, eagle rewinds the hit by one hash unit + instead of a whole cache block: the tail block's KV is append-only, so it + still covers the reduced length and stays in the hit as a partial block.""" + from vllm.v1.core.block_pool import BlockPool + from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager + + hash_block_size = 2 + block_size = 4 + pool = BlockPool( + num_gpu_blocks=10, enable_caching=True, hash_block_size=hash_block_size + ) + spec = FullAttentionSpec( + block_size=block_size, num_kv_heads=1, head_size=1, dtype=torch.float32 + ) + req = make_request("0", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256) + + def find(drop_eagle_block): + return FullAttentionManager.find_longest_cache_hit( + block_hashes=req.block_hashes, + max_length=8, + kv_cache_group_ids=[0], + block_pool=pool, + kv_cache_spec=spec, + drop_eagle_block=drop_eagle_block, + alignment_tokens=hash_block_size, + ) + + # Two full cached blocks (hit 8): eagle rewinds to 6, keeping the last + # block as a partial hit instead of dropping it to 4. + blocks = pool.get_new_blocks(2) + pool.cache_full_blocks( + request=req, + blocks=blocks, + num_cached_blocks=0, + num_full_blocks=2, + block_size=block_size, + kv_cache_group_id=0, + ) + hit_blocks, hit_length = find(drop_eagle_block=False) + assert (hit_length, len(hit_blocks[0])) == (8, 2) + hit_blocks, hit_length = find(drop_eagle_block=True) + assert (hit_length, len(hit_blocks[0])) == (6, 2) + + # A partial tail at 6 (block 1 not fully cached): eagle rewinds to the + # block boundary and trims the tail block. + pool2 = BlockPool( + num_gpu_blocks=10, enable_caching=True, hash_block_size=hash_block_size + ) + pool = pool2 + blocks = pool.get_new_blocks(2) + pool.cache_full_blocks( + request=req, + blocks=blocks[:1], + num_cached_blocks=0, + num_full_blocks=1, + block_size=block_size, + kv_cache_group_id=0, + ) + assert ( + pool.cache_partial_block( + request=req, + block=blocks[1], + num_tokens=6, + kv_cache_group_id=0, + block_size=block_size, + ) + is not None + ) + hit_blocks, hit_length = find(drop_eagle_block=False) + assert (hit_length, len(hit_blocks[0])) == (6, 2) + hit_blocks, hit_length = find(drop_eagle_block=True) + assert (hit_length, len(hit_blocks[0])) == (4, 1) + + +def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks(): + """Regression: with eagle, the mamba group must not receive the eagle + lookup margin — its finder never applies the drop, so it could return a + hit past the blocks the (dropped) full-attention group covers, crashing + the consumer's CoW with block_idx >= len(req_blocks).""" + hash_block_size = 2 + block_size = 2 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + use_eagle=True, + ) + + # The owner prefills in scheduler-split style: stop at the block boundary + # (4), then at the prompt's last hash boundary (6, partial entries). + req0 = make_request("0", [7] * 6, hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 4 + manager.new_step_starts() + assert manager.allocate_slots(req0, 2) is not None + req0.num_computed_tokens = 6 + manager.new_step_starts() + + # A longer request with eagle: full attention drops the partial tail, so + # the joint hit must fall back to the block boundary the FA blocks cover. + req1 = make_request("1", [7] * 6 + [9] * 2, hash_block_size, sha256) + computed_blocks, num_computed = manager.get_computed_blocks(req1) + assert num_computed == 4 + assert all( + len(group) * block_size >= num_computed for group in computed_blocks.blocks + ) + assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None diff --git a/tests/v1/core/test_deferred_block_free.py b/tests/v1/core/test_deferred_block_free.py index 8cab620f0e3..13789a39601 100644 --- a/tests/v1/core/test_deferred_block_free.py +++ b/tests/v1/core/test_deferred_block_free.py @@ -18,6 +18,7 @@ from unittest.mock import PropertyMock, patch import pytest from vllm.config import VllmConfig +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import RequestStatus @@ -412,3 +413,64 @@ def test_non_async_abort_defers_via_last_sched_seq(): scheduler.update_from_output(out0, _make_model_runner_output(out0)) assert not scheduler.deferred_frees assert pool.get_num_free_blocks() == num_free_initially + + +def test_cow_retentions_deferred_until_copy_step_processed(): + """The endpoints of a queued KV block copy must stay out of the free + pool until the step that runs the copy has been processed. Freed + earlier, an endpoint can be reallocated (e.g. as a PD KV-load + destination) and overwritten by a transfer that is not ordered against + the copy still pending in the in-flight step. + """ + scheduler = _create_deferring_scheduler() + pool = scheduler.kv_cache_manager.block_pool + manager = scheduler.kv_cache_manager.coordinator.single_type_managers[0] + + request = create_requests( + num_requests=1, + num_tokens=NUM_PROMPT_TOKENS, + max_tokens=5, + stop_token_ids=[STOP_TOKEN_ID], + )[0] + scheduler.add_request(request) + + # Simulate a partial-hit CoW performed while scheduling step 1, whose + # hitting request was freed within the same step: the copy rides out0 + # and each endpoint stays alive only through its copy retention. + src_block, dst_block = pool.get_new_blocks(2) + block_copy = KVCacheBlockCopy( + src_block_id=src_block.block_id, dst_block_id=dst_block.block_id + ) + manager._pending_cow_copies.append((src_block, dst_block)) + out0 = scheduler.schedule() + assert out0.kv_cache_block_copies == [block_copy] + + # Exhaust the rest of the pool so the copy endpoints are the only blocks + # a new request could receive, then add one that fits exactly in them. + pool.get_new_blocks(pool.get_num_free_blocks()) + late_request = create_requests( + num_requests=1, + num_tokens=2 * scheduler.block_size, + max_tokens=5, + req_ids=["late"], + )[0] + scheduler.add_request(late_request) + + # Step 2 is scheduled while step 1 (which runs the copy) is still in + # flight: the retentions are released against the copy's fence, so the + # endpoints must not reach the free pool -- the late request must not be + # scheduled onto them. + out1 = scheduler.schedule() + assert src_block.ref_cnt == 1 + assert dst_block.ref_cnt == 1 + assert scheduler.deferred_frees + assert not out1.scheduled_new_reqs + + # Step 1's output is processed: the copy has run, endpoints return to + # the pool and the late request can be scheduled onto them safely. + scheduler.update_from_output(out0, _make_model_runner_output(out0)) + assert src_block.ref_cnt == 0 + assert dst_block.ref_cnt == 0 + assert not scheduler.deferred_frees + out2 = scheduler.schedule() + assert [r.req_id for r in out2.scheduled_new_reqs] == ["late"] diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 947672c48af..0c212cb089a 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -2597,3 +2597,63 @@ def test_hma_not_disabled_when_kv_events_enabled(): assert vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is False, ( "kv_events_config must not force-disable the hybrid KV cache manager." ) + + +def test_resolve_block_hashes_gate(): + # Resolve symbols through the module so they stay consistent with each other + # even after other tests reload ``kv_cache_utils``. + resolve_block_hashes = kv_cache_utils.resolve_block_hashes + BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize + # Raw, hash_block_size-granularity hashes (contents are opaque here). + raw = [BlockHash(bytes([i])) for i in range(8)] + + # block_size == hash_block_size: always reuse the raw hashes. + assert resolve_block_hashes(raw, 2, 2, alignment_tokens=2) is raw + assert ( + resolve_block_hashes( + raw, 2, 2, supports_fine_grained_hash_lookup=True, alignment_tokens=2 + ) + is raw + ) + + # Fine-grained manager, partial hits ON (alignment_tokens == hash_block_size + # < block_size): keep raw hashes so the manager can scan at hash granularity. + assert ( + resolve_block_hashes( + raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=2 + ) + is raw + ) + + # Fine-grained manager, partial hits OFF (alignment_tokens == + # scheduler_block_size >= block_size): must fall back to a block-size view, + # exactly like the pre-refactor coordinator did. + off = resolve_block_hashes( + raw, 2, 4, supports_fine_grained_hash_lookup=True, alignment_tokens=4 + ) + assert isinstance(off, BlockHashListWithBlockSize) + assert off.scale_factor == 2 + + # Non-fine-grained manager (e.g. sliding window): always a block-size view + # when block_size != hash_block_size, regardless of alignment_tokens. + swa = resolve_block_hashes( + raw, 2, 4, supports_fine_grained_hash_lookup=False, alignment_tokens=2 + ) + assert isinstance(swa, BlockHashListWithBlockSize) + assert swa.scale_factor == 2 + + +def test_resolve_block_hashes_rejects_mismatched_view(): + resolve_block_hashes = kv_cache_utils.resolve_block_hashes + BlockHashListWithBlockSize = kv_cache_utils.BlockHashListWithBlockSize + raw = [BlockHash(bytes([i])) for i in range(8)] + + # A view built at this block_size is returned as-is (idempotent). + view = BlockHashListWithBlockSize(raw, 2, 4) + assert resolve_block_hashes(view, 2, 4) is view + + # A view built at a different block_size must fail loudly rather than be + # silently reinterpreted at the wrong granularity. + mismatched = BlockHashListWithBlockSize(raw, 2, 8) + with pytest.raises(AssertionError): + resolve_block_hashes(mismatched, 2, 4) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 03d61ec299f..d0dc8a97d63 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1072,7 +1072,10 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): # Next, validate scheduler logic for num_uncached_common_prefix_tokens > 0 # Create minimal mock with just the needed attributes mock = SimpleNamespace( - cache_config=SimpleNamespace(block_size=block_size), use_eagle=False + cache_config=SimpleNamespace(block_size=block_size), + use_eagle=False, + hash_block_size=block_size, + mamba_partial_cache_hit=False, ) num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split( self=mock, diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 609c1428d19..7300d45fbfc 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -93,7 +93,7 @@ def test_chunked_local_attention_possible_cached_prefix(): kv_cache_spec=chunked_local_attention_spec, drop_eagle_block=False, alignment_tokens=block_size, - )[0] + )[0][0] assert len(computed_blocks) == expect_length assert all( @@ -164,7 +164,7 @@ def test_sliding_window_possible_cached_prefix(): kv_cache_spec=sliding_window_spec, drop_eagle_block=False, alignment_tokens=block_size, - )[0] + )[0][0] assert len(computed_blocks) == expect_length assert all( @@ -398,13 +398,13 @@ def test_get_num_blocks_to_allocate(): assert ( manager.get_num_blocks_to_allocate( - "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size + "1", 20 * block_size, cached_blocks_1, 0, 0, 20 * block_size ) == 20 ) assert ( manager.get_num_blocks_to_allocate( - "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size + "2", 20 * block_size, cached_blocks_2, 0, 0, 20 * block_size ) == 15 ) @@ -434,6 +434,7 @@ def test_evictable_cached_blocks_not_double_allocated(): num_tokens=2 * block_size, new_computed_blocks=[evictable_block], total_computed_tokens=block_size, + num_local_computed_tokens=block_size, num_tokens_main_model=2 * block_size, ) # Free capacity check should count evictable cached blocks, but allocation @@ -474,13 +475,13 @@ def test_chunked_local_attention_get_num_blocks_to_allocate(): assert ( manager.get_num_blocks_to_allocate( - "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size + "1", 20 * block_size, cached_blocks_1, 0, 0, 20 * block_size ) == 20 ) assert ( manager.get_num_blocks_to_allocate( - "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size + "2", 20 * block_size, cached_blocks_2, 0, 0, 20 * block_size ) == 15 ) @@ -524,6 +525,7 @@ def test_predictor_matches_allocator_blocks_calculation_with_admission_cap(): num_tokens=num_tokens, new_computed_blocks=[], total_computed_tokens=total_computed, + num_local_computed_tokens=0, num_tokens_main_model=num_tokens, ) new_blocks = manager.allocate_new_blocks( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index 2ad4b79164a..6e003798c7a 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -37,7 +37,7 @@ def _make_coord(groups, hash_block_size, use_eagle=False, retention_interval=Non def test_external_cached_block_pool_tautological_returns_present_for_any_hash(): - cmap = ExternalCachedBlockPool() + cmap = ExternalCachedBlockPool(16) h = BlockHash(b"\xaa" * 4) res = cmap.get_cached_block(h, [0, 1]) assert res is not None @@ -48,7 +48,7 @@ def test_external_cached_block_pool_tautological_returns_present_for_any_hash(): def test_external_cached_block_pool_hit_all_groups(): h = BlockHash(b"\x11\x22\x33\x44") - cmap = ExternalCachedBlockPool({(0, bytes(h)), (1, bytes(h))}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)), (1, bytes(h))}) res = cmap.get_cached_block(h, [0, 1]) assert res is not None assert len(res) == 2 @@ -58,14 +58,14 @@ def test_external_cached_block_pool_hit_all_groups(): def test_external_cached_block_pool_miss_one_group(): h = BlockHash(b"\x11\x22\x33\x44") - cmap = ExternalCachedBlockPool({(0, bytes(h))}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(h))}) assert cmap.get_cached_block(h, [0, 1]) is None def test_external_cached_block_pool_unknown_hash(): h_known = BlockHash(b"\x01" * 4) h_unknown = BlockHash(b"\x02" * 4) - cmap = ExternalCachedBlockPool({(0, bytes(h_known))}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(h_known))}) assert cmap.get_cached_block(h_unknown, [0]) is None @@ -103,7 +103,7 @@ def test_coordinator_single_full_attention_all_hits(): groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs}) masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap) assert hit == 64 assert masks[0] == [True, True, True, True] @@ -113,7 +113,7 @@ def test_coordinator_single_full_attention_partial_prefix(): groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) - cmap = ExternalCachedBlockPool({(0, bytes(hs[0])), (0, bytes(hs[1]))}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(hs[0])), (0, bytes(hs[1]))}) masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap) assert hit == 32 assert masks[0] == [True, True] @@ -123,7 +123,7 @@ def test_coordinator_single_full_attention_no_hits(): groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) - cmap = ExternalCachedBlockPool(set()) + cmap = ExternalCachedBlockPool(16, set()) masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap) assert hit == 0 assert masks[0] == [] @@ -135,7 +135,7 @@ def test_coordinator_single_swa_tautological_pool_masks_pre_window(): groups = [KVCacheGroupSpec(["L0"], _swa(block_size=16, sliding_window=32))] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) # 4 chunks * 16 tokens - cmap = ExternalCachedBlockPool() + cmap = ExternalCachedBlockPool(16) masks, hit = coord.find_longest_cache_hit(hs, max_length=64, cached_block_pool=cmap) assert hit == 64 # ceil((sw-1)/block_size) = ceil(31/16) = 2 tail blocks. @@ -153,7 +153,7 @@ def test_coordinator_hybrid_full_plus_swa_all_hit(): ] coord = _make_coord(groups, hash_block_size=16) hs = _hashes(4) - cmap = ExternalCachedBlockPool({(g, bytes(h)) for g in (0, 1) for h in hs}) + cmap = ExternalCachedBlockPool(16, {(g, bytes(h)) for g in (0, 1) for h in hs}) _masks, hit = coord.find_longest_cache_hit( hs, max_length=64, cached_block_pool=cmap ) @@ -169,7 +169,7 @@ def test_coordinator_hybrid_hole_in_full_clips_both(): hs = _hashes(4) exists = {(0, bytes(hs[0])), (0, bytes(hs[2])), (0, bytes(hs[3]))} exists |= {(1, bytes(h)) for h in hs} - cmap = ExternalCachedBlockPool(exists) + cmap = ExternalCachedBlockPool(16, exists) _masks, hit = coord.find_longest_cache_hit( hs, max_length=64, cached_block_pool=cmap ) @@ -188,7 +188,7 @@ def test_coordinator_group_block_size_double_hash(): big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32)) exists = {(0, bytes(h)) for h in hs} exists |= {(1, bytes(bh)) for bh in big_hashes} - cmap = ExternalCachedBlockPool(exists) + cmap = ExternalCachedBlockPool(16, exists) _masks, hit = coord.find_longest_cache_hit( hs, max_length=64, cached_block_pool=cmap ) @@ -405,7 +405,7 @@ def test_lookup_with_eagle_pops_last_full_attention_block(): groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16, use_eagle=True) hs = _hashes(4) - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs}) _masks, hit = coord.find_longest_cache_hit( hs, max_length=64, cached_block_pool=cmap ) @@ -426,7 +426,7 @@ def test_load_mask_with_eagle_does_not_double_prune_full_attention(): groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16, use_eagle=True) hs = _hashes(4) - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs}) _masks, hit = coord.find_longest_cache_hit( hs, max_length=64, cached_block_pool=cmap ) @@ -450,7 +450,7 @@ def test_load_mask_with_eagle_hybrid_full_plus_swa(): coord = _make_coord(groups, hash_block_size=16, use_eagle=True) hs = _hashes(4) exists = {(g, bytes(h)) for g in (0, 1) for h in hs} - cmap = ExternalCachedBlockPool(exists) + cmap = ExternalCachedBlockPool(16, exists) _masks, hit = coord.find_longest_cache_hit( hs, max_length=64, cached_block_pool=cmap ) @@ -469,7 +469,7 @@ def test_load_mask_without_eagle_unchanged(): groups = [KVCacheGroupSpec(["L0"], _full(16))] coord = _make_coord(groups, hash_block_size=16, use_eagle=False) hs = _hashes(4) - cmap = ExternalCachedBlockPool({(0, bytes(h)) for h in hs}) + cmap = ExternalCachedBlockPool(16, {(0, bytes(h)) for h in hs}) _masks, hit = coord.find_longest_cache_hit( hs, max_length=64, cached_block_pool=cmap ) diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index 6dcb3c914ba..d0a50f46579 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py @@ -65,6 +65,7 @@ def _minimal_vllm_config(cache_block_size=16): cfg.cache_config.block_size = cache_block_size cfg.cache_config.num_gpu_blocks = 4 cfg.cache_config.hash_block_size = None + cfg.cache_config.prefix_match_unit = None cfg.cache_config.enable_prefix_caching = True cfg.parallel_config.prefill_context_parallel_size = 1 cfg.parallel_config.decode_context_parallel_size = 1 diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 1091ec5f505..a628e7d7cdd 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -53,17 +53,17 @@ class CacheConfig: """Whether block_size was explicitly provided. Derived automatically.""" user_specified_mamba_block_size: bool = field(default=False, init=False) """Whether mamba_block_size was explicitly provided. Derived automatically.""" - hash_block_size: int | None = Field(default=None, gt=0) - """Block size (in tokens) used for computing Request's block_hashes. + prefix_match_unit: int | None = Field(default=None, gt=0) + """The finest token boundary (in tokens) a prefix-cache hit can land on. - This can be set to a finer granularity than the physical KV cache block - sizes (e.g. 8) as long as every KV cache group's `block_size` is divisible - by it. This enables prefix-caching keys to be computed at the finest common - granularity and then merged for larger physical block sizes. + Prefix-cache keys are computed every `prefix_match_unit` tokens. It can + be set finer than the physical KV cache block sizes (e.g. 32 vs a + 1024-token hybrid-model block) as long as every KV cache group's + `block_size` is divisible by it, enabling cache hits at boundaries + inside a physical block. It controls matching granularity only, not how + often states are stored. - This config is not static default. If left unspecified, vLLM will choose a - default based on the resolved KV cache groups (typically the smallest KV - cache block size when there are multiple groups). + This equals to the `hash_block_size` used throughout the KV cache code. """ gpu_memory_utilization: float = Field(default=0.92, gt=0, le=1) """The fraction of GPU memory to be used for the model executor, which can @@ -209,7 +209,7 @@ class CacheConfig: "enable_prefix_caching", "prefix_caching_hash_algo", # Prefix-caching implementation detail (doesn't affect compiled graph). - "hash_block_size", + "prefix_match_unit", "mamba_page_size_padded", "skip_page_size_padded", "user_specified_block_size", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 6923ceb24df..b08f451d2df 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -29,10 +29,15 @@ from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry class ExternalCachedBlockPool: """Duck-typed BlockPool backed by a ``(group_id, hash)`` exists set.""" - def __init__(self, exists: set[tuple[int, bytes]] | None = None) -> None: + def __init__( + self, + hash_block_size: int, + exists: set[tuple[int, bytes]] | None = None, + ) -> None: # ``exists=None`` is used on the recv side where hit_length is already # determined and we just want each spec's manager to apply its own mask. self._exists = exists + self.hash_block_size = hash_block_size self.null_block = KVCacheBlock(block_id=0) # Dummy ID 1 for present block for duck-typing. self._present_block = KVCacheBlock(block_id=1) @@ -166,7 +171,7 @@ class MooncakeStoreCoordinator: masks, _ = self.find_longest_cache_hit( block_hashes, token_len, - ExternalCachedBlockPool(), + ExternalCachedBlockPool(self.hash_block_size), apply_eagle=False, ) return masks @@ -270,7 +275,7 @@ class MooncakeStoreCoordinator: if len(self.attention_groups) == 1: spec, group_ids, manager_cls = self.attention_groups[0] hashes = self.block_hashes_for_spec(block_hashes, spec) - hit_blocks = manager_cls.find_longest_cache_hit( + hit_blocks, hit_length = manager_cls.find_longest_cache_hit( block_hashes=hashes, # type: ignore[arg-type] max_length=max_length, kv_cache_group_ids=group_ids, @@ -283,11 +288,12 @@ class MooncakeStoreCoordinator: blocks_by_group: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)] for gid, blks in zip(group_ids, hit_blocks, strict=True): blocks_by_group[gid] = blks - return tuple(blocks_by_group), len(hit_blocks[0]) * spec.block_size + return tuple(blocks_by_group), hit_length num_groups = len(self.kv_cache_groups) hit_length = max_length hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups + hit_length_by_group: list[int] = [0] * num_groups is_simple_hybrid = len(self.attention_groups) == 2 and isinstance( self.attention_groups[0][0], FullAttentionSpec @@ -298,10 +304,11 @@ class MooncakeStoreCoordinator: curr_hit_length = hit_length for idx, (spec, group_ids, manager_cls) in enumerate(self.attention_groups): - cached = hit_blocks_by_group[group_ids[0]] + first_group_id = group_ids[0] + cached = hit_blocks_by_group[first_group_id] if isinstance(spec, FullAttentionSpec) and cached is not None: - curr_hit_length = ( - curr_hit_length // spec.block_size * spec.block_size + curr_hit_length = min( + curr_hit_length, hit_length_by_group[first_group_id] ) continue @@ -310,7 +317,7 @@ class MooncakeStoreCoordinator: if drop_eagle_block: _max_length = min(curr_hit_length + spec.block_size, max_length) hashes = self.block_hashes_for_spec(block_hashes, spec) - hit_blocks = manager_cls.find_longest_cache_hit( + hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( block_hashes=hashes, # type: ignore[arg-type] max_length=_max_length, kv_cache_group_ids=group_ids, @@ -319,7 +326,6 @@ class MooncakeStoreCoordinator: drop_eagle_block=drop_eagle_block, alignment_tokens=self.lcm_block_size, ) - _new_hit_length = len(hit_blocks[0]) * spec.block_size if drop_eagle_block: eagle_verified.add(idx) elif _new_hit_length < curr_hit_length: @@ -327,6 +333,7 @@ class MooncakeStoreCoordinator: curr_hit_length = _new_hit_length for gid, blocks in zip(group_ids, hit_blocks, strict=True): hit_blocks_by_group[gid] = blocks + hit_length_by_group[gid] = _new_hit_length if curr_hit_length >= hit_length: break @@ -343,6 +350,7 @@ class MooncakeStoreCoordinator: full_blks = hit_blocks_by_group[gid] assert full_blks is not None del full_blks[num_blocks:] + hit_length_by_group[gid] = hit_length return ( tuple(blks if blks is not None else [] for blks in hit_blocks_by_group), diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index aa48d5123ef..38f4cb0c3a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1523,7 +1523,9 @@ class MooncakeStoreWorker: } _masks, hit_length = self.coord.find_longest_cache_hit( - block_hashes, token_len, ExternalCachedBlockPool(exists_set) + block_hashes, + token_len, + ExternalCachedBlockPool(self.hash_block_size, exists_set), ) return hit_length diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index c602eaf8f2e..742d62ac369 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -693,6 +693,7 @@ class EngineArgs: mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") + prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON @@ -1194,6 +1195,9 @@ class EngineArgs: cache_group.add_argument( "--mamba-block-size", **cache_kwargs["mamba_block_size"] ) + cache_group.add_argument( + "--prefix-match-unit", **cache_kwargs["prefix_match_unit"] + ) cache_group.add_argument( "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] ) @@ -1899,6 +1903,7 @@ class EngineArgs: mamba_cache_dtype=self.mamba_cache_dtype, mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, + prefix_match_unit=self.prefix_match_unit, mamba_cache_mode=self.mamba_cache_mode, kv_offloading_size=self.kv_offloading_size, kv_offloading_backend=self.kv_offloading_backend, diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index bc8f2d87adb..b42c7662d2b 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -260,7 +260,9 @@ class BlockPool: return new_full_blocks = blocks[num_cached_blocks:num_full_blocks] assert block_mask is None or len(block_mask) == len(new_full_blocks) - block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size) + block_hashes = resolve_block_hashes( + request.block_hashes, self.hash_block_size, block_size + ) new_block_hashes = block_hashes[num_cached_blocks:] new_hashes: list[ExternalBlockHash] | None = ( @@ -390,7 +392,9 @@ class BlockPool: if not self.enable_kv_cache_events or num_cached_blocks == 0: return - block_hashes = resolve_block_hashes(request, self.hash_block_size, block_size) + block_hashes = resolve_block_hashes( + request.block_hashes, self.hash_block_size, block_size + ) # Collect external hashes and extra_keys for cached blocks. cached_hashes: list[ExternalBlockHash] = [] @@ -622,6 +626,24 @@ class BlockPool: ) self.cached_block_hash_to_block.insert(block_hash_with_group_id, block) + def move_block_hashes( + self, + src_block: KVCacheBlock, + dst_block: KVCacheBlock, + ) -> None: + """Re-point ``src_block``'s prefix-cache entries to ``dst_block``. + + Used when the request owning ``src_block`` keeps writing into it + : the prefix cache holds a private copy (``dst_block``) + under the same hashes instead. Entries stay live; no events emitted. + """ + assert dst_block.block_hash is None + assert dst_block.block_id not in self.cached_block_hashes_by_block + num_tokens = src_block.block_hash_num_tokens + for block_hash in self._remove_cached_block_hashes(src_block): + # `num_tokens` only applies to the first (primary) insertion. + self._insert_block_hash(block_hash, dst_block, num_tokens=num_tokens) + def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: """Get new blocks from the free block pool. diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 2e978451b88..bf88712aed6 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -5,12 +5,11 @@ from collections.abc import Sequence from typing import NamedTuple from vllm import envs +from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector from vllm.v1.core.kv_cache_utils import ( BlockHash, - BlockHashList, - BlockHashListWithBlockSize, KVCacheBlock, ) from vllm.v1.core.single_type_kv_cache_manager import ( @@ -135,6 +134,7 @@ class KVCacheCoordinator(ABC): new_computed_blocks: tuple[Sequence[KVCacheBlock], ...], num_encoder_tokens: int, total_computed_tokens: int, + num_local_computed_tokens: int, num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: @@ -150,6 +150,8 @@ class KVCacheCoordinator(ABC): num_encoder_tokens: The number of encoder tokens for allocating blocks for cross-attention. total_computed_tokens: Include both local and external tokens. + num_local_computed_tokens: The number of local prefix-cache computed + tokens. num_tokens_main_model: The number of tokens for the main model (aka target model in spec decode). w/o spec decode, it is num_tokens; with spec decode, it is num_tokens - num_lookahead_tokens. @@ -171,6 +173,7 @@ class KVCacheCoordinator(ABC): num_encoder_tokens, [], 0, + 0, num_encoder_tokens, apply_admission_cap=apply_admission_cap, ) @@ -180,6 +183,7 @@ class KVCacheCoordinator(ABC): num_tokens, new_computed_blocks[i], total_computed_tokens, + num_local_computed_tokens, num_tokens_main_model, apply_admission_cap=apply_admission_cap, ) @@ -370,7 +374,7 @@ class KVCacheCoordinator(ABC): pass def new_step_starts(self) -> None: - """Called when a new step is started.""" + """Notify each manager that a new step is starting.""" for manager in self.single_type_managers: manager.new_step_starts() @@ -483,7 +487,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): block_hashes: list[BlockHash], max_cache_hit_length: int, ) -> tuple[tuple[list[KVCacheBlock], ...], int]: - hit_blocks = self.single_type_managers[0].find_longest_cache_hit( + hit_blocks, hit_length = self.single_type_managers[0].find_longest_cache_hit( block_hashes=block_hashes, max_length=max_cache_hit_length, kv_cache_group_ids=[0], @@ -494,7 +498,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): dcp_world_size=self.dcp_world_size, pcp_world_size=self.pcp_world_size, ) - return hit_blocks, len(hit_blocks[0]) * self.block_size + return hit_blocks, hit_length class SpecGroup(NamedTuple): @@ -572,8 +576,26 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): "full-attention and Mamba groups, got: " f"{type(g.kv_cache_spec).__name__}." ) + # Partial hash hits are limited to full-attention + mamba ("align") + # without context parallelism. + self.enable_partial_hash_hits = dcp_world_size == 1 and any( + isinstance(g.kv_cache_spec, MambaSpec) + and g.kv_cache_spec.mamba_cache_mode == "align" + and g.kv_cache_spec.block_size > hash_block_size + for g in kv_cache_config.kv_cache_groups + ) self.verify_and_split_kv_cache_groups() + @property + def _cache_hit_alignment_tokens(self) -> int: + # Fine-grained partial hits may return hash-block-aligned lengths; + # otherwise it must stay scheduler-block-aligned. + return ( + self.hash_block_size + if self.enable_partial_hash_hits + else self.scheduler_block_size + ) + def verify_and_split_kv_cache_groups(self) -> None: """ Groups KV cache groups by their spec type for efficient batch processing @@ -617,14 +639,19 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): self.single_type_managers[gid].use_eagle = True def cache_blocks(self, request: Request, num_computed_tokens: int) -> None: - # Cache hits in this coordinator are always a multiple of - # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``). - # Within an aligned region, SWA groups may only consult a subset of blocks - # per ``scheduler_block_size``-segment so the unused blocks also stay - # out of the prefix-cache hash map. - aligned_num_computed_tokens = ( - num_computed_tokens // self.scheduler_block_size * self.scheduler_block_size - ) + if self.enable_partial_hash_hits: + aligned_num_computed_tokens = num_computed_tokens + else: + # Cache hits in this coordinator are always a multiple of + # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``). + # Within an aligned region, SWA groups may only consult a subset of + # blocks per ``scheduler_block_size``-segment so the unused blocks + # also stay out of the prefix-cache hash map. + aligned_num_computed_tokens = ( + num_computed_tokens + // self.scheduler_block_size + * self.scheduler_block_size + ) for manager in self.single_type_managers: num_tokens_to_cache = aligned_num_computed_tokens # EAGLE groups match one block past each aligned boundary and drop @@ -667,17 +694,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): - The number of tokens of the longest cache hit. """ - def _get_block_hashes(block_size: int) -> BlockHashList: - if block_size == self.hash_block_size: - return block_hashes - return BlockHashListWithBlockSize( - block_hashes, self.hash_block_size, block_size - ) - num_groups = len(self.kv_cache_config.kv_cache_groups) hit_length = max_cache_hit_length longest_hit_length = 0 hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups + hit_length_by_group: list[int] = [0] * num_groups # Simple hybrid (1 full attn + 1 other): one iteration suffices. # Full attn is always first if it exists. @@ -696,40 +717,53 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): for idx, (spec, group_ids, manager_cls, use_eagle) in enumerate( self.attention_groups ): - group_block_size = self.single_type_managers[group_ids[0]].block_size - cached_blocks = hit_blocks_by_group[group_ids[0]] + first_group_id = group_ids[0] + # DCP/PCP shard each block's KV across ranks, so the manager's + # effective block size may exceed the spec's. + group_block_size = self.single_type_managers[first_group_id].block_size + cached_blocks = hit_blocks_by_group[first_group_id] if isinstance(spec, FullAttentionSpec) and cached_blocks is not None: # Full attention is downward-closed: we only need to look # up cached blocks once; on subsequent iterations just trim # to the (reduced) current hit length. - curr_hit_length = ( - curr_hit_length // group_block_size * group_block_size + curr_hit_length = min( + curr_hit_length, hit_length_by_group[first_group_id] ) continue drop_eagle_block = use_eagle and idx not in eagle_verified _max_length = curr_hit_length - if drop_eagle_block: - # Eagle needs to match one more block and then pop the last. - _max_length = min( - curr_hit_length + group_block_size, max_cache_hit_length + # Eagle matches one extra drop unit (one hash unit for + # fine-grained managers, else one cache block) and then drops + # it, landing back at the candidate length. No margin for + # mamba: its finder never drops (draft models have no mamba + # layers), so the hit would grow past the candidate. + if drop_eagle_block and not isinstance(spec, MambaSpec): + eagle_margin = ( + self.hash_block_size + if self.enable_partial_hash_hits + and manager_cls.supports_fine_grained_hash_lookup + and group_block_size > self.hash_block_size + else group_block_size ) - hit_blocks = manager_cls.find_longest_cache_hit( - block_hashes=_get_block_hashes(group_block_size), + _max_length = min( + curr_hit_length + eagle_margin, max_cache_hit_length + ) + hit_blocks, _new_hit_length = manager_cls.find_longest_cache_hit( + block_hashes=block_hashes, max_length=_max_length, kv_cache_group_ids=group_ids, block_pool=self.block_pool, kv_cache_spec=spec, drop_eagle_block=drop_eagle_block, - alignment_tokens=self.scheduler_block_size, + alignment_tokens=self._cache_hit_alignment_tokens, dcp_world_size=( self.dcp_world_size if isinstance(spec, FullAttentionSpec) else 1 ), ) - _new_hit_length = len(hit_blocks[0]) * group_block_size if drop_eagle_block: eagle_verified.add(idx) elif _new_hit_length < curr_hit_length: @@ -738,6 +772,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): curr_hit_length = _new_hit_length for group_id, blocks in zip(group_ids, hit_blocks): hit_blocks_by_group[group_id] = blocks + hit_length_by_group[group_id] = _new_hit_length longest_hit_length = max(longest_hit_length, curr_hit_length) @@ -753,10 +788,11 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): group_block_size = self.single_type_managers[ first_group.group_ids[0] ].block_size - num_blocks = hit_length // group_block_size + num_blocks = cdiv(hit_length, group_block_size) for group_id in first_group.group_ids: if (blks := hit_blocks_by_group[group_id]) is not None: del blks[num_blocks:] + hit_length_by_group[group_id] = hit_length # Uncached shared prefix detection: If any attn. group cached a longer prefix # than the current prefix, it is an uncached common prefix across requests: @@ -776,28 +812,20 @@ class HybridKVCacheCoordinator(KVCacheCoordinator): (blocks_per_group, hit_lengths_per_group) """ - def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList: - if kv_cache_spec.block_size == self.hash_block_size: - return block_hashes - return BlockHashListWithBlockSize( - block_hashes, self.hash_block_size, kv_cache_spec.block_size - ) - num_groups = len(self.kv_cache_config.kv_cache_groups) hit_blocks: list[list[KVCacheBlock]] = [[] for _ in range(num_groups)] hit_lengths: list[int] = [0] * num_groups for spec, group_ids, manager_cls, use_eagle in self.attention_groups: - blocks = manager_cls.find_longest_cache_hit( - block_hashes=_get_block_hashes(spec), + blocks, group_hit = manager_cls.find_longest_cache_hit( + block_hashes=block_hashes, max_length=max_cache_hit_length, kv_cache_group_ids=group_ids, block_pool=self.block_pool, kv_cache_spec=spec, drop_eagle_block=use_eagle, - alignment_tokens=self.scheduler_block_size, + alignment_tokens=self._cache_hit_alignment_tokens, ) - group_hit = len(blocks[0]) * spec.block_size for gid, blks in zip(group_ids, blocks): hit_blocks[gid] = blks hit_lengths[gid] = group_hit diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 4d5d613866a..4f06b58e309 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/vllm/v1/core/kv_cache_manager.py @@ -11,7 +11,7 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector -from vllm.v1.core.kv_cache_utils import KVCacheBlock +from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy from vllm.v1.kv_cache_interface import ( AttentionSpec, CrossAttentionSpec, @@ -404,6 +404,7 @@ class KVCacheManager: new_computed_blocks=new_computed_block_list, num_encoder_tokens=num_encoder_tokens, total_computed_tokens=total_computed_tokens, + num_local_computed_tokens=num_local_computed_tokens, num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) @@ -438,6 +439,7 @@ class KVCacheManager: num_encoder_tokens=num_encoder_tokens, total_computed_tokens=num_local_computed_tokens + num_external_computed_tokens, + num_local_computed_tokens=num_local_computed_tokens, num_tokens_main_model=num_tokens_main_model, ) @@ -665,6 +667,23 @@ class KVCacheManager: ids.extend(mgr.take_new_block_ids()) return ids + def take_kv_cache_block_copies( + self, + ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]: + """Drain pending copies and return their retained endpoints.""" + pending_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] + for mgr in self.coordinator.single_type_managers: + pending_copies.extend(mgr.take_pending_cow_copies()) + copies = [ + KVCacheBlockCopy( + src_block_id=source_block.block_id, + dst_block_id=cow_block.block_id, + ) + for source_block, cow_block in pending_copies + ] + retained_blocks = [block for pair in pending_copies for block in pair] + return copies, retained_blocks + def new_step_starts(self) -> None: - """Called when a new step is started.""" + """Notify the coordinator that a new step is starting.""" self.coordinator.new_step_starts() diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 27af83c9a51..8adc4726727 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -10,7 +10,7 @@ from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, replace from functools import partial -from typing import Any, NewType, TypeAlias, cast, overload +from typing import Any, NamedTuple, NewType, TypeAlias, cast, overload from vllm import envs from vllm.config import VllmConfig @@ -126,7 +126,7 @@ class KVCacheBlock: # when the block is full and cached. _block_hash: BlockHashWithGroupId | None = None # Number of prefix tokens covered by _block_hash. For full blocks this is - # the full block boundary; partial aliases can end inside a cache block. + # the full block boundary; partial entries can end inside a cache block. _block_hash_num_tokens: int | None = None # Used to construct a doubly linked list for free blocks. @@ -176,6 +176,11 @@ class KVCacheBlock: ) +class KVCacheBlockCopy(NamedTuple): + src_block_id: int + dst_block_id: int + + class FreeKVCacheBlockQueue: """This class organizes a list of KVCacheBlock objects to a doubly linked list of free blocks. We implement this class instead of using Python @@ -631,11 +636,11 @@ def resolve_kv_cache_block_sizes( Mamba groups keep their full per-rank state and are not scaled. - ``hash_block_size`` is the granularity at which ``Request.block_hashes`` is computed. Single group: equals scheduler block size. Multiple groups: - ``cache_config.hash_block_size`` override if set, else the GCD of group - block sizes; every group's block size must be divisible by it. Returns - the scheduler block size (i.e. disables finer hashing) if block hashing - is inactive or a mamba group's block size diverges from the cache - block size (mamba_cache_mode != "align"). + ``cache_config.prefix_match_unit`` override if set, else the GCD of + group block sizes; every group's block size must be divisible by it. + Returns the scheduler block size (i.e. disables finer hashing) if block + hashing is inactive or a mamba group's block size diverges from the + cache block size (mamba_cache_mode != "align"). """ cache_config = vllm_config.cache_config dcp = vllm_config.parallel_config.decode_context_parallel_size @@ -671,14 +676,14 @@ def resolve_kv_cache_block_sizes( ): return scheduler_block_size, scheduler_block_size - requested = cache_config.hash_block_size + requested = cache_config.prefix_match_unit hash_block_size = ( requested if requested is not None else math.gcd(*group_block_sizes) ) if any(bs % hash_block_size != 0 for bs in group_block_sizes): raise ValueError( - f"Invalid hash_block_size={hash_block_size}; all KV cache group " - f"block sizes must be divisible by hash_block_size. " + f"Invalid prefix_match_unit={hash_block_size}; all KV cache group " + f"block sizes must be divisible by prefix_match_unit. " f"Got group block sizes={group_block_sizes}." ) return scheduler_block_size, hash_block_size @@ -2252,19 +2257,33 @@ BlockHashList = list[BlockHash] | BlockHashListWithBlockSize def resolve_block_hashes( - request: Request, + block_hashes: BlockHashList, hash_block_size: int, block_size: int, + *, + supports_fine_grained_hash_lookup: bool = False, + alignment_tokens: int | None = None, ) -> BlockHashList: - """Resolve the block-hash view for ``request`` at ``block_size``. + """Resolve the block-hash view at ``block_size``. - When ``block_size`` equals ``hash_block_size``, reuse the request's - precomputed ``block_hashes`` directly; otherwise recalculate at - ``block_size`` granularity (``block_size`` must be a multiple of - ``hash_block_size``, which happens when KV cache groups differ in - block size). + When ``block_size`` equals ``hash_block_size``, reuse the precomputed block + hashes directly; otherwise view them at ``block_size`` granularity. + Fine-grained lookup keeps the original hashes for partial cache hits. """ if block_size == hash_block_size: - return request.block_hashes + return block_hashes + if isinstance(block_hashes, BlockHashListWithBlockSize): + # Already a block-size view + assert block_hashes.scale_factor == block_size // hash_block_size + return block_hashes + # Fine-grained partial hits keep the raw hashes. The caller passes + # alignment_tokens = hash_block_size to enable them, else >= block_size. + if ( + supports_fine_grained_hash_lookup + and alignment_tokens is not None + and alignment_tokens < block_size + and block_size % alignment_tokens == 0 + ): + return block_hashes assert block_size % hash_block_size == 0 - return BlockHashListWithBlockSize(request.block_hashes, hash_block_size, block_size) + return BlockHashListWithBlockSize(block_hashes, hash_block_size, block_size) diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 291e73bc64b..4401fb050b3 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -16,10 +16,12 @@ if TYPE_CHECKING: from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams + from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.request import Request else: ECConnectorMetadata = object KVConnectorMetadata = object + KVCacheBlockCopy = object LoRARequest = object MultiModalFeatureSpec = object PoolingParams = object @@ -240,6 +242,9 @@ class SchedulerOutput: # preventing stale NaN/data from corrupting attention or SSM computation. new_block_ids_to_zero: list[int] | None = None + # CoW copies to apply after zeroing new blocks and before forward. + kv_cache_block_copies: list[KVCacheBlockCopy] | None = None + # Dynamic speculative decoding: optimal K chosen by scheduler. # Number of spec tokens to schedule for the next step. num_spec_tokens_to_schedule: int = 0 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 95071408876..47d3a53d20a 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -259,6 +259,7 @@ class Scheduler(SchedulerInterface): # Create the KV cache manager. if hash_block_size is None: hash_block_size = block_size + self.hash_block_size = hash_block_size self.kv_cache_manager = KVCacheManager( kv_cache_config=kv_cache_config, max_model_len=self.max_model_len, @@ -297,6 +298,13 @@ class Scheduler(SchedulerInterface): self.need_mamba_block_aligned_split = ( self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" ) + # A finer prefix_match_unit is configured: a mamba partial tail entry + # can only be registered by a step ending exactly at the prompt's last + # hash boundary, so the split adds that stop. + self.mamba_partial_cache_hit = ( + self.need_mamba_block_aligned_split + and self.hash_block_size < self.block_size + ) # Counts of non-empty steps scheduled / processed. update_from_output # is called once per scheduled step in FIFO order, so these stay in sync. @@ -368,9 +376,23 @@ class Scheduler(SchedulerInterface): if self.use_eagle: last_cache_position = max(last_cache_position - block_size, 0) num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens - if num_computed_tokens_after_sched < last_cache_position: - # align to block_size - num_new_tokens = num_new_tokens // block_size * block_size + next_boundary = (num_computed_tokens // block_size + 1) * block_size + if ( + num_computed_tokens % block_size != 0 + and next_boundary <= last_cache_position + and num_computed_tokens_after_sched > next_boundary + ): + # Resumed mid-block (partial hash hit): stop the first chunk + # at the next block boundary so chunk ends re-align to blocks + # and the boundary state is materialized before running past. + num_new_tokens = next_boundary - num_computed_tokens + elif num_computed_tokens_after_sched < last_cache_position: + # Align the chunk END (not its length) to block_size; + # identical to flooring the length when the start is + # block-aligned. May yield 0 (insufficient budget to reach + # the next boundary); the caller then skips the request. + aligned_end = num_computed_tokens_after_sched // block_size * block_size + num_new_tokens = max(aligned_end - num_computed_tokens, 0) elif ( num_computed_tokens < last_cache_position @@ -378,9 +400,23 @@ class Scheduler(SchedulerInterface): ): # force to cache the last chunk num_new_tokens = last_cache_position - num_computed_tokens - else: - # prefill the last few tokens - pass + elif self.mamba_partial_cache_hit: + # Prefill of the final partial block: stop once at the + # prompt's last hash boundary so the mamba partial tail entry + # can be registered. + tail_boundary = ( + request.num_prompt_tokens + // self.hash_block_size + * self.hash_block_size + ) + if ( + num_computed_tokens + < tail_boundary + < num_computed_tokens_after_sched + and tail_boundary < request.num_prompt_tokens + and tail_boundary > last_cache_position + ): + num_new_tokens = tail_boundary - num_computed_tokens # Marconi cache admission optimization: # cache common prefixes by scheduling num_new_tokens = common prefix length @@ -1084,6 +1120,16 @@ class Scheduler(SchedulerInterface): new_block_ids_to_zero = ( (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None ) + kv_cache_block_copies, cow_retained_blocks = ( + self.kv_cache_manager.take_kv_cache_block_copies() + ) + if kv_cache_block_copies: + # The copies run with this step's execution; the first non-empty + # step at or after it gets seq `sched_step_seq + 1` (0-token steps + # do not advance the seq), and its completion implies the copies + # have run. + self._free_cow_retained_blocks(cow_retained_blocks, self.sched_step_seq + 1) + pending_kv_cache_block_copies = kv_cache_block_copies or None # Dynamic speculative decoding: compute optimal K num_spec_tokens_to_schedule = self.num_spec_tokens @@ -1108,6 +1154,7 @@ class Scheduler(SchedulerInterface): finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=new_block_ids_to_zero, + kv_cache_block_copies=pending_kv_cache_block_copies, num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, ) @@ -2148,11 +2195,23 @@ class Scheduler(SchedulerInterface): if blocks: self.deferred_frees.append((self.sched_step_seq, blocks)) + def _free_cow_retained_blocks( + self, blocks: list[KVCacheBlock], fence_seq: int + ) -> None: + """Release CoW copy retentions, deferring their return to the block + pool while the step that runs the copy may still be in flight. + """ + if not self.defer_block_free or fence_seq <= self.processed_step_seq: + self.kv_cache_manager.block_pool.free_blocks(blocks) + return + self.deferred_frees.append((fence_seq, blocks[::-1])) + def _drain_deferred_frees(self): """Return deferred blocks whose fence step has completed. - Entries are appended with monotonically non-decreasing fences, so - stop at the first one that is still pending. + Fences are appended in near-monotonic order (a CoW retention fence + can lead request-free fences by one step), so stop at the first + pending one; any satisfied entry behind it is merely freed later. """ while self.deferred_frees: fence, _ = self.deferred_frees[0] @@ -2401,6 +2460,7 @@ class Scheduler(SchedulerInterface): new_computed_blocks=self.kv_cache_manager.empty_kv_cache_blocks.blocks, num_encoder_tokens=0, total_computed_tokens=request.num_computed_tokens, + num_local_computed_tokens=request.num_computed_tokens, num_tokens_main_model=full_num_tokens, apply_admission_cap=True, ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 324d0016cd0..21e3ad8ffda 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -4,13 +4,16 @@ import itertools from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Sequence +from typing import ClassVar from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import ( BlockHashList, + BlockHashListWithBlockSize, BlockHashWithGroupId, KVCacheBlock, + resolve_block_hashes, ) from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, @@ -36,6 +39,8 @@ class SingleTypeKVCacheManager(ABC): logic of one specific type of attention layer. """ + supports_fine_grained_hash_lookup: ClassVar[bool] = False + def __init__( self, kv_cache_spec: KVCacheSpec, @@ -106,16 +111,34 @@ class SingleTypeKVCacheManager(ABC): # determining the attention groups. self.use_eagle = False + # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained + # managers (full attention, mamba "align"); harmlessly empty elsewhere. + self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {} + self._pending_cow_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = [] + @classmethod def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]): return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks) + def _has_partial_local_hit( + self, + new_computed_blocks: Sequence[KVCacheBlock], + num_local_computed_tokens: int, + ) -> bool: + # The local prefix-cache hit ends inside one of this manager's + # blocks: the shared tail block needs CoW. + return ( + len(new_computed_blocks) > 0 + and num_local_computed_tokens % self.block_size != 0 + ) + def get_num_blocks_to_allocate( self, request_id: str, num_tokens: int, new_computed_blocks: Sequence[KVCacheBlock], total_computed_tokens: int, + num_local_computed_tokens: int, num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: @@ -130,6 +153,8 @@ class SingleTypeKVCacheManager(ABC): prefix caching. total_computed_tokens: Include both local and external computed tokens. + num_local_computed_tokens: The number of local prefix-cache computed + tokens. num_tokens_main_model: The number of tokens for the main model (aka target model in spec decode). w/o spec decode, it is num_tokens; with spec decode, it is num_tokens - num_lookahead_tokens. @@ -189,6 +214,10 @@ class SingleTypeKVCacheManager(ABC): num_evictable_blocks = self._get_num_evictable_blocks( new_computed_blocks[num_skipped_new_computed_blocks:] ) + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens): + # Reserve the extra block that allocate_new_blocks pulls for the + # partial-hit CoW redirect. + num_new_blocks += 1 return num_new_blocks + num_evictable_blocks def add_local_computed_blocks( @@ -242,6 +271,13 @@ class SingleTypeKVCacheManager(ABC): # them so cache_blocks() will not try to re-cache blocks that already # have a block_hash set. self.num_cached_block[request_id] = len(req_blocks) + if self._has_partial_local_hit(new_computed_blocks, num_local_computed_tokens): + # Record the partial tail for the CoW redirect in + # allocate_new_blocks; cap the cached count at the full blocks so + # cache_blocks() re-caches the private copy once full. + block_idx = num_local_computed_tokens // self.block_size + self._partial_hit_reqs[request_id] = (block_idx, new_computed_blocks[-1]) + self.num_cached_block[request_id] = block_idx def allocate_external_computed_blocks( self, @@ -299,17 +335,29 @@ class SingleTypeKVCacheManager(ABC): Returns: The new allocated blocks. """ + cow_blocks: list[KVCacheBlock] = [] + if request_id in self._partial_hit_reqs: + # Partial hit: redirect the shared tail to a private CoW block. + # Replacing in place keeps the length-based allocation below + # correct; the extra block was reserved by + # get_num_blocks_to_allocate. + block_idx, source_block = self._partial_hit_reqs.pop(request_id) + cow_block = self.block_pool.get_new_blocks(1)[0] + self._apply_cow(request_id, block_idx, source_block, cow_block) + self.new_block_ids.append(cow_block.block_id) + cow_blocks.append(cow_block) + req_blocks = self.req_to_blocks[request_id] num_required_blocks = cdiv(num_tokens, self.block_size) num_new_blocks = num_required_blocks - len(req_blocks) if num_new_blocks <= 0: - return [] + return cow_blocks else: new_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(new_blocks) if self._record_new_block_ids: self.new_block_ids.extend(b.block_id for b in new_blocks) - return new_blocks + return cow_blocks + new_blocks def take_new_block_ids(self) -> list[int]: """Drain and return block IDs allocated since the last call.""" @@ -317,6 +365,36 @@ class SingleTypeKVCacheManager(ABC): self.new_block_ids = [] return ids + def take_pending_cow_copies( + self, + ) -> list[tuple[KVCacheBlock, KVCacheBlock]]: + """Drain pending CoW source and destination block pairs.""" + pending_copies = self._pending_cow_copies + self._pending_cow_copies = [] + return pending_copies + + def _apply_cow( + self, + request_id: str, + block_idx: int, + source_block: KVCacheBlock, + cow_block: KVCacheBlock, + ) -> None: + """Redirect a partial prefix-cache hit to a private CoW block. + + Both copy endpoints stay retained until the copy has run on the worker, + so a same-step free cannot recycle them: ``source_block`` keeps its + hit-ref, ``cow_block`` takes an extra ref beyond the one handed to the + request. + """ + req_blocks = self.req_to_blocks[request_id] + assert block_idx < len(req_blocks) + assert req_blocks[block_idx] is source_block + assert not source_block.is_null and source_block.ref_cnt > 0 + req_blocks[block_idx] = cow_block + self._pending_cow_copies.append((source_block, cow_block)) + cow_block.ref_cnt += 1 + def cache_blocks( self, request: Request, @@ -398,6 +476,7 @@ class SingleTypeKVCacheManager(ABC): # Default to [] in case a request is freed (aborted) before alloc. req_blocks = self.req_to_blocks.pop(request_id, []) self.num_cached_block.pop(request_id, None) + self._partial_hit_reqs.pop(request_id, None) return req_blocks def free(self, request_id: str) -> None: @@ -439,7 +518,7 @@ class SingleTypeKVCacheManager(ABC): alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: """ Get the longest cache hit prefix of the blocks that is not longer than `max_length`. The prefix should be a common prefix hit for all the @@ -466,11 +545,9 @@ class SingleTypeKVCacheManager(ABC): pcp_world_size: The world size of prefill context parallelism. Returns: - A list of cached blocks with skipped blocks replaced by null block - for each kv cache group in `kv_cache_group_ids`. - Return a list of length `len(kv_cache_group_ids)`, where the i-th - element is a list of cached blocks for the i-th kv cache group - in `kv_cache_group_ids`. + A tuple containing cached blocks and the exact cache-hit length in + tokens. The cached block tuple has skipped blocks replaced by null + blocks for each kv cache group in `kv_cache_group_ids`. For example, sliding window manager should return a list like ([NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]) for block size 4 and sliding window 8 and len(kv_cache_group_ids) = 1. @@ -558,11 +635,12 @@ class SingleTypeKVCacheManager(ABC): return 0 def new_step_starts(self) -> None: - # do nothing by default return None class FullAttentionManager(SingleTypeKVCacheManager): + supports_fine_grained_hash_lookup: ClassVar[bool] = True + @classmethod def find_longest_cache_hit( cls, @@ -575,42 +653,131 @@ class FullAttentionManager(SingleTypeKVCacheManager): alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance( kv_cache_spec, FullAttentionSpec | ChunkedLocalAttentionSpec ), ( "FullAttentionManager can only be used for full attention " "and chunked local attention groups" ) + block_size = kv_cache_spec.block_size + if dcp_world_size * pcp_world_size > 1: + # DCP/PCP shard each block's KV across ranks; hashes must be + # viewed at the sharded (scaled) block size. + block_size *= dcp_world_size * pcp_world_size + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) + + # Fine-grained mode (alignment_tokens == hash_block_size < + # block_size): resolve_block_hashes kept the raw hash-granularity + # list so interior boundaries can be probed. + fine_grained = ( + alignment_tokens < block_size and block_size % alignment_tokens == 0 + ) + if fine_grained: + assert isinstance(block_hashes, list) + full_block_hashes: BlockHashList = BlockHashListWithBlockSize( + block_hashes, alignment_tokens, block_size + ) + else: + full_block_hashes = block_hashes + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( [] for _ in range(len(kv_cache_group_ids)) ) - block_size = kv_cache_spec.block_size - if dcp_world_size * pcp_world_size > 1: - block_size *= dcp_world_size * pcp_world_size - max_num_blocks = max_length // block_size - for block_hash in itertools.islice(block_hashes, max_num_blocks): - # block_hashes is a chain of block hashes. If a block hash is not - # in the cached_block_hash_to_id, the following block hashes are - # not computed yet for sure. - if cached_block := block_pool.get_cached_block( - block_hash, kv_cache_group_ids - ): - for computed, cached in zip(computed_blocks, cached_block): - computed.append(cached) - else: + # Phase 1: longest run of cached full blocks from the start. A missing + # block implies every later block misses too (chained hashes). + for block_hash in itertools.islice(full_block_hashes, max_length // block_size): + cached_block = block_pool.get_cached_block(block_hash, kv_cache_group_ids) + if not cached_block: break - if drop_eagle_block and computed_blocks[0]: - # Need to drop the last matched block if eagle is enabled. - for computed in computed_blocks: - computed.pop() - while ( - block_size != alignment_tokens # Faster for common case. - and len(computed_blocks[0]) * block_size % alignment_tokens != 0 - ): - for computed in computed_blocks: - computed.pop() - return computed_blocks + for computed, cached in zip(computed_blocks, cached_block): + computed.append(cached) + hit_length = len(computed_blocks[0]) * block_size + + # Phase 2 (fine-grained only): extend into the first non-full block by + # probing its interior hash boundaries high-to-low (longest hit first). + if fine_grained: + assert isinstance(block_hashes, list) + scale_factor = block_size // alignment_tokens + first_partial_idx = len(computed_blocks[0]) * scale_factor + max_partial_idx = min( + first_partial_idx + scale_factor - 1, + max_length // alignment_tokens, + len(block_hashes), + ) + for fine_idx in range(max_partial_idx - 1, first_partial_idx - 1, -1): + cached_tail = block_pool.get_cached_block( + block_hashes[fine_idx], kv_cache_group_ids + ) + if not cached_tail: + continue + for computed, cached in zip(computed_blocks, cached_tail): + computed.append(cached) + hit_length = (fine_idx + 1) * alignment_tokens + break + + # Eagle needs the tokens right before the generation point recomputed: + # drop one hash unit when fine-grained (the tail block's KV is + # append-only, so it still covers the reduced length), else one cache + # block. + if drop_eagle_block and hit_length > 0: + hit_length -= min(alignment_tokens, block_size) + # Round down to the alignment; a no-op when fine-grained (hits land on + # hash boundaries by construction) and when alignment_tokens == + # block_size. Then trim blocks past the new tail. + hit_length -= hit_length % alignment_tokens + num_blocks = cdiv(hit_length, block_size) + for computed in computed_blocks: + del computed[num_blocks:] + return computed_blocks, hit_length + + def cache_blocks( + self, + request: Request, + num_tokens: int, + retention_interval: int | None = None, + ) -> None: + super().cache_blocks(request, num_tokens, retention_interval=retention_interval) + hash_block_size = self.block_pool.hash_block_size + if self.block_size == hash_block_size: + return + self._cache_partial_tail_block(request, num_tokens) + + def _cache_partial_tail_block( + self, + request: Request, + num_tokens: int, + ) -> None: + """Cache the prompt tail when it ends inside a cache block. + + Only the final prompt hash boundary is registered as a partial + prefix-cache entry; intermediate hash boundaries inside the same cache + block are intentionally skipped. + """ + hash_block_size = self.block_pool.hash_block_size + boundary_tokens = request.num_prompt_tokens // hash_block_size * hash_block_size + if boundary_tokens == 0 or boundary_tokens > num_tokens: + return + if boundary_tokens % self.block_size == 0: + return + + blocks = self.req_to_blocks[request.request_id] + block_idx = boundary_tokens // self.block_size + if block_idx >= len(blocks): + return + self.block_pool.cache_partial_block( + request=request, + block=blocks[block_idx], + num_tokens=boundary_tokens, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + ) def get_num_common_prefix_blocks(self, running_request_id: str) -> int: blocks = self.req_to_blocks[running_request_id] @@ -699,12 +866,23 @@ class SlidingWindowManager(SingleTypeKVCacheManager): alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance(kv_cache_spec, SlidingWindowSpec), ( "SlidingWindowManager can only be used for sliding window groups" ) assert dcp_world_size == 1, "DCP not support sliding window attn now." assert pcp_world_size == 1, "PCP not support sliding window attn now." + # Fine-grained partial hits are not supported for sliding window now + assert alignment_tokens % kv_cache_spec.block_size == 0, ( + "SlidingWindowManager does not support fine-grained (partial) cache hits" + ) + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) # The number of contiguous blocks needed for a prefix cache hit. sliding_window_contiguous_blocks = cls._contiguous_blocks_for_hit( @@ -717,7 +895,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager): # sliding_window_contiguous_blocks), # which is good for low cache hit rate scenarios. max_num_blocks = max_length // kv_cache_spec.block_size - computed_blocks = tuple( + computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( [block_pool.null_block] * max_num_blocks for _ in range(len(kv_cache_group_ids)) ) @@ -772,7 +950,8 @@ class SlidingWindowManager(SingleTypeKVCacheManager): ): for computed in computed_blocks: computed.pop() - return computed_blocks + hit_length = len(computed_blocks[0]) * block_size + return computed_blocks, hit_length @classmethod def reachable_block_mask( @@ -893,7 +1072,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: """ For chunked local attention, we need to find the longest cache hit prefix of the blocks that is not longer than `max_length`. The prefix @@ -942,6 +1121,13 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): "KV cache groups with different block sizes are not compatible with " "chunked local attention now" ) + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) max_num_blocks = max_length // kv_cache_spec.block_size if max_length > 0: local_attention_start_idx = ( @@ -971,7 +1157,8 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): computed.append(cached) else: break - return computed_blocks + hit_length = len(computed_blocks[0]) * kv_cache_spec.block_size + return computed_blocks, hit_length def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: """ @@ -1027,6 +1214,8 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager): class MambaManager(SingleTypeKVCacheManager): + supports_fine_grained_hash_lookup: ClassVar[bool] = True + def __init__( self, kv_cache_spec: MambaSpec, block_pool: BlockPool, **kwargs ) -> None: @@ -1035,9 +1224,9 @@ class MambaManager(SingleTypeKVCacheManager): # recurrent state. Undo the DCP/PCP block_size scaling that the base # class applies for attention groups whose KV cache is partitioned. self.block_size = kv_cache_spec.block_size - self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks + self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() if self.mamba_cache_mode == "align": # Mapping from request ID to the index of the block # allocated in the previous step @@ -1057,17 +1246,46 @@ class MambaManager(SingleTypeKVCacheManager): alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance(kv_cache_spec, MambaSpec), ( "MambaManager can only be used for mamba groups" ) assert dcp_world_size == 1, "DCP not support mamba now." assert pcp_world_size == 1, "PCP not support mamba now." + block_hashes = resolve_block_hashes( + block_hashes, + block_pool.hash_block_size, + kv_cache_spec.block_size, + supports_fine_grained_hash_lookup=cls.supports_fine_grained_hash_lookup, + alignment_tokens=alignment_tokens, + ) computed_blocks: tuple[list[KVCacheBlock], ...] = tuple( [] for _ in range(len(kv_cache_group_ids)) ) + hit_length = 0 block_size = kv_cache_spec.block_size + if alignment_tokens < block_size and block_size % alignment_tokens == 0: + assert isinstance(block_hashes, list) + hash_block_size = alignment_tokens + scale_factor = block_size // hash_block_size + max_num_partial_units = min( + max_length // hash_block_size, len(block_hashes) + ) + for fine_idx in range(max_num_partial_units - 1, -1, -1): + num_tokens = (fine_idx + 1) * hash_block_size + block_hash = block_hashes[fine_idx] + if cached_block := block_pool.get_cached_block( + block_hash, kv_cache_group_ids + ): + block_idx = fine_idx // scale_factor + for computed, cached in zip(computed_blocks, cached_block): + computed.extend([block_pool.null_block] * block_idx) + computed.append(cached) + hit_length = num_tokens + break + return computed_blocks, hit_length + max_num_blocks = max_length // block_size # Search from right to left and early stop when a match is found. for i in range(max_num_blocks - 1, -1, -1): @@ -1089,9 +1307,10 @@ class MambaManager(SingleTypeKVCacheManager): # so we insert dummy blocks at the beginning: computed.extend([block_pool.null_block] * i) computed.append(cached) + hit_length = (i + 1) * block_size break # we just need the last match - early stopping - return computed_blocks + return computed_blocks, hit_length @classmethod def reachable_block_mask( @@ -1189,6 +1408,7 @@ class MambaManager(SingleTypeKVCacheManager): num_tokens: int, new_computed_blocks: Sequence[KVCacheBlock], total_computed_tokens: int, + num_local_computed_tokens: int, num_tokens_main_model: int, apply_admission_cap: bool = False, ) -> int: @@ -1214,6 +1434,7 @@ class MambaManager(SingleTypeKVCacheManager): num_tokens, new_computed_blocks, total_computed_tokens, + num_local_computed_tokens, num_tokens_main_model, apply_admission_cap=apply_admission_cap, ) @@ -1235,15 +1456,26 @@ class MambaManager(SingleTypeKVCacheManager): - len(new_computed_blocks) - len(self.req_to_blocks[request_id]) ) + has_partial_hit = ( + self._has_partial_local_hit( + new_computed_blocks, num_local_computed_tokens + ) + or request_id in self._partial_hit_reqs + ) + if has_partial_hit: + num_new_blocks = max(num_new_blocks, 0) + 1 if num_new_blocks > 0: if request_id in self._allocated_block_reqs: # Old request. Needs at most 1 more blocks as we can reuse the # speculative blocks in previous step. - num_new_blocks = 1 + num_new_blocks = 1 + int(has_partial_hit) else: - # First prefill. Allocate 1 block for running state and the - # speculative blocks. - num_new_blocks = 1 + self.num_speculative_blocks + # First prefill. Allocate 1 block for running state, the + # speculative blocks, and one extra block if a partial cache + # hit must be copy-on-written before the new tokens run. + num_new_blocks = ( + 1 + self.num_speculative_blocks + int(has_partial_hit) + ) num_evictable_computed_blocks = self._get_num_evictable_blocks( new_computed_blocks @@ -1275,9 +1507,11 @@ class MambaManager(SingleTypeKVCacheManager): num_required_blocks = ( cdiv(num_tokens, self.block_size) + self.num_speculative_blocks ) + partial_hit = self._partial_hit_reqs.get(request_id) + has_partial_hit = partial_hit is not None # `num_required_blocks` might be less than `len(req_blocks)` if blocks are # over-allocated at last round. - if num_required_blocks <= len(req_blocks): + if num_required_blocks <= len(req_blocks) and not has_partial_hit: return [] else: prev_block_len = len(req_blocks) @@ -1317,14 +1551,42 @@ class MambaManager(SingleTypeKVCacheManager): else: break num_new_blocks = num_required_blocks - len(req_blocks) + if has_partial_hit: + num_new_blocks = max(num_new_blocks, 0) + 1 if blocks_allocated: - assert num_new_blocks <= 1 + assert num_new_blocks <= 1 + int(has_partial_hit) else: - assert num_new_blocks <= self.num_speculative_blocks + 1 + assert num_new_blocks <= self.num_speculative_blocks + 1 + int( + has_partial_hit + ) new_blocks = self.block_pool.get_new_blocks(num_new_blocks) + returned_blocks = req_blocks[prev_block_len:] + if partial_hit is not None: + block_idx, source_block = partial_hit + cow_block = new_blocks[0] + new_blocks = new_blocks[1:] + if blocks_allocated: + # The worker block table of a running request is + # append-only, so the request must stay on + # source_block. Move the cache entry to cow_block + # instead; the queued copy fills it before forward + # overwrites source_block. + assert req_blocks[block_idx] is source_block + self.block_pool.move_block_hashes(source_block, cow_block) + self._pending_cow_copies.append((source_block, cow_block)) + source_block.ref_cnt += 1 + if cow_block.block_hash is not None: + # The moved entry is only filled by this step's + # copy, so defer same-step hits on it. + self.cached_blocks_this_step.add(cow_block.block_hash) + else: + self._apply_cow(request_id, block_idx, source_block, cow_block) + returned_blocks = [cow_block] + returned_blocks req_blocks.extend(new_blocks) self._allocated_block_reqs.add(request_id) - return req_blocks[prev_block_len:] + self._partial_hit_reqs.pop(request_id, None) + returned_blocks.extend(new_blocks) + return returned_blocks def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]: if self.mamba_cache_mode == "align": @@ -1349,6 +1611,10 @@ class MambaManager(SingleTypeKVCacheManager): num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0) super().cache_blocks(request, num_tokens, retention_interval=retention_interval) num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0) + if self.mamba_cache_mode == "align": + partial_hash = self._cache_partial_tail_block(request, num_tokens) + if partial_hash is not None: + self.cached_blocks_this_step.add(partial_hash) if num_cached_blocks_after > num_cached_blocks_before: for block in self.req_to_blocks[request.request_id][ num_cached_blocks_before:num_cached_blocks_after @@ -1364,6 +1630,44 @@ class MambaManager(SingleTypeKVCacheManager): def new_step_starts(self) -> None: self.cached_blocks_this_step.clear() + def _cache_partial_tail_block( + self, + request: Request, + num_tokens: int, + ) -> BlockHashWithGroupId | None: + hash_block_size = self.block_pool.hash_block_size + if self.block_size == hash_block_size: + return None + if num_tokens % self.block_size == 0: + return None + if num_tokens % hash_block_size != 0: + return None + latest_prompt_hash_boundary = ( + request.num_prompt_tokens // hash_block_size + ) * hash_block_size + if num_tokens != latest_prompt_hash_boundary: + return None + + block_idx = num_tokens // self.block_size + blocks = self.req_to_blocks[request.request_id] + if block_idx >= len(blocks): + return None + source_block = blocks[block_idx] + if source_block.is_null: + return None + + partial_hash = self.block_pool.cache_partial_block( + request=request, + block=source_block, + num_tokens=num_tokens, + kv_cache_group_id=self.kv_cache_group_id, + block_size=self.block_size, + ) + if partial_hash is not None: + self._partial_hit_reqs[request.request_id] = (block_idx, source_block) + self.num_cached_block[request.request_id] = block_idx + return partial_hash + class CrossAttentionManager(SingleTypeKVCacheManager): """Manager for cross-attention KV cache in encoder-decoder models.""" @@ -1415,7 +1719,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager): alignment_tokens: int, dcp_world_size: int = 1, pcp_world_size: int = 1, - ) -> tuple[list[KVCacheBlock], ...]: + ) -> tuple[tuple[list[KVCacheBlock], ...], int]: assert isinstance(kv_cache_spec, CrossAttentionSpec), ( "CrossAttentionManager can only be used for cross-attention groups" ) @@ -1435,16 +1739,18 @@ class SinkFullAttentionManager(FullAttentionManager): block_pool: BlockPool, enable_caching: bool, kv_cache_group_id: int, + scheduler_block_size: int, dcp_world_size: int = 1, pcp_world_size: int = 1, ): super().__init__( - kv_cache_spec, - block_pool, - enable_caching, - kv_cache_group_id, - dcp_world_size, - pcp_world_size, + kv_cache_spec=kv_cache_spec, + block_pool=block_pool, + enable_caching=enable_caching, + kv_cache_group_id=kv_cache_group_id, + scheduler_block_size=scheduler_block_size, + dcp_world_size=dcp_world_size, + pcp_world_size=pcp_world_size, ) sink_len = kv_cache_spec.sink_len assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1f5dd4d2fba..0a62295629f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -112,7 +112,7 @@ from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler from vllm.v1.worker.gpu.states import RequestState from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin -from vllm.v1.worker.utils import KVBlockZeroer +from vllm.v1.worker.utils import KVBlockZeroer, copy_kv_cache_blocks_inplace logger = init_logger(__name__) @@ -838,6 +838,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): assert self.kv_block_zeroer is not None self.kv_block_zeroer.zero_block_ids(scheduler_output.new_block_ids_to_zero) + # Apply copy-on-write block copies for partial prefix-cache hits, after + # zeroing new blocks and before the forward pass reads them. + if scheduler_output.kv_cache_block_copies: + copy_kv_cache_blocks_inplace( + self.kv_caches, + self.kv_cache_config.num_blocks, + scheduler_output.kv_cache_block_copies, + ) + def prepare_inputs( self, scheduler_output: SchedulerOutput, batch_desc: BatchExecutionDescriptor ) -> InputBatch: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index efb302b90f4..cb8ed5f367c 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -230,6 +230,7 @@ from .utils import ( KVBlockZeroer, add_kv_sharing_layers_to_kv_cache_groups, bind_kv_cache, + copy_kv_cache_blocks_inplace, prepare_kernel_block_sizes, sanity_check_mm_encoder_outputs, ) @@ -1188,6 +1189,12 @@ class GPUModelRunner( # stale NaN/data from corrupting attention or SSM computation. if scheduler_output.new_block_ids_to_zero: self._zero_block_ids(scheduler_output.new_block_ids_to_zero) + if scheduler_output.kv_cache_block_copies: + copy_kv_cache_blocks_inplace( + self.kv_caches, + self.kv_cache_config.num_blocks, + scheduler_output.kv_cache_block_copies, + ) # Free the cached encoder outputs. for mm_hash in scheduler_output.free_encoder_mm_hashes: diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 2c2f930001b..9e60d2ae7fa 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -2,11 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from itertools import product as iprod from typing import Any +import numpy as np import torch from vllm.config import CacheConfig, VllmConfig @@ -18,11 +19,13 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import largest_power_of_2_divisor from vllm.utils.mem_utils import MemorySnapshot, format_gib +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, AttentionMetadataBuilder, MultipleOf, ) +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.kv_cache_interface import ( AttentionSpec, EncoderOnlyAttentionSpec, @@ -535,6 +538,45 @@ def bind_kv_cache( forward_context[layer_name].kv_cache = kv_cache +def copy_kv_cache_blocks_inplace( + kv_caches: Iterable[torch.Tensor | list[torch.Tensor]], + num_blocks: int, + kv_cache_block_copies: Sequence[KVCacheBlockCopy], +) -> None: + if not kv_cache_block_copies: + return + + storage_tensors: list[torch.Tensor] = [] + seen_storage: set[int] = set() + for entry in kv_caches: + # Mamba layers hold a list of state tensors; attention layers a single + # tensor. Both alias the shared block-major backing storage. + tensors = entry if isinstance(entry, (list, tuple)) else (entry,) + for tensor in tensors: + ptr = tensor.untyped_storage().data_ptr() + if ptr in seen_storage: + continue + seen_storage.add(ptr) + storage_tensors.append(tensor) + + if not storage_tensors: + return + device = storage_tensors[0].device + indices_np = np.array(kv_cache_block_copies, dtype=np.int64) + indices = async_tensor_h2d(indices_np, device=device) + src_indices, dst_indices = indices.unbind(dim=1) + + for tensor in storage_tensors: + assert tensor.device == device + blocks = torch.empty(0, dtype=torch.uint8, device=device) + blocks.set_(tensor.untyped_storage()) + # Block-major backing storage: block i owns the contiguous byte range + # [i * page_size, (i + 1) * page_size). + assert blocks.numel() % num_blocks == 0 + blocks = blocks.view(num_blocks, -1) + blocks[dst_indices] = blocks[src_indices] + + def is_residual_scattered_for_sp( vllm_config: VllmConfig, num_input_tokens: int ) -> bool: