# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Compare the with and without prefix caching."""

import copy
from collections.abc import Callable
from dataclasses import replace
from math import lcm
from types import MethodType, SimpleNamespace
from unittest.mock import MagicMock

import pytest
import torch

import vllm.v1.core.kv_cache_manager as kv_cache_manager
import vllm.v1.core.kv_cache_utils as kv_cache_utils
from vllm.config import VllmConfig
from vllm.distributed.kv_events import (
    MEDIUM_CPU,
    MEDIUM_GPU,
    AllBlocksCleared,
    BlockRemoved,
    BlockStored,
)
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (
    HiSparseConnector,
    HiSparseConnectorScheduler,
    HiSparseConnectorWorkerMetadata,
)
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector
from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
    OffloadingConnector,
)
from vllm.lora.request import LoRARequest
from vllm.model_executor.layers.mamba.checkpoint import (
    compute_mamba_prefill_checkpoints,
)
from vllm.multimodal.inputs import (
    MultiModalFeatureSpec,
    MultiModalKwargsItem,
    PlaceholderRange,
)
from vllm.sampling_params import SamplingParams
from vllm.utils.hashing import sha256, sha256_cbor
from vllm.v1.core.block_pool import BlockHashToBlockMap, BlockPool
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager, Request
from vllm.v1.core.kv_cache_utils import (
    BlockHash,
    BlockHashWithGroupId,
    KVCacheBlock,
    get_block_hash,
    get_group_id,
    get_request_block_hasher,
    hash_block_tokens,
    init_none_hash,
    make_block_hash_with_group_id,
)
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.core.sched.scheduler import Scheduler
from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    HiSparseHotSpec,
    HiSparseResidentSpec,
    KVCacheConfig,
    KVCacheGroupRole,
    KVCacheGroupSpec,
    KVCacheSpecKind,
    MambaSpec,
    MLAAttentionSpec,
    SlidingWindowSpec,
)
from vllm.v1.outputs import KVConnectorOutput
from vllm.v1.request import RequestStatus

pytestmark = pytest.mark.cpu_test


@pytest.fixture(autouse=True)
def _auto_init_hash_fn(request):
    hash_fn: Callable
    if "hash_fn" in request.fixturenames:
        hash_fn = request.getfixturevalue("hash_fn")
    else:
        hash_fn = sha256
    init_none_hash(hash_fn)


def make_request(
    request_id: str,
    prompt_token_ids: list[int],
    block_size: int,
    hash_fn: Callable,
    mm_positions: list[PlaceholderRange] | None = None,
    mm_hashes: list[str] | None = None,
    prompt_logprobs: int | None = None,
    cache_salt: str | None = None,
    lora_request: LoRARequest | None = None,
    num_hash_lookahead_tokens: int = 0,
    session_id: str | None = None,
    sampling_params: SamplingParams | None = None,
):
    mm_features = []
    if mm_positions is not None:
        for j, position in enumerate(mm_positions):
            identifier = mm_hashes[j] if mm_hashes else f"hash_{j}"
            mm_feature = MultiModalFeatureSpec(
                data=MultiModalKwargsItem.dummy(),
                mm_position=position,
                identifier=identifier,
                modality="image",
            )
            mm_features.append(mm_feature)

    if sampling_params is None:
        sampling_params = SamplingParams(max_tokens=17, prompt_logprobs=prompt_logprobs)
    sampling_params.update_from_generation_config({}, eos_token_id=100)

    return Request(
        request_id=request_id,
        prompt_token_ids=prompt_token_ids,
        mm_features=mm_features if mm_features else None,
        sampling_params=sampling_params,
        pooling_params=None,
        lora_request=lora_request,
        cache_salt=cache_salt,
        block_hasher=get_request_block_hasher(
            block_size, hash_fn, num_hash_lookahead_tokens
        ),
        session_id=session_id,
    )


def make_kv_cache_manager(
    kv_cache_config: KVCacheConfig, vllm_config: VllmConfig | None = None, **kwargs
) -> KVCacheManager:
    """Build a ``KVCacheManager``, deriving ``scheduler_block_size`` from the
    config (LCM of group block sizes) unless explicitly provided. This mirrors
    ``resolve_kv_cache_block_sizes`` for the non-context-parallel case used by
    these tests, so callers don't have to pass it at every site. With
    ``vllm_config``, a HiSparseConnector builds the HiSparse coordinator."""
    kwargs.setdefault(
        "scheduler_block_size",
        lcm(*(g.kv_cache_spec.block_size for g in kv_cache_config.kv_cache_groups)),
    )
    if "retention_interval" in kwargs:
        kv_cache_config = replace(
            kv_cache_config,
            prefix_cache_retention_interval=kwargs.pop("retention_interval"),
        )
    manager = KVCacheManager(kv_cache_config, **kwargs)
    if kv_cache_config.hisparse_host_num_blocks is not None:
        if vllm_config is not None:
            HiSparseConnector(
                vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
            ).bind_kv_cache_manager(manager)
        else:
            get_hisparse_coordinator(manager)
    return manager


def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
    return KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            )
        ],
    )


HISPARSE_BLOCK_SIZE = 16


def make_hisparse_kv_cache_config(
    num_blocks: int,
    host_num_blocks: int,
    *,
    transfer_device_cache: bool = False,
    num_resident_groups: int = 1,
) -> KVCacheConfig:
    source_spec = FullAttentionSpec(
        block_size=HISPARSE_BLOCK_SIZE,
        num_kv_heads=1,
        head_size=1,
        dtype=torch.float32,
    )
    groups = [
        KVCacheGroupSpec(
            ["source"],
            source_spec,
            role=KVCacheGroupRole.HISPARSE_SOURCE,
            host_resident=True,
        ),
        KVCacheGroupSpec(
            ["indexer"],
            source_spec,
            enable_kv_transfer=transfer_device_cache,
            role=KVCacheGroupRole.HISPARSE_INDEXER,
        ),
    ]
    groups.append(
        KVCacheGroupSpec(
            ["resident"],
            HiSparseResidentSpec(
                block_size=HISPARSE_BLOCK_SIZE,
                page_size=HISPARSE_BLOCK_SIZE * 4,
            ),
            enable_kv_transfer=transfer_device_cache,
        )
    )
    groups.append(
        KVCacheGroupSpec(
            ["hot"],
            HiSparseHotSpec(
                block_size=HISPARSE_BLOCK_SIZE,
                page_size=HISPARSE_BLOCK_SIZE * 4,
                blocks_per_request=2,
            ),
            enable_kv_transfer=False,
        )
    )
    for group_idx in range(1, num_resident_groups):
        groups.extend(
            replace(
                group, layer_names=[f"{name}_{group_idx}" for name in group.layer_names]
            )
            for group in groups[2:4]
        )
    return KVCacheConfig(
        num_blocks=num_blocks,
        hisparse_host_num_blocks=host_num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=groups,
    )


def make_hisparse_kv_cache_manager(
    num_blocks: int,
    host_num_blocks: int,
    *,
    max_model_len: int = 128,
    max_in_flight_tokens: int | None = None,
    enable_caching: bool = False,
    num_prefill_lookahead: int = 0,
    vllm_config: VllmConfig | None = None,
    **config_kwargs,
) -> KVCacheManager:
    config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs)
    return make_kv_cache_manager(
        config,
        vllm_config,
        max_model_len=max_model_len,
        max_in_flight_tokens=max_in_flight_tokens,
        enable_caching=enable_caching,
        hash_block_size=HISPARSE_BLOCK_SIZE,
        num_prefill_lookahead=num_prefill_lookahead,
    )


def _allocate_scheduled(
    manager: KVCacheManager,
    request: Request,
    num_new_tokens: int,
    *,
    num_new_computed_tokens: int = 0,
    num_external_computed_tokens: int = 0,
    **kwargs,
) -> KVCacheBlocks | None:
    """``allocate_slots`` plus the per-step residency work the connector runs."""
    blocks = manager.allocate_slots(
        request,
        num_new_tokens,
        num_new_computed_tokens=num_new_computed_tokens,
        num_external_computed_tokens=num_external_computed_tokens,
        **kwargs,
    )
    if blocks is None or kwargs.get("delay_cache_blocks"):
        return blocks
    num_computed_tokens = min(
        request.num_computed_tokens
        + num_new_computed_tokens
        + num_external_computed_tokens,
        manager.max_model_len,
    )
    num_tokens = min(num_computed_tokens + num_new_tokens, request.num_tokens)
    get_hisparse_coordinator(manager).advance_scheduled(
        [(request.request_id, num_tokens)]
    )
    return blocks


def _resident_block_ids(manager: KVCacheManager, request_id: str) -> list[int]:
    """The GPU block HiSparse reads each resident page from; 0 if host-only."""
    resident = get_hisparse_coordinator(manager).resident_managers[0]
    return [block.block_id for block in resident.get_residency_row(request_id)]


def test_hisparse_does_not_write_back_reprefillable_tokens():
    """Multi-module MTP re-prefills the last tokens, so they are not final."""
    manager = make_hisparse_kv_cache_manager(32, 16, num_prefill_lookahead=3)
    request = make_request(
        "request", list(range(2 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256
    )
    assert _allocate_scheduled(manager, request, 2 * HISPARSE_BLOCK_SIZE)
    transfers = get_hisparse_coordinator(manager).build_offload_command().page_transfers
    assert len(transfers) == 1


def test_hisparse_builds_dma_row_mirrors_across_pages():
    manager = make_hisparse_kv_cache_manager(32, 16)
    request = make_request(
        "request",
        list(range(2 * HISPARSE_BLOCK_SIZE)),
        HISPARSE_BLOCK_SIZE,
        sha256,
    )
    assert _allocate_scheduled(manager, request, num_new_tokens=32) is not None
    coordinator = get_hisparse_coordinator(manager)
    resident_blocks = coordinator.resident_managers[0].req_to_blocks[request.request_id]
    host_blocks = coordinator.host_manager.req_to_blocks[request.request_id]

    mirrors = coordinator.build_row_mirrors([(request.request_id, 14, 4)])

    assert [mirror.source_starts for mirror in mirrors] == [
        (resident_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,),
        (resident_blocks[1].block_id * HISPARSE_BLOCK_SIZE,),
    ]
    assert [mirror.destination_start for mirror in mirrors] == [
        host_blocks[0].block_id * HISPARSE_BLOCK_SIZE + 14,
        host_blocks[1].block_id * HISPARSE_BLOCK_SIZE,
    ]
    assert [mirror.num_rows for mirror in mirrors] == [2, 2]


def test_hisparse_async_speculation_mirrors_uncertain_position_range():
    """Unresolved drafts must not leave gaps in the eager host mirror."""
    coordinator = MagicMock(host_group_id=0)
    coordinator.take_residency_updates.return_value = {}
    coordinator.build_offload_command.return_value = None
    coordinator.build_row_mirrors.return_value = ()
    scheduler = HiSparseConnectorScheduler(
        async_speculative=True,
        draft_kv_lookahead=4,
    )
    scheduler.bind_coordinator(coordinator)
    request = SimpleNamespace(request_id="request", num_output_placeholders=3)
    connector = object.__new__(HiSparseConnector)
    connector.connector_scheduler = scheduler
    connector.update_state_after_alloc(request, None, 0)
    scheduler_output = SimpleNamespace(
        kv_cache_block_copies=None,
        scheduled_new_reqs=[],
        scheduled_cached_reqs=SimpleNamespace(
            new_block_ids=[],
            req_ids=["request"],
            num_computed_tokens=[103],
        ),
        num_scheduled_tokens={"request": 4},
    )

    scheduler.build_connector_meta(scheduler_output)

    coordinator.build_row_mirrors.assert_called_once_with((("request", 100, 11),))

    # Acceptance arrives between steps; the connector must read the live count.
    request.num_output_placeholders = 1
    scheduler.build_connector_meta(scheduler_output)
    coordinator.build_row_mirrors.assert_called_with((("request", 102, 9),))

    connector.request_finished_all_groups(request, ())
    assert not scheduler.requests


def test_hisparse_reports_when_context_is_fully_resident():
    """A page the pool reuses under a host-reading request stops being resident."""
    manager = make_hisparse_kv_cache_manager(32, 16)
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None
    coordinator = get_hisparse_coordinator(manager)
    scheduled = ((request.request_id, len(tokens) - 1, 1),)
    assert coordinator.all_context_pages_resident(scheduled)

    for hot_manager in coordinator.hot_managers:
        hot_manager.require_hot(request.request_id)
    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None
    _publish_hisparse_pages(manager)
    coordinator.take_residency_updates([request.request_id])
    core_row = manager.get_block_ids(request.request_id)[2]
    pool = manager.block_pool
    pool.get_new_blocks(pool.get_num_free_blocks())

    assert not coordinator.all_context_pages_resident(scheduled)
    # The block-table row stays append-only; the worker learns of the lost
    # pages from the residency update instead.
    assert manager.get_block_ids(request.request_id)[2] == core_row
    resident_ids = _resident_block_ids(manager, request.request_id)
    lost = [page for page, block_id in enumerate(resident_ids) if block_id == 0]
    assert lost
    update = coordinator.take_residency_updates([request.request_id])
    assert update[request.request_id].pages == lost
    assert update[request.request_id].block_ids[0] == [0] * len(lost)


def test_hisparse_residency_update_sends_only_changed_pages():
    """Losing one early page must not resend the pages after it: per-step
    residency updates would otherwise grow with the context length."""
    manager = make_hisparse_kv_cache_manager(32, 16)
    tokens = list(range(6 * HISPARSE_BLOCK_SIZE))
    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens))
    coordinator = get_hisparse_coordinator(manager)
    for hot_manager in coordinator.hot_managers:
        hot_manager.require_hot(request.request_id)
    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens))
    coordinator.take_residency_updates([request.request_id])
    pool = manager.block_pool
    num_never_used = pool.get_num_free_blocks()
    # Pages unpin as their host copies complete, so page 0 is reused first.
    _publish_hisparse_pages(manager)
    pool.get_new_blocks(num_never_used + 1)

    update = coordinator.take_residency_updates([request.request_id])

    assert update[request.request_id].pages == [0]
    assert update[request.request_id].block_ids == ([0],)


def test_hisparse_readmitted_request_resends_its_full_residency():
    """A preempted request may resume in a worker state row holding another
    request's residency, so its first update must cover every page again."""
    manager = make_hisparse_kv_cache_manager(32, 16)
    coordinator = get_hisparse_coordinator(manager)
    tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, len(tokens)) is not None
    assert request.request_id in coordinator.take_residency_updates(
        [request.request_id]
    )
    manager.free(request)

    assert _allocate_scheduled(manager, request, len(tokens)) is not None
    update = coordinator.take_residency_updates([request.request_id])
    assert update[request.request_id].pages == [0, 1]
    assert update[request.request_id].block_ids[0] == _resident_block_ids(
        manager, request.request_id
    )


def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():
    """Restore the final host page when completing an indexer-only import."""
    manager = make_hisparse_kv_cache_manager(
        32,
        16,
        enable_caching=True,
    )
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None
    )
    spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
    spill_counts = {spill.transfer_id: 1 for spill in spills}
    get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)
    _, indexer_blocks, _, _ = manager.get_blocks(original.request_id).blocks
    evicted_indexer_id = indexer_blocks[2].block_id
    manager.free(original)
    manager.block_pool.evict_blocks({evicted_indexer_id})

    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
    blocks, num_local, _ = manager.get_computed_blocks(resumed)
    max_completion = HISPARSE_BLOCK_SIZE
    assert num_local == 2 * HISPARSE_BLOCK_SIZE
    assert [len(group_blocks) for group_blocks in blocks.blocks] == [3, 2, 0, 0]

    allocated = _allocate_scheduled(
        manager,
        resumed,
        num_new_tokens=1,
        num_new_computed_tokens=num_local,
        new_computed_blocks=blocks,
        num_external_computed_tokens=max_completion,
    )

    assert allocated is not None
    source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks
    assert len(source) == len(indexer) == len(resident) == 4
    assert len(hot) == 2
    assert 0 not in _resident_block_ids(manager, resumed.request_id)
    coordinator = get_hisparse_coordinator(manager)
    assert not coordinator.build_offload_command().page_transfers
    coordinator.finish_host_import(resumed.request_id, failed=False)
    transfers = coordinator.build_offload_command().page_transfers
    assert len(transfers) == 1 and transfers[0].restore
    assert transfers[0].resident_block_ids == (resident[2].block_id,)


def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix():
    manager = make_hisparse_kv_cache_manager(
        32,
        16,
        enable_caching=True,
    )
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None
    )
    spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
    spill_counts = {spill.transfer_id: 1 for spill in spills}
    get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)
    host_blocks, _, _, _ = manager.get_blocks(original.request_id).blocks
    evicted_host_id = host_blocks[0].block_id
    manager.free(original)
    get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
        {evicted_host_id}
    )

    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
    _, num_local, _ = manager.get_computed_blocks(resumed)
    connector = object.__new__(OffloadingConnector)
    connector._kv_cache_config = manager.kv_cache_config
    connector.bind_kv_cache_manager(manager)
    assert num_local == 0
    assert connector._max_loadable_tokens(resumed, num_local) == 0


@pytest.mark.parametrize("nested", [False, True])
@pytest.mark.parametrize(
    "available_tokens,imported_tokens,host_blocks",
    [(0, 0, 2), (16, 16, 2), (64, 32, 3)],
)
def test_connector_completes_partial_prefix_without_importing_missing_host_kv(
    nested, available_tokens, imported_tokens, host_blocks
):
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    tokens = list(range(64))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, original, num_new_tokens=64) is not None
    _publish_hisparse_pages(manager)
    host, indexer, _, _ = manager.get_blocks(original.request_id).blocks
    missing_indexer = indexer[1].block_id
    missing_host = host[0].block_id
    manager.free(original)
    manager.block_pool.evict_blocks({missing_indexer})

    defer = True

    def lookup(request, num_computed_tokens, max_num_new_tokens=None):
        nonlocal defer
        if defer:
            defer = False
            return None, False
        available = available_tokens
        if max_num_new_tokens is not None:
            available = min(available, max_num_new_tokens)
        return available, available > 0

    connector = object.__new__(OffloadingConnector)
    connector._kv_cache_config = manager.kv_cache_config
    connector.connector_scheduler = SimpleNamespace(get_num_new_matched_tokens=lookup)
    if nested:
        for _ in range(2):
            outer = object.__new__(MultiConnector)
            outer._connectors = (connector,)
            outer._requests_to_connector = {}
            connector = outer
    connector.bind_kv_cache_manager(manager)
    scheduler = SimpleNamespace(connector=connector, kv_cache_manager=manager)
    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)

    blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit(
        scheduler, resumed
    )
    assert local == 16 and not diverged
    assert connector.get_num_new_matched_tokens(resumed, local) == (None, False)

    blocks, local, _, diverged = Scheduler._get_local_prefix_cache_hit(
        scheduler, resumed
    )
    external, asynchronous = connector.get_num_new_matched_tokens(resumed, local)
    assert external == imported_tokens
    assert asynchronous == (external > 0)
    assert [len(group) for group in blocks.blocks] == [3, 1, 0, 0]
    assert (
        _allocate_scheduled(
            manager,
            resumed,
            int(external == 0),
            delay_cache_blocks=True,
            num_new_computed_tokens=local,
            new_computed_blocks=blocks,
            num_external_computed_tokens=external,
        )
        is not None
    )
    allocated = manager.get_blocks(resumed.request_id).blocks
    assert len(allocated[0]) == host_blocks
    assert allocated[2][1].is_null == (host_blocks > 2)
    assert not allocated[2][-1].is_null
    manager.free(resumed)

    # A surviving indexer offload must not make absent host KV look computed.
    get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
        {missing_host}
    )
    retry = make_request("retry", tokens, HISPARSE_BLOCK_SIZE, sha256)
    _, local, _, _ = Scheduler._get_local_prefix_cache_hit(scheduler, retry)
    assert local == 0
    assert connector.get_num_new_matched_tokens(retry, local) == (0, False)


def allocate_external_prefix(
    manager: KVCacheManager, request: Request, num_tokens: int
) -> KVCacheBlocks | None:
    return _allocate_scheduled(
        manager,
        request,
        num_new_tokens=0,
        num_external_computed_tokens=num_tokens,
        delay_cache_blocks=True,
        full_sequence_must_fit=True,
    )


@pytest.mark.parametrize("num_tokens", [1, 15, 16, 17, 31, 32, 33, 127])
def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens):
    """Import every real token and restore the page used by last-token replay."""
    from tests.v1.kv_connector.unit.utils import create_vllm_config
    from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (
        NixlPullConnectorScheduler,
    )

    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    config = create_vllm_config(block_size=HISPARSE_BLOCK_SIZE)
    connector = NixlPullConnectorScheduler(config, "test", manager.kv_cache_config)
    request = make_request(
        "import", list(range(num_tokens)), HISPARSE_BLOCK_SIZE, sha256
    )
    request.kv_transfer_params = {"do_remote_prefill": True}
    count, is_async = connector.get_num_new_matched_tokens(request, 0)
    assert count == num_tokens and is_async
    assert allocate_external_prefix(manager, request, count) is not None
    source, _, resident, _ = manager.get_blocks(request.request_id).blocks
    assert len(source) == len(resident) == (num_tokens + 15) // 16
    assert all(block.is_null for block in resident[:-1])
    assert not resident[-1].is_null
    coordinator = get_hisparse_coordinator(manager)
    assert not coordinator.build_offload_command().page_transfers

    coordinator.finish_host_import(request.request_id, failed=False)
    transfers = coordinator.build_offload_command().page_transfers
    assert len(transfers) == 1
    restore = transfers[0]
    assert restore.restore and not restore.after_forward
    assert restore.host_block_id == source[-1].block_id
    assert restore.resident_block_ids == (resident[-1].block_id,)
    request.num_computed_tokens = count
    scheduler = SimpleNamespace(
        connector=connector,
        kv_cache_manager=manager,
        failed_recving_kv_req_ids=set(),
        finished_recving_kv_req_ids={request.request_id},
        prefix_replay_tokens=0,
        prefix_replay_group_ids=(),
    )
    scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler)
    scheduler._load_restores_replay_window = MethodType(
        Scheduler._load_restores_replay_window, scheduler
    )
    Scheduler._update_waiting_for_remote_kv(scheduler, request)
    assert request.num_tokens - request.num_computed_tokens == 1
    assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None
    assert coordinator.build_row_mirrors([(request.request_id, num_tokens - 1, 1)])
    # A pending restore must not expose uninitialized GPU copies to prefix hits.
    assert all(resident[-1] not in copies for copies in coordinator.copies.values())

    coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1})
    assert resident[-1].ref_cnt == 2
    assert all(resident[-1] not in copies for copies in coordinator.copies.values())
    coordinator.update_spills({}, {restore.transfer_id: 1})
    assert resident[-1].ref_cnt == 1
    assert coordinator.request_states[request.request_id].valid_pages == set(
        range(num_tokens // HISPARSE_BLOCK_SIZE)
    )
    if num_tokens % HISPARSE_BLOCK_SIZE:
        assert source[-1].block_hash is None
    else:
        assert coordinator.copies[source[-1].block_hash] == (resident[-1],)


def test_hisparse_aborted_tail_restore_retains_both_endpoints():
    """An abort cannot recycle storage still used by a queued tail restore."""
    manager = make_hisparse_kv_cache_manager(16, 16)
    request = make_request("import", list(range(17)), HISPARSE_BLOCK_SIZE, sha256)
    assert allocate_external_prefix(manager, request, 17) is not None
    source, _, resident, _ = manager.get_blocks(request.request_id).blocks
    host_tail, gpu_tail = source[-1], resident[-1]
    coordinator = get_hisparse_coordinator(manager)
    coordinator.finish_host_import(request.request_id, failed=False)
    restore = coordinator.build_offload_command().page_transfers[0]
    manager.free(request)
    assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1
    coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1})
    assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1
    coordinator.update_spills({}, {restore.transfer_id: 1})
    assert host_tail.ref_cnt == gpu_tail.ref_cnt == 0
    assert not coordinator.has_pending_work()


@pytest.mark.parametrize("enable_caching", [False, True])
def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_caching):
    """Later admissions must not remove a scheduled request's only readable tier."""
    manager = make_hisparse_kv_cache_manager(
        32, 32, max_model_len=256, enable_caching=enable_caching
    )
    first = make_request("first", list(range(160)), HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, first, num_new_tokens=128) is not None
    _publish_hisparse_pages(manager)
    first.num_computed_tokens = 128
    pool = manager.block_pool
    held = pool.get_new_blocks(pool.get_num_free_blocks() - 3)

    assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None
    second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)
    _allocate_scheduled(manager, second, num_new_tokens=16)
    coordinator = get_hisparse_coordinator(manager)
    scheduled = (("first", 128, 1),)
    assert coordinator.all_context_pages_resident(scheduled)
    assert all(not m.has_hot("first") for m in coordinator.hot_managers)

    # Once another allocation finishes, A can acquire its hot buffer and
    # safely release its clean pages for other requests to reuse.
    pool.free_blocks(held)
    assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None
    assert all(m.has_hot("first") for m in coordinator.hot_managers)
    coordinator.take_residency_updates(["first"])
    pool.get_new_blocks(pool.get_num_free_blocks())
    assert not coordinator.all_context_pages_resident(scheduled)
    assert "first" in coordinator.take_residency_updates(["first"])


def test_hisparse_full_pool_keeps_pages_pinned_until_preemption():
    """Without room for a hot buffer, admission must defer instead of losing KV."""
    manager = make_hisparse_kv_cache_manager(18, 18, max_model_len=160)
    first = make_request("first", list(range(128)), HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, first, num_new_tokens=128) is not None
    _publish_hisparse_pages(manager)
    assert manager.block_pool.get_num_free_blocks() == 1
    first.num_computed_tokens = 128
    assert _allocate_scheduled(manager, first, num_new_tokens=1) is None
    assert get_hisparse_coordinator(manager).all_context_pages_resident(
        (("first", 127, 1),)
    )

    manager.free(first)
    second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, second, num_new_tokens=16) is not None


@pytest.mark.parametrize("num_resident_groups", [1, 3])
def test_hisparse_materializes_prefix_without_allocating_hot_blocks(
    num_resident_groups,
):
    """A host prefix becomes visible only when every page is durable.

    Completing a later page first must not expose a prefix with a hole.
    """
    manager = make_hisparse_kv_cache_manager(
        32,
        16,
        enable_caching=True,
        num_resident_groups=num_resident_groups,
    )
    tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
    request = make_request("resident", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None
    blocks = manager.get_block_ids(request.request_id)
    assert len(blocks[2]) == 2
    assert blocks[3] == []

    spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers
    assert len(spills) == 2
    assert all(spill.after_forward for spill in spills)
    resident_blocks = manager.get_blocks(request.request_id).blocks[2]
    assert [block.ref_cnt for block in resident_blocks] == [2, 2]
    spill_counts = {spill.transfer_id: 1 for spill in spills}
    duplicate = make_request(
        "duplicate",
        list(range(3 * HISPARSE_BLOCK_SIZE)),
        HISPARSE_BLOCK_SIZE,
        sha256,
    )
    _, num_computed, _ = manager.get_computed_blocks(duplicate)
    assert num_computed == 0

    get_hisparse_coordinator(manager).update_spills(spill_counts, {})
    assert [block.ref_cnt for block in resident_blocks] == [2, 2]
    _, num_computed, _ = manager.get_computed_blocks(duplicate)
    assert num_computed == 0

    get_hisparse_coordinator(manager).update_spills({}, {spills[1].transfer_id: 1})
    assert [block.ref_cnt for block in resident_blocks] == [2, 1]
    _, num_computed, _ = manager.get_computed_blocks(duplicate)
    assert num_computed == 0

    get_hisparse_coordinator(manager).update_spills({}, {spills[0].transfer_id: 1})
    _, num_computed, _ = manager.get_computed_blocks(duplicate)
    assert num_computed == 2 * HISPARSE_BLOCK_SIZE
    blocks = manager.get_block_ids(request.request_id)
    assert len(blocks[2]) == 2
    assert blocks[3] == []


@pytest.mark.parametrize("num_resident_groups", [1, 3])
def test_hisparse_materialization_respects_per_step_spill_budget(num_resident_groups):
    """Prefix publication must not bypass the configured spill batch limit."""
    manager = make_hisparse_kv_cache_manager(
        32,
        16,
        enable_caching=True,
        num_resident_groups=num_resident_groups,
    )
    coordinator = get_hisparse_coordinator(manager)
    coordinator.max_spill_pages = 1
    tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
    request = make_request("bounded", tokens, HISPARSE_BLOCK_SIZE, sha256)

    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None
    first = coordinator.build_offload_command().page_transfers
    assert len(first) == 1

    coordinator.plan_prefix_materialization(request.request_id, len(tokens))
    second = coordinator.build_offload_command().page_transfers
    assert len(second) == 1
    assert first[0].transfer_id != second[0].transfer_id


def test_hisparse_deferred_free_retains_host_blocks_until_fence():
    """Detaching a request must not make in-flight host pages reusable."""
    manager = make_hisparse_kv_cache_manager(16, 16)
    host_pool = get_hisparse_coordinator(manager).get_host_block_pool()
    assert host_pool is not None
    host_free = host_pool.get_num_free_blocks()
    device_free = manager.block_pool.get_num_free_blocks()
    request = make_request("free", list(range(16)), HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, request, 16, delay_cache_blocks=True) is not None
    )
    blocks = manager.pop_blocks_for_free(request)
    assert host_pool.get_num_free_blocks() < host_free
    assert blocks and any(host_pool.blocks[b.block_id] is b for b in blocks)
    manager.block_pool.free_blocks(reversed(blocks))
    assert host_pool.get_num_free_blocks() == host_free
    assert manager.block_pool.get_num_free_blocks() == device_free


def test_hisparse_host_cow_copy_is_drained_without_a_gpu_pool():
    """Host-only copy-on-write work must reach the worker copy queue."""
    manager = make_hisparse_kv_cache_manager(16, 16)
    source_manager = manager.coordinator.single_type_managers[0]
    source_block = source_manager.block_pool.get_new_blocks(1)[0]
    source_manager.req_to_blocks["cow"] = [source_block]
    source_manager._partial_hit_reqs["cow"] = (0, source_block)

    new_blocks = source_manager.allocate_new_blocks(
        "cow", HISPARSE_BLOCK_SIZE, HISPARSE_BLOCK_SIZE
    )
    new_block_ids = manager.take_new_block_ids()
    coordinator = get_hisparse_coordinator(manager)
    assert manager.take_kv_cache_block_copies() == ([], [])
    connector = HiSparseConnectorScheduler(async_speculative=False)
    connector.bind_coordinator(coordinator)
    output = SchedulerOutput.make_empty()
    copies = connector.build_connector_meta(output).host_block_copies
    assert output.has_sync_kv_loads

    assert new_blocks and new_block_ids == []
    assert len(copies) == 1
    assert copies[0].src_block_id == source_block.block_id
    assert copies[0].dst_block_id == new_blocks[0].block_id
    source_manager.free("cow")
    assert source_block.ref_cnt > 0 and new_blocks[0].ref_cnt > 0
    assert coordinator.has_pending_work()
    connector.update_connector_output(
        KVConnectorOutput(
            kv_connector_worker_meta=HiSparseConnectorWorkerMetadata(
                {}, {}, completed_host_copy_dst_ids=(copies[0].dst_block_id,)
            )
        )
    )
    assert source_block.ref_cnt == new_blocks[0].ref_cnt == 0
    assert not coordinator.has_pending_work()


def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages():
    """An in-flight host import must reserve its unwritten resident pages."""
    manager = make_hisparse_kv_cache_manager(
        12,
        16,
        transfer_device_cache=True,
    )
    request = make_request(
        "partial-import",
        list(range(6 * HISPARSE_BLOCK_SIZE)),
        HISPARSE_BLOCK_SIZE,
        sha256,
    )
    imported_tokens = 4 * HISPARSE_BLOCK_SIZE

    assert allocate_external_prefix(manager, request, imported_tokens) is not None
    required = manager.coordinator.get_num_blocks_to_allocate(
        request_id=request.request_id,
        num_tokens=request.num_tokens,
        new_computed_blocks=manager.empty_kv_cache_blocks.blocks,
        num_encoder_tokens=0,
        total_computed_tokens=imported_tokens,
        num_local_computed_tokens=0,
        num_tokens_main_model=request.num_tokens,
        apply_admission_cap=True,
    )

    assert required == 4


def _make_hisparse_scheduler(
    tmp_path, host_num_blocks: int, matched_tokens: int, **kwargs
) -> tuple[Scheduler, KVCacheManager]:
    """A scheduler over a HiSparse manager whose async mock connector matches
    ``matched_tokens``."""
    from .utils import create_scheduler, mock_kv

    (tmp_path / "config.json").write_text(
        '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}'
    )
    scheduler = create_scheduler(
        model=str(tmp_path),
        skip_tokenizer_init=True,
        max_model_len=128,
        use_kv_connector=mock_kv(matched_tokens=matched_tokens, is_async=True),
        **kwargs,
    )
    manager = make_hisparse_kv_cache_manager(
        32, host_num_blocks, vllm_config=scheduler.vllm_config
    )
    scheduler.kv_cache_manager = manager
    scheduler.kv_cache_config = manager.kv_cache_config
    assert scheduler.scheduler_reserve_full_isl
    return scheduler, manager


def _add_hisparse_request(
    scheduler: Scheduler, request_id: str, num_pages: int, offset: int = 0
) -> Request:
    request = make_request(
        request_id,
        list(range(offset, offset + num_pages * HISPARSE_BLOCK_SIZE)),
        HISPARSE_BLOCK_SIZE,
        sha256,
    )
    scheduler.add_request(request)
    scheduler.schedule()
    return request


def test_hisparse_async_admission_fits_beside_inflight_import(tmp_path):
    """The in-flight import's three prompt pages stay reserved, so the waiting
    request's four prompt pages fit in eight host blocks including the null
    block."""
    scheduler, manager = _make_hisparse_scheduler(tmp_path, 8, HISPARSE_BLOCK_SIZE)
    inflight = _add_hisparse_request(scheduler, "inflight", 3)
    assert inflight in scheduler._inflight_prefills

    request = _add_hisparse_request(scheduler, "waiting", 4)
    assert request in scheduler._inflight_prefills
    assert manager.get_blocks(request.request_id).blocks[0]


@pytest.mark.parametrize("release,second_pages", [("free", 6), ("resume", 4)])
def test_hisparse_import_reservation_released(release, second_pages, tmp_path):
    """An import stops reserving host pages once it allocates them or is freed.

    Seven usable host blocks: the first import holds one page and reserves two.
    Once it is released, the second import fits only if nothing is still
    reserved for the first.
    """
    scheduler, manager = _make_hisparse_scheduler(tmp_path, 8, HISPARSE_BLOCK_SIZE)
    first = _add_hisparse_request(scheduler, "first", 3)
    assert first in scheduler._inflight_prefills
    if release == "free":
        # What abort and preemption run; the mock connector cannot finish a
        # request over several cache groups.
        manager.free(first)
    else:
        scheduler.finished_recving_kv_req_ids.add(first.request_id)
        assert scheduler.schedule().num_scheduled_tokens[first.request_id] > 0

    second = _add_hisparse_request(scheduler, "second", second_pages, offset=1000)
    assert second in scheduler._inflight_prefills


@pytest.mark.parametrize(
    "host_num_blocks,first_pages,matched_tokens,num_speculative_tokens",
    [
        # Once first imports one page, three are free: enough for second's
        # prompt but not beside first's remaining three.
        (5, 4, HISPARSE_BLOCK_SIZE, None),
        # Both prompts are fully imported; first's draft tokens need a fourth
        # page on resume.
        (7, 3, 3 * HISPARSE_BLOCK_SIZE - 1, 4),
    ],
)
def test_hisparse_admitted_async_loads_can_finish_with_nothing_running(
    tmp_path, host_num_blocks, first_pages, matched_tokens, num_speculative_tokens
):
    """Admitted async loads must not strand each other on host capacity.

    Waiting requests keep their blocks and are never preempted, so two imports
    that each hold host pages while together needing more than the pool has
    left would never be scheduled again. A load is admitted only if its whole
    sequence, including the spec decode step, fits beside the pages earlier
    imports still need.
    """
    scheduler, _ = _make_hisparse_scheduler(
        tmp_path,
        host_num_blocks,
        matched_tokens,
        num_speculative_tokens=num_speculative_tokens,
    )
    _add_hisparse_request(scheduler, "first", first_pages)
    _add_hisparse_request(scheduler, "second", 3, offset=1000)
    admitted = {r.request_id for r in scheduler._inflight_prefills}
    assert admitted == {"first"}

    scheduler.finished_recving_kv_req_ids.update(admitted)
    output = scheduler.schedule()
    assert output.num_scheduled_tokens.get("first", 0) > 0


def test_hisparse_import_reservation_does_not_gate_local_admissions(tmp_path):
    """Like the GPU pool's, the host reservation only gates async loads; a
    local admission runs and can be preempted, so it cannot strand an import."""
    # Seven usable host blocks: the import holds one and reserves five, leaving
    # one, so gating the local request's two pages on the reservation fails.
    scheduler, _ = _make_hisparse_scheduler(tmp_path, 8, HISPARSE_BLOCK_SIZE)
    _add_hisparse_request(scheduler, "import", 6)
    match_external = scheduler.connector.get_num_new_matched_tokens
    scheduler.connector.get_num_new_matched_tokens = lambda request, num_computed: (
        (0, False)
        if request.request_id == "local"
        else match_external(request, num_computed)
    )
    local = _add_hisparse_request(scheduler, "local", 2, offset=1000)
    assert local in scheduler.running


@pytest.mark.parametrize(
    "evictable,local_tokens,free_blocks,admitted",
    [
        (False, HISPARSE_BLOCK_SIZE, 1, True),
        (True, HISPARSE_BLOCK_SIZE, 1, False),
        (False, HISPARSE_BLOCK_SIZE - 1, 1, False),
        (False, HISPARSE_BLOCK_SIZE - 1, 2, True),
    ],
)
def test_hisparse_import_capacity_includes_hits_and_cow(
    evictable, local_tokens, free_blocks, admitted
):
    """Touching an evictable hit and copying a partial hit consume capacity."""
    manager = make_hisparse_kv_cache_manager(32, free_blocks + 2 - int(evictable))
    coordinator = get_hisparse_coordinator(manager)
    host_pool = coordinator.get_host_block_pool()
    assert host_pool is not None
    hit = host_pool.get_new_blocks(1)[0]
    if evictable:
        host_pool.free_blocks([hit])
    assert host_pool.get_num_free_blocks() == free_blocks
    assert (
        coordinator.host_manager.get_num_blocks_to_allocate(
            "waiting", 32, [hit], 32, local_tokens, 32
        )
        <= host_pool.get_num_free_blocks()
    ) == admitted


@pytest.mark.parametrize("enable_caching", [False, True])
def test_hisparse_host_exhaustion_defers_allocation(enable_caching):
    """Pages are only allocated with host backing, so every one can be reclaimed."""
    manager = make_hisparse_kv_cache_manager(32, 5, enable_caching=enable_caching)
    donor = make_request("donor", [99] * 16, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, donor, 16) is not None
    request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
    free_gpu = manager.block_pool.get_num_free_blocks()
    assert manager.allocate_slots(request, 64) is None
    assert manager.block_pool.get_num_free_blocks() == free_gpu

    # The donor's in-flight write-back holds its host block until acked.
    manager.free(donor)
    coordinator = get_hisparse_coordinator(manager)
    assert coordinator.has_pending_reclamation()
    assert manager.allocate_slots(request, 64) is None
    _publish_hisparse_pages(manager)
    assert not coordinator.has_pending_reclamation()
    assert _allocate_scheduled(manager, request, 64) is not None
    host = coordinator.host_manager
    assert host is not None
    assert not any(b.is_null for b in host.req_to_blocks[request.request_id])
    _publish_hisparse_pages(manager)
    for hot in coordinator.hot_managers:
        hot.require_hot(request.request_id)
        hot.allocate_new_blocks(request.request_id, 64, 64)
    coordinator.update_residency(request.request_id)
    assert coordinator.request_states[request.request_id].unpinned_pages == {0, 1}
    for resident in coordinator.resident_managers:
        assert [b.ref_cnt for b in resident.req_to_blocks[request.request_id]] == [
            0,
            0,
            1,
            1,
        ]


@pytest.mark.parametrize("failed", [False, True])
def test_host_receive_completion_without_spill_metadata(failed):
    """Only successful external receives make host pages readable."""
    manager = make_hisparse_kv_cache_manager(32, 16)
    request = make_request("import", list(range(32)), HISPARSE_BLOCK_SIZE, sha256)
    assert allocate_external_prefix(manager, request, 32) is not None
    request.num_computed_tokens = 32
    request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
    scheduler = HiSparseConnectorScheduler(async_speculative=False)
    scheduler.bind_coordinator(get_hisparse_coordinator(manager))
    scheduler.requests[request.request_id] = request
    connector = object.__new__(HiSparseConnector)
    connector.connector_scheduler = scheduler
    composite = object.__new__(MultiConnector)
    composite._connectors = (connector,)
    composite.update_connector_output(
        KVConnectorOutput(
            finished_recving={request.request_id},
            failed_recving={request.request_id} if failed else set(),
        )
    )
    state = get_hisparse_coordinator(manager).request_states.get(request.request_id)
    if failed:
        assert state is None
    else:
        assert state is not None and state.valid_pages == {0, 1}


def test_hisparse_recomputation_does_not_complete_failed_host_import():
    """Allocating recomputation is not evidence that imported host KV is valid."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    request = make_request(
        "partial-failure", list(range(48)), HISPARSE_BLOCK_SIZE, sha256
    )
    assert allocate_external_prefix(manager, request, 32) is not None
    # A partial load failure leaves only the first page computed. The next
    # allocation reaches the original import boundary before the worker runs.
    request.num_computed_tokens = 16
    manager.cache_blocks(request, 16)
    assert _allocate_scheduled(manager, request, 16) is not None
    state = get_hisparse_coordinator(manager).request_states.get(request.request_id)
    assert state is None or 1 not in state.valid_pages


def test_hisparse_host_import_ignores_unsealed_tail():
    """A partial imported page must not become readable from stale host data."""
    manager = make_hisparse_kv_cache_manager(16, 16)
    coordinator = get_hisparse_coordinator(manager)

    coordinator.complete_host_import("partial", HISPARSE_BLOCK_SIZE + 1)

    state = coordinator.request_states["partial"]
    assert state.valid_pages == {0}
    assert state.ready_prefix_pages == 1


def test_hisparse_resident_request_can_grow_without_hot_capacity():
    """Running resident history must not be mistaken for a new host-prefix hit."""
    manager = make_hisparse_kv_cache_manager(8, 16)
    request = make_request("resident", list(range(48)), HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, num_new_tokens=32) is not None
    request.num_computed_tokens = 32

    # Three free blocks fit the next indexer/resident pages, but not a hot region.
    assert manager.block_pool.get_num_free_blocks() == 3
    assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None
    assert manager.get_block_ids(request.request_id)[3] == []


def test_hisparse_capacity_query_does_not_require_hot_blocks():
    """A read-only capacity query must not mutate hot-block requirements."""
    manager = make_hisparse_kv_cache_manager(16, 16)
    coordinator = get_hisparse_coordinator(manager)
    host_pool = coordinator.get_host_block_pool()
    assert host_pool is not None
    computed: tuple[list[KVCacheBlock], ...] = (
        [host_pool.blocks[1]],
        [],
        [],
        [],
    )

    manager.coordinator.get_num_blocks_to_allocate(
        request_id="query-only",
        num_tokens=HISPARSE_BLOCK_SIZE,
        new_computed_blocks=computed,
        num_encoder_tokens=0,
        total_computed_tokens=HISPARSE_BLOCK_SIZE,
        num_local_computed_tokens=HISPARSE_BLOCK_SIZE,
        num_tokens_main_model=HISPARSE_BLOCK_SIZE,
    )

    assert all(
        "query-only" not in hot_manager.hot_required
        for hot_manager in coordinator.hot_managers
    )


def _publish_hisparse_pages(manager: KVCacheManager) -> None:
    """Ack all planned spills so host pages publish to the prefix cache."""
    command = get_hisparse_coordinator(manager).build_offload_command()
    counts = {spill.transfer_id: 1 for spill in command.page_transfers}
    get_hisparse_coordinator(manager).update_spills(counts, counts)


def test_hisparse_events_report_host_and_device_placement():
    """Host publication and eviction must not advertise GPU-resident KV."""
    manager = make_kv_cache_manager(
        make_hisparse_kv_cache_config(32, 16),
        max_model_len=128,
        enable_caching=True,
        hash_block_size=HISPARSE_BLOCK_SIZE,
        enable_kv_cache_events=True,
    )
    request = make_request("request", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, 64) is not None
    _publish_hisparse_pages(manager)

    stored = [
        event for event in manager.take_events() if isinstance(event, BlockStored)
    ]
    assert {(event.group_idx, event.medium) for event in stored} == {
        (0, MEDIUM_CPU),
        (1, MEDIUM_GPU),
    }

    host_block = manager.get_blocks(request.request_id).blocks[0][0]
    get_hisparse_coordinator(manager).host_manager.block_pool.evict_blocks(
        {host_block.block_id}
    )
    [removed] = manager.take_events()
    assert isinstance(removed, BlockRemoved)
    assert removed.medium == MEDIUM_CPU
    assert removed.group_idx == 0


def test_mamba_boundary_handoffs_do_not_pin_obsolete_blocks():
    """Draining unused connector offers must not retain old recurrent states."""
    manager = make_kv_cache_manager(
        KVCacheConfig(
            num_blocks=16,
            kv_cache_tensors=[],
            kv_cache_groups=[
                KVCacheGroupSpec(
                    ["mamba"],
                    MambaSpec(
                        block_size=16,
                        shapes=((1,),),
                        dtypes=(torch.float32,),
                        mamba_cache_mode="align",
                    ),
                )
            ],
        ),
        max_model_len=128,
        enable_caching=True,
        hash_block_size=16,
    )
    request = make_request("request", list(range(64)), 16, sha256)
    old_blocks = []
    for end in (16, 32, 48):
        assert manager.allocate_slots(request, 16) is not None
        request.num_computed_tokens = end
        offers = manager.take_boundary_state_offloads()
        [(_, block_id, boundary)] = offers[request.request_id]
        assert boundary == end
        old_blocks.append(manager.block_pool.blocks[block_id])
        manager.remove_skipped_blocks(request.request_id, end, 64)
        manager.new_step_starts()

    assert all(block.ref_cnt == 0 for block in old_blocks[:-1])


def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():
    """A host prefix hit must come back GPU-resident while shadows survive."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None
    )
    _publish_hisparse_pages(manager)
    original_resident_ids = [
        block.block_id for block in manager.get_blocks("original").blocks[2]
    ]
    manager.free(original)

    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
    computed, num_computed, _ = manager.get_computed_blocks(resumed)
    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
    assert _allocate_scheduled(
        manager,
        resumed,
        num_new_tokens=len(tokens) - num_computed,
        num_new_computed_tokens=num_computed,
        new_computed_blocks=computed,
    )
    assert _resident_block_ids(manager, "resumed")[:3] == original_resident_ids[:3]
    assert get_hisparse_coordinator(manager).all_context_pages_resident(
        ((resumed.request_id, num_computed, len(tokens) - num_computed),)
    )


@pytest.mark.parametrize("free_blocks,adopted_pages", [(7, []), (10, [0, 1])])
def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies(
    free_blocks, adopted_pages
):
    """A host prefix hit must not starve its own allocation of free blocks.

    Admission counts free GPU copies as free, but adopting a copy pins it. When
    adoption ran before the hit's allocation, it took blocks the allocation was
    promised and the pool ran dry (``Cannot get N free blocks``). The allocation
    now goes first, evicting a prefix's copies tail first; the copies it leaves
    are adopted.
    """
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None
    )
    _publish_hisparse_pages(manager)
    copy_ids = [block.block_id for block in manager.get_blocks("original").blocks[2]]
    manager.free(original)
    pool = manager.block_pool
    pool.get_new_blocks(pool.get_num_free_blocks() - free_blocks)

    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
    computed, num_computed, _ = manager.get_computed_blocks(resumed)
    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
    assert _allocate_scheduled(
        manager,
        resumed,
        num_new_tokens=len(tokens) - num_computed,
        num_new_computed_tokens=num_computed,
        new_computed_blocks=computed,
    )

    resident_ids = _resident_block_ids(manager, "resumed")
    assert resident_ids[:3] == [
        copy_ids[page] if page in adopted_pages else pool.null_block.block_id
        for page in range(3)
    ]
    _, indexer, _, hot = manager.get_block_ids("resumed")
    gpu_ids = [block_id for block_id in indexer + resident_ids + hot if block_id]
    assert len(gpu_ids) == len(set(gpu_ids))


def test_hisparse_host_backed_request_accepts_local_prefix_hit():
    """Local group-completion hits need no external host import allocation."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None
    )
    _publish_hisparse_pages(manager)
    manager.free(original)

    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
    computed, num_computed, _ = manager.get_computed_blocks(resumed)
    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
    assert _allocate_scheduled(
        manager,
        resumed,
        num_new_tokens=len(tokens) - num_computed,
        num_new_computed_tokens=num_computed,
        new_computed_blocks=computed,
        num_external_computed_tokens=0,
    )


def test_hisparse_finished_request_leaves_free_copies():
    """Clean pages of a finished request are free-queue blocks, not pins."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None
    )
    _publish_hisparse_pages(manager)
    manager.free(original)

    coordinator = get_hisparse_coordinator(manager)
    assert coordinator.copies
    pool = manager.block_pool
    assert pool.get_num_free_blocks() == pool.num_gpu_blocks - 1
    pool.get_new_blocks(pool.get_num_free_blocks())
    assert not coordinator.copies
    assert not coordinator.spills_to_send


@pytest.mark.parametrize(
    "status,expected_cached_pages",
    [
        (RequestStatus.FINISHED_STOPPED, 2),
        (RequestStatus.FINISHED_ABORTED, 2),
        (RequestStatus.RUNNING, 0),
    ],
    ids=["stopped", "aborted", "running"],
)
def test_hisparse_cleanup_publishes_only_finalized_terminal_pages(
    status, expected_cached_pages
):
    """Only terminal cleanup may publish completed, finalized pages."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    coordinator = get_hisparse_coordinator(manager)
    request = make_request("terminal", list(range(33)), HISPARSE_BLOCK_SIZE, sha256)
    request.spec_token_ids = list(range(33, 49))
    assert (
        _allocate_scheduled(manager, request, 49, num_lookahead_tokens=16) is not None
    )
    host_blocks = list(manager.get_blocks(request.request_id).blocks[0])
    command = coordinator.build_offload_command()
    counts = {transfer.transfer_id: 1 for transfer in command.page_transfers}
    assert len(counts) == 2
    request.status = status
    manager.free(request)
    coordinator.update_spills(counts, counts)
    assert all(
        block.block_hash is not None for block in host_blocks[:expected_cached_pages]
    )
    assert all(
        block.block_hash is None for block in host_blocks[expected_cached_pages:]
    )
    assert all(block.ref_cnt == 0 for block in host_blocks)
    assert not coordinator.has_pending_work()


@pytest.mark.parametrize("num_pages,reads_host", [(4, True), (3, False)])
def test_hisparse_prefill_reads_host_once_it_fills_admission_window(
    num_pages, reads_host
):
    """A prefill switches to host reads once it fills its admission window.

    Admission caps each request's resident pages at the in-flight window,
    assuming older pages move to host, but pinned pages were released only
    once free blocks fell below a fixed watermark. A chunked prefill larger
    than the watermark outgrew the pool first and was preempted.
    """
    window_pages = 4
    manager = make_hisparse_kv_cache_manager(
        32,
        16,
        enable_caching=True,
        max_in_flight_tokens=window_pages * HISPARSE_BLOCK_SIZE,
    )
    coordinator = get_hisparse_coordinator(manager)
    tokens = list(range(2 * window_pages * HISPARSE_BLOCK_SIZE))
    request = make_request("prefill", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, request, num_pages * HISPARSE_BLOCK_SIZE)
        is not None
    )
    pool = coordinator.gpu_pool
    assert pool is not None
    assert pool.get_num_free_blocks() >= coordinator.transition_watermark

    hot_manager = coordinator.hot_managers[0]
    assert (request.request_id in hot_manager.hot_required) == reads_host


def test_hisparse_preempted_prefill_resumes_from_durable_prefix():
    """A prefill preempted mid-way must resume from the pages already on host.

    Host writes trail a chunked prefill by a chunk, and host pages were only
    published once the whole computed prefix was durable. Each chunk moved
    that target past the finished writes, so nothing was published before the
    prefill ended: a prefill preempted each time it outgrew the GPU pool
    recomputed from the start forever.
    """
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    coordinator = get_hisparse_coordinator(manager)
    chunk = 2 * HISPARSE_BLOCK_SIZE
    tokens = list(range(3 * chunk))
    request = make_request("long", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, chunk) is not None
    request.num_computed_tokens = chunk
    first_chunk_writes = coordinator.build_offload_command().page_transfers
    assert _allocate_scheduled(manager, request, chunk) is not None
    request.num_computed_tokens = 2 * chunk
    acks = {transfer.transfer_id: 1 for transfer in first_chunk_writes}
    coordinator.update_spills(acks, acks)
    manager.free(request)

    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert manager.get_computed_blocks(resumed)[1] == chunk


def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity():
    """Old completions retain their pages through pressure and request-ID reuse."""
    manager = make_hisparse_kv_cache_manager(32, 8, enable_caching=True)
    coordinator = get_hisparse_coordinator(manager)
    host_pool = coordinator.get_host_block_pool()
    tokens = list(range(2 * HISPARSE_BLOCK_SIZE + 1))
    request = make_request("reused", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, len(tokens)) is not None
    old_blocks = list(manager.get_blocks(request.request_id).blocks[0])
    command = coordinator.build_offload_command()
    counts = {transfer.transfer_id: 2 for transfer in command.page_transfers}
    partial = dict.fromkeys(counts, 1)
    first_completions = partial.copy()
    first_completions[next(iter(counts))] = 2
    coordinator.update_spills(counts, first_completions)
    request.status = RequestStatus.FINISHED_LENGTH_CAPPED
    manager.free(request)
    replacement = make_request(
        "reused", list(range(100, 117)), HISPARSE_BLOCK_SIZE, sha256
    )
    assert _allocate_scheduled(manager, replacement, 17) is not None
    new_blocks = list(manager.get_blocks(replacement.request_id).blocks[0])
    pressure = host_pool.get_new_blocks(host_pool.get_num_free_blocks())
    assert not {block.block_id for block in old_blocks[:2]} & {
        block.block_id for block in pressure + new_blocks
    }
    # Only the page every worker finished writing is published.
    assert old_blocks[0].block_hash is not None
    assert all(block.block_hash is None for block in old_blocks[1:])
    coordinator.update_spills({}, partial)
    repeated = make_request("probe", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert manager.get_computed_blocks(repeated)[1] == 2 * HISPARSE_BLOCK_SIZE
    assert all(block.block_hash is None for block in new_blocks)
    assert all(block.ref_cnt == 0 for block in old_blocks[:2])
    coordinator.update_spills({}, partial)
    assert all(block.ref_cnt == 0 for block in old_blocks[:2])
    host_pool.free_blocks(pressure)
    manager.free(replacement)
    _publish_hisparse_pages(manager)
    assert not coordinator.has_pending_work()


def test_hisparse_reset_prefix_cache_drops_copies():
    """Reset must forget GPU copies along with the host hashes they mirror."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None
    _publish_hisparse_pages(manager)
    manager.free(request)

    coordinator = get_hisparse_coordinator(manager)
    assert coordinator.copies
    assert manager.reset_prefix_cache()
    assert not coordinator.copies


def test_hisparse_reused_copy_is_not_adopted():
    """A host prefix hit must not adopt a GPU copy the pool has handed out."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)
    assert (
        _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None
    )
    _publish_hisparse_pages(manager)
    indexer_blocks = list(manager.get_blocks("original").blocks[1])
    manager.free(original)
    pool = manager.block_pool
    pool.touch(indexer_blocks)
    reused = pool.get_new_blocks(pool.get_num_free_blocks())
    pool.free_blocks(reused)
    pool.free_blocks(indexer_blocks)

    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
    computed, num_computed, _ = manager.get_computed_blocks(resumed)
    assert num_computed == 3 * HISPARSE_BLOCK_SIZE
    assert _allocate_scheduled(
        manager,
        resumed,
        num_new_tokens=len(tokens) - num_computed,
        num_new_computed_tokens=num_computed,
        new_computed_blocks=computed,
    )
    assert _resident_block_ids(manager, "resumed")[:3] == [0, 0, 0]


def test_hisparse_external_import_uses_hard_gpu_footprint():
    """An external prefix larger than resident capacity must remain admissible.

    This guards the admission boundary where indexer + one writable resident
    page + the fixed hot allocation fit, but indexer + full resident history do
    not. Requiring one more resident page would silently strand long P/D inputs.
    """
    num_prompt_blocks = 4
    num_prompt_tokens = num_prompt_blocks * HISPARSE_BLOCK_SIZE
    manager = make_hisparse_kv_cache_manager(
        8,
        9,
        transfer_device_cache=True,
    )
    request = make_request(
        "host-import",
        list(range(num_prompt_tokens)),
        HISPARSE_BLOCK_SIZE,
        sha256,
    )

    allocated = allocate_external_prefix(manager, request, num_prompt_tokens)

    assert allocated is not None
    source, indexer, resident, hot = manager.get_blocks(request.request_id).blocks
    assert len(source) == num_prompt_blocks
    assert len(indexer) == num_prompt_blocks
    assert len(resident) == num_prompt_blocks
    assert all(block.is_null for block in resident[:-1])
    assert not resident[-1].is_null
    assert len(hot) == 2
    assert manager.block_pool.get_num_free_blocks() == 0


def test_hisparse_external_import_survives_capacity_retry():
    manager = make_hisparse_kv_cache_manager(
        9,
        7,
        transfer_device_cache=True,
    )
    tokens = list(range(4 * HISPARSE_BLOCK_SIZE))
    first = make_request("first", tokens, HISPARSE_BLOCK_SIZE, sha256)
    second = make_request("second", tokens, HISPARSE_BLOCK_SIZE, sha256)

    assert allocate_external_prefix(manager, first, len(tokens)) is not None

    assert allocate_external_prefix(manager, second, len(tokens)) is None

    manager.free(first)
    assert allocate_external_prefix(manager, second, len(tokens)) is not None
    _, _, resident, hot = manager.get_blocks(second.request_id).blocks
    assert all(block.is_null for block in resident[:-1])
    assert not resident[-1].is_null
    assert len(hot) == 2


def make_kv_cache_config_hybrid_model(
    block_size: int,
    num_blocks: int,
    sliding_window_blocks: int,
    second_spec_type: str = "sliding_window",
) -> KVCacheConfig:
    if second_spec_type == "sliding_window":
        second_spec = SlidingWindowSpec(
            block_size=block_size,
            num_kv_heads=1,
            head_size=1,
            dtype=torch.float32,
            sliding_window=sliding_window_blocks * block_size,
        )
    elif second_spec_type == "mamba":
        second_spec = MambaSpec(
            block_size=block_size,
            shapes=((1, 1),),
            dtypes=(torch.float32,),
            mamba_cache_mode="align",
        )

    return KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                second_spec,
            ),
            KVCacheGroupSpec(
                ["layer3"],
                second_spec,
            ),
        ],
    )


def make_kv_cache_config_three_types(
    block_size: int, num_blocks: int, third_spec_type: str = "mamba"
) -> KVCacheConfig:
    if third_spec_type == "mamba":
        third_spec = MambaSpec(
            block_size=block_size,
            shapes=((1, 1),),
            dtypes=(torch.float32,),
            mamba_cache_mode="align",
        )
    elif third_spec_type == "sliding_window":
        third_spec = SlidingWindowSpec(
            block_size=block_size,
            num_kv_heads=1,
            head_size=1,
            dtype=torch.float32,
            sliding_window=4 * block_size,
        )

    return KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=2 * block_size,
                ),
            ),
            KVCacheGroupSpec(
                ["layer3"],
                third_spec,
            ),
        ],
    )


@pytest.mark.parametrize("draft_sharded", [True, False])
def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded):
    """Prefix lookup must use each group's DCP size, not the process-wide one.

    Target MLA is sharded; draft MLA can opt out and Mamba stays replicated.
    Hits must align to the target block in either case.
    """
    block_size = 16
    dcp = 8
    sharded_block = block_size * dcp
    kv_cache_config = KVCacheConfig(
        num_blocks=64,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["target_mla"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            ),
            KVCacheGroupSpec(
                ["draft_mla"],
                MLAAttentionSpec(
                    dcp_sharded=draft_sharded,
                    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 = KVCacheManager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        scheduler_block_size=sharded_block,
        dcp_world_size=dcp,
    )
    target_mgr, draft_mgr, mamba_mgr = manager.coordinator.single_type_managers
    assert target_mgr.dcp_world_size == dcp
    assert draft_mgr.dcp_world_size == (dcp if draft_sharded else 1)
    assert mamba_mgr.dcp_world_size == 1
    assert target_mgr.block_size == sharded_block
    assert draft_mgr.block_size == (sharded_block if draft_sharded else block_size)
    assert mamba_mgr.block_size == block_size

    hash_fn = sha256
    common_token_ids = [i for i in range(2) for _ in range(sharded_block)]
    req0 = make_request("0", common_token_ids + [99] * 7, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert num_computed_tokens == 0
    # Align-mode Mamba keeps a state only where a chunk ends, so end the first
    # chunk on the shared prefix.
    for num_new_tokens in (len(common_token_ids), 7):
        blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks)
        assert blocks is not None
        req0.num_computed_tokens += num_new_tokens
        computed_blocks = None

    req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert num_computed_tokens == 2 * sharded_block
    assert [len(group) for group in computed_blocks.blocks] == [
        2,
        2 if draft_sharded else 16,
        16,
    ]

    manager.free(req0)
    manager.free(req1)


@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
def test_prefill(hash_fn):
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # Complete 3 blocks (48 tokens)
    common_token_ids = [i for i in range(3) for _ in range(16)]

    # Fully cache miss
    # Incomplete 1 block (7 tokens)
    unique_token_ids = [3] * 7
    all_token_ids = common_token_ids + unique_token_ids
    req0 = make_request("0", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert len(req0.block_hashes) == 3
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, 55, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4],)

    # Check full block metadata
    parent_block_hash = None
    for block_id in (1, 2, 3):
        block_tokens = tuple(all_token_ids[(block_id - 1) * 16 : block_id * 16])
        block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_tokens)
        blk_hash = manager.block_pool.blocks[block_id].block_hash
        assert blk_hash is not None
        assert get_block_hash(blk_hash) == block_hash
        assert get_group_id(blk_hash) == 0
        assert manager.block_pool.blocks[block_id].ref_cnt == 1
        parent_block_hash = block_hash

    # Check partial block metadata
    for block_id in (4,):
        assert manager.block_pool.blocks[block_id].block_hash is None
        assert manager.block_pool.blocks[block_id].ref_cnt == 1

    # Cache hit in the common prefix when the original block is still in use.
    # Incomplete 1 block (5 tokens)
    unique_token_ids = [3] * 5
    req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert len(req1.block_hashes) == 3
    assert computed_blocks.get_block_ids() == ([1, 2, 3],)
    assert num_computed_tokens == 3 * 16
    num_new_tokens = 53 - 3 * 16
    blocks = manager.allocate_slots(
        req1, num_new_tokens, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([5],)
    for block in computed_blocks.blocks[0]:
        assert block.ref_cnt == 2

    # At this point, we should have 5 free blocks left.
    free_block_queue = manager.block_pool.free_block_queue
    assert free_block_queue.num_free_blocks == 5

    manager.free(req0)
    manager.free(req1)

    # All blocks should be available.
    assert free_block_queue.num_free_blocks == 10
    # The order should be
    # [partial without hashes from req1 and req0 (5, 4) - prepended for immediate reuse]
    # [unallocated (6, 7, 8, 9, 10)]
    # [common (3, 2, 1)]
    assert [
        b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()
    ] == [5, 4, 6, 7, 8, 9, 10, 3, 2, 1]

    # Cache hit in the common prefix when the original block is already free.
    # Incomplete 1 block (6 tokens)
    unique_token_ids = [3] * 6
    req2 = make_request("2", common_token_ids + unique_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert len(req2.block_hashes) == 3
    assert computed_blocks.get_block_ids() == ([1, 2, 3],)
    assert num_computed_tokens == 3 * 16
    num_new_tokens = 53 - 3 * 16
    blocks = manager.allocate_slots(
        req2, num_new_tokens, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([5],)  # reuse partial [5]

    # Although we only have 6 free blocks, we have 8 blocks in
    # the free block queue due to lazy removal.
    assert free_block_queue.num_free_blocks == 6
    assert all([b.ref_cnt == 0 for b in free_block_queue.get_all_free_blocks()])
    assert len([b for b in free_block_queue.get_all_free_blocks()]) == 6

    manager.free(req2)

    # Cache miss and eviction.
    req3 = make_request("3", [99] * (16 * 10), block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req3)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req3, 16 * 10, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    # This block ID order also checks the eviction order.
    assert blocks is not None and blocks.get_block_ids() == (
        [5, 4, 6, 7, 8, 9, 10, 3, 2, 1],
    )

    assert free_block_queue.num_free_blocks == 0
    assert (
        free_block_queue.fake_free_list_head.next_free_block
        is free_block_queue.fake_free_list_tail
    )
    assert (
        free_block_queue.fake_free_list_tail.prev_free_block
        is free_block_queue.fake_free_list_head
    )


def test_prefill_hybrid_model():
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config_hybrid_model(block_size, 21, 2),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    hash_fn = sha256

    # Complete 3 blocks (48 tokens)
    num_full_blocks = 3
    common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]

    # Fully cache miss
    # Incomplete 1 block (7 tokens)
    unique_token_ids = [3] * 7
    all_token_ids = common_token_ids + unique_token_ids
    req0 = make_request("0", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert len(req0.block_hashes) == 3
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, 55, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == (
        [1, 2, 3, 4],
        [5, 6, 7, 8],
        [9, 10, 11, 12],
    )

    # Check full block metadata
    parent_block_hash = None
    for length, block_ids in zip((1, 2, 3), ((1, 5, 9), (2, 6, 10), (3, 7, 11))):
        block_tokens = tuple(all_token_ids[(length - 1) * 16 : length * 16])
        block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_tokens)
        for group_id, block_id in enumerate(block_ids):
            blk_hash = manager.block_pool.blocks[block_id].block_hash
            assert blk_hash is not None
            assert get_block_hash(blk_hash) == block_hash
            assert get_group_id(blk_hash) == group_id
            assert manager.block_pool.blocks[block_id].ref_cnt == 1
        parent_block_hash = block_hash

    # Check partial block metadata
    for block_id in (4, 8, 12):
        assert manager.block_pool.blocks[block_id].block_hash is None
        assert manager.block_pool.blocks[block_id].ref_cnt == 1

    # Cache hit in the common prefix
    # Incomplete 1 block (5 tokens)
    unique_token_ids = [3] * 5
    all_token_ids = common_token_ids + unique_token_ids
    req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert len(req1.block_hashes) == 3
    assert computed_blocks.get_block_ids() == ([1, 2, 3], [0, 6, 7], [0, 10, 11])
    assert num_computed_tokens == 3 * 16
    num_new_tokens = 53 - 3 * 16
    blocks = manager.allocate_slots(
        req1, num_new_tokens, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([13], [14], [15])
    for block_per_group in computed_blocks.blocks:
        for block in block_per_group:
            if block != manager.block_pool.null_block:
                assert block.ref_cnt == 2

    block_hashes = req1.block_hashes
    manager.free(req0)
    manager.free(req1)

    # Evict the blocks outside sliding window, does not affect the hit length.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "2",
        all_token_ids,
        [
            make_block_hash_with_group_id(block_hashes[0], 1),
            make_block_hash_with_group_id(block_hashes[0], 2),
        ],
        3,
    )

    # Evict the first block of full attention, makes total cache miss.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "3",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[0], 0)],
        0,
    )

    # Evict the last block of all layers, reduces the hit length to 2.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "4",
        all_token_ids,
        [
            make_block_hash_with_group_id(block_hashes[2], 0),
            make_block_hash_with_group_id(block_hashes[2], 1),
            make_block_hash_with_group_id(block_hashes[2], 2),
        ],
        2,
    )

    # Evict the last block of full attention, reduces the hit length to 2.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "5",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[2], 0)],
        2,
    )

    # Evict the last block of sliding window, reduces the hit length to 2.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "6",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[2], 1)],
        2,
    )

    # Evict the last block of sliding window, reduces the hit length to 2.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "7",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[2], 2)],
        2,
    )

    # Evict different set of blocks for full attention and sliding window makes
    # total cache miss.
    # The cache hit length of full attention is 1 * block_size.
    # The cache hit length of sliding window is 2 * block_size.
    # Then it is cache miss as the two type of layers
    # have different hit length.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "8",
        all_token_ids,
        [
            make_block_hash_with_group_id(block_hashes[2], 0),
            make_block_hash_with_group_id(block_hashes[0], 1),
            make_block_hash_with_group_id(block_hashes[0], 2),
        ],
        0,
    )


def test_prefill_hybrid_model_eagle():
    block_size = 16
    kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 3)
    manager = make_kv_cache_manager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        use_eagle=True,
    )

    hash_fn = sha256

    # Complete 6 blocks (96 tokens)
    num_full_blocks = 6
    common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]

    # Fully cache miss
    # Incomplete 1 block (7 tokens)
    unique_token_ids = [6] * 7
    all_token_ids = common_token_ids + unique_token_ids
    req0 = make_request("0", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert len(req0.block_hashes) == len(all_token_ids) // block_size
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, len(all_token_ids), num_computed_tokens, computed_blocks
    )
    block_ids = (
        [1, 2, 3, 4, 5, 6, 7],
        [8, 9, 10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19, 20, 21],
    )
    assert blocks is not None and blocks.get_block_ids() == block_ids

    # Check full block metadata
    parent_block_hash = None
    for i, full_block_ids in enumerate(zip(*(row[:-1] for row in block_ids))):
        block_tokens = tuple(all_token_ids[i * block_size : (i + 1) * block_size])
        block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_tokens)
        for group_id, block_id in enumerate(full_block_ids):
            blk_hash = manager.block_pool.blocks[block_id].block_hash
            assert blk_hash is not None
            assert get_block_hash(blk_hash) == block_hash
            assert get_group_id(blk_hash) == group_id
            assert manager.block_pool.blocks[block_id].ref_cnt == 1
        parent_block_hash = block_hash

    # Check partial block metadata
    for partial_block_id in (row[-1] for row in block_ids):
        assert manager.block_pool.blocks[partial_block_id].block_hash is None
        assert manager.block_pool.blocks[partial_block_id].ref_cnt == 1

    # Cache hit in the common prefix
    # Incomplete 1 block (5 tokens)
    unique_token_ids = [6] * 5
    all_token_ids = common_token_ids + unique_token_ids
    req1 = make_request("1", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert len(req1.block_hashes) == num_full_blocks
    assert computed_blocks.get_block_ids() == (
        [1, 2, 3, 4, 5],
        [0, 0, 10, 11, 12],
        [0, 0, 17, 18, 19],
    )
    assert num_computed_tokens == 5 * block_size
    num_new_tokens = len(all_token_ids) - num_computed_tokens
    blocks = manager.allocate_slots(
        req1, num_new_tokens, num_computed_tokens, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == (
        [22, 23],
        [24, 25],
        [26, 27],
    )
    for block_per_group in computed_blocks.blocks:
        for block in block_per_group:
            if block != manager.block_pool.null_block:
                assert block.ref_cnt == 2

    block_hashes = req1.block_hashes
    manager.free(req0)
    manager.free(req1)

    # Evict the blocks outside sliding window, does not affect the hit length.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "2",
        all_token_ids,
        [
            make_block_hash_with_group_id(block_hashes[0], 1),
            make_block_hash_with_group_id(block_hashes[0], 2),
        ],
        5,
    )

    # Evict the first block of full attention, makes total cache miss.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "3",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[0], 0)],
        0,
    )

    # Evict the last block of all layers, reduces the hit length to 4.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "4",
        all_token_ids,
        [
            make_block_hash_with_group_id(block_hashes[-1], 0),
            make_block_hash_with_group_id(block_hashes[-1], 1),
            make_block_hash_with_group_id(block_hashes[-1], 2),
        ],
        4,
    )

    # Evict the last block of full attention, reduces the hit length to 4.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "5",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[-1], 0)],
        4,
    )

    # Since the last block of full attention is dropped for eagle, evict
    # the second last block of sliding window, reduces the hit length to 3.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "6",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[-2], 1)],
        3,
    )

    # Since the last block of full attention is dropped for eagle, evict
    # the second last block of sliding window, reduces the hit length to 3.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "7",
        all_token_ids,
        [make_block_hash_with_group_id(block_hashes[-2], 2)],
        3,
    )

    # Evict different set of blocks for full attention and sliding window.
    # Full loses its last block so it drops to 4 full blocks after the eagle
    # pop; SWA lost block 0 (outside the sliding window of the final hit),
    # which is not required for the K+1 anchor at position 4. Coordinated
    # single-drop aligns both groups at hit=4.
    _test_partial_request_hit(
        manager,
        block_size,
        num_full_blocks,
        "8",
        all_token_ids,
        [
            make_block_hash_with_group_id(block_hashes[-1], 0),
            make_block_hash_with_group_id(block_hashes[0], 1),
            make_block_hash_with_group_id(block_hashes[0], 2),
        ],
        4,
    )


def _test_partial_request_hit(
    manager: KVCacheManager,
    block_size: int,
    num_full_blocks,
    request_id: str,
    prompt_token_ids: list[int],
    hash_to_evict: list[BlockHashWithGroupId],
    expect_hit_length: int,
):
    cached_block_hash_to_block_bak = copy.copy(
        manager.block_pool.cached_block_hash_to_block._cache
    )
    req = make_request(request_id, prompt_token_ids, block_size, sha256)
    for hash_with_group_id in hash_to_evict:
        manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)
    assert len(req.block_hashes) == num_full_blocks
    assert num_computed_tokens == expect_hit_length * block_size
    for block_per_group in computed_blocks.blocks:
        assert len(block_per_group) == num_computed_tokens // block_size
    for hash_with_group_id in hash_to_evict:
        manager.block_pool.cached_block_hash_to_block._cache[hash_with_group_id] = (
            cached_block_hash_to_block_bak[hash_with_group_id]
        )
    manager.free(req)


def _make_hybrid_kv_cache_config(
    block_size: int, num_blocks: int, spec_types: list[str]
) -> KVCacheConfig:
    """Create a KVCacheConfig with the specified spec types.

    Args:
        block_size: The block size for KV cache.
        num_blocks: The number of blocks in the KV cache.
        spec_types: List of spec type strings. Supported types:
            - "full": FullAttentionSpec
            - "sliding_window": SlidingWindowSpec with window=2*block_size
            - "sliding_window_large": SlidingWindowSpec with window=4*block_size
            - "mamba": MambaSpec in prefix-caching ("align") mode

    """
    spec_map = {
        "full": lambda: FullAttentionSpec(
            block_size=block_size,
            num_kv_heads=1,
            head_size=1,
            dtype=torch.float32,
        ),
        "sliding_window": lambda: SlidingWindowSpec(
            block_size=block_size,
            num_kv_heads=1,
            head_size=1,
            dtype=torch.float32,
            sliding_window=2 * block_size,
        ),
        "sliding_window_large": lambda: SlidingWindowSpec(
            block_size=block_size,
            num_kv_heads=1,
            head_size=1,
            dtype=torch.float32,
            sliding_window=4 * block_size,
        ),
        "mamba": lambda: MambaSpec(
            block_size=block_size,
            shapes=((1, 1),),
            dtypes=(torch.float32,),
            mamba_cache_mode="align",
        ),
    }

    kv_cache_groups = [
        KVCacheGroupSpec([f"layer{i}"], spec_map[spec_type]())
        for i, spec_type in enumerate(spec_types)
    ]

    return KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=kv_cache_groups,
    )


# Test cases covering various combinations of KV cache spec types:
# - Varying number of groups (2, 3, or 4)
# - 0, 1, or 2 full attention groups
# - Sliding window with different window sizes
# - Interleaved group IDs (full attn and other types mixed)
# - Mamba spec combinations
_HYBRID_MODEL_TEST_CASES = [
    # 2 groups: 1 full + 1 other
    pytest.param(["full", "sliding_window"], id="2g-full+sw"),
    pytest.param(["full", "mamba"], id="2g-full+mamba"),
    # 2 groups: 0 full (all other types)
    pytest.param(["sliding_window", "mamba"], id="2g-sw+mamba"),
    pytest.param(["sliding_window", "sliding_window_large"], id="2g-sw+sw_large"),
    # 3 groups: 1 full + 2 others (same type)
    pytest.param(["full", "sliding_window", "sliding_window"], id="3g-full+2sw"),
    pytest.param(["full", "mamba", "mamba"], id="3g-full+2mamba"),
    # 3 groups: 1 full + 2 others (different types)
    pytest.param(["full", "sliding_window", "mamba"], id="3g-full+sw+mamba"),
    pytest.param(
        ["full", "sliding_window", "sliding_window_large"],
        id="3g-full+sw+sw_large",
    ),
    # 3 groups: 2 full + 1 other
    pytest.param(["full", "full", "sliding_window"], id="3g-2full+sw"),
    pytest.param(["full", "full", "mamba"], id="3g-2full+mamba"),
    # 4 groups: interleaved (full, other, full, other)
    pytest.param(
        ["full", "sliding_window", "full", "sliding_window_large"],
        id="4g-interleaved-full+sw+sw_large",
    ),
    pytest.param(
        ["full", "mamba", "full", "mamba"],
        id="4g-interleaved-full+mamba",
    ),
    # 4 groups: interleaved with different sliding windows
    pytest.param(
        ["full", "sliding_window", "full", "sliding_window_large"],
        id="4g-interleaved-full+sw_mixed",
    ),
    # 4 groups: 0 full (all other types)
    pytest.param(
        ["sliding_window", "mamba", "sliding_window_large", "mamba"],
        id="4g-sw+mamba+sw_large+mamba",
    ),
    # 4 groups: 2 full + 2 others (grouped)
    pytest.param(
        ["full", "full", "sliding_window", "mamba"],
        id="4g-2full+sw+mamba",
    ),
]


@pytest.mark.parametrize("spec_types", _HYBRID_MODEL_TEST_CASES)
def test_prefill_hybrid_model_combinations(spec_types: list[str]):
    """Test prefix caching with hybrid models containing various combinations of
    KV cache spec types.

    This unified test covers:
    - Various combinations (full attn + other attn types)
    - Varying number of groups (2, 3, or 4)
    - 0, 1, or 2 full attention groups in the combination
    - Two sliding_window attn groups with different window sizes
    - Interleaved group IDs (full attn and other types alternating)
    - Mamba spec with other attention types
    """
    block_size = 16
    num_groups = len(spec_types)
    # Allocate enough blocks for all groups
    num_blocks = 10 * num_groups

    kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types)
    manager = make_kv_cache_manager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    hash_fn = sha256

    # Complete 3 blocks (48 tokens)
    common_token_ids = [i for i in range(3) for _ in range(block_size)]
    unique_token_ids = [3] * 7
    all_token_ids = common_token_ids + unique_token_ids

    # First request: no cache hit initially
    req0 = make_request("0", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

    assert len(req0.block_hashes) == 3
    assert not computed_blocks.blocks[0]  # No cache hit initially
    assert num_computed_tokens == 0

    # Align-mode Mamba keeps a state only where a chunk ends, so end the first
    # chunk on the common prefix.
    for num_new_tokens in (len(common_token_ids), len(unique_token_ids)):
        blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks)
        assert blocks is not None
        # Should have blocks for all groups
        assert len(blocks.get_block_ids()) == num_groups
        req0.num_computed_tokens += num_new_tokens
        computed_blocks = None
        manager.new_step_starts()

    # Second request: should hit cached blocks for common prefix
    req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

    # Should hit cached blocks for all groups
    assert num_computed_tokens == 3 * block_size
    assert len(computed_blocks.blocks) == num_groups

    # Allocate and verify blocks for second request
    blocks = manager.allocate_slots(
        req1,
        len(common_token_ids) + 5 - num_computed_tokens,
        num_computed_tokens,
        computed_blocks,
    )
    assert blocks is not None
    assert len(blocks.get_block_ids()) == num_groups

    manager.free(req0)
    manager.free(req1)


# Test cases with eagle enabled: Only test a single simple case for now.
# - 2 groups: 1 full + 1 other
_EAGLE_HYBRID_MODEL_TEST_CASES = [
    # 2 groups: 1 full + 1 other
    pytest.param(["full", "sliding_window"], 3, id="2g-full+sw"),
]


@pytest.mark.parametrize("spec_types,expect_hit_length", _EAGLE_HYBRID_MODEL_TEST_CASES)
def test_prefill_hybrid_model_combinations_eagle(
    spec_types: list[str], expect_hit_length: int
):
    """Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE.
    More complex hybrid models with EAGLE are not yet supported (see issue #32802).
    """
    block_size = 16
    num_groups = len(spec_types)
    # Allocate enough blocks for all groups
    num_blocks = 10 * num_groups

    kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types)
    manager = make_kv_cache_manager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        use_eagle=True,
    )

    hash_fn = sha256

    # Complete 3 blocks (48 tokens)
    num_full_blocks = 4
    common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]
    unique_token_ids = [4] * 7
    all_token_ids = common_token_ids + unique_token_ids

    # First request: no cache hit initially
    req0 = make_request("0", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

    assert len(req0.block_hashes) == num_full_blocks
    assert not computed_blocks.blocks[0]  # No cache hit initially
    assert num_computed_tokens == 0

    blocks = manager.allocate_slots(
        req0, len(all_token_ids), num_computed_tokens, computed_blocks
    )
    assert blocks is not None
    # Should have blocks for all groups
    assert len(blocks.get_block_ids()) == num_groups

    # Second request: should hit cached blocks for common prefix
    all_token_ids = common_token_ids + [6] * 5
    req1 = make_request("1", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)

    # Should hit cached blocks for all groups
    assert num_computed_tokens == expect_hit_length * block_size
    assert len(computed_blocks.blocks) == num_groups
    # Verify each group has the correct number of computed blocks
    for block_per_group in computed_blocks.blocks:
        assert len(block_per_group) == expect_hit_length

    # Allocate and verify blocks for second request
    blocks = manager.allocate_slots(
        req1,
        len(all_token_ids) - num_computed_tokens,
        num_computed_tokens,
        computed_blocks,
    )
    assert blocks is not None
    assert len(blocks.get_block_ids()) == num_groups

    manager.free(req0)
    manager.free(req1)


def test_prefill_hybrid_model_mamba_align():
    """Test that MambaManager.cache_blocks() handles null blocks in align mode.

    Regression test for https://github.com/vllm-project/vllm/issues/34361.
    In mamba_cache_mode="align", allocate_new_blocks() pads req_to_blocks with
    null blocks. cache_full_blocks() correctly skips them, but
    MambaManager.cache_blocks() must also skip null blocks when tracking
    cached_blocks_this_step.
    """
    block_size = 16
    num_blocks = 30

    kv_cache_config = _make_hybrid_kv_cache_config(
        block_size, num_blocks, ["full", "mamba"]
    )
    manager = make_kv_cache_manager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    hash_fn = sha256

    # 3 full blocks (48 tokens) + 7 partial tokens = 55 tokens total
    all_token_ids = [i for i in range(3) for _ in range(block_size)] + [3] * 7

    # First request: allocate_slots should not crash with the assertion error
    # in MambaManager.cache_blocks() when null blocks are present.
    req0 = make_request("0", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert num_computed_tokens == 0

    blocks = manager.allocate_slots(req0, 55, num_computed_tokens, computed_blocks)
    assert blocks is not None
    assert len(blocks.get_block_ids()) == 2  # full_attn + mamba groups

    manager.free(req0)


def test_hybrid_cache_mamba_align_shared_prefix_detection():
    """Test shared prefix detection heuristic for mamba align cache mode

    HybridKVCacheCoordinator returns num_uncached_common > 0 when a shared
    uncached prefix is detected. With mamba_align cache, _mamba_block_aligned_split
    enforces scheduling aligned with the common prefix.
    """
    block_size = 16
    manager = make_kv_cache_manager(
        _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba"]),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )
    hash_fn = sha256

    # Request: 3 blocks
    prefix = [i for i in range(3) for _ in range(block_size)]
    req_0 = make_request("0", prefix, block_size, hash_fn)
    computed_blocks, num_computed, shared_prefix_boundary = manager.get_computed_blocks(
        req_0
    )
    assert num_computed == 0  # nothing cached yet
    assert shared_prefix_boundary == 0
    manager.allocate_slots(req_0, 3 * block_size, 0, computed_blocks)

    # Request: 3 blocks (shared with above) + 7 different tokens
    req_1 = make_request("1", prefix + [100] * 7, block_size, hash_fn)
    computed_blocks, num_computed, shared_prefix_boundary = manager.get_computed_blocks(
        req_1
    )
    assert num_computed == 3 * block_size  # we should observe a 3-block cache hit
    assert shared_prefix_boundary == 0
    manager.allocate_slots(req_1, 7, 3 * block_size, computed_blocks)

    # Request: 3 blocks, but only 2 blocks shared (replace the last token in 3rd block):
    req_2 = make_request("2", prefix[:-1] + [101], block_size, hash_fn)
    computed_blocks, num_computed, shared_prefix_boundary = manager.get_computed_blocks(
        req_2
    )
    assert num_computed == 0  # mamba_align doesn't cache intermediate blocks
    # boundary == junction: 2 shared blocks (num_computed 0 + 2 uncached blocks).
    assert shared_prefix_boundary == 2 * block_size

    # Next, validate the align-mode Marconi split. It takes the shared-prefix
    # junction (absolute) and stops the chunk there (req_2 has num_computed 0).
    # Create minimal mock with just the needed attributes
    mock = SimpleNamespace(
        block_size=block_size,
        cache_config=SimpleNamespace(block_size=block_size),
        max_num_scheduled_tokens=3 * block_size,
        scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
        use_eagle_block_drop=False,
        num_hash_lookahead_tokens=0,
        hash_block_size=block_size,
        mamba_partial_cache_hit=False,
        mamba_shared_prefix_checkpoint=False,
        mamba_has_prefill_checkpoint_blocks=False,
    )
    req_2.shared_prefix_boundary = shared_prefix_boundary
    num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split(
        self=mock,
        request=req_2,
        num_new_tokens=3 * block_size,
    )
    assert num_new_tokens_adjusted == 2 * block_size  # adjust to the common prefix

    manager.allocate_slots(req_2, 3 * block_size, 0, computed_blocks)
    # Cleanup
    manager.free(req_0)
    manager.free(req_1)
    manager.free(req_2)


def test_hybrid_model_mamba_align_with_dynamic_draft_tokens():
    """Regression test for https://github.com/vllm-project/vllm/issues/39271.

    With suffix decoding enabled, the number of proposed draft token may
    change dynamically each round, causing the MambaManager to crash during
    allocate_slots() as it originally assumes the `num_blocks` to increase.
    """
    block_size = 16
    num_blocks = 30

    kv_cache_config = _make_hybrid_kv_cache_config(
        block_size, num_blocks, ["full", "mamba"]
    )
    manager = KVCacheManager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        scheduler_block_size=block_size,
    )

    # the default hash function is sha256
    hash_fn = sha256

    all_token_ids = [i for i in range(3) for _ in range(block_size)] + [3] * 7
    req0 = make_request("0", all_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, len(all_token_ids), num_computed_tokens, computed_blocks
    )
    assert blocks is not None

    # prefill forward finished
    req0.append_output_token_ids([1])
    req0.num_computed_tokens = len(all_token_ids)

    # Round1: propose 16 draft tokens, accept only one
    req0.spec_token_ids = [4] * 16
    blocks = manager.allocate_slots(req0, num_new_tokens=16, num_new_computed_tokens=0)
    assert blocks is not None
    req0.append_output_token_ids([4])
    req0.num_computed_tokens += 1

    # Round2: propose only one token, allocate should not crash
    req0.spec_token_ids = [5] * 1
    blocks = manager.allocate_slots(req0, num_new_tokens=1, num_new_computed_tokens=0)
    assert blocks is not None and all(len(group) == 0 for group in blocks.blocks)

    manager.free(req0)


def test_prefill_plp():
    """Test prefill with APC and some prompt logprobs (plp) requests.

    1. Schedule plp request and validate APC block allocation
    2. Schedule non-plp request and validate blocks
    3. Schedule plp request; no hit should occur; validate blocks
    """
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )
    # the default hash function is sha256
    hash_fn = sha256

    # Complete 3 blocks (48 tokens)
    common_token_ids = [i for i in range(3) for _ in range(16)]

    # Request #0 is a prompt logprobs request
    # Fully cache miss
    # Incomplete 1 block (7 tokens)
    unique_token_ids = [3] * 7
    all_token_ids = common_token_ids + unique_token_ids
    req0 = make_request("0", all_token_ids, block_size, hash_fn, prompt_logprobs=5)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert len(req0.block_hashes) == 3
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, 55, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4],)
    req0_block_hashes = [b.block_hash for b in blocks.blocks[0]]

    # Check full block metadata
    parent_block_hash = None
    for block_id in (1, 2, 3):
        block_tokens = tuple(all_token_ids[(block_id - 1) * 16 : block_id * 16])
        block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_tokens)
        blk_hash = manager.block_pool.blocks[block_id].block_hash
        assert blk_hash is not None
        assert get_block_hash(blk_hash) == block_hash
        assert get_group_id(blk_hash) == 0
        assert manager.block_pool.blocks[block_id].ref_cnt == 1
        parent_block_hash = block_hash

    # Check partial block metadata
    for block_id in (4,):
        assert manager.block_pool.blocks[block_id].block_hash is None
        assert manager.block_pool.blocks[block_id].ref_cnt == 1

    # Request #1 is a non-prompt-logprobs request:
    # Cache hit in the common prefix when the original block is still in use.
    # Incomplete 1 block (5 tokens)
    unique_token_ids = [3] * 5
    req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert len(req1.block_hashes) == 3
    assert computed_blocks.get_block_ids() == ([1, 2, 3],)
    assert num_computed_tokens == 3 * 16
    num_new_tokens = 53 - 3 * 16
    blocks = manager.allocate_slots(
        req1, num_new_tokens, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([5],)
    for block in computed_blocks.blocks[0]:
        assert block.ref_cnt == 2

    # At this point, we should have 5 free blocks left.
    assert manager.block_pool.free_block_queue.num_free_blocks == 5

    manager.free(req0)
    manager.free(req1)

    # All blocks should be available.
    assert manager.block_pool.free_block_queue.num_free_blocks == 10
    # The order should be
    # [partial without hashes from req1 and req0 (5, 4) - prepended for immediate reuse]
    # [unallocated (6, 7, 8, 9, 10)]
    # [common (3, 2, 1)]
    assert [
        b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()
    ] == [5, 4, 6, 7, 8, 9, 10, 3, 2, 1]

    # Request #2 is a prompt-logprobs request:
    # NO cache hit in the common prefix; duplicates request #0 cached blocks
    unique_token_ids = [3] * 6
    req2 = make_request(
        "2", common_token_ids + unique_token_ids, block_size, hash_fn, prompt_logprobs=5
    )
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert len(req2.block_hashes) == 3
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req2, 55, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None
    block_ids = blocks.get_block_ids()
    # Duplicate cached blocks have different ids but same hashes vs request #0
    assert [b.block_hash for b in blocks.blocks[0]] == req0_block_hashes
    assert block_ids != ([1, 2, 3, 4],)

    # Request #2 block hashes are valid since request #0 hashes are.
    # Check block reference counts.
    for block_id in block_ids[0]:
        assert manager.block_pool.blocks[block_id].ref_cnt == 1

    manager.free(req2)


def test_decode():
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # Complete 3 blocks (48 tokens)
    common_token_ids = [i for i in range(3) for _ in range(16)]

    # Fully cache miss
    # Incomplete 1 block (7 tokens)
    unique_token_ids = [3] * 7
    req0 = make_request("0", common_token_ids + unique_token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, 55, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4],)

    # Append slots without allocating a new block.
    req0.num_computed_tokens = 55
    for _ in range(4):
        req0.append_output_token_ids(8)
    new_blocks = manager.allocate_slots(
        req0, 4, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
    assert (
        manager.coordinator.single_type_managers[0]
        .req_to_blocks[req0.request_id][-1]
        .block_hash
        is None
    )

    # Append slots with allocating a new block.
    req0.num_computed_tokens = 59
    # 9 tokens to fill the previous block, and 10 tokens to fill
    # the preallocated block.
    for _ in range(9 + 10):
        req0.append_output_token_ids(7)
    new_blocks = manager.allocate_slots(
        req0, 19, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert new_blocks is not None and len(new_blocks.blocks[0]) == 1
    assert (
        manager.coordinator.single_type_managers[0]
        .req_to_blocks[req0.request_id][-2]
        .block_hash
        is not None
    )
    assert (
        manager.coordinator.single_type_managers[0]
        .req_to_blocks[req0.request_id][-1]
        .block_hash
        is None
    )


def test_evict():
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    last_token_id = 5 * 16 + 7
    req0 = make_request("0", list(range(last_token_id)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, 5 * 16 + 7, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    # 5 full + 1 partial
    assert blocks is not None and len(blocks.blocks[0]) == 6

    # 3 blocks.
    req1 = make_request(
        "1", list(range(last_token_id, last_token_id + 3 * 16)), block_size, sha256
    )
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req1, 3 * 16, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and len(blocks.blocks[0]) == 3  # 3 full blocks
    last_token_id += 3 * 16

    # 10 - (6 + 3) == 1
    assert manager.block_pool.free_block_queue.num_free_blocks == 1

    manager.free(req0)
    # partial blocks (without hash) at head, other at tail (LRU policy):
    assert [
        b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()
    ] == [6, 10, 5, 4, 3, 2, 1]
    manager.free(req1)
    assert manager.block_pool.free_block_queue.num_free_blocks == 10
    assert [
        b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()
    ] == [6, 10, 5, 4, 3, 2, 1, 9, 8, 7]

    # Touch the first 2 blocks.
    req2 = make_request("2", list(range(2 * 16 + 3)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert computed_blocks.get_block_ids() == ([1, 2],)
    assert num_computed_tokens == 2 * 16
    blocks = manager.allocate_slots(
        req2, 3, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([6],)
    assert manager.block_pool.free_block_queue.num_free_blocks == 7


def test_hash_block_correct_reuse():
    """This tests when a previously cached block is reused as a new block,
    its hash metadata should be correctly reset.
    """
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(16, 2),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # Allocate 1 block and cache it.
    num_tokens = block_size * 1
    req = make_request("0", list(range(num_tokens)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req, num_tokens, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and len(blocks.blocks[0]) == 1

    # Deallocate the block.
    manager.free(req)

    # Allocate a new block that's not full, make sure hash info on the
    # block is cleared.
    req = make_request("1", list(range(num_tokens - 1)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req, num_tokens - 1, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and len(blocks.blocks[0]) == 1

    assert manager.block_pool.blocks[blocks.blocks[0][0].block_id].block_hash is None


def test_computed_blocks_not_evicted():
    """Test that the computed blocks are not evicted when getting new blocks
    for a request if there are any other free blocks.
    """
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 3),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # Allocate a block and cache it.
    num_tokens = block_size * 1
    req0 = make_request("0", list(range(num_tokens)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, num_tokens, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and len(blocks.blocks[0]) == 1
    assert blocks.blocks[0][0].block_id == 1

    # Allocate another block.
    req1 = make_request(
        "1", list(range(num_tokens, num_tokens * 2)), block_size, sha256
    )
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req1, num_tokens, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and len(blocks.blocks[0]) == 1
    assert blocks.blocks[0][0].block_id == 2

    # Free the blocks.
    manager.free(req0)
    manager.free(req1)

    # Now if we have a cache hit on the first block, we should evict the second
    # cached block rather than the first one.
    req2 = make_request("2", list(range(num_tokens * 2)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert len(computed_blocks.blocks[0]) == 1
    assert computed_blocks.blocks[0][0].block_id == 1
    assert num_computed_tokens == block_size

    blocks = manager.allocate_slots(
        req2,
        num_tokens * 2 - num_tokens,
        len(computed_blocks.blocks[0]) * 16,
        computed_blocks,
    )
    assert blocks is not None and len(blocks.blocks[0]) == 1
    assert blocks.blocks[0][0].block_id == 2


@pytest.mark.parametrize(
    "start,prompt_logprobs,for_connector,expected_hit",
    [(0, None, False, 0), (5, None, False, 4), (5, 0, False, 4), (5, None, True, 4)],
)
def test_fixed_token_scores_reuse_only_unscored_prefix(
    start, prompt_logprobs, for_connector, expected_hit
):
    block_size = 4
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 8),
        max_model_len=32,
        enable_caching=True,
        hash_block_size=block_size,
    )
    tokens = list(range(16))
    warm = make_request("warm", tokens, block_size, sha256)
    assert manager.allocate_slots(warm, len(tokens)) is not None
    manager.free(warm)
    scoring = make_request(
        "scoring",
        tokens,
        block_size,
        sha256,
        sampling_params=SamplingParams(
            prompt_logprob_token_ids=[[1]] * (len(tokens) - 1 - start),
            prompt_logprob_start=start,
            prompt_logprobs=prompt_logprobs,
            skip_reading_prefix_cache=False,
        ),
    )
    lookup = (
        manager.get_computed_blocks_for_connector
        if for_connector
        else manager.get_computed_blocks
    )
    assert lookup(scoring)[1] == expected_hit
    if expected_hit:
        assert lookup(scoring, max_cache_hit_length=len(tokens))[1] == expected_hit
        assert lookup(scoring, max_cache_hit_length=block_size - 1)[1] == 0


@pytest.mark.parametrize(
    "kwargs,expected_limit",
    [
        ({}, None),
        ({"skip_reading_prefix_cache": True}, 0),
        # Prompt scoring skips cache reads by default.
        ({"prompt_logprobs": 1, "prompt_logprob_start": 5}, 0),
        ({"prompt_logprobs": 1, "skip_reading_prefix_cache": False}, 0),
        (
            {
                "prompt_logprobs": 1,
                "prompt_logprob_start": 5,
                "skip_reading_prefix_cache": False,
            },
            5,
        ),
    ],
)
def test_prefix_cache_hit_limit(kwargs, expected_limit):
    request = make_request(
        "req", list(range(16)), 4, sha256, sampling_params=SamplingParams(**kwargs)
    )
    assert request.prefix_cache_hit_limit == expected_limit
    assert request.skip_reading_external_kv == (expected_limit is not None)


def test_basic_prefix_caching_disabled():
    """This tests that the prefix caching is disabled."""
    block_size = 4
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 5),
        max_model_len=8192,
        enable_caching=False,
        hash_block_size=block_size,
    )

    req1 = make_request(
        "1", list(range(10)), block_size, sha256
    )  # 2 blocks and some more

    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req1, 10, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and len(blocks.blocks[0]) == 3

    # Free the blocks.
    manager.free(req1)

    # No caching.
    req2 = make_request("2", list(range(16)), block_size, sha256)  # shared prefix
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req2, 16, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and len(blocks.blocks[0]) == 4

    # New requests should not have any blocks.
    req3 = make_request("3", list(range(4)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req3)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req3, 4, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert not blocks


@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
def test_cache_blocks(hash_fn):
    """This is a unit test that tests the correctness of the _cache_full_blocks
    function of KVCacheManager.
    """
    block_size = 4
    block_pool = BlockPool(
        num_gpu_blocks=5,
        enable_caching=True,
        hash_block_size=block_size,
    )
    # Req:
    #  Block 0: [0, 1, 2, 3]
    #  Block 1: [4, 5, 6, 7]
    #  Block 2: [8, 9, 10, 11]
    #  Block 3: [12, 13]
    req = make_request("0", list(range(14)), block_size, hash_fn)

    # Test that blocks are cached correctly for 2 full blocks from the start.
    blocks = [KVCacheBlock(block_id=i) for i in range(2)]

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

    assert len(block_pool.cached_block_hash_to_block) == 2
    assert all([block.block_hash is not None for block in blocks])

    # Test that blocks that don't start from the beginning are cached
    # correctly.
    blocks += [KVCacheBlock(block_id=2)]
    block_pool.cache_full_blocks(
        request=req,
        blocks=blocks,
        num_cached_blocks=2,
        num_full_blocks=3,
        block_size=block_size,
        kv_cache_group_id=0,
    )
    assert len(block_pool.cached_block_hash_to_block) == 3
    assert blocks[0].block_hash is not None


def test_cache_blocks_multi_group():
    """This tests that blocks are cached correctly for different kv cache groups."""
    block_size = 4
    block_pool = BlockPool(
        num_gpu_blocks=10, enable_caching=True, hash_block_size=block_size
    )

    # Req:
    #  Block 0/4: [0, 1, 2, 3]
    #  Block 1/5: [4, 5, 6, 7]
    #  Block 2/6: [8, 9, 10, 11]
    #  Block 3/7: [12, 13]
    req = make_request("0", list(range(14)), block_size, sha256)

    # Cache the blocks for group 0.
    blocks = [KVCacheBlock(block_id=i) for i in range(2)]
    block_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,
    )
    assert len(block_pool.cached_block_hash_to_block) == 2
    assert len(req.block_hashes) == 3
    assert all([block.block_hash is not None for block in blocks])

    # Cache the blocks for group 1.
    blocks = [KVCacheBlock(block_id=i) for i in range(3)]
    block_pool.cache_full_blocks(
        request=req,
        blocks=blocks,
        num_cached_blocks=0,
        num_full_blocks=3,
        block_size=block_size,
        kv_cache_group_id=1,
    )
    assert len(block_pool.cached_block_hash_to_block) == 5
    assert len(req.block_hashes) == 3
    assert all([block.block_hash is not None for block in blocks])

    # Block hash 0: hit for group 0 and 1
    # Block hash 1: hit for group 0 and 1
    # Block hash 2: hit for group 1

    assert (
        block_pool.get_cached_block(req.block_hashes[0], kv_cache_group_ids=[0])
        is not None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[1], kv_cache_group_ids=[0])
        is not None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[2], kv_cache_group_ids=[0]) is None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[0], kv_cache_group_ids=[1])
        is not None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[1], kv_cache_group_ids=[1])
        is not None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[2], kv_cache_group_ids=[1])
        is not None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[0], kv_cache_group_ids=[0, 1])
        is not None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[1], kv_cache_group_ids=[0, 1])
        is not None
    )
    assert (
        block_pool.get_cached_block(req.block_hashes[2], kv_cache_group_ids=[0, 1])
        is None
    )


def test_mm_prefix_caching():
    """This tests that the multi-modal prefix caching is correct."""
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # Common prompt tokens (T is text tokens and P is image placeholder tokens)
    # [T,...,T, P0,...,P0], [P0,...,P0,T,...,T,P1,...,P1], [P1,...,P1]
    common_token_ids = list(range(10)) + [-1] * 6
    common_token_ids += [-1] * 4 + list(range(10, 20)) + [-1] * 2
    common_token_ids += [-1] * 16

    common_mm_positions = [
        PlaceholderRange(offset=11, length=10),
        PlaceholderRange(offset=30, length=18),
    ]
    common_mm_hashes = ["aaa", "bbb"]

    # A unique image plus some text tokens.
    unique_token_ids = [-1] * 7 + [100] * 4
    all_token_ids = common_token_ids + unique_token_ids
    mm_positions = common_mm_positions + [PlaceholderRange(offset=48, length=7)]
    mm_hashes = common_mm_hashes + ["ccc"]
    req0 = make_request(
        "0",
        all_token_ids,
        block_size,
        sha256,
        mm_positions=mm_positions,
        mm_hashes=mm_hashes,
    )
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

    # Completed block should have hashes
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    block_hashes = req0.block_hashes
    assert len(block_hashes) == 3
    assert block_hashes[0] == sha256(
        (
            kv_cache_utils.NONE_HASH,
            tuple(all_token_ids[:block_size]),
            (("mm", "aaa", 11),),
        )
    )
    assert block_hashes[1] == sha256(
        (
            block_hashes[0],
            tuple(all_token_ids[block_size : block_size * 2]),
            (("mm", "aaa", -5), ("mm", "bbb", 14)),
        )
    )
    assert block_hashes[2] == sha256(
        (
            block_hashes[1],
            tuple(all_token_ids[block_size * 2 : block_size * 3]),
            (("mm", "bbb", -2),),
        )
    )

    blocks = manager.allocate_slots(
        req0, 59, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None
    assert blocks.get_block_ids() == ([1, 2, 3, 4],)
    req0.num_computed_tokens = 59

    # Append slots without allocating a new block.
    for _ in range(5):
        req0.append_output_token_ids(8)
    new_blocks = manager.allocate_slots(
        req0, 5, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
    assert len(block_hashes) == 4
    assert block_hashes[3] == sha256(
        (
            block_hashes[2],
            tuple(all_token_ids[3 * block_size :] + [8] * 5),
            (("mm", "ccc", 0),),
        )
    )

    # Cache hit.
    unique_token_ids = [-1] * 7 + [200] * 5
    all_token_ids = common_token_ids + unique_token_ids
    mm_positions = common_mm_positions + [PlaceholderRange(offset=48, length=7)]
    mm_hashes = common_mm_hashes + ["ccc"]
    req1 = make_request(
        "1",
        all_token_ids,
        block_size,
        sha256,
        mm_positions=mm_positions,
        mm_hashes=mm_hashes,
    )
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert len(computed_blocks.blocks[0]) == 3
    assert num_computed_tokens == 3 * 16


def test_cache_key_salting():
    """This tests that cache salts are applied during hashing and the cache
    is separated cache as expected.
    """
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # 3 complete blocks and an incomplete block with 11 tokens.
    common_token_ids = [i for i in range(3) for _ in range(block_size)]
    token_ids = common_token_ids + [3] * 11
    req0 = make_request("0", token_ids, block_size, sha256, cache_salt="salt1")
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)

    # Completed block should have hashes
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    block_hashes = req0.block_hashes
    assert len(block_hashes) == 3
    assert block_hashes[0] == sha256(
        (
            kv_cache_utils.NONE_HASH,
            tuple(token_ids[:block_size]),
            (("cache_salt", "salt1"),),
        )
    )
    assert block_hashes[1] == sha256(
        (block_hashes[0], tuple(token_ids[block_size : block_size * 2]), None)
    )
    assert block_hashes[2] == sha256(
        (block_hashes[1], tuple(token_ids[block_size * 2 : block_size * 3]), None)
    )

    blocks = manager.allocate_slots(
        req0, 59, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None
    assert blocks.get_block_ids() == ([1, 2, 3, 4],)
    req0.num_computed_tokens = 59

    # Append slots without allocating a new block.
    for _ in range(5):
        req0.append_output_token_ids(8)
    new_blocks = manager.allocate_slots(
        req0, 5, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
    assert len(block_hashes) == 4
    assert block_hashes[3] == sha256(
        (block_hashes[2], tuple(token_ids[3 * block_size :] + [8] * 5), None)
    )

    # Test cache hit with a new request that has the same salt.
    token_ids = common_token_ids + [4] * 11
    req1 = make_request("1", token_ids, block_size, sha256, cache_salt="salt1")
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    # Should match only a prefix of 3 blocks.
    assert len(computed_blocks.blocks[0]) == 3
    assert num_computed_tokens == 3 * block_size

    # Test cache miss with same content but different salt.
    token_ids = common_token_ids + [4] * 11
    req2 = make_request("2", token_ids, block_size, sha256, cache_salt="salt2")
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert len(computed_blocks.blocks[0]) == 0
    assert num_computed_tokens == 0
    block_hashes = req2.block_hashes
    assert len(block_hashes) == 3
    assert block_hashes[0] == sha256(
        (
            kv_cache_utils.NONE_HASH,
            tuple(token_ids[:block_size]),
            (("cache_salt", "salt2"),),
        )
    )
    assert block_hashes[1] == sha256(
        (block_hashes[0], tuple(token_ids[block_size : block_size * 2]), None)
    )
    assert block_hashes[2] == sha256(
        (block_hashes[1], tuple(token_ids[block_size * 2 : block_size * 3]), None)
    )


def test_prefill_not_enough_free_blocks_with_computed_blocks():
    """This is a unit test that tests the correctness of the allocate_slots
    when there is not enough free blocks. Specifically, when a request
    has computed blocks but cannot be allocated due to not enough free blocks,
    the computed blocks should not be touched.
    """
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )
    # Complete 3 blocks (48 tokens)
    # | Common-0 | Common-1 | Common-2 | ... |
    common_token_ids = [i for i in range(3) for _ in range(16)]
    req0 = make_request("0", common_token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    manager.allocate_slots(
        req0, 48, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    block_part0 = manager.coordinator.single_type_managers[0].req_to_blocks[
        req0.request_id
    ]

    # | Common-0 | Common-1 | Common-2 | Req1-3 | Req1-4 | Req1-5 | ... |
    req1 = make_request("1", common_token_ids * 2, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert computed_blocks.blocks[0] == block_part0
    assert num_computed_tokens == 3 * 16
    manager.allocate_slots(
        req1, 48, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    block_part1 = manager.coordinator.single_type_managers[0].req_to_blocks[
        req1.request_id
    ]
    # | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
    # | Req1-5(F)| ... |
    manager.free(req1)
    assert {block.ref_cnt for block in block_part1[:3]} == {1}
    assert {block.ref_cnt for block in block_part1[3:]} == {0}

    # | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
    # | Req1-5(F)| Req2-0   | Req2-1   | ... |
    req2 = make_request("2", [7] * block_size * 2, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    manager.allocate_slots(
        req2,
        block_size * 2,
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )

    # Req3 is Req2 + 3 new blocks, so the first 6 blocks are computed,
    # but it cannot be allocated due to insufficient free blocks (2).
    # In this case, the ref_cnt of the computed blocks should not be changed.
    assert manager.block_pool.free_block_queue.num_free_blocks == 5
    req3 = make_request("3", common_token_ids * 3, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req3)
    assert computed_blocks.blocks[0] == block_part1
    assert num_computed_tokens == 6 * 16
    # Req3 cannot be allocated.
    assert (
        manager.allocate_slots(
            req3, 48, len(computed_blocks.blocks[0]) * 16, computed_blocks
        )
        is None
    )
    # Block 0-2 are used by Req 1.
    assert {block.ref_cnt for block in block_part1[:3]} == {1}
    # Block 3-5 are free.
    assert {block.ref_cnt for block in block_part1[3:]} == {0}


def test_reset_prefix_cache():
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    full_block_token_ids = [i for i in range(3) for _ in range(16)]
    unique_token_ids = [3] * 7
    all_token_ids = full_block_token_ids + unique_token_ids
    req0 = make_request("0", all_token_ids, block_size, sha256)
    blocks = manager.allocate_slots(req0, 55)
    assert blocks is not None and blocks.get_block_ids() == ([1, 2, 3, 4],)

    unique_token_ids = [4] * 7
    all_token_ids = full_block_token_ids + unique_token_ids
    req1 = make_request("1", all_token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req1)
    assert len(req1.block_hashes) == 3
    assert len(computed_blocks.blocks[0]) == 3
    blocks = manager.allocate_slots(
        req1, 7, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks is not None and blocks.get_block_ids() == ([5],)

    # Failed to reset prefix cache because some blocks are not freed yet.
    assert not manager.reset_prefix_cache()
    assert manager.block_pool.cached_block_hash_to_block

    # Free the blocks.
    manager.free(req0)
    manager.free(req1)

    assert manager.reset_prefix_cache()
    assert not manager.block_pool.cached_block_hash_to_block
    assert all([blk.block_hash is None for blk in manager.block_pool.blocks])


def test_prefix_cache_stats_disabled():
    """Test that prefix_cache_stats is None when log_stats is False."""
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, 11),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        log_stats=False,  # Disable logging stats
    )
    assert manager.prefix_cache_stats is None

    # Call all functions that check whether log_stats is disabled.
    req = make_request("0", list(range(16)), block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req)
    assert not computed_blocks.blocks[0]
    assert num_computed_tokens == 0
    manager.allocate_slots(
        req, 16, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    manager.reset_prefix_cache()

    # Ensure prefix_cache_stats remains None
    assert manager.prefix_cache_stats is None


def test_maybe_evict_cached_block():
    pool = BlockPool(num_gpu_blocks=4, enable_caching=True, hash_block_size=16)
    block_hash0 = make_block_hash_with_group_id(BlockHash(b"10"), 1000)
    block_hash1 = make_block_hash_with_group_id(BlockHash(b"20"), 2000)
    block_hash2 = make_block_hash_with_group_id(BlockHash(b"30"), 3000)
    block_hashes = [
        block_hash0,
        block_hash1,
        block_hash2,
        # block3 had the exact same block_hash as the first block
        block_hash0,
    ]
    assert len(pool.blocks) == len(block_hashes)
    # Manually add all blocks to cached_blocks
    for block, block_hash in zip(pool.blocks, block_hashes):
        block.set_block_hash(block_hash)
        pool.cached_block_hash_to_block.insert(block_hash, block)

    block0, block1, block2, block3 = pool.blocks
    assert pool.cached_block_hash_to_block._cache == {
        block_hash0: {
            block0.block_id: block0,
            block3.block_id: block3,
        },
        block_hash1: block1,
        block_hash2: block2,
    }
    # Evict block1
    pool._maybe_evict_cached_block(block1)
    assert pool.cached_block_hash_to_block._cache == {
        block_hash0: {block0.block_id: block0, block3.block_id: block3},
        block_hash2: block2,
    }
    # Evict block0: block_hash0 entry should NOT be removed, as block3
    # also use the same hash
    pool._maybe_evict_cached_block(block0)
    assert pool.cached_block_hash_to_block._cache == {
        block_hash0: {block3.block_id: block3},
        block_hash2: block2,
    }
    # Evict block2
    pool._maybe_evict_cached_block(block2)
    assert pool.cached_block_hash_to_block._cache == {block_hash0: {3: block3}}
    # Evict block3
    pool._maybe_evict_cached_block(block3)
    assert pool.cached_block_hash_to_block._cache == {}


@pytest.mark.parametrize("blocks_to_cache", [2, 3, 10])
def test_kv_cache_events(blocks_to_cache: int):
    block_size = 16
    num_blocks = blocks_to_cache + 1

    # Allocate Blocks
    # Should see a single block stored event with a blocks_to_cache number of
    # block hashes
    # take_events should reset the kv_event_queue
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, num_blocks),
        max_model_len=8192,
        enable_caching=True,
        enable_kv_cache_events=True,
        hash_block_size=block_size,
    )

    num_tokens = block_size * blocks_to_cache
    req0 = make_request("0", list(range(num_tokens)), block_size, sha256)
    _ = manager.allocate_slots(req0, num_tokens)
    events = manager.take_events()

    block = events[-1]
    assert (
        len(block.block_hashes)
        == blocks_to_cache
        == len(manager.block_pool.cached_block_hash_to_block)
    )
    assert len(block.token_ids) == block.block_size * len(block.block_hashes)
    assert block.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
    assert len(manager.block_pool.kv_event_queue) == 0

    stored_block_hash = block.block_hashes

    # Remove blocks and send another request
    # Should see block_to_cache number of removed block events and a new block
    # stored event
    manager.free(req0)
    req1 = make_request("1", list(range(num_tokens)), block_size, sha256)
    _ = manager.allocate_slots(req1, num_tokens)
    events = manager.take_events()

    for blocks in events[:-1]:
        assert isinstance(blocks, BlockRemoved)
        assert blocks.block_hashes[0] in stored_block_hash
    assert len(events) == blocks_to_cache + 1
    assert isinstance(events[-2], BlockRemoved)
    assert (
        len(events[-1].block_hashes)
        == blocks_to_cache
        == len(manager.block_pool.cached_block_hash_to_block)
    )

    # All Blocks Cleared
    # Should see a single all blocks cleared event
    manager.free(req1)
    manager.reset_prefix_cache()
    events = manager.take_events()

    assert isinstance(events[-1], AllBlocksCleared)
    assert len(manager.block_pool.cached_block_hash_to_block) == 0


def test_null_parent_block_hash():
    block_size = 1
    num_cached_blocks = 2
    num_full_blocks = 4
    kv_cache_group_id = 0

    pool = BlockPool(
        num_gpu_blocks=8,
        enable_caching=True,
        hash_block_size=block_size,
        enable_kv_cache_events=True,
    )

    req = make_request(
        "req_null_parent",
        prompt_token_ids=[10, 11, 12, 13],
        block_size=block_size,
        hash_fn=sha256,
    )
    assert len(req.block_hashes) == num_full_blocks

    # Physical parent is `null_block` (no hash), while the logical parent hash
    # still exists in `request.block_hashes[num_cached_blocks - 1]`.
    assert pool.null_block.block_hash is None
    new_blocks = pool.get_new_blocks(num_full_blocks - 1)
    blocks = [
        new_blocks[: num_cached_blocks - 1],
        pool.null_block,  # physical parent
        *new_blocks[num_cached_blocks - 1 :],
    ]

    pool.cache_full_blocks(
        request=req,
        blocks=blocks,
        num_cached_blocks=num_cached_blocks,
        num_full_blocks=num_full_blocks,
        block_size=block_size,
        kv_cache_group_id=kv_cache_group_id,
    )

    events = pool.take_events()
    assert len(events) == 1
    event = events[0]
    assert isinstance(event, BlockStored)

    expected_parent = kv_cache_utils.maybe_convert_block_hash(
        req.block_hashes[num_cached_blocks - 1]
    )
    assert event.parent_block_hash == expected_parent
    assert event.parent_block_hash is not None

    expected_new_hashes = [
        kv_cache_utils.maybe_convert_block_hash(h)
        for h in req.block_hashes[num_cached_blocks:num_full_blocks]
    ]
    assert event.block_hashes == expected_new_hashes
    assert event.group_idx == kv_cache_group_id
    assert event.kv_cache_spec_kind is None
    assert event.kv_cache_spec_sliding_window is None

    # Ensure we didn't accidentally assign a hash to the null block.
    assert pool.null_block.block_hash is None
    # Sanity check: newly cached physical blocks should have hashes assigned.
    assert blocks[num_cached_blocks].block_hash is not None
    assert blocks[num_full_blocks - 1].block_hash is not None


@pytest.mark.parametrize("blocks_to_cache", [2, 3, 10])
def test_kv_cache_events_with_lora(blocks_to_cache: int):
    """Test BlockStored events contain correct lora_id when using LoRA requests."""
    block_size = 16
    num_blocks = blocks_to_cache + 1

    # Create KVCacheManager with events enabled
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, num_blocks),
        max_model_len=8192,
        enable_caching=True,
        enable_kv_cache_events=True,
        hash_block_size=block_size,
    )

    # Test with LoRA request
    lora_request = LoRARequest(
        lora_name="test_lora", lora_int_id=42, lora_path="/test/path"
    )

    num_tokens = block_size * blocks_to_cache
    req_with_lora = make_request(
        "lora_req",
        list(range(num_tokens)),
        block_size,
        sha256,
        lora_request=lora_request,
    )

    # Allocate slots and get events
    _ = manager.allocate_slots(req_with_lora, num_tokens)
    events = manager.take_events()

    # Verify BlockStored event contains correct lora_id
    block_stored_event = events[-1]
    assert isinstance(block_stored_event, BlockStored)
    assert block_stored_event.lora_id == 42  # Should match lora_request.adapter_id
    assert len(block_stored_event.block_hashes) == blocks_to_cache
    assert block_stored_event.block_size == block_size

    # Clean up
    manager.free(req_with_lora)

    # Test without LoRA request (should have lora_id=None)
    req_without_lora = make_request(
        "no_lora_req", list(range(num_tokens)), block_size, sha256
    )

    _ = manager.allocate_slots(req_without_lora, num_tokens)
    events = manager.take_events()

    block_stored_event = events[-1]
    assert isinstance(block_stored_event, BlockStored)
    assert block_stored_event.lora_id is None  # Should be None when no LoRA request
    assert len(block_stored_event.block_hashes) == blocks_to_cache
    assert block_stored_event.block_size == block_size


@pytest.mark.parametrize("group_id", [0, 1, 2])
def test_block_stored_event_group_idx(group_id: int):
    """Test BlockStored events emitted by cache_full_blocks carry the correct
    group_idx."""
    block_size = 4
    num_tokens = block_size * 2

    manager = make_kv_cache_manager(
        make_kv_cache_config_three_types(block_size, num_blocks=5),
        max_model_len=8192,
        enable_caching=True,
        enable_kv_cache_events=True,
        hash_block_size=block_size,
    )
    pool = manager.block_pool

    req = make_request(
        "req_grp_idx",
        prompt_token_ids=list(range(num_tokens)),
        block_size=block_size,
        hash_fn=sha256,
    )

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

    events = manager.take_events()
    assert len(events) == 1
    assert isinstance(events[0], BlockStored)
    assert events[0].group_idx == group_id
    assert (
        events[0].kv_cache_spec_kind
        == [
            KVCacheSpecKind.FULL_ATTENTION.value,
            KVCacheSpecKind.SLIDING_WINDOW.value,
            KVCacheSpecKind.MAMBA.value,
        ][group_id]
    )
    assert (
        events[0].kv_cache_spec_sliding_window
        == [
            None,
            2 * block_size,
            None,
        ][group_id]
    )


def test_block_stored_event_session_id():
    block_size = 4
    session_id = "agent-session-1"
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, num_blocks=5),
        max_model_len=8192,
        enable_caching=True,
        enable_kv_cache_events=True,
        hash_block_size=block_size,
    )
    req = make_request(
        "req_session_id",
        prompt_token_ids=list(range(block_size * 2)),
        block_size=block_size,
        hash_fn=sha256,
        session_id=session_id,
    )

    manager.allocate_slots(req, req.num_tokens)
    events = manager.take_events()

    assert len(events) == 1
    assert isinstance(events[0], BlockStored)
    assert events[0].session_id == session_id


def test_session_id_does_not_affect_prefix_cache_identity():
    block_size = 4
    prompt_token_ids = list(range(block_size * 2 + 1))
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, num_blocks=8),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )
    first_req = make_request(
        "req_session_a",
        prompt_token_ids,
        block_size,
        sha256,
        session_id="session-a",
    )
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(first_req)
    allocated_blocks = manager.allocate_slots(
        first_req,
        first_req.num_tokens,
        num_computed_tokens,
        computed_blocks,
    )
    assert allocated_blocks is not None

    second_req = make_request(
        "req_session_b",
        prompt_token_ids,
        block_size,
        sha256,
        session_id="session-b",
    )
    assert second_req.block_hashes == first_req.block_hashes

    cached_blocks, num_cached_tokens, _ = manager.get_computed_blocks(second_req)
    assert num_cached_tokens == block_size * 2
    assert cached_blocks.get_block_ids() == (allocated_blocks.get_block_ids()[0][:2],)


def test_block_stored_event_group_idx_multiple_groups():
    """Test BlockStored events for separate HMA groups that each carry the
    correct group_idx.

    Simulates the HMA scenario where full-attention blocks (group 0) and
    sliding-window blocks (group 1) are cached independently and must be
    distinguishable by consumers doing HMA-aware prefix-cache routing.
    """
    block_size = 4
    num_tokens = block_size * 2

    manager = make_kv_cache_manager(
        KVCacheConfig(
            num_blocks=5,
            kv_cache_tensors=[],
            kv_cache_groups=[
                KVCacheGroupSpec(
                    ["layer1"],
                    FullAttentionSpec(
                        block_size=block_size,
                        num_kv_heads=1,
                        head_size=1,
                        dtype=torch.float32,
                    ),
                ),
                KVCacheGroupSpec(
                    ["layer2"],
                    SlidingWindowSpec(
                        block_size=block_size,
                        num_kv_heads=1,
                        head_size=1,
                        dtype=torch.float32,
                        sliding_window=128,
                    ),
                ),
            ],
        ),
        max_model_len=8192,
        enable_caching=True,
        enable_kv_cache_events=True,
        hash_block_size=block_size,
    )
    pool = manager.block_pool

    req = make_request(
        "req_multi_grp",
        prompt_token_ids=list(range(num_tokens)),
        block_size=block_size,
        hash_fn=sha256,
    )

    # Cache blocks for group 0 (full-attention)
    blocks_grp0 = pool.get_new_blocks(2)
    pool.cache_full_blocks(
        request=req,
        blocks=blocks_grp0,
        num_cached_blocks=0,
        num_full_blocks=2,
        block_size=block_size,
        kv_cache_group_id=0,
    )

    # Cache blocks for group 1 (sliding-window)
    blocks_grp1 = pool.get_new_blocks(2)
    pool.cache_full_blocks(
        request=req,
        blocks=blocks_grp1,
        num_cached_blocks=0,
        num_full_blocks=2,
        block_size=block_size,
        kv_cache_group_id=1,
    )

    events = manager.take_events()
    assert len(events) == 2
    assert isinstance(events[0], BlockStored)
    assert events[0].group_idx == 0
    assert events[0].kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
    assert events[0].kv_cache_spec_sliding_window is None
    assert isinstance(events[1], BlockStored)
    assert events[1].group_idx == 1
    assert events[1].kv_cache_spec_kind == KVCacheSpecKind.SLIDING_WINDOW.value
    assert events[1].kv_cache_spec_sliding_window == 128


def test_block_stored_event_group_idx_out_of_bounds(monkeypatch):
    """Out-of-range group_idx events are returned without metadata annotation."""
    block_size = 4
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, num_blocks=5),
        max_model_len=8192,
        enable_caching=True,
        enable_kv_cache_events=True,
        hash_block_size=block_size,
    )
    event = BlockStored(
        block_hashes=[1],
        parent_block_hash=None,
        token_ids=list(range(block_size)),
        block_size=block_size,
        lora_id=None,
        medium=None,
        lora_name=None,
        group_idx=1,
    )
    manager.block_pool.kv_event_queue.append(event)
    warnings = []

    def collect_warning(message, *args, **kwargs):
        del kwargs
        warnings.append(message % args if args else message)

    monkeypatch.setattr(kv_cache_manager.logger, "warning", collect_warning)
    events = manager.take_events()

    assert events == [event]
    assert event.kv_cache_spec_kind is None
    assert event.kv_cache_spec_sliding_window is None
    assert warnings == ["Group index `1` not in KV cache metadata"]


@pytest.mark.parametrize("group_id", [0, 1, 2])
def test_block_removed_event_group_idx(group_id: int):
    """Test BlockRemoved events emitted on eviction carry the group_idx extracted
    from the evicted block's BlockHashWithGroupId via get_group_id().
    """
    block_size = 4
    num_tokens = block_size * 2

    # null block + 4 usable; allocate all 4, cache 2, free all, re-allocate
    # all 4 so the 2 cached blocks are forced through _maybe_evict_cached_block.
    pool = BlockPool(
        num_gpu_blocks=5,
        enable_caching=True,
        hash_block_size=block_size,
        enable_kv_cache_events=True,
    )

    req = make_request(
        "req_evict_grp",
        prompt_token_ids=list(range(num_tokens)),
        block_size=block_size,
        hash_fn=sha256,
    )

    # Allocate all usable blocks and cache the first two for the target group.
    all_blocks = pool.get_new_blocks(4)
    pool.cache_full_blocks(
        request=req,
        blocks=all_blocks,
        num_cached_blocks=0,
        num_full_blocks=2,
        block_size=block_size,
        kv_cache_group_id=group_id,
    )

    # Drain the BlockStored events so only eviction events remain later.
    pool.take_events()

    # Return all blocks to the free queue so they become eviction candidates.
    pool.free_blocks(all_blocks)

    # Re-allocate all blocks; the two with hashes trigger BlockRemoved events.
    pool.get_new_blocks(4)

    events = pool.take_events()
    removed_events = [e for e in events if isinstance(e, BlockRemoved)]

    assert len(removed_events) == 2
    for event in removed_events:
        assert event.group_idx == group_id


def test_emit_cached_block_events():
    """emit_cached_block_events emits one BlockStored for already-cached
    (reused) prefix blocks, carrying the correct group_idx /
    parent_block_hash / token_ids, and without mutating block state."""
    block_size = 4
    num_cached_blocks = 3
    kv_cache_group_id = 1
    num_tokens = block_size * 4  # 4 full blocks; reuse the first 3

    pool = BlockPool(
        num_gpu_blocks=8,
        enable_caching=True,
        hash_block_size=block_size,
        enable_kv_cache_events=True,
    )

    req = make_request(
        "req_emit_cached",
        prompt_token_ids=list(range(num_tokens)),
        block_size=block_size,
        hash_fn=sha256,
        session_id="agent-session-reuse",
    )
    assert len(req.block_hashes) >= num_cached_blocks

    # Snapshot block state to prove emit_cached_block_events does not mutate it.
    free_before = pool.get_num_free_blocks()
    assert len(pool.cached_block_hash_to_block) == 0

    pool.emit_cached_block_events(
        request=req,
        num_cached_blocks=num_cached_blocks,
        block_size=block_size,
        kv_cache_group_id=kv_cache_group_id,
    )

    # No block-state mutation: nothing allocated, nothing inserted into the
    # prefix-cache map.
    assert pool.get_num_free_blocks() == free_before
    assert len(pool.cached_block_hash_to_block) == 0

    events = pool.take_events()
    assert len(events) == 1
    event = events[0]
    assert isinstance(event, BlockStored)

    expected_hashes = [
        kv_cache_utils.maybe_convert_block_hash(req.block_hashes[i])
        for i in range(num_cached_blocks)
    ]
    assert event.block_hashes == expected_hashes
    # Reused blocks start from block 0, so there is no parent block hash.
    assert event.parent_block_hash is None
    assert event.token_ids == list(req.all_token_ids[: num_cached_blocks * block_size])
    assert event.group_idx == kv_cache_group_id
    assert event.block_size == block_size
    assert event.medium == MEDIUM_GPU
    assert event.lora_id is None
    assert event.lora_name is None
    assert event.session_id == "agent-session-reuse"


def test_emit_cached_block_events_disabled():
    """No events are emitted when enable_kv_cache_events is False."""
    block_size = 4
    pool = BlockPool(
        num_gpu_blocks=8,
        enable_caching=True,
        hash_block_size=block_size,
        enable_kv_cache_events=False,
    )
    req = make_request(
        "req_emit_disabled",
        prompt_token_ids=list(range(block_size * 4)),
        block_size=block_size,
        hash_fn=sha256,
    )

    pool.emit_cached_block_events(
        request=req,
        num_cached_blocks=3,
        block_size=block_size,
        kv_cache_group_id=0,
    )

    assert pool.take_events() == []


def test_emit_cached_block_events_zero_cached():
    """No events are emitted when num_cached_blocks == 0."""
    block_size = 4
    pool = BlockPool(
        num_gpu_blocks=8,
        enable_caching=True,
        hash_block_size=block_size,
        enable_kv_cache_events=True,
    )
    req = make_request(
        "req_emit_zero",
        prompt_token_ids=list(range(block_size * 4)),
        block_size=block_size,
        hash_fn=sha256,
    )

    pool.emit_cached_block_events(
        request=req,
        num_cached_blocks=0,
        block_size=block_size,
        kv_cache_group_id=0,
    )

    assert pool.take_events() == []


@pytest.mark.parametrize(
    ("lookahead", "num_prompt_tokens", "expected_blocks"),
    [
        # Legacy drop: pop the logits block, then drop one more.
        (False, 48, 1),
        (False, 37, 1),
        # Lookahead hashes prove the boundary, so no extra block is dropped.
        (True, 48, 2),
        (True, 37, 2),
    ],
)
def test_eagle_identical_prompt_hit(
    lookahead: bool, num_prompt_tokens: int, expected_blocks: int
):
    block_size = 16
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, num_blocks=10),
        max_model_len=8192,
        enable_caching=True,
        use_eagle=not lookahead,
        num_hash_lookahead_tokens=lookahead,
        hash_block_size=block_size,
    )
    token_ids = [0] * num_prompt_tokens

    req = make_request(
        "prime", token_ids, block_size, sha256, num_hash_lookahead_tokens=lookahead
    )
    assert manager.allocate_slots(req, len(token_ids)) is not None
    manager.free(req)

    req_eagle = make_request(
        "eagle", token_ids, block_size, sha256, num_hash_lookahead_tokens=lookahead
    )
    computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_eagle)
    assert len(computed_blocks.blocks[0]) == expected_blocks
    assert num_tokens == expected_blocks * block_size


@pytest.mark.parametrize("lookahead", [1, 2])
def test_kv_events_reconstruct_lookahead_block_hashes(lookahead: int):
    block_size = 2
    manager = make_kv_cache_manager(
        make_kv_cache_config(block_size, num_blocks=20),
        max_model_len=8192,
        enable_caching=True,
        use_eagle=False,
        num_hash_lookahead_tokens=lookahead,
        hash_block_size=block_size,
        enable_kv_cache_events=True,
    )
    token_ids = list(range(2 * block_size + lookahead))
    request = make_request(
        "request",
        token_ids,
        block_size,
        sha256,
        num_hash_lookahead_tokens=lookahead,
    )
    assert manager.allocate_slots(request, request.num_tokens) is not None

    (event,) = manager.take_events()
    assert isinstance(event, BlockStored)
    assert event.block_hashes == [
        kv_cache_utils.maybe_convert_block_hash(block_hash)
        for block_hash in request.block_hashes
    ]
    assert event.token_ids == [0, 1, 2, 3]
    assert event.lookahead_token_ids == [
        token_ids[end : end + lookahead] for end in (2, 4)
    ]
    reconstructed_hashes = []
    parent_hash = event.parent_block_hash
    assert event.extra_keys is not None
    for block_start, extra_keys, lookahead_token_ids in zip(
        range(0, len(event.token_ids), block_size),
        event.extra_keys,
        event.lookahead_token_ids,
        strict=True,
    ):
        parent_hash = kv_cache_utils.hash_block_tokens(
            sha256,
            parent_hash,
            event.token_ids[block_start : block_start + block_size],
            (extra_keys, *lookahead_token_ids),
        )
        reconstructed_hashes.append(parent_hash)
    assert reconstructed_hashes == request.block_hashes


@pytest.mark.parametrize(
    "lookahead,prompt_len,expected_boundary",
    [
        (0, 1089, 1088),
        (1, 1089, 1088),
        (2, 1089, 1024),
        (2, 1090, 1088),
        (3, 1089, 1024),
        (3, 1091, 1088),
    ],
)
@pytest.mark.parametrize("checkpoint_blocks", [0, 1])
@pytest.mark.parametrize("full_block_size", [64, 128])
def test_lookahead_checkpoint_has_available_hash(
    lookahead, prompt_len, expected_boundary, checkpoint_blocks, full_block_size
):
    """Prefill checkpoints must be reusable at the boundary their hash proves."""
    config = _make_hybrid_kv_cache_config(64, 100, ["full", "mamba", "full"])
    config.kv_cache_groups[0].kv_cache_spec = replace(
        config.kv_cache_groups[0].kv_cache_spec, block_size=full_block_size
    )
    config.kv_cache_groups[1].kv_cache_spec = replace(
        config.kv_cache_groups[1].kv_cache_spec,
        block_size=512,
        num_prefill_checkpoint_blocks=checkpoint_blocks,
        prefill_checkpoint_alignment=64 if checkpoint_blocks else None,
    )
    config.kv_cache_groups[2].is_eagle_group = bool(lookahead)
    manager = make_kv_cache_manager(
        config,
        max_model_len=4096,
        enable_caching=True,
        hash_block_size=64,
        num_prefill_lookahead=lookahead,
        num_hash_lookahead_tokens=lookahead,
    )
    tokens = list(range(prompt_len))
    request = make_request(
        "producer", tokens, 64, sha256, num_hash_lookahead_tokens=lookahead
    )
    if not checkpoint_blocks:
        assert manager.allocate_slots(request, expected_boundary) is not None
        request.num_computed_tokens = expected_boundary
        manager.new_step_starts()
    assert manager.allocate_slots(request, len(tokens) - request.num_computed_tokens)
    checkpoint_hash = request.block_hashes[expected_boundary // 64 - 1]
    blocks = manager.block_pool.get_cached_block(checkpoint_hash, [1])
    assert blocks is not None
    assert blocks[0].block_hash_num_tokens == expected_boundary
    if checkpoint_blocks:
        offsets, cols = compute_mamba_prefill_checkpoints(
            [prompt_len],
            [prompt_len],
            hash_block_size=64,
            mamba_block_size=512,
            checkpoint_alignment=64,
            drop_eagle_block=False,
            cache_hit_alignment_tokens=64,
            num_hash_lookahead_tokens=lookahead,
        )
        assert offsets == [expected_boundary]
        assert (
            blocks[0].block_id == manager.get_block_ids(request.request_id)[1][cols[0]]
        )

    request.num_computed_tokens = len(tokens)
    manager.free(request)
    manager.new_step_starts()
    replay = make_request(
        "replay", tokens, 64, sha256, num_hash_lookahead_tokens=lookahead
    )
    _, hit_tokens, _ = manager.get_computed_blocks(replay)
    assert hit_tokens == expected_boundary


def test_eagle_hybrid_mamba_hits_partial_prompt_boundary():
    hash_block_size = 16
    mamba_block_size = 544
    token_ids = list(range(1249))
    config = _make_hybrid_kv_cache_config(hash_block_size, 200, ["full", "mamba"])
    config.kv_cache_groups[1].kv_cache_spec = replace(
        config.kv_cache_groups[1].kv_cache_spec, block_size=mamba_block_size
    )
    manager = make_kv_cache_manager(
        config,
        max_model_len=8192,
        enable_caching=True,
        use_eagle=False,
        num_hash_lookahead_tokens=1,
        hash_block_size=hash_block_size,
    )

    first = make_request(
        "first",
        token_ids,
        hash_block_size,
        sha256,
        num_hash_lookahead_tokens=1,
    )
    assert manager.allocate_slots(first, 1248) is not None
    first.num_computed_tokens = 1248
    manager.new_step_starts()
    assert manager.allocate_slots(first, 1) is not None
    first.num_computed_tokens = 1249
    manager.new_step_starts()
    manager.free(first)

    second = make_request(
        "second",
        token_ids,
        hash_block_size,
        sha256,
        num_hash_lookahead_tokens=1,
    )
    _, num_computed, _ = manager.get_computed_blocks(second)

    assert num_computed == 1248
    # Per-group lookup must not add the legacy EAGLE drop margin.
    _, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
        second.block_hashes, 2 * hash_block_size
    )
    assert per_group_hits[0] == 2 * hash_block_size


@pytest.mark.parametrize("lookahead", [0, 1, 2])
@pytest.mark.parametrize("diverge_offset", [0, 1, 2])
def test_mamba_partial_hit_requires_matching_lookahead_tokens(
    lookahead: int, diverge_offset: int
):
    """A partial Mamba hit at boundary b also needs tokens [b, b + lookahead)."""
    hash_block_size = 16
    mamba_block_size = 544
    prompt_len = 1249
    config = _make_hybrid_kv_cache_config(hash_block_size, 200, ["full", "mamba"])
    config.kv_cache_groups[1].kv_cache_spec = replace(
        config.kv_cache_groups[1].kv_cache_spec, block_size=mamba_block_size
    )
    manager = make_kv_cache_manager(
        config,
        max_model_len=8192,
        enable_caching=True,
        use_eagle=False,
        num_hash_lookahead_tokens=lookahead,
        hash_block_size=hash_block_size,
    )
    boundary = (prompt_len - max(1, lookahead)) // hash_block_size * hash_block_size
    assert boundary % mamba_block_size != 0

    token_ids = list(range(prompt_len))
    producer = make_request(
        "producer",
        token_ids,
        hash_block_size,
        sha256,
        num_hash_lookahead_tokens=lookahead,
    )
    for chunk_end in (boundary, prompt_len):
        assert (
            manager.allocate_slots(producer, chunk_end - producer.num_computed_tokens)
            is not None
        )
        producer.num_computed_tokens = chunk_end
        manager.new_step_starts()
    manager.free(producer)

    shared_len = boundary + diverge_offset
    consumer_ids = token_ids[:shared_len] + [
        100_000 + i for i in range(prompt_len + hash_block_size - shared_len)
    ]
    consumer = make_request(
        "consumer",
        consumer_ids,
        hash_block_size,
        sha256,
        num_hash_lookahead_tokens=lookahead,
    )
    _, num_computed, _ = manager.get_computed_blocks(consumer)
    # No Mamba state is retained below the prompt's partial boundary.
    assert num_computed == (boundary if diverge_offset >= lookahead else 0)


def test_eagle_with_sliding_window():
    """Test Eagle behavior with sliding window."""
    block_size = 16
    sliding_window_spec = SlidingWindowSpec(
        block_size=block_size,
        num_kv_heads=1,
        head_size=1,
        dtype=torch.float32,
        sliding_window=block_size,
    )
    manager = make_kv_cache_manager(
        KVCacheConfig(
            num_blocks=10,
            kv_cache_tensors=[],
            kv_cache_groups=[KVCacheGroupSpec(["layer"], sliding_window_spec)],
        ),
        max_model_len=8192,
        enable_caching=True,
        use_eagle=True,
        hash_block_size=block_size,
    )

    # 2 full blocks + 5 tokens (non-divisible length)
    token_ids = [0] * (2 * block_size + 5)
    req = make_request("partial_block_test", token_ids, block_size, sha256)

    # Prime the cache
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    manager.allocate_slots(
        req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    # record the block hash of the first block in the request for later use
    block_hash_first_block = req.block_hashes[0]
    assert block_hash_first_block is not None
    manager.free(req)

    # New request with Eagle enabled
    req_eagle = make_request("partial_eagle", token_ids, block_size, sha256)
    computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_eagle)
    # Original match: 2 full blocks → Eagle removes 1 → 1 remaining
    assert len(computed_blocks.blocks[0]) == 1
    assert num_tokens == 1 * block_size

    # Evict the first block in the request
    assert (
        manager.block_pool.get_cached_block(
            block_hash_first_block, kv_cache_group_ids=[0]
        )
        is not None
    )
    manager.block_pool.cached_block_hash_to_block._cache.pop(
        make_block_hash_with_group_id(block_hash_first_block, 0)
    )

    # New request
    req_after_evict = make_request(
        "partial_eagle_after_evict", token_ids, block_size, sha256
    )
    computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_after_evict)
    # Cache miss. The only hit prefix is [NULL_BLOCK, BLOCK_2] if eagle is
    # not considered. But after dropping the last matched block due to eagle,
    # there will be no matched prefix.
    assert len(computed_blocks.blocks[0]) == 0
    assert num_tokens == 0


def test_eagle_swa_alignment_caches_extra_block():
    """Regression: SWA + EAGLE with `sliding_window <= alignment_tokens`.

    When the cache-hit alignment (lcm of per-group block sizes) is larger than
    the SWA window, the SWA mask only kept the last block of each aligned
    segment. EAGLE/MTP lookup needs ``tail + 1`` contiguous cached blocks and
    that +1 block lives at the next segment's first position, which was left
    uncached. The fix caches that extra block when ``use_eagle=True``.
    """
    block_size = 8
    # Full group uses 4 * block_size, so lcm/alignment is 4 * block_size.
    # SWA group has sliding_window = block_size (i.e., tail = 1 block).
    # Without the fix, the second cached block needed for the EAGLE 2-block
    # match never exists -> EAGLE cache hit fails entirely.
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["full"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["swa_mtp"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
                is_eagle_group=True,
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        use_eagle=True,
    )

    # Prime the cache with a long prompt (16 swa blocks = 4 aligned segments).
    token_ids = [i for i in range(16) for _ in range(block_size)]
    req0 = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req0)
    blocks = manager.allocate_slots(
        req0,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None
    manager.free(req0)

    # Second request with identical prompt should find an EAGLE cache hit.
    # Without the fix, ``num_computed_tokens`` is 0; with the fix, it lands at
    # an alignment boundary (multiple of 32 tokens, minus the EAGLE drop).
    req1 = make_request("1", token_ids, block_size, sha256)
    _, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert num_computed_tokens > 0, (
        "EAGLE + SWA with sliding_window <= alignment failed to find any "
        "cache hit; the +1 block past each segment boundary must be cached."
    )
    # Each aligned segment contributes 4 * block_size = 32 tokens; EAGLE drops
    # the last block (block_size tokens) from the hit.
    assert num_computed_tokens % (4 * block_size) == 0


def test_eagle_swa_boundary_caches_post_boundary_block():
    """EAGLE + SWA must cache the first block after an alignment boundary.

    A 40-token computed prefix with 8-token SWA blocks and 32-token hybrid
    alignment needs SWA blocks 3 and 4 cached to reuse a 32-token prefix:
    block 3 is the segment tail, and block 4 is the EAGLE lookahead block
    that gets dropped after lookup.
    """
    block_size = 8
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["full"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["swa_mtp"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
                is_eagle_group=True,
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        use_eagle=True,
    )

    token_ids = [i for i in range(5) for _ in range(block_size)]
    req0 = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req0)
    blocks = manager.allocate_slots(
        req0,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    assert pool.get_cached_block(req0.block_hashes[3], kv_cache_group_ids=[1])
    assert pool.get_cached_block(req0.block_hashes[4], kv_cache_group_ids=[1])
    manager.free(req0)

    req1 = make_request("1", token_ids + [999], block_size, sha256)
    _, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert num_computed_tokens == 4 * block_size


def test_eagle_grouped_swa_siblings_use_same_cache_mask():
    """Grouped SWA siblings must cache the EAGLE lookahead block together."""
    block_size = 8
    swa_spec = SlidingWindowSpec(
        block_size=block_size,
        num_kv_heads=1,
        head_size=1,
        dtype=torch.float32,
        sliding_window=block_size,
    )
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["full"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(["swa_main"], swa_spec),
            KVCacheGroupSpec(["swa_mtp"], swa_spec, is_eagle_group=True),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        use_eagle=True,
    )

    token_ids = [i for i in range(9) for _ in range(block_size)]
    req0 = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req0)
    blocks = manager.allocate_slots(
        req0,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    assert pool.get_cached_block(req0.block_hashes[4], kv_cache_group_ids=[1, 2])
    assert pool.get_cached_block(req0.block_hashes[8], kv_cache_group_ids=[1, 2])
    manager.free(req0)

    req1 = make_request("1", token_ids + [999], block_size, sha256)
    _, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert num_computed_tokens == 8 * block_size


def test_different_block_size():
    block_size = 16
    # full attention and sliding window attention layers have the same page size:
    # (32 tokens/block * float16 token, vs. 16 tokens/block * float32 token)
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=block_size * 2,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=2 * block_size,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # 10 blocks of 16 tokens each. Token ids are not strictly aligned for each block.
    common_token_ids = [i for i in range(10) for _ in range(block_size)]

    req0 = make_request("0", common_token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert not computed_blocks.blocks[0]
    assert not computed_blocks.blocks[1]
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0, 7 * block_size, len(computed_blocks.blocks[0]) * 16, computed_blocks
    )
    assert blocks.get_block_ids() == ([1, 2, 3, 4], [5, 6, 7, 8, 9, 10, 11])
    req1 = make_request("1", common_token_ids[: 7 * block_size + 1], block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert len(computed_blocks.blocks[0]) == 3
    assert len(computed_blocks.blocks[1]) == 6
    assert num_computed_tokens == 6 * 16

    req2 = make_request("2", common_token_ids[: 6 * block_size + 1], block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req2)
    assert len(computed_blocks.blocks[0]) == 3
    assert len(computed_blocks.blocks[1]) == 6
    assert num_computed_tokens == 6 * 16

    # Evict some blocks to make sliding window cache hit length 5*16
    # But should return 4 * 16 because full attention cache hit length must be
    # a multiple of 32
    manager.block_pool.cached_block_hash_to_block.pop(
        make_block_hash_with_group_id(req1.block_hashes[6], 1), 11
    )
    manager.block_pool.cached_block_hash_to_block.pop(
        make_block_hash_with_group_id(req1.block_hashes[5], 1), 10
    )
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert len(computed_blocks.blocks[0]) == 2
    assert len(computed_blocks.blocks[1]) == 4
    assert num_computed_tokens == 4 * 16


def test_hybrid_cache_blocks_swa_tail_window_only():
    """Within each lcm-aligned segment, SWA's ``find_longest_cache_hit`` only
    returns the trailing ``ceil((sliding_window - 1) / block_size)`` blocks
    (its right-to-left scan stops once a contiguous match is found). Blocks
    earlier in the segment can never serve a hit, so
    ``HybridKVCacheCoordinator.cache_blocks`` should skip them rather than
    polluting the prefix-cache hash map."""
    block_size = 8
    # Full attn block_size=32, SWA block_size=8, sw=8 -> lcm=32.
    # tail = ceil(7/8) = 1; per_segment = 32/8 = 4.
    # Per-segment template = [F, F, F, T]; only the last SWA block in each
    # 32-token segment ends up in the prefix-cache hash map.
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # 8 hash-blocks of 8 tokens (64 tokens, two lcm-aligned segments).
    token_ids = [i for i in range(8) for _ in range(block_size)]
    req = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    blocks = manager.allocate_slots(
        req,
        8 * block_size,
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None
    assert len(req.block_hashes) == 8

    pool = manager.block_pool
    # SWA group_id=1: only hash 3 and hash 7 (the last block of each
    # 32-token segment) should be cached. Hashes 0,1,2,4,5,6 cannot serve
    # a hit at any lcm-aligned length, so they must NOT be cached.
    expected_cached = {3, 7}
    for i in range(8):
        cached = pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1])
        if i in expected_cached:
            assert cached is not None, f"SWA hash {i} should be cached"
        else:
            assert cached is None, (
                f"SWA hash {i} cannot serve any lcm-aligned hit; should not be cached"
            )


def test_hybrid_cache_blocks_clamped_to_lcm():
    """HybridKVCacheCoordinator.cache_blocks() clamps to scheduler_block_size.
    Chunks past the last lcm-aligned boundary can never participate in a
    cache hit (find_longest_cache_hit always returns lcm-aligned hits), so
    caching them only pollutes the prefix-cache hash map and keeps blocks
    on the LRU list that could otherwise return to the free pool."""
    block_size = 16
    # Full attn block_size=32, SWA block_size=16 -> lcm=32.
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=block_size * 2,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=2 * block_size,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # 7 hash-blocks of 16 tokens (112 tokens). With lcm=32 the clamp truncates
    # to 96 tokens — SWA caches 6 hashes, full-attn caches 3.
    token_ids = [i for i in range(7) for _ in range(block_size)]
    req = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    blocks = manager.allocate_slots(
        req,
        7 * block_size,
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None
    assert len(req.block_hashes) == 7

    pool = manager.block_pool
    # SWA group_id=1: hashes 0..5 cached (6 blocks * 16 tokens = 96), hash 6
    # spans tokens [96, 112) past the lcm boundary and must NOT be cached.
    for i in range(6):
        assert (
            pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1])
            is not None
        ), f"SWA hash {i} should be cached"
    assert pool.get_cached_block(req.block_hashes[6], kv_cache_group_ids=[1]) is None, (
        "SWA hash 6 spans tokens past the lcm boundary; should not be cached"
    )


def test_hybrid_local_kv_retention_interval_aligns_in_manager():
    """Verify fixed intervals retain sparse tails plus the latest replay tail."""
    block_size = 8
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        retention_interval=64,
    )

    # The SWA manager uses the configured 64-token interval (a multiple of the
    # 32-token lcm_block_size) as its retention segment. For this 128-token
    # prompt, the retained SWA tails are the 64-token interval boundary, the
    # 96-token replay boundary, and the 128-token interval boundary.
    token_ids = [i for i in range(16) for _ in range(block_size)]
    req = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    blocks = manager.allocate_slots(
        req,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    expected_swa_cached = {7, 11, 15}
    for i in range(16):
        cached = pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1])
        if i in expected_swa_cached:
            assert cached is not None, f"SWA hash {i} should be cached"
        else:
            assert cached is None, f"SWA hash {i} should not be cached"


@pytest.mark.parametrize(
    "interval, expected_match",
    [
        # scheduler_block_size is 32 (= lcm(4*8, 8)); 33 is not a multiple of it.
        (33, "multiple of scheduler_block_size"),
        # A negative multiple (-32 % 32 == 0) must still be rejected explicitly,
        # otherwise it would pass the modulo check and silently degrade to dense.
        (-32, "non-negative"),
    ],
)
def test_hybrid_local_kv_retention_interval_rejects_invalid(interval, expected_match):
    """A retention interval that is negative or not a multiple of
    scheduler_block_size errors out at construction time."""
    block_size = 8
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
            ),
        ],
    )
    with pytest.raises(ValueError, match=expected_match):
        make_kv_cache_manager(
            kv_cache_config=kv_cache_config,
            max_model_len=8192,
            enable_caching=True,
            hash_block_size=block_size,
            retention_interval=interval,
        )


def test_zero_retention_is_ignored_for_full_attention():
    kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10)
    manager = make_kv_cache_manager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=16,
        retention_interval=0,
    )
    assert manager.coordinator.retention_interval == 0


def test_positive_retention_rejects_full_attention():
    kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=10)
    with pytest.raises(ValueError, match="no sliding-window or Mamba"):
        make_kv_cache_manager(
            kv_cache_config,
            max_model_len=8192,
            enable_caching=True,
            hash_block_size=16,
            retention_interval=16,
        )


def test_hybrid_local_kv_retention_interval_survives_recycling():
    """Verify retained local checkpoints are reused after block recycling."""
    hash_block_size = 4
    kv_cache_config = KVCacheConfig(
        num_blocks=800,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["full"],
                MLAAttentionSpec(
                    block_size=64 * hash_block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.uint8,
                    tokens_per_state=4,
                ),
            ),
            KVCacheGroupSpec(
                ["swa"],
                SlidingWindowSpec(
                    block_size=16 * hash_block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=512,
                ),
            ),
            KVCacheGroupSpec(
                ["c128"],
                SlidingWindowSpec(
                    block_size=2 * hash_block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=128,
                ),
            ),
            KVCacheGroupSpec(
                ["c4"],
                SlidingWindowSpec(
                    block_size=hash_block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=8,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=4096,
        enable_caching=True,
        hash_block_size=hash_block_size,
        retention_interval=1024,
    )

    def fill_request(request_id: str, token_offset: int) -> list[int]:
        token_ids = [
            token_offset + i for i in range(1024) for _ in range(hash_block_size)
        ]
        fill_req = make_request(request_id, token_ids, hash_block_size, sha256)
        while fill_req.num_computed_tokens < len(token_ids):
            num_new_tokens = min(512, len(token_ids) - fill_req.num_computed_tokens)
            blocks = manager.allocate_slots(fill_req, num_new_tokens)
            assert blocks is not None
            fill_req.num_computed_tokens += num_new_tokens
        manager.free(fill_req)
        return token_ids

    token_ids = fill_request("fill_0", 0)
    replay_req = make_request("replay", token_ids[:1800], hash_block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(replay_req)
    assert num_computed_tokens == 1024
    assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256]

    fill_request("fill_1", 100_000)
    replay_req = make_request("replay_again", token_ids[:1800], hash_block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(replay_req)
    assert num_computed_tokens == 1024
    assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256]


def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary():
    """Verify latest-only retention reuses only the replayable prompt boundary."""
    block_size = 8
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        retention_interval=0,
    )

    token_ids = [i for i in range(16) for _ in range(block_size)]
    req0 = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req0)
    blocks = manager.allocate_slots(
        req0,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    expected_swa_cached = {11}
    for i in range(16):
        cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
        if i in expected_swa_cached:
            assert cached is not None, f"SWA hash {i} should be cached"
        else:
            assert cached is None, f"SWA hash {i} should not be cached"

    manager.free(req0)
    retained_swa_block = pool.get_cached_block(req0.block_hashes[11], [1])
    assert retained_swa_block is not None
    assert retained_swa_block[0].ref_cnt == 0

    req1 = make_request("1", token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    # Full prompt hits intentionally recompute the final block for logits, so
    # the longest usable hit is the previous LCM boundary: 96 tokens.
    assert num_computed_tokens == 12 * block_size
    assert len(computed_blocks.blocks[1]) == 12

    shorter_req = make_request("2", token_ids[: 12 * block_size], block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(shorter_req)
    assert num_computed_tokens == 0
    assert len(computed_blocks.blocks[1]) == 0


def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary():
    """Verify MTP/EAGLE SWA retention keeps the extra proof block.

    EAGLE/MTP lookup matches one additional local block after the returned
    prefix and then drops it. Sparse retention must therefore cache the normal
    local tail at the latest replay boundary plus one extra SWA block.
    """
    block_size = 8
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["full"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["swa_mtp"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
                is_eagle_group=True,
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        retention_interval=0,
        use_eagle=True,
    )

    # 127 tokens: eagle proves a boundary only if an aligned unit exists past
    # it, so the replay boundary rewinds one unit to
    # floor(127 / 32) * 32 - 32 = 64. The SWA tail plus its proof block then
    # ends at 72, two 8-token blocks wide: hashes 7 and 8.
    token_ids = [i for i in range(15) for _ in range(block_size)] + [15] * 7
    req0 = make_request("0", token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert num_computed_tokens == 0
    blocks = manager.allocate_slots(
        req0,
        len(token_ids),
        num_computed_tokens,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    expected_swa_cached = {7, 8}
    for i in range(15):
        cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
        if i in expected_swa_cached:
            assert cached is not None, f"SWA hash {i} should be cached"
        else:
            assert cached is None, f"SWA hash {i} should not be cached"

    manager.free(req0)

    req1 = make_request("1", token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert num_computed_tokens == 8 * block_size
    assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 8]


def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt():
    """EAGLE resend and extension reuse the same retained checkpoint.

    Attention may inspect the prompt's final block before dropping it, so an
    aligned resend must not lose another block to the P - 1 lookup limit.
    """
    block_size = 32
    num_spec = 3
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["full"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["mamba_mtp"],
                MambaSpec(
                    block_size=block_size,
                    shapes=((1, 1),),
                    dtypes=(torch.float32,),
                    mamba_cache_mode="align",
                    num_speculative_blocks=num_spec,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        retention_interval=0,
        use_eagle=True,
    )

    # Both requests prove 128 attention tokens, drop one block, and reuse 96.
    token_ids = [i for i in range(4) for _ in range(block_size)]
    req0 = make_request("0", token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
    assert num_computed_tokens == 0
    # Materialize the checkpoint through normal block-aligned prefill chunks.
    for chunk_end in (32, 64, 96, 128):
        blocks = manager.allocate_slots(
            req0,
            chunk_end - req0.num_computed_tokens,
            num_computed_tokens,
            computed_blocks,
            num_lookahead_tokens=num_spec,
        )
        assert blocks is not None
        req0.num_computed_tokens = chunk_end

    # Both materialized replay checkpoints survive; resend and extension use 96.
    pool = manager.block_pool
    expected_mamba_cached = {1, 2}
    for i in range(4):
        cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
        if i in expected_mamba_cached:
            assert cached is not None, f"mamba hash {i} should be cached"
        else:
            assert cached is None, f"mamba hash {i} should not be cached"
    manager.free(req0)

    # The identical resend must reuse the retained checkpoint, not miss to zero.
    req1 = make_request("1", token_ids, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert num_computed_tokens == 3 * block_size
    assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 3]

    # The longer sibling reuses that same checkpoint.
    longer = make_request("2", token_ids + [9] * block_size, block_size, sha256)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(longer)
    assert num_computed_tokens == 3 * block_size


def test_hybrid_mamba_retention_lookahead_keeps_last_prompt_boundary():
    """Sparse retention keeps the last boundary covered by lookahead hashes."""
    block_size = 32
    kv_cache_config = _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba"])
    kv_cache_config.kv_cache_groups[1].kv_cache_spec = replace(
        kv_cache_config.kv_cache_groups[1].kv_cache_spec,
        num_speculative_blocks=1,
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        retention_interval=0,
        use_eagle=False,
        num_hash_lookahead_tokens=1,
    )

    # 129 tokens: the last boundary below the logits token is 128.
    token_ids = [i for i in range(4) for _ in range(block_size)] + [4]
    req0 = make_request("0", token_ids, block_size, sha256, num_hash_lookahead_tokens=1)
    for chunk_end in (32, 64, 96, 128, 129):
        blocks = manager.allocate_slots(
            req0,
            chunk_end - req0.num_computed_tokens,
            num_lookahead_tokens=1,
        )
        assert blocks is not None
        req0.num_computed_tokens = chunk_end

    pool = manager.block_pool
    for i in range(4):
        cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
        if i == 3:
            assert cached is not None, "mamba state at 128 should be retained"
        else:
            assert cached is None, f"mamba hash {i} should not be cached"
    manager.free(req0)

    req1 = make_request("1", token_ids, block_size, sha256, num_hash_lookahead_tokens=1)
    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
    assert num_computed_tokens == 4 * block_size
    assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 4]


def test_block_lookup_cache_single_block_per_key():
    cache = BlockHashToBlockMap()
    key0 = BlockHashWithGroupId(b"hash0")
    key1 = BlockHashWithGroupId(b"hash1")
    key2 = BlockHashWithGroupId(b"hash2")
    block0 = KVCacheBlock(0)
    block1 = KVCacheBlock(1)

    assert cache.get_one_block(key0) is None
    assert cache.get_one_block(key1) is None
    assert cache.get_one_block(key2) is None
    # key0 inserted
    cache.insert(key0, block0)
    assert cache.get_one_block(key0) is block0
    assert cache.get_one_block(key1) is None
    assert cache.get_one_block(key2) is None
    # key1 inserted
    cache.insert(key1, block1)
    assert cache.get_one_block(key0) is block0
    assert cache.get_one_block(key1) is block1
    assert cache.get_one_block(key2) is None
    # No block popped due to block_id mismatch
    assert cache.pop(key0, 100) is None
    assert cache.get_one_block(key0) is block0
    assert cache.get_one_block(key1) is block1
    assert cache.get_one_block(key2) is None
    # block popped with (key0, block ID 0)
    assert cache.pop(key0, 0) is block0
    assert cache.get_one_block(key0) is None
    assert cache.get_one_block(key1) is block1
    assert cache.get_one_block(key2) is None
    # No block popped due to block_id mismatch
    assert cache.pop(key0, 1) is None
    assert cache.get_one_block(key0) is None
    assert cache.get_one_block(key1) is block1
    assert cache.get_one_block(key2) is None
    # block popped with (key1, block ID 1)
    assert cache.pop(key1, 1) is block1
    assert cache.get_one_block(key0) is None
    assert cache.get_one_block(key1) is None
    assert cache.get_one_block(key2) is None


def test_block_lookup_cache_multi_blocks_per_key():
    cache = BlockHashToBlockMap()
    key0 = BlockHashWithGroupId(b"hash0")
    key1 = BlockHashWithGroupId(b"hash1")
    block00 = KVCacheBlock(0)
    block01 = KVCacheBlock(1)
    block10 = KVCacheBlock(10)
    block11 = KVCacheBlock(11)

    assert cache.get_one_block(key0) is None
    assert cache.get_one_block(key1) is None

    cache.insert(key0, block00)
    cache.insert(key0, block01)
    cache.insert(key1, block10)
    cache.insert(key1, block11)

    assert cache.get_one_block(key0) is block00
    assert cache.pop(key0, 0) is block00
    assert cache.get_one_block(key0) is block01
    assert cache.pop(key0, 1) is block01
    assert cache.get_one_block(key0) is None
    assert cache.pop(key0, 2) is None

    assert cache.get_one_block(key1) is block10
    assert cache.pop(key1, 10) is block10
    assert cache.get_one_block(key1) is block11
    assert cache.pop(key1, 11) is block11
    assert cache.get_one_block(key1) is None
    assert cache.pop(key1, 12) is None


def test_can_fit_full_sequence_swa_cap_admits_long_prompt():
    """Hybrid full+SWA model with a pool sized at the startup minimum should
    admit a prompt longer than the SWA cap, because SlidingWindowManager
    recycles blocks during chunked prefill (issue #39734)."""
    block_size = 16
    sliding_window = 4 * block_size  # 64 tokens
    max_num_batched_tokens = 8 * block_size  # 128 tokens
    max_model_len = 64 * block_size  # 1024 tokens — much larger than the SWA cap
    # Startup pool sizing: full demands cdiv(max_model_len, bs) = 64 blocks,
    # SWA demands cdiv(SW-1+max_batched, bs) + 1 = cdiv(191, 16) + 1 = 13.
    # Pool minimum = 64 + 13 = 77; +1 for the null block.
    num_blocks = 64 + 13 + 1

    config = KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer_full"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            ),
            KVCacheGroupSpec(
                ["layer_swa"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=sliding_window,
                ),
            ),
        ],
    )

    manager = make_kv_cache_manager(
        config,
        max_model_len=max_model_len,
        # Single (sync) batch in flight, so in-flight tokens == batched tokens.
        max_in_flight_tokens=max_num_batched_tokens,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # A prompt that is shorter than max_model_len but longer than SW + chunk:
    # cdiv(prompt_len, bs) = 32 blocks. Without the cap, admission would
    # demand 32 (full) + 32 (SWA) = 64 blocks. With the cap, SWA contributes
    # only 13, so total = 32 + 13 = 45 ≤ pool size.
    prompt_len = 32 * block_size
    req = make_request("long", list(range(prompt_len)), block_size, sha256)

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


@pytest.mark.parametrize(("reserved_blocks", "fits"), [(4, True), (10, False)])
def test_full_sequence_admission_reserves_device_blocks_for_immediate_chunk(
    reserved_blocks, fits
):
    """Reservations constrain the current allocation, not the entire prompt."""
    block_size = 16
    config = KVCacheConfig(
        num_blocks=11,  # Ten usable blocks after the null block.
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            )
        ],
    )
    manager = make_kv_cache_manager(
        config, max_model_len=128, enable_caching=True, hash_block_size=block_size
    )
    request = make_request("chunk", list(range(128)), block_size, sha256)

    blocks = manager.allocate_slots(
        request,
        block_size,
        full_sequence_must_fit=True,
        reserved_blocks=reserved_blocks,
    )

    assert (blocks is not None) is fits


def test_can_fit_full_sequence_full_attention_still_gates_oversized():
    """The cap only loosens the SWA group; a prompt that exceeds the
    full-attention pool capacity must still be rejected."""
    block_size = 16
    sliding_window = 4 * block_size
    max_num_batched_tokens = 8 * block_size
    max_model_len = 64 * block_size
    # Provide a tiny pool — even a small prompt should be rejected.
    num_blocks = 5

    config = KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer_full"],
                FullAttentionSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                ),
            ),
            KVCacheGroupSpec(
                ["layer_swa"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=sliding_window,
                ),
            ),
        ],
    )

    manager = make_kv_cache_manager(
        config,
        max_model_len=max_model_len,
        # Single (sync) batch in flight, so in-flight tokens == batched tokens.
        max_in_flight_tokens=max_num_batched_tokens,
        enable_caching=True,
        hash_block_size=block_size,
    )

    # 16 blocks of full attention demand alone exceeds the 5-block pool.
    prompt_len = 16 * block_size
    req = make_request("oversized", list(range(prompt_len)), block_size, sha256)

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


def test_can_fit_full_sequence_hisparse_caps_resident_pages():
    manager = make_hisparse_kv_cache_manager(
        num_blocks=11,
        host_num_blocks=9,
        max_model_len=8 * HISPARSE_BLOCK_SIZE,
        max_in_flight_tokens=2 * HISPARSE_BLOCK_SIZE,
    )
    request = make_request(
        "long-hisparse",
        list(range(8 * HISPARSE_BLOCK_SIZE)),
        HISPARSE_BLOCK_SIZE,
        sha256,
    )

    assert (
        _allocate_scheduled(
            manager,
            request,
            num_new_tokens=HISPARSE_BLOCK_SIZE,
            full_sequence_must_fit=True,
        )
        is not None
    )


def test_cache_hit_local_and_external():
    # Regression test for #33775: when a request hits the local prefix cache
    # in one KV cache group and needs external (connector) blocks in another,
    # the external allocation of an earlier group must not evict the local
    # cache-hit blocks of a later group. Otherwise the same physical block can
    # be handed out twice, producing duplicate block IDs / ref_cnt corruption.
    block_size = 16
    kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 100)
    del kv_cache_config.kv_cache_groups[2:]
    req_id = "test"
    manager = make_kv_cache_manager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        use_eagle=True,
    )

    top_blocks = []
    head = manager.block_pool.free_block_queue.fake_free_list_head
    for _ in range(10):
        top_blocks.append(head.next_free_block)
        head = head.next_free_block
    cache_hit = KVCacheBlocks((top_blocks[:5], top_blocks[5:]))

    manager.allocate_slots(
        make_request(req_id, [0] * (8 * block_size), block_size, sha256),
        16,
        5 * block_size,
        cache_hit,
        0,
        2 * block_size,
    )

    req_blocks = manager.get_blocks(req_id)
    req_block_ids = req_blocks.get_block_ids()
    all_block_ids = req_block_ids[0] + req_block_ids[1]
    assert len(set(all_block_ids)) == len(all_block_ids), "Block IDs are not unique"


def _take_free_blocks(manager: KVCacheManager, num_blocks: int) -> list[KVCacheBlock]:
    """Grab the first ``num_blocks`` blocks at the head of the free queue
    without removing them. These ref_cnt==0 blocks stand in for evictable
    cache-hit blocks left behind by a previous (e.g. preempted) request, and
    sitting at the head guarantees a later group's external ``get_new_blocks``
    would contend for them on unpatched code (issue #33775)."""
    blocks: list[KVCacheBlock] = []
    head = manager.block_pool.free_block_queue.fake_free_list_head
    for _ in range(num_blocks):
        head = head.next_free_block
        blocks.append(head)
    return blocks


def _assert_no_double_allocation(manager: KVCacheManager, req_id: str) -> None:
    """No physical block may be handed out twice across groups, and every
    block referenced by the request must have a live ref_cnt."""
    block_ids = manager.get_blocks(req_id).get_block_ids()
    flat = [block_id for group in block_ids for block_id in group]
    assert len(set(flat)) == len(flat), "Block IDs are not unique across groups"
    null_id = manager.block_pool.null_block.block_id
    for block_id in flat:
        if block_id == null_id:
            continue
        assert manager.block_pool.blocks[block_id].ref_cnt >= 1, (
            f"block {block_id} referenced by the request has ref_cnt 0"
        )


def _two_phase_block_size(manager: KVCacheManager) -> int:
    return manager.kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size


def _cross_group_cache_hit(
    manager: KVCacheManager,
    req_id: str,
    num_groups: int,
    local_blocks_per_group: int = 5,
    num_external_blocks: int = 2,
    num_new_blocks: int = 1,
) -> Request:
    """Allocate ``req_id`` with a per-group local prefix hit plus external
    (connector) computed tokens, driving the coordinator's two-phase path.
    Returns the allocated request so callers can free it (e.g. to preempt)."""
    block_size = _two_phase_block_size(manager)
    hit_blocks = _take_free_blocks(manager, num_groups * local_blocks_per_group)
    cache_hit = KVCacheBlocks(
        tuple(
            hit_blocks[i * local_blocks_per_group : (i + 1) * local_blocks_per_group]
            for i in range(num_groups)
        )
    )
    prompt_blocks = local_blocks_per_group + num_external_blocks + num_new_blocks
    request = make_request(
        req_id, [0] * (prompt_blocks * block_size), block_size, sha256
    )
    manager.allocate_slots(
        request,
        num_new_blocks * block_size,
        local_blocks_per_group * block_size,
        cache_hit,
        0,
        num_external_blocks * block_size,
    )
    return request


def _make_two_phase_manager(num_groups: int) -> KVCacheManager:
    assert num_groups in (2, 3)
    block_size = 16
    kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 100)
    del kv_cache_config.kv_cache_groups[num_groups:]
    return make_kv_cache_manager(
        kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        use_eagle=True,
    )


def test_cache_hit_local_and_external_three_groups():
    # Scenario 1 (issue #33775): SWA + full attention with *three* KV cache
    # groups (1 full + 2 sliding-window). A local prefix hit in some groups
    # combined with external (connector) blocks in others must not let one
    # group's external `get_new_blocks` evict another group's not-yet-touched
    # cache-hit blocks, which would hand the same physical block out twice.
    manager = _make_two_phase_manager(num_groups=3)
    _cross_group_cache_hit(manager, "test", num_groups=3)
    _assert_no_double_allocation(manager, "test")


def test_cache_hit_local_and_external_three_groups_preempt_and_reallocate():
    # Scenario 2: the same 3-group hybrid config, but the request is preempted
    # (freed) and then reallocated. After the free, the coordinator must treat
    # the request as new again so external blocks are re-allocated, and the
    # two-phase ordering must still prevent cross-group double allocation when
    # reallocating against the now-evictable cache-hit blocks.
    manager = _make_two_phase_manager(num_groups=3)

    request = _cross_group_cache_hit(manager, "test", num_groups=3)
    _assert_no_double_allocation(manager, "test")

    # Preempt: free the request; its blocks return to the pool (full ones stay
    # cached/evictable) and the coordinator forgets it.
    manager.free(request)
    assert manager.get_blocks("test").get_block_ids() == ([], [], [])

    # Reallocate the same request id against fresh cache-hit blocks taken from
    # the current free-queue head, mirroring a preempted request being
    # scheduled again. Because the request is no longer known, the coordinator
    # re-arms `is_new_request` and re-runs external allocation, which must still
    # not double-allocate across groups.
    _cross_group_cache_hit(manager, "test", num_groups=3)
    _assert_no_double_allocation(manager, "test")
    assert manager.get_blocks("test").get_block_ids() != ([], [], [])


def test_cache_hit_local_and_external_two_groups_preempt_and_reallocate():
    # Scenario 3: the minimal 2-group hybrid config (1 full + 1 sliding-window)
    # exercised through the same preempt -> reallocate cycle as scenario 2.
    manager = _make_two_phase_manager(num_groups=2)

    request = _cross_group_cache_hit(manager, "test", num_groups=2)
    _assert_no_double_allocation(manager, "test")

    manager.free(request)
    assert manager.get_blocks("test").get_block_ids() == ([], [])

    _cross_group_cache_hit(manager, "test", num_groups=2)
    _assert_no_double_allocation(manager, "test")
    assert manager.get_blocks("test").get_block_ids() != ([], [])


def test_swa_free_split_keeps_cached_tail_ahead_of_scratch():
    """Dense retention: freeing an SWA request must place its
    uncached scratch blocks at the front of the free queue (recycled first)
    and keep its cached checkpoint blocks at the back (retained for prefix
    hits). This split is always-on, independent of the retention interval."""
    block_size = 8
    kv_cache_config = KVCacheConfig(
        num_blocks=100,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer1"],
                FullAttentionSpec(
                    block_size=4 * block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float16,
                ),
            ),
            KVCacheGroupSpec(
                ["layer2"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=block_size,
                ),
            ),
        ],
    )
    manager = make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )

    token_ids = [i for i in range(16) for _ in range(block_size)]
    req = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    blocks = manager.allocate_slots(
        req,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    swa_manager = manager.coordinator.single_type_managers[1]
    null_block = manager.block_pool.null_block
    cached_ids: set[int] = set()
    uncached_ids: set[int] = set()
    cached_hash_indices: list[int] = []
    for i, block in enumerate(swa_manager.req_to_blocks[req.request_id]):
        if block is null_block:
            continue
        if block.block_hash is None:
            uncached_ids.add(block.block_id)
        else:
            cached_ids.add(block.block_id)
            cached_hash_indices.append(i)
    # The dense default mask caches only the per-segment tails, so a 16-block
    # SWA prompt must produce a mix of retained and scratch blocks.
    assert cached_ids, "expected some retained (cached) SWA tail blocks"
    assert uncached_ids, "expected some scratch (uncached) SWA blocks"

    manager.free(req)

    order = [
        b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()
    ]
    pos = {bid: i for i, bid in enumerate(order)}
    # Every scratch block is recycled before every retained block.
    assert max(pos[bid] for bid in uncached_ids) < min(pos[bid] for bid in cached_ids)
    # The retained tails survive the free and still serve a prefix-cache hit.
    for i in cached_hash_indices:
        assert (
            manager.block_pool.get_cached_block(
                req.block_hashes[i], kv_cache_group_ids=[1]
            )
            is not None
        )


def _make_pure_swa_manager(block_size, sliding_window, num_blocks=100, **kwargs):
    """Single sliding-window group (UnitaryKVCacheCoordinator)."""
    kv_cache_config = KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=[
            KVCacheGroupSpec(
                ["layer"],
                SlidingWindowSpec(
                    block_size=block_size,
                    num_kv_heads=1,
                    head_size=1,
                    dtype=torch.float32,
                    sliding_window=sliding_window,
                ),
            ),
        ],
    )
    return make_kv_cache_manager(
        kv_cache_config=kv_cache_config,
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
        **kwargs,
    )


def test_pure_swa_retention_interval_caches_sparse_tails():
    """Sparse retention must work for a pure-SWA single-group model, not just
    hybrid models: only the per-interval tails plus the latest replay tail are
    cached, and a replay still hits the latest replayable boundary."""
    block_size = 16
    manager = _make_pure_swa_manager(
        block_size, sliding_window=block_size, retention_interval=64
    )
    assert type(manager.coordinator).__name__ == "UnitaryKVCacheCoordinator"

    token_ids = [i for i in range(16) for _ in range(block_size)]
    req = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    blocks = manager.allocate_slots(
        req,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    cached = {
        i
        for i in range(16)
        if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[0])
        is not None
    }
    # per_segment = 64 / 16 = 4, need = cdiv(16-1, 16) = 1 -> segment tails at
    # i%4==3 -> {3,7,11,15}; latest replay boundary (255//16*16 = 240) -> tail
    # block 14. Crucially this is a strict subset of all 16 blocks: retention
    # is actually sparse for pure SWA (not silently dense).
    assert cached == {3, 7, 11, 14, 15}

    # A replay of the same prompt hits the latest replayable boundary (240).
    replay = make_request("1", token_ids, block_size, sha256)
    _, num_computed, _ = manager.get_computed_blocks(replay)
    assert num_computed == 240


def test_pure_swa_retention_latest_only():
    """`=0` on a pure-SWA model keeps only the latest replay tail."""
    block_size = 16
    manager = _make_pure_swa_manager(
        block_size, sliding_window=block_size, retention_interval=0
    )

    token_ids = [i for i in range(16) for _ in range(block_size)]
    req = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    blocks = manager.allocate_slots(
        req,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    cached = {
        i
        for i in range(16)
        if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[0])
        is not None
    }
    # No segment tails (interval 0); only the latest replay tail (block 14).
    assert cached == {14}

    replay = make_request("1", token_ids, block_size, sha256)
    _, num_computed, _ = manager.get_computed_blocks(replay)
    assert num_computed == 240


def test_pure_swa_dense_retention_caches_all():
    """With retention set to ``None``, a pure-SWA model keeps dense behavior:
    every block boundary is a potential hit, so all blocks are cached."""
    block_size = 16
    manager = _make_pure_swa_manager(block_size, sliding_window=block_size)

    token_ids = [i for i in range(16) for _ in range(block_size)]
    req = make_request("0", token_ids, block_size, sha256)
    computed_blocks, _, _ = manager.get_computed_blocks(req)
    blocks = manager.allocate_slots(
        req,
        len(token_ids),
        len(computed_blocks.blocks[0]) * block_size,
        computed_blocks,
    )
    assert blocks is not None

    pool = manager.block_pool
    cached = {
        i
        for i in range(16)
        if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[0])
        is not None
    }
    assert cached == set(range(16))


def test_mamba_reachable_block_mask_sparsifies_retention():
    """Mamba state-snapshot retention: with a configured retention interval,
    the manager keeps one cached state per interval-sized segment (plus the
    latest replay boundary) instead of a snapshot per block, which is what
    lets a small attention block_size avoid Mamba dominating the KV pool."""
    from vllm.v1.core.single_type_kv_cache_manager import MambaManager

    block_size = 16
    spec = MambaSpec(
        block_size=block_size,
        shapes=(1, 1),
        dtypes=(torch.float32,),
        mamba_cache_mode="align",
    )

    def retained(retention_interval, num_prompt_tokens=256, end_block=16):
        m = MambaManager.reachable_block_mask(
            start_block=0,
            end_block=end_block,
            alignment_tokens=block_size,
            kv_cache_spec=spec,
            use_eagle=False,
            retention_interval=retention_interval,
            reachable_boundaries=(num_prompt_tokens - 1,),
        )
        return None if m is None else {i for i, v in enumerate(m) if v}

    # Dense retention (None) -> no mask, every block cached.
    assert retained(None) is None
    # interval == block_size -> every block is a boundary -> stays dense.
    assert retained(block_size) is None
    # interval 64 = 4 blocks: segment tails at i%4==3 -> {3,7,11,15}; latest
    # replay boundary 240//16 - 1 = 14. Sparse subset of the 16 blocks.
    assert retained(64) == {3, 7, 11, 14, 15}
    # interval 0 -> only the latest replay boundary (block 14).
    assert retained(0) == {14}


def test_mamba_reachable_block_mask_pins_shared_prefix():
    """A Marconi-detected shared prefix (``shared_prefix_boundary``) lands before
    ``num_prompt`` so the replay-boundary rule alone would drop it. The mask must
    pin the single state block ending on that boundary so sparse retention does
    not defeat cross-request shared-prefix reuse."""
    from vllm.v1.core.single_type_kv_cache_manager import MambaManager

    block_size = 16
    spec = MambaSpec(
        block_size=block_size,
        shapes=(1, 1),
        dtypes=(torch.float32,),
        mamba_cache_mode="align",
    )

    def retained(retention_interval, shared_prefix_boundary, end_block=16):
        boundaries = [255]  # replay boundary (num_prompt 256 - 1)
        if shared_prefix_boundary:
            boundaries.append(shared_prefix_boundary)
        m = MambaManager.reachable_block_mask(
            start_block=0,
            end_block=end_block,
            alignment_tokens=block_size,
            kv_cache_spec=spec,
            use_eagle=False,
            retention_interval=retention_interval,
            reachable_boundaries=boundaries,
        )
        return None if m is None else {i for i, v in enumerate(m) if v}

    # interval 0 keeps only replay boundary 14; the shared prefix at token 96
    # (state block 96//16 - 1 = 5) is now pinned too.
    assert retained(0, 96) == {5, 14}
    # A non-aligned boundary floors to the enclosing aligned boundary.
    assert retained(0, 100) == {5, 14}
    # Coexists with segment tails (interval 64 -> {3,7,11,15} + replay 14).
    assert retained(64, 96) == {3, 5, 7, 11, 14, 15}
    # Dense retention ignores the hint (nothing to sparsify).
    assert retained(None, 96) is None
    # Out-of-range boundary is a no-op (only replay 14 remains).
    assert retained(0, 16 * block_size * 2) == {14}
    # No boundary given -> unchanged replay-only behavior.
    assert retained(0, 0) == {14}
    assert retained(0, None) == {14}


def test_mamba_align_siblings_hit_committed_boundary_in_one_step():
    """Requests that hit a Mamba state committed in an earlier step must not
    defer each other. Re-caching that boundary through the prompt-tail path
    used to mark its hash as cached in the current step, so every sibling
    sharing the boundary was refused until the next step and siblings ran one
    per engine step (e.g. many reads over one shared prompt)."""
    block_size = 16
    manager = make_kv_cache_manager(
        _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba"]),
        max_model_len=8192,
        enable_caching=True,
        hash_block_size=block_size,
    )
    shared = [7] * (2 * block_size)

    # Step 1: the producer prefills up to the end of the shared prefix, which
    # commits the Mamba state at that block boundary.
    manager.new_step_starts()
    producer = make_request("p", shared + [50] * 8, block_size, sha256)
    blocks, num_computed, _ = manager.get_computed_blocks(producer)
    assert num_computed == 0
    assert manager.allocate_slots(producer, 2 * block_size, num_computed, blocks)
    producer.num_computed_tokens = 2 * block_size

    # Step 2: the producer finishes its prompt, and siblings that share the
    # committed boundary are all admitted in the same step.
    manager.new_step_starts()
    assert manager.allocate_slots(producer, 8) is not None
    for i, tail in enumerate((60, 70, 80)):
        sibling = make_request(f"s{i}", shared + [tail] * 8, block_size, sha256)
        blocks, num_computed, _ = manager.get_computed_blocks(sibling)
        assert num_computed == 2 * block_size
        new_blocks = manager.allocate_slots(
            sibling, len(sibling.all_token_ids) - num_computed, num_computed, blocks
        )
        assert new_blocks is not None, f"sibling {i} was deferred"


def test_mamba_shared_prefix_reuse_under_zero_retention():
    """Full cross-request Marconi flow: a partial shared prefix cached by the
    detecting request must stay reusable by a later request under
    ``prefix_cache_retention_interval=0``. Without the pin the junction is
    masked out and the later request misses; with it (and under dense) the reuse
    is preserved."""
    block_size = 16

    def last_req_hit(retention, pin):
        manager = make_kv_cache_manager(
            _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba"]),
            max_model_len=8192,
            enable_caching=True,
            hash_block_size=block_size,
            retention_interval=retention,
        )
        shared = [7 for _ in range(2 * block_size)]  # 2-block shared prefix

        def distinct(v):
            return [v for _ in range(2 * block_size)]

        # req0 primes the shared prefix's (dense) attention cache; align-mode
        # Mamba keeps only its own tail, not the shared-prefix junction.
        req0 = make_request("0", shared + distinct(50), block_size, sha256)
        cb, nc, _ = manager.get_computed_blocks(req0)
        manager.allocate_slots(req0, len(req0.all_token_ids), nc, cb)

        # req1 detects the shared prefix (attn hit, Mamba lag) and, via the
        # Marconi chunk, caches the junction's Mamba state.
        req1 = make_request("1", shared + distinct(60), block_size, sha256)
        cb, nc, boundary = manager.get_computed_blocks(req1)
        assert boundary == 2 * block_size  # junction detected at 2 shared blocks
        if pin:
            req1.shared_prefix_boundary = boundary
        # Simulate the Marconi chunk: schedule up to the junction (boundary - nc).
        manager.allocate_slots(req1, boundary - nc, nc, cb)

        # req2 shares the same prefix: does it reuse the cached junction?
        req2 = make_request("2", shared + distinct(70), block_size, sha256)
        _, nc2, _ = manager.get_computed_blocks(req2)
        return nc2

    # Dense retains the junction -> reuse works (baseline ceiling).
    assert last_req_hit(retention=None, pin=False) == 2 * block_size
    # retention=0 without the pin masks the junction out -> reuse lost (the bug).
    assert last_req_hit(retention=0, pin=False) == 0
    # retention=0 with the pin keeps the junction -> reuse restored.
    assert last_req_hit(retention=0, pin=True) == 2 * block_size


@pytest.mark.skip_global_cleanup
def test_swa_reachable_block_mask_final_partial_segment():
    """Only a final SWA horizon makes its incomplete segment tail reachable."""
    from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager

    spec = SlidingWindowSpec(
        block_size=4,
        num_kv_heads=1,
        head_size=1,
        dtype=torch.float32,
        sliding_window=16,
    )
    start_block = 243
    end_block = 250

    def reachable_blocks(
        final_segment_end_block: int, *, use_eagle: bool = False
    ) -> set[int]:
        mask = SlidingWindowManager.reachable_block_mask(
            start_block=start_block,
            end_block=end_block,
            alignment_tokens=36,
            kv_cache_spec=spec,
            use_eagle=use_eagle,
            final_segment_end_block=final_segment_end_block,
        )
        assert mask is not None
        return {start_block + i for i, reachable in enumerate(mask) if reachable}

    assert reachable_blocks(300) == {248, 249}
    assert reachable_blocks(250) == {246, 247, 248, 249}
    assert reachable_blocks(300, use_eagle=True) == {243, 248, 249}
    assert reachable_blocks(250, use_eagle=True) == {243, 248, 249}


def test_swa_reachable_block_mask_pins_shared_prefix():
    """SWA analog of the Mamba pin: the shared-prefix junction must keep the
    ``need``-block sliding-window tail ending on that boundary (not a single
    block), so a windowed hit can land there under sparse retention."""
    from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager

    block_size = 16

    def retained(retention, boundary, window, end_block=16):
        spec = SlidingWindowSpec(
            block_size=block_size,
            num_kv_heads=1,
            head_size=1,
            dtype=torch.float32,
            sliding_window=window,
        )
        boundaries = [255]  # replay boundary (num_prompt 256 - 1)
        if boundary:
            boundaries.append(boundary)
        m = SlidingWindowManager.reachable_block_mask(
            start_block=0,
            end_block=end_block,
            alignment_tokens=block_size,
            kv_cache_spec=spec,
            use_eagle=False,
            retention_interval=retention,
            reachable_boundaries=boundaries,
        )
        return None if m is None else {i for i, v in enumerate(m) if v}

    # window == block_size -> need = cdiv(15, 16) = 1: single-block tail (like
    # Mamba). Junction at token 96 -> block 5; replay boundary -> block 14.
    assert retained(0, 96, block_size) == {5, 14}
    # window == 3 * block_size -> need = cdiv(47, 16) = 3: the junction keeps a
    # 3-block WINDOW {3,4,5} (the SWA distinction), plus replay window {12,13,14}.
    assert retained(0, 96, 3 * block_size) == {3, 4, 5, 12, 13, 14}
    # Coexists with segment tails (interval 64, need=1): {3,7,11,15} + replay 14
    # + junction 5.
    assert retained(64, 96, block_size) == {3, 5, 7, 11, 14, 15}
    # Dense (all blocks reachable) ignores the hint.
    assert retained(None, 96, block_size) is None
    # No boundary -> unchanged replay-only behavior.
    assert retained(0, 0, block_size) == {14}


def test_swa_reachable_block_mask_with_dcp_scaling():
    """DCP shards each block's KV across ranks, scaling the effective block size.
    Verify that dcp_world_size > 1 scales block size in reachability calculations,
    producing different masks than dcp_world_size=1."""
    from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager

    block_size = 16
    sliding_window = 32  # need = cdiv(31, 16) = 2

    def get_mask(dcp_world_size):
        spec = SlidingWindowSpec(
            block_size=block_size,
            num_kv_heads=1,
            head_size=1,
            dtype=torch.float32,
            sliding_window=sliding_window,
        )
        m = SlidingWindowManager.reachable_block_mask(
            start_block=0,
            end_block=16,
            alignment_tokens=64,  # alignment constraint
            kv_cache_spec=spec,
            use_eagle=False,
            retention_interval=64,  # sparse: 1 state per 64-token segment
            reachable_boundaries=(255,),
            dcp_world_size=dcp_world_size,
        )
        return None if m is None else {i for i, v in enumerate(m) if v}

    # With dcp_world_size=1: effective block_size = 16
    # per_segment = 64 // 16 = 4 blocks per segment
    # need = 2, so 2 < 4 -> sparse retention with coarser granularity
    mask_no_dcp = get_mask(dcp_world_size=1)

    # With dcp_world_size=2: effective block_size = 32
    # per_segment = 64 // 32 = 2 blocks per segment
    # need = 2, so 2 >= 2 -> different retention granularity
    mask_with_dcp = get_mask(dcp_world_size=2)

    # DCP scaling changes block_size, altering segment granularity
    assert mask_no_dcp is not None
    assert mask_with_dcp is not None
    # Different DCP world sizes should produce different masks
    assert mask_no_dcp != mask_with_dcp, (
        "DCP scaling should change retention granularity"
    )


@pytest.mark.parametrize(
    "block_size,dcp_world_size,alignment_tokens",
    [
        (16, 1, 24),  # 24 % 16 != 0
        (16, 3, 64),  # effective block_size 48; 64 % 48 != 0
    ],
)
def test_swa_reachable_block_mask_sub_block_alignment_is_dense(
    block_size, dcp_world_size, alignment_tokens
):
    """The mask is block-granular, so a sub-block alignment_tokens (not a
    multiple of the DCP-scaled block size) cannot be represented exactly. This
    happens for hybrid offloading (e.g. Gemma), where alignment_tokens is the
    full-attention chunk size and need not divide the SWA block size. The mask
    must fall back to dense (None) rather than raise."""
    from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager

    spec = SlidingWindowSpec(
        block_size=block_size,
        num_kv_heads=1,
        head_size=1,
        dtype=torch.float32,
        sliding_window=32,
    )

    mask = SlidingWindowManager.reachable_block_mask(
        start_block=0,
        end_block=16,
        alignment_tokens=alignment_tokens,
        kv_cache_spec=spec,
        use_eagle=False,
        retention_interval=alignment_tokens,
        reachable_boundaries=(),
        dcp_world_size=dcp_world_size,
    )
    assert mask is None


def test_mamba_reachable_block_mask_ignores_dcp():
    """Mamba uses TP, not DCP: each rank holds the full recurrent state, so a
    state block spans kv_cache_spec.block_size tokens regardless of DCP. The
    mask must not scale by dcp_world_size, so retention granularity is
    identical for any DCP world size."""
    from vllm.v1.core.single_type_kv_cache_manager import MambaManager

    spec = MambaSpec(
        block_size=16,
        shapes=(1, 1),
        dtypes=(torch.float32,),
        mamba_cache_mode="align",
    )

    def get_mask(dcp_world_size):
        m = MambaManager.reachable_block_mask(
            start_block=0,
            end_block=16,
            alignment_tokens=64,
            kv_cache_spec=spec,
            use_eagle=False,
            retention_interval=64,  # sparse: one state per 64-token segment
            reachable_boundaries=(255,),
            dcp_world_size=dcp_world_size,
        )
        return None if m is None else {i for i, v in enumerate(m) if v}

    # block_size stays 16 regardless of DCP: per_segment = 64 // 16 = 4.
    mask_no_dcp = get_mask(dcp_world_size=1)
    mask_with_dcp = get_mask(dcp_world_size=2)

    assert mask_no_dcp is not None
    assert mask_with_dcp is not None
    assert mask_no_dcp == mask_with_dcp, "DCP must not change Mamba retention"


def test_mamba_reachable_block_mask_large_dcp_stays_sparse():
    """A large DCP world size must not scale the Mamba block size, so it can
    never collapse sparse retention into dense caching."""
    from vllm.v1.core.single_type_kv_cache_manager import MambaManager

    spec = MambaSpec(
        block_size=8,
        shapes=(1, 1),
        dtypes=(torch.float32,),
        mamba_cache_mode="align",
    )

    def get_mask(dcp_world_size):
        return MambaManager.reachable_block_mask(
            start_block=0,
            end_block=16,
            alignment_tokens=64,
            kv_cache_spec=spec,
            use_eagle=False,
            retention_interval=64,
            reachable_boundaries=(255,),
            dcp_world_size=dcp_world_size,
        )

    # block_size stays 8 regardless of DCP: per_segment = 64 // 8 = 8, sparse.
    mask_no_dcp = get_mask(dcp_world_size=1)
    mask_large_dcp = get_mask(dcp_world_size=16)

    assert mask_no_dcp is not None, "Sparse retention expected"
    assert mask_large_dcp is not None, "Large DCP must not force dense caching"
    assert mask_no_dcp == mask_large_dcp, "DCP must not change Mamba retention"


@pytest.mark.parametrize("use_eagle", [False, True])
def test_swa_shared_prefix_reuse_under_zero_retention(use_eagle):
    """SWA cross-request analog: a partial shared prefix's sliding-window tail
    must stay reusable under ``prefix_cache_retention_interval=0``. Without
    the pin the junction window is masked out and a later request misses; with
    it (and under dense) reuse is preserved. With only the SWA group drafting
    under EAGLE, the full group reports the junction at the shared prefix, but
    the SWA hit lands one block below it: the block past the junction is not
    shared and cannot serve as the EAGLE peek."""
    block_size = 16

    def last_req_hit(retention, pin):
        config = _make_hybrid_kv_cache_config(
            block_size, 200, ["full", "sliding_window"]
        )
        config.kv_cache_groups[1].is_eagle_group = use_eagle
        manager = make_kv_cache_manager(
            config,
            max_model_len=8192,
            enable_caching=True,
            use_eagle=use_eagle,
            hash_block_size=block_size,
            retention_interval=retention,
        )
        shared = [7 for _ in range(4 * block_size)]  # 4-block shared prefix

        def distinct(v):
            return [v for _ in range(4 * block_size)]

        # req0 primes the (dense) full-attention cache; SWA under RET=0 keeps
        # only its own replay window, not the shared-prefix boundary.
        req0 = make_request("0", shared + distinct(50), block_size, sha256)
        cb, nc, _ = manager.get_computed_blocks(req0)
        manager.allocate_slots(req0, len(req0.all_token_ids), nc, cb)

        # req1 detects the shared prefix (full-attn hit, SWA lag). No chunk for
        # SWA -- the pin fires during normal prefill.
        req1 = make_request("1", shared + distinct(60), block_size, sha256)
        cb, nc, boundary = manager.get_computed_blocks(req1)
        if retention == 0:
            assert boundary == 4 * block_size  # junction detected when SWA lags
        if pin and boundary:
            req1.shared_prefix_boundary = boundary
        manager.allocate_slots(req1, len(req1.all_token_ids), nc, cb)

        # req2 shares the same prefix: does its SWA window reuse the junction?
        req2 = make_request("2", shared + distinct(70), block_size, sha256)
        _, nc2, _ = manager.get_computed_blocks(req2)
        return nc2

    expected = (3 if use_eagle else 4) * block_size
    # Dense keeps every SWA tail -> reuse works (ceiling).
    assert last_req_hit(retention=None, pin=False) == expected
    # retention=0 without the pin drops the junction window -> reuse lost (bug).
    assert last_req_hit(retention=0, pin=False) == 0
    # retention=0 with the pin keeps the junction window -> reuse restored.
    assert last_req_hit(retention=0, pin=True) == expected


def test_device_eviction_keeps_host_prefix():
    """Device eviction must not target a same-numbered host block."""
    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
    request = make_request("producer", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)
    assert _allocate_scheduled(manager, request, 64) is not None
    _publish_hisparse_pages(manager)
    host, indexer, _, _ = manager.get_blocks(request.request_id).blocks
    device_block = indexer[0]
    host_block = host[0]
    assert device_block.block_id == host_block.block_id
    manager.free(request)
    assert host_block.block_hash is not None and device_block.block_hash is not None
    manager.evict_blocks({device_block.block_id})
    assert device_block.block_hash is None, (
        "Device eviction left the indexer prefix cached"
    )
    assert host_block.block_hash is not None, (
        "Device eviction invalidated unrelated host KV"
    )


def test_get_unhashed_block_ids_all_groups():
    """Unhashed, non-null block ids are reported per KV cache group.
    A group with no unhashed blocks reports an empty list, not a missing entry.
    """

    def hashed(block_id: int, group_id: int) -> KVCacheBlock:
        block = KVCacheBlock(block_id=block_id)
        block.set_block_hash(make_block_hash_with_group_id(BlockHash(b"h"), group_id))
        return block

    blocks = KVCacheBlocks(
        (
            [
                KVCacheBlock(block_id=1),  # unhashed
                hashed(2, group_id=0),  # cached
                KVCacheBlock(block_id=3, is_null=True),  # null padding
                KVCacheBlock(block_id=4),  # unhashed
            ],
            # No unhashed blocks in this group.
            [hashed(5, group_id=1), KVCacheBlock(block_id=6, is_null=True)],
        )
    )

    assert blocks.get_unhashed_block_ids_all_groups() == [[1, 4], []]