[2/N][Core] support partial prefix cache hit for hybrid model (#46384)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

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

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"]

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)

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,

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(

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
     )

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

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",

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),

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
 

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,

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.
 

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

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()

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)

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

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,
         )

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

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:

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:

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: