# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for SimpleCPUOffloadScheduler."""

from __future__ import annotations

from dataclasses import dataclass

import pytest
import torch

from vllm import SamplingParams
from vllm.config import (
    CacheConfig,
    DeviceConfig,
    KVTransferConfig,
    ModelConfig,
    SchedulerConfig,
    VllmConfig,
)
from vllm.utils.hashing import sha256
from vllm.v1.core.block_pool import BlockPool
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.core.kv_cache_utils import (
    get_request_block_hasher,
    init_none_hash,
    make_block_hash_with_group_id,
)
from vllm.v1.core.sched.output import (
    CachedRequestData,
    NewRequestData,
    SchedulerOutput,
)
from vllm.v1.core.single_type_kv_cache_manager import (
    register_all_kvcache_specs,
)
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheConfig,
    KVCacheGroupSpec,
    KVCacheTensor,
)
from vllm.v1.outputs import KVConnectorOutput
from vllm.v1.request import Request
from vllm.v1.simple_kv_offload.manager import SimpleCPUOffloadScheduler
from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadWorkerMetadata

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
BLOCK_SIZE = 16
HEAD_SIZE = 16
NUM_KV_HEADS = 1
DTYPE = torch.float16
# bytes per block per tensor:
# block_size * num_kv_heads * head_size * 2 (K+V) * element_size
_BYTES_PER_BLOCK = BLOCK_SIZE * NUM_KV_HEADS * HEAD_SIZE * 2 * DTYPE.itemsize

# Ensure none_hash is initialized once
init_none_hash(sha256)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_kv_cache_config(
    num_blocks: int,
    num_groups: int = 1,
) -> KVCacheConfig:
    """Build a KVCacheConfig with non-empty kv_cache_tensors."""
    groups = []
    tensors = []
    register_all_kvcache_specs(
        vllm_config=None
    )  # Ensure specs are registered for tests
    for g in range(num_groups):
        layer_names = [f"layer_{g}"]
        groups.append(
            KVCacheGroupSpec(
                layer_names,
                FullAttentionSpec(
                    block_size=BLOCK_SIZE,
                    num_kv_heads=NUM_KV_HEADS,
                    head_size=HEAD_SIZE,
                    dtype=DTYPE,
                ),
            )
        )
        tensors.append(
            KVCacheTensor(
                size=_BYTES_PER_BLOCK * num_blocks,
                shared_by=layer_names,
            )
        )
    return KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=tensors,
        kv_cache_groups=groups,
    )


def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig:
    """Minimal VllmConfig for scheduler tests (no GPU)."""
    model_config = ModelConfig(
        model="facebook/opt-125m",
        trust_remote_code=True,
        dtype="float16",
        seed=42,
    )
    scheduler_config = SchedulerConfig(
        max_num_seqs=16,
        max_num_batched_tokens=64,
        max_model_len=10000,
        enable_chunked_prefill=True,
        is_encoder_decoder=False,
    )
    cache_config = CacheConfig(
        block_size=block_size,
        gpu_memory_utilization=0.9,
        enable_prefix_caching=True,
    )
    kv_transfer_config = KVTransferConfig(
        kv_connector="SimpleCPUOffloadConnector",
        kv_role="kv_both",
    )
    return VllmConfig(
        scheduler_config=scheduler_config,
        model_config=model_config,
        cache_config=cache_config,
        kv_transfer_config=kv_transfer_config,
        device_config=DeviceConfig("cpu"),
    )


@dataclass
class SchedulerFixture:
    """Bundle returned by make_scheduler for convenient access."""

    scheduler: SimpleCPUOffloadScheduler
    gpu_block_pool: BlockPool
    vllm_config: VllmConfig
    kv_cache_config: KVCacheConfig
    num_groups: int = 1


def make_scheduler(
    num_cpu_blocks: int = 8,
    num_gpu_blocks: int = 16,
    num_groups: int = 1,
    lazy: bool = False,
) -> SchedulerFixture:
    """Build a SimpleCPUOffloadScheduler with small block pools."""
    kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)
    vllm_config = _make_vllm_config()
    cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups

    sched = SimpleCPUOffloadScheduler(
        vllm_config=vllm_config,
        kv_cache_config=kv_cache_config,
        cpu_capacity_bytes=cpu_capacity_bytes,
        scheduler_block_size=BLOCK_SIZE,
        hash_block_size=BLOCK_SIZE,
        lazy_offload=lazy,
    )

    # Build a real GPU block pool and bind it
    gpu_block_pool = BlockPool(
        num_gpu_blocks=num_gpu_blocks,
        enable_caching=True,
        hash_block_size=BLOCK_SIZE,
    )
    sched.bind_gpu_block_pool(gpu_block_pool)

    return SchedulerFixture(
        scheduler=sched,
        gpu_block_pool=gpu_block_pool,
        vllm_config=vllm_config,
        kv_cache_config=kv_cache_config,
        num_groups=num_groups,
    )


_req_counter = 0


def make_request(
    num_blocks: int = 2,
    request_id: str | None = None,
    extra_tokens: int = 1,
) -> Request:
    """Create a Request with deterministic block hashes."""
    global _req_counter
    _req_counter += 1
    if request_id is None:
        request_id = f"req-{_req_counter}"

    num_tokens = num_blocks * BLOCK_SIZE + extra_tokens
    start = _req_counter * 10000
    prompt_token_ids = list(range(start, start + num_tokens))
    sampling_params = SamplingParams(max_tokens=1)

    req = Request(
        request_id=request_id,
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256),
    )
    return req


def make_scheduler_output(
    req_id_to_num_tokens: dict[str, int],
    *,
    new_reqs: dict[str, tuple[list[int], ...]] | None = None,
    cached_req_new_blocks: dict[str, tuple[list[int], ...] | None] | None = None,
) -> SchedulerOutput:
    """Build a minimal SchedulerOutput with num_scheduled_tokens.

    Args:
        new_reqs: For first-time requests, maps req_id -> block_ids tuple.
            These are placed into ``scheduled_new_reqs`` as ``NewRequestData``.
        cached_req_new_blocks: For returning (cached) requests, maps
            req_id -> new_block_ids (incremental) or None.
            These are placed into ``scheduled_cached_reqs``.
    """
    scheduled_new_reqs: list[NewRequestData] = []
    if new_reqs:
        for req_id, block_ids in new_reqs.items():
            scheduled_new_reqs.append(
                NewRequestData(
                    req_id=req_id,
                    prompt_token_ids=None,
                    mm_features=[],
                    sampling_params=None,
                    pooling_params=None,
                    block_ids=block_ids,
                    num_computed_tokens=0,
                    lora_request=None,
                )
            )

    if cached_req_new_blocks:
        cached_req_ids = list(cached_req_new_blocks.keys())
        cached_new_block_ids = [cached_req_new_blocks[rid] for rid in cached_req_ids]
        cached_reqs = CachedRequestData(
            req_ids=cached_req_ids,
            resumed_req_ids=set(),
            new_token_ids=[[] for _ in cached_req_ids],
            all_token_ids={},
            new_block_ids=cached_new_block_ids,
            num_computed_tokens=[0] * len(cached_req_ids),
            num_output_tokens=[0] * len(cached_req_ids),
        )
    else:
        cached_reqs = CachedRequestData.make_empty()

    return SchedulerOutput(
        scheduled_new_reqs=scheduled_new_reqs,
        scheduled_cached_reqs=cached_reqs,
        num_scheduled_tokens=req_id_to_num_tokens,
        total_num_scheduled_tokens=sum(req_id_to_num_tokens.values()),
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        num_common_prefix_blocks=[],
        preempted_req_ids=set(),
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
    )


def simulate_store_completion(
    scheduler: SimpleCPUOffloadScheduler,
    event_idx: int,
) -> None:
    """Simulate worker reporting a store event completion."""
    output = KVConnectorOutput(
        finished_recving=set(),
        kv_connector_worker_meta=SimpleCPUOffloadWorkerMetadata(
            completed_store_events={event_idx: scheduler._expected_worker_count},
        ),
    )
    scheduler.update_connector_output(output)


def simulate_load_completion(
    scheduler: SimpleCPUOffloadScheduler,
    req_ids: set[str],
) -> None:
    """Simulate worker reporting load completions for requests."""
    output = KVConnectorOutput(
        finished_sending=set(),
        finished_recving=req_ids,
    )
    scheduler.update_connector_output(output)


def get_cpu_free_blocks(scheduler: SimpleCPUOffloadScheduler) -> int:
    """Return number of free CPU blocks."""
    return scheduler.cpu_block_pool.get_num_free_blocks()


def _allocate_gpu_blocks(
    gpu_block_pool: BlockPool,
    request: Request,
    num_blocks: int,
    group_id: int = 0,
) -> list:
    """Allocate GPU blocks, cache them with hashes, return block list.

    Mimics what KVCacheManager does: allocate blocks from pool, then
    register them in the prefix cache via cache_full_blocks so that
    re-allocation properly evicts stale hashes.
    """
    blocks = gpu_block_pool.get_new_blocks(num_blocks)
    num_full = min(num_blocks, len(request.block_hashes))
    if num_full > 0:
        gpu_block_pool.cache_full_blocks(
            request=request,
            blocks=blocks,
            num_cached_blocks=0,
            num_full_blocks=num_full,
            block_size=BLOCK_SIZE,
            kv_cache_group_id=group_id,
        )
    return blocks


def _alloc_and_register(
    fix: SchedulerFixture,
    request: Request,
    num_blocks: int,
    *,
    confirmed: bool = True,
    group_id: int = 0,
) -> KVCacheBlocks:
    """Allocate GPU blocks and return KVCacheBlocks.

    Block IDs are no longer registered in a mock KVCacheManager; instead
    tests pass them through ``make_scheduler_output`` so that
    ``yield_req_data`` can pick them up.

    If ``confirmed`` is True, advance ``request.num_computed_tokens`` to simulate
    the scheduler's ``_update_after_schedule`` from a prior step.
    """
    gpu_blocks = _allocate_gpu_blocks(
        fix.gpu_block_pool, request, num_blocks, group_id=group_id
    )
    kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
    if confirmed:
        request.num_computed_tokens = num_blocks * BLOCK_SIZE
    return kv_blocks


# ---------------------------------------------------------------------------
# Test 1a: Eager store-and-load roundtrip
# ---------------------------------------------------------------------------
def test_eager_store_and_load_roundtrip() -> None:
    """Eager mode: store blocks on compute, complete store, verify cache hit."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)

    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )

    meta = sched.build_connector_meta(sched_out)
    assert meta.store_event >= 0, "Expected a store event to be scheduled"
    assert len(meta.store_gpu_blocks) > 0
    assert len(meta.store_cpu_blocks) == len(meta.store_gpu_blocks)
    simulate_store_completion(sched, meta.store_event)

    # New request with same tokens should get CPU cache hit
    req2 = Request(
        request_id="req-eager-load",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    # make_request pads num_tokens by +1 beyond the last full block, so the
    # manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.
    assert hit_tokens == num_blocks * BLOCK_SIZE
    assert is_async is True

    gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

    block_ids2 = kv_blocks2.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req2.request_id: 1},
        new_reqs={req2.request_id: block_ids2},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.load_event >= 0, "Expected a load event to be assigned"
    assert len(meta2.load_gpu_blocks) > 0
    assert len(meta2.load_cpu_blocks) == len(meta2.load_gpu_blocks)


# ---------------------------------------------------------------------------
# Test 1b: Boundary — max_hit_len cap drops the last full block when the
# prompt is an exact multiple of BLOCK_SIZE.
# ---------------------------------------------------------------------------
def test_max_hit_len_cap_drops_last_full_block() -> None:
    """When num_tokens is an exact multiple of BLOCK_SIZE, the manager's
    ``max_hit_len = num_tokens - 1`` cap forces ``find_longest_cache_hit`` to
    drop the final block (since ``max_length // block_size`` rounds down).
    """
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 2
    req = make_request(num_blocks=num_blocks, extra_tokens=0)
    assert req.num_tokens == num_blocks * BLOCK_SIZE

    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: kv_blocks.get_block_ids()},
    )
    meta = sched.build_connector_meta(sched_out)
    simulate_store_completion(sched, meta.store_event)

    req2 = Request(
        request_id="req-cap-boundary",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    assert hit_tokens == (num_blocks - 1) * BLOCK_SIZE


# ---------------------------------------------------------------------------
# Test 1c: Lazy store-and-load roundtrip
# ---------------------------------------------------------------------------
def _flush_old_blocks_to_lru_head(
    gpu_pool: BlockPool,
    num_filler_blocks: int,
) -> list:
    """Allocate filler blocks so that previously-freed (hashed) blocks migrate
    to the LRU head of the free queue.  Returns the filler blocks (caller must
    free them later to restore pool capacity).

    In a real engine the same thing happens naturally: after one request
    finishes and frees its blocks, subsequent requests allocate from the LRU
    head, consuming the unhashed blocks and leaving the old hashed blocks at
    the front of the queue.
    """
    fillers = gpu_pool.get_new_blocks(num_filler_blocks)
    return fillers


def test_lazy_store_and_load_roundtrip() -> None:
    """Lazy mode: schedule a request, finish it so its hashed blocks are freed,
    then schedule new requests so the old blocks migrate to the LRU head.
    The lazy scanner offloads them to CPU.  Re-scheduling the old request
    triggers a CPU cache hit + load.

    GPU pool: 8 blocks (7 usable).  _target_free = ceil(64/16) = 4.
    """
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True)
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool

    num_blocks = 2

    # --- Step 1: Schedule req_old, compute, and finish ---
    req_old = make_request(num_blocks=num_blocks)
    gpu_blocks_old = _allocate_gpu_blocks(gpu_pool, req_old, num_blocks, group_id=0)
    gpu_pool.free_blocks(gpu_blocks_old)

    # Allocate filler blocks so req_old's hashed blocks move to LRU head.
    # 7 usable - 2 (req_old freed) = 5 other free blocks to consume.
    fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)

    # --- Step 2: Lazy scanner should offload req_old's blocks ---
    sched_out = make_scheduler_output({})
    meta = sched.build_connector_meta(sched_out)
    assert meta.store_event >= 0, "Expected lazy store to offload old blocks"
    assert len(meta.store_gpu_blocks) == num_blocks
    simulate_store_completion(sched, meta.store_event)

    # Free fillers to restore pool capacity.
    gpu_pool.free_blocks(fillers)

    # --- Step 3: Re-schedule req_old — should get CPU cache hit ---
    req_old2 = Request(
        request_id="req-old-reload",
        prompt_token_ids=req_old.prompt_token_ids,
        sampling_params=req_old.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req_old._block_hasher,
    )
    hit_tokens, is_async = sched.get_num_new_matched_tokens(
        req_old2, num_computed_tokens=0
    )
    # make_request pads num_tokens by +1 beyond the last full block, so the
    # manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.
    expected_hit = num_blocks * BLOCK_SIZE
    assert hit_tokens == expected_hit, (
        f"Expected {expected_hit} hit tokens, got {hit_tokens}"
    )
    assert is_async is True

    # Allocate fresh GPU blocks for the load.
    gpu_blocks_load = gpu_pool.get_new_blocks(num_blocks)
    kv_blocks_load = KVCacheBlocks(blocks=(gpu_blocks_load,))
    sched.update_state_after_alloc(
        req_old2, kv_blocks_load, num_external_tokens=hit_tokens
    )

    sched_out2 = make_scheduler_output({req_old2.request_id: 1})
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.load_event >= 0, "Expected a load event to be assigned"
    assert len(meta2.load_gpu_blocks) > 0


# ---------------------------------------------------------------------------
# Test 2a: Eager duplicate store is skipped
# ---------------------------------------------------------------------------
def test_eager_duplicate_store_skipped() -> None:
    """Eager: storing the same block hashes twice should not allocate new CPU blocks."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)

    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )

    meta1 = sched.build_connector_meta(sched_out)
    assert meta1.store_event >= 0
    simulate_store_completion(sched, meta1.store_event)
    cpu_free_after_first = get_cpu_free_blocks(sched)

    # Second request with identical hashes — should skip store
    req2 = Request(
        request_id="req-dup-eager",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    kv_blocks2 = _alloc_and_register(fix, req2, num_blocks)
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=0)
    block_ids2 = kv_blocks2.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req2.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req2.request_id: block_ids2},
    )

    meta2 = sched.build_connector_meta(sched_out2)
    if meta2.store_event >= 0:
        assert len(meta2.store_cpu_blocks) == 0, (
            "Expected no new CPU blocks for duplicate hashes"
        )
    assert get_cpu_free_blocks(sched) == cpu_free_after_first


# ---------------------------------------------------------------------------
# Test 2b: Eager dedup of in-flight stores across consecutive steps
# ---------------------------------------------------------------------------
def test_eager_in_flight_store_dedup_across_steps() -> None:
    """Eager: a second request sharing a prefix with an in-flight store
    must not re-offload the same GPU blocks before completion lands.

    Simulates a GPU prefix-cache hit by reusing the first request's
    GPU block IDs in the second scheduler step, which is the path the
    real scheduler takes when two requests share a prefix.
    """
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)

    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )

    meta1 = sched.build_connector_meta(sched_out)
    assert meta1.store_event >= 0
    assert len(meta1.store_cpu_blocks) == num_blocks
    # In-flight set tracks the scheduled GPU blocks until completion.
    assert sched._in_flight_store_gpu_blocks == set(meta1.store_gpu_blocks)
    cpu_free_after_first = get_cpu_free_blocks(sched)

    # Second request shares the prefix and reuses the same GPU block IDs
    # (the real scheduler path: GPU prefix cache returns the same blocks).
    # Do NOT simulate completion — the first store is still in-flight.
    req2 = Request(
        request_id="req-dup-eager-inflight",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    req2.num_computed_tokens = num_blocks * BLOCK_SIZE
    sched.update_state_after_alloc(req2, kv_blocks, num_external_tokens=0)
    sched_out2 = make_scheduler_output(
        {req2.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req2.request_id: block_ids},
    )

    meta2 = sched.build_connector_meta(sched_out2)
    if meta2.store_event >= 0:
        assert len(meta2.store_cpu_blocks) == 0, (
            "Expected no new CPU blocks for in-flight duplicate hashes"
        )
    assert get_cpu_free_blocks(sched) == cpu_free_after_first, (
        "Second request should not consume CPU blocks while the first "
        "store is still in-flight"
    )

    # After completion, the in-flight set is cleared.
    simulate_store_completion(sched, meta1.store_event)
    assert sched._in_flight_store_gpu_blocks == set()


# ---------------------------------------------------------------------------
# Test 2c: Lazy duplicate store is skipped
# ---------------------------------------------------------------------------
def test_lazy_duplicate_store_skipped() -> None:
    """Lazy: blocks already offloaded to CPU should not be offloaded again.

    Same pattern as the lazy roundtrip: flush old blocks to LRU head, offload,
    then repeat with the same hashes and verify no new CPU allocation.
    """
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True)
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)

    # Schedule + finish → hashed blocks in free queue
    gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
    gpu_pool.free_blocks(gpu_blocks)

    # Flush old blocks to LRU head, then trigger lazy offload.
    fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)
    meta1 = sched.build_connector_meta(make_scheduler_output({}))
    assert meta1.store_event >= 0
    simulate_store_completion(sched, meta1.store_event)
    gpu_pool.free_blocks(fillers)
    cpu_free_after_first = get_cpu_free_blocks(sched)

    # Allocate blocks with the same hashes and free them again.
    # The scanner should see they are already in CPU cache and skip them.
    req2 = Request(
        request_id="req-dup-lazy",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    gpu_blocks2 = _allocate_gpu_blocks(gpu_pool, req2, num_blocks, group_id=0)
    gpu_pool.free_blocks(gpu_blocks2)

    # Flush again so the hashed blocks are at LRU head for the scanner.
    fillers2 = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)
    meta2 = sched.build_connector_meta(make_scheduler_output({}))
    gpu_pool.free_blocks(fillers2)

    # Either no store event, or zero new CPU blocks (already cached).
    if meta2.store_event >= 0:
        assert len(meta2.store_cpu_blocks) == 0, (
            "Expected no new CPU blocks for duplicate hashes"
        )
    assert get_cpu_free_blocks(sched) == cpu_free_after_first


# ---------------------------------------------------------------------------
# Test 3: LRU eviction order
# ---------------------------------------------------------------------------
def test_lru_eviction_order() -> None:
    """With limited CPU space, oldest blocks should be evicted first.

    CPU block pool: num_cpu_blocks=5 -> 4 free usable blocks (1 taken by null_block).
    After storing 4 blocks (2 req_a + 2 req_b), all free slots are occupied by
    cached blocks (ref_cnt=0, in hash map).  When 2 more are stored (req_c),
    2 LRU blocks from req_a get evicted from the cache to make room.
    """
    # 5 total = 4 usable (null_block takes 1), filling exactly with 4 blocks
    fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    # Fill CPU with 4 blocks: 2 requests x 2 blocks (in LRU insertion order)
    req_a = make_request(num_blocks=2)
    req_b = make_request(num_blocks=2)

    kv_a = _alloc_and_register(fix, req_a, 2)
    kv_b = _alloc_and_register(fix, req_b, 2)
    sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)
    sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0)

    ids_a = kv_a.get_block_ids()
    ids_b = kv_b.get_block_ids()
    sched_out = make_scheduler_output(
        {
            req_a.request_id: 2 * BLOCK_SIZE,
            req_b.request_id: 2 * BLOCK_SIZE,
        },
        new_reqs={
            req_a.request_id: ids_a,
            req_b.request_id: ids_b,
        },
    )
    meta = sched.build_connector_meta(sched_out)
    assert meta.store_event >= 0
    simulate_store_completion(sched, meta.store_event)

    # Verify all 4 blocks are cached in CPU hash map
    for i, bhash in enumerate(req_a.block_hashes[:2]):
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        assert (
            sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
                bhash_with_group
            )
            is not None
        ), f"req_a block {i} should be cached after store"
    for i, bhash in enumerate(req_b.block_hashes[:2]):
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        assert (
            sched.cpu_block_pool.cached_block_hash_to_block.get_one_block(
                bhash_with_group
            )
            is not None
        ), f"req_b block {i} should be cached after store"

    # Store 2 more blocks from a new request - must evict 2 LRU blocks (req_a)
    req_c = make_request(num_blocks=2)
    kv_c = _alloc_and_register(fix, req_c, 2)
    sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)

    ids_c = kv_c.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req_c.request_id: 2 * BLOCK_SIZE},
        new_reqs={req_c.request_id: ids_c},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.store_event >= 0
    simulate_store_completion(sched, meta2.store_event)

    # req_a hashes should be evicted from CPU (they were LRU)
    for i, bhash in enumerate(req_a.block_hashes[:2]):
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        cache_map = sched.cpu_block_pool.cached_block_hash_to_block
        cached = cache_map.get_one_block(bhash_with_group)
        assert cached is None, f"req_a block {i} should have been evicted"

    # req_b and req_c hashes should be present
    for i, bhash in enumerate(req_b.block_hashes[:2]):
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        cache_map = sched.cpu_block_pool.cached_block_hash_to_block
        cached = cache_map.get_one_block(bhash_with_group)
        assert cached is not None, f"req_b block {i} should still be cached"

    for i, bhash in enumerate(req_c.block_hashes[:2]):
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        cache_map = sched.cpu_block_pool.cached_block_hash_to_block
        cached = cache_map.get_one_block(bhash_with_group)
        assert cached is not None, f"req_c block {i} should still be cached"


# ---------------------------------------------------------------------------
# Test 4: Touched blocks survive eviction
# ---------------------------------------------------------------------------
def test_touched_blocks_survive_eviction() -> None:
    """Touching CPU blocks updates their LRU position, protecting them from eviction."""
    # 5 total = 4 usable (null_block takes 1)
    fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    # Fill CPU with 4 blocks (req_a: 2, req_b: 2) in LRU order
    req_a = make_request(num_blocks=2)
    req_b = make_request(num_blocks=2)

    kv_a = _alloc_and_register(fix, req_a, 2)
    kv_b = _alloc_and_register(fix, req_b, 2)
    sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)
    sched.update_state_after_alloc(req_b, kv_b, num_external_tokens=0)

    ids_a = kv_a.get_block_ids()
    ids_b = kv_b.get_block_ids()
    sched_out = make_scheduler_output(
        {
            req_a.request_id: 2 * BLOCK_SIZE,
            req_b.request_id: 2 * BLOCK_SIZE,
        },
        new_reqs={
            req_a.request_id: ids_a,
            req_b.request_id: ids_b,
        },
    )
    meta = sched.build_connector_meta(sched_out)
    simulate_store_completion(sched, meta.store_event)

    # Touch req_a's CPU blocks to make them most-recently-used
    cpu_pool = sched.cpu_block_pool
    for bhash in req_a.block_hashes[:2]:
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        cached_blk = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group)
        assert cached_blk is not None
        cpu_pool.touch([cached_blk])
        # Undo touch to return ref_cnt to 0
        # (so it's a free candidate but at MRU position)
        cpu_pool.free_blocks([cached_blk])

    # Now store 2 more blocks; req_b (LRU front) should be evicted, not req_a
    req_c = make_request(num_blocks=2)
    kv_c = _alloc_and_register(fix, req_c, 2)
    sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)

    ids_c = kv_c.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req_c.request_id: 2 * BLOCK_SIZE},
        new_reqs={req_c.request_id: ids_c},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    simulate_store_completion(sched, meta2.store_event)

    # req_b should be evicted (LRU), req_a and req_c should survive
    for i, bhash in enumerate(req_b.block_hashes[:2]):
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group)
        assert cached is None, f"req_b block {i} should have been evicted (it was LRU)"

    for i, bhash in enumerate(req_a.block_hashes[:2]):
        bhash_with_group = make_block_hash_with_group_id(bhash, 0)
        cached = cpu_pool.cached_block_hash_to_block.get_one_block(bhash_with_group)
        assert cached is not None, f"req_a block {i} should survive (was touched/MRU)"


# ---------------------------------------------------------------------------
# Test 5: Preemption no CPU block leak
# ---------------------------------------------------------------------------
def test_preemption_no_cpu_block_leak() -> None:
    """request_finished during in-flight load defers cleanup;
    completes after load done."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 2

    # First: store blocks to CPU
    req = make_request(num_blocks=num_blocks)
    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )
    meta = sched.build_connector_meta(sched_out)
    simulate_store_completion(sched, meta.store_event)

    # Create new request with same tokens, check hit
    req2 = Request(
        request_id="req-preempt-load",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    assert hit_tokens > 0

    gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

    # Assign load_event via build_connector_meta
    block_ids2 = kv_blocks2.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req2.request_id: 1},
        new_reqs={req2.request_id: block_ids2},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.load_event >= 0

    # Request finishes BEFORE load completes -> deferred
    sched.request_finished(req2, block_ids=[])
    assert req2.request_id in sched._reqs_to_load
    assert sched._reqs_to_load[req2.request_id].finished is True

    # Now simulate load completion -> cleanup fires
    simulate_load_completion(sched, {req2.request_id})
    assert req2.request_id not in sched._reqs_to_load


# ---------------------------------------------------------------------------
# Test 6: Eager store preemption cleanup
# ---------------------------------------------------------------------------
def test_eager_store_preemption_cleanup() -> None:
    """In eager mode, finishing a request during in-flight store defers cleanup."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)
    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )
    meta = sched.build_connector_meta(sched_out)
    store_event = meta.store_event
    assert store_event >= 0

    # The request gets store_events populated
    assert req.request_id in sched._reqs_to_store
    store_state = sched._reqs_to_store[req.request_id]
    assert store_event in store_state.store_events

    # Finish request while store still in-flight -> deferred
    sched.request_finished(req, block_ids=[])
    assert req.request_id in sched._reqs_to_store
    assert sched._reqs_to_store[req.request_id].finished is True

    # Simulate store completion -> deferred cleanup fires
    simulate_store_completion(sched, store_event)
    assert req.request_id not in sched._reqs_to_store


# ---------------------------------------------------------------------------
# Test 7: In-flight finish deferred cleanup (load variant)
# ---------------------------------------------------------------------------
def test_inflight_finish_deferred_cleanup() -> None:
    """Store, then start a load, request_finished defers,
    load completion fires cleanup."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 2

    # Store
    req = make_request(num_blocks=num_blocks)
    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )
    meta = sched.build_connector_meta(sched_out)
    simulate_store_completion(sched, meta.store_event)

    # Load
    req2 = Request(
        request_id="req-inflight-load",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    assert hit_tokens > 0

    gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

    block_ids2 = kv_blocks2.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req2.request_id: 1},
        new_reqs={req2.request_id: block_ids2},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.load_event >= 0

    # Finish before load completes
    sched.request_finished(req2, block_ids=[])
    assert req2.request_id in sched._reqs_to_load

    # Simulate load completion -> request removed
    simulate_load_completion(sched, {req2.request_id})
    assert req2.request_id not in sched._reqs_to_load


# ---------------------------------------------------------------------------
# Test 8: Null GPU blocks are skipped in store and load transfer pairs
# ---------------------------------------------------------------------------
def test_multi_group_null_blocks_skipped() -> None:
    """Null GPU blocks (no block_hash) must not appear in store or load pairs.

    In eager store mode, _prepare_eager_store_specs skips blocks whose
    block_hash is None (null blocks have no hash). We verify this by mixing
    real hashed blocks with unhashed (null-like) blocks in a single group and
    checking that only real blocks appear in the store list.
    """
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, num_groups=1, lazy=False)
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)

    # Allocate real blocks (with hashes) and use the null_block as a placeholder
    gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
    null_block = gpu_pool.null_block

    # Mix: [real_block, null_block] — null_block has no hash, should be skipped
    mixed_blocks = [gpu_blocks[0], null_block]
    kv_blocks = KVCacheBlocks(blocks=(mixed_blocks,))
    req.num_computed_tokens = num_blocks * BLOCK_SIZE
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )
    meta = sched.build_connector_meta(sched_out)

    # Null block's ID should NOT appear in store_gpu_blocks
    null_block_id = null_block.block_id
    assert null_block_id not in meta.store_gpu_blocks, (
        f"Null block id {null_block_id} should not appear in store transfer pairs"
    )

    # Only real block should be scheduled for store
    assert len(meta.store_gpu_blocks) == 1
    assert gpu_blocks[0].block_id in meta.store_gpu_blocks

    # Complete the store
    assert meta.store_event >= 0
    simulate_store_completion(sched, meta.store_event)

    # Create matching request and get load hit
    req2 = Request(
        request_id="req-null-load",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    # Only 1 block was stored (the real one)
    assert hit_tokens == BLOCK_SIZE
    assert is_async is True

    # Allocate new GPU blocks for the load
    gpu_blocks2 = gpu_pool.get_new_blocks(1)
    kv_blocks2 = KVCacheBlocks(blocks=([gpu_blocks2[0], null_block],))
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

    sched_out2 = make_scheduler_output({req2.request_id: 1})
    meta2 = sched.build_connector_meta(sched_out2)

    # Null block's ID should NOT appear in load_gpu_blocks
    assert null_block_id not in meta2.load_gpu_blocks, (
        f"Null block id {null_block_id} should not appear in load transfer pairs"
    )


# ---------------------------------------------------------------------------
# Test 9: Chunked prefill accumulates block_ids across steps
# ---------------------------------------------------------------------------
def test_chunked_prefill_reads_live_block_ids() -> None:
    """With chunked prefill, block IDs accumulate across scheduler steps.
    _prepare_eager_store_specs reads block IDs from scheduler_output via
    yield_req_data, so the store should reflect the updated (larger) block
    list, not a stale snapshot."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    num_blocks = 4
    req = make_request(num_blocks=num_blocks)

    # First chunk: allocate 2 blocks
    kv_blocks_first = _alloc_and_register(fix, req, 2)
    sched.update_state_after_alloc(req, kv_blocks_first, num_external_tokens=0)

    assert req.request_id in sched._reqs_to_store
    # Should still be exactly 1 entry in _reqs_to_store
    assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1

    # Build connector meta with 2 blocks — stores the first 2
    ids_first = kv_blocks_first.get_block_ids()
    sched_out1 = make_scheduler_output(
        {req.request_id: 2 * BLOCK_SIZE},
        new_reqs={req.request_id: ids_first},
    )
    meta1 = sched.build_connector_meta(sched_out1)
    assert meta1.store_event >= 0
    assert len(meta1.store_gpu_blocks) == 2
    simulate_store_completion(sched, meta1.store_event)

    # Second chunk: allocate 4 blocks total (2 new ones)
    kv_blocks_second = _alloc_and_register(fix, req, num_blocks)
    # update_state_after_alloc is idempotent for store registration
    sched.update_state_after_alloc(req, kv_blocks_second, num_external_tokens=0)

    # Still exactly 1 entry
    assert list(sched._reqs_to_store.keys()).count(req.request_id) == 1

    # The second chunk's NEW block IDs (positions 2,3) are passed as
    # cached_req_new_blocks. The full block_ids include both old and new,
    # but yield_req_data only appends the new_block_ids for cached reqs.
    ids_second_full = kv_blocks_second.get_block_ids()
    # New blocks are those beyond the first chunk
    new_block_ids = tuple(ids_second_full[g][2:] for g in range(len(ids_second_full)))
    sched_out2 = make_scheduler_output(
        {req.request_id: 2 * BLOCK_SIZE},
        cached_req_new_blocks={req.request_id: new_block_ids},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.store_event >= 0
    # Only the 2 NEW blocks should be stored (first 2 already done)
    assert len(meta2.store_gpu_blocks) == 2


# ---------------------------------------------------------------------------
# Test 10: Partial GPU prefix hit + CPU load + new compute blocks
# ---------------------------------------------------------------------------
def test_partial_gpu_prefix_plus_cpu_load() -> None:
    """When GPU has a prefix cache hit for the first N blocks, CPU has a
    hit for the next M blocks, and there are P new blocks needing fresh
    compute, the block layout is:

        | comp (N) | ext_comp (M) | new (P) |

    External blocks sit in the middle — not at the beginning or end.
    The load path must target hashes at positions [N, N+M).

    Request: 6 blocks (0..5).
    - Store all 6 to CPU.
    - New request: GPU prefix cache hits blocks 0,1 (hashed).
      CPU hits blocks 2,3. Blocks 4,5 are new (need compute).
    - update_state_after_alloc receives 6 GPU blocks:
      [0,1] hashed (comp), [2,3] unhashed (ext_comp), [4,5] unhashed (new).
    - Load must target hash positions 2,3.
    """
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool

    num_blocks = 6
    req = make_request(num_blocks=num_blocks)

    # Store all 6 blocks to CPU via eager store.
    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )
    meta = sched.build_connector_meta(sched_out)
    assert meta.store_event >= 0
    simulate_store_completion(sched, meta.store_event)

    # New request with same tokens — but only partial GPU prefix hit.
    req2 = Request(
        request_id="req-partial-gpu",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )

    # GPU prefix cache hits the first 2 blocks.
    gpu_local_computed = 2 * BLOCK_SIZE
    hit_tokens, is_async = sched.get_num_new_matched_tokens(
        req2, num_computed_tokens=gpu_local_computed
    )
    # CPU has all 6 blocks stored. make_request pads num_tokens by +1, so
    # the manager's num_tokens - 1 cap leaves all full blocks intact:
    # remaining hashable range = 6 - 2 = 4 blocks, all hit.
    num_cpu_hit_blocks = 4
    assert hit_tokens == num_cpu_hit_blocks * BLOCK_SIZE, (
        f"Expected {num_cpu_hit_blocks * BLOCK_SIZE} CPU hit tokens, got {hit_tokens}"
    )
    assert is_async is True

    # Simulate what the real scheduler does: only accept 2 of the 4 CPU hit
    # blocks as external (e.g. due to budget constraints), leaving 2 new
    # blocks for fresh compute.
    num_ext_blocks = 2
    num_new_blocks = 2
    external_tokens = num_ext_blocks * BLOCK_SIZE

    # Build block list matching real layout: | comp(2) | ext_comp(2) | new(2) |
    # comp: GPU prefix cache hit — blocks with hashes
    gpu_comp = _allocate_gpu_blocks(gpu_pool, req2, 2, group_id=0)
    # ext_comp + new: freshly allocated, no hashes
    gpu_ext_and_new = gpu_pool.get_new_blocks(num_ext_blocks + num_new_blocks)
    all_gpu_blocks = gpu_comp + gpu_ext_and_new
    kv_blocks2 = KVCacheBlocks(blocks=(all_gpu_blocks,))

    # Critical call: with 2 hashed comp blocks and 2 external tokens worth
    # of blocks, the manager must derive skipped=2 and load hashes [2,3].
    sched.update_state_after_alloc(
        req2, kv_blocks2, num_external_tokens=external_tokens
    )

    block_ids2 = kv_blocks2.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req2.request_id: num_new_blocks * BLOCK_SIZE},
        new_reqs={req2.request_id: block_ids2},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.load_event >= 0, "Expected a load event for partial GPU + CPU hit"
    assert len(meta2.load_gpu_blocks) == num_ext_blocks
    assert len(meta2.load_cpu_blocks) == num_ext_blocks

    # Verify the load targets the ext_comp GPU blocks (positions 2,3),
    # not the comp blocks (0,1) or new blocks (4,5).
    ext_block_ids = [b.block_id for b in gpu_ext_and_new[:num_ext_blocks]]
    for bid in meta2.load_gpu_blocks:
        assert bid in ext_block_ids, (
            f"Load GPU block {bid} should be an ext_comp block, not a comp or new block"
        )


# ---------------------------------------------------------------------------
# Test 11: TOCTOU between Phase A and Phase B (regression for #39702)
# ---------------------------------------------------------------------------
def test_toctou_cpu_hit_evicted_between_phases_no_crash() -> None:
    """Regression for vllm-project/vllm#39702.

    When ``get_num_new_matched_tokens`` (Phase A) reports a CPU cache hit
    of ``N`` tokens but ``update_state_after_alloc`` (Phase B) runs after
    other requests have caused LRU eviction of those exact blocks, the
    second ``find_longest_cache_hit`` call returns 0 while
    ``num_external_tokens`` is still ``N``, triggering
    ``AssertionError: Expected N hit tokens, got 0``.

    Setup: ``num_cpu_blocks=5`` (4 usable; null_block takes 1).
        1. Store req_a's 2 blocks to CPU. CPU: [a0, a1, _, _].
        2. Phase A on req_b (same prompt as req_a) reports a 2-block hit.
           Without the fix this does NOT pin a0/a1 — they remain at LRU front.
        3. Store req_c (2) + req_d (2) — 4 more blocks into 4 slots. With
           a0/a1 unpinned and at LRU front, they get evicted. CPU: [c0, c1,
           d0, d1].
        4. Phase B on req_b: re-searches the CPU coordinator, finds 0
           cached hashes, asserts.

    After the fix, Phase A pins the hit blocks so step 3 evicts req_c's
    blocks instead, and Phase B reads the cached ``(cpu_hit_blocks,
    hit_length)`` tuple without re-searching.
    """
    # 5 total = 4 usable (null_block takes 1)
    fix = make_scheduler(num_cpu_blocks=5, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler

    # --- Step 1: Store req_a's 2 blocks to CPU cache ---
    req_a = make_request(num_blocks=2)
    kv_a = _alloc_and_register(fix, req_a, 2)
    sched.update_state_after_alloc(req_a, kv_a, num_external_tokens=0)
    sched_out_a = make_scheduler_output(
        {req_a.request_id: 2 * BLOCK_SIZE},
        new_reqs={req_a.request_id: kv_a.get_block_ids()},
    )
    meta_a = sched.build_connector_meta(sched_out_a)
    assert meta_a.store_event >= 0
    simulate_store_completion(sched, meta_a.store_event)

    # --- Step 2: Phase A — req_b (same prompt as req_a) reports a CPU hit ---
    req_b = Request(
        request_id="req-b-toctou",
        prompt_token_ids=req_a.prompt_token_ids,
        sampling_params=req_a.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req_a._block_hasher,
    )
    hit_tokens, is_async = sched.get_num_new_matched_tokens(
        req_b, num_computed_tokens=0
    )
    assert hit_tokens == 2 * BLOCK_SIZE, (
        f"Phase A should report 2 blocks of CPU hit, got {hit_tokens}"
    )
    assert is_async is True

    # --- Step 3: TOCTOU window — fill CPU cache so LRU evicts req_a's blocks
    # (in production this corresponds to other concurrent requests landing
    # between Phase A and Phase B for req_b). 4 usable slots, req_a occupies 2;
    # req_c (2) + req_d (2) require evicting 2 LRU blocks. ---
    req_c = make_request(num_blocks=2)
    req_d = make_request(num_blocks=2)
    kv_c = _alloc_and_register(fix, req_c, 2)
    kv_d = _alloc_and_register(fix, req_d, 2)
    sched.update_state_after_alloc(req_c, kv_c, num_external_tokens=0)
    sched.update_state_after_alloc(req_d, kv_d, num_external_tokens=0)
    sched_out_pressure = make_scheduler_output(
        {
            req_c.request_id: 2 * BLOCK_SIZE,
            req_d.request_id: 2 * BLOCK_SIZE,
        },
        new_reqs={
            req_c.request_id: kv_c.get_block_ids(),
            req_d.request_id: kv_d.get_block_ids(),
        },
    )
    meta_pressure = sched.build_connector_meta(sched_out_pressure)
    assert meta_pressure.store_event >= 0
    simulate_store_completion(sched, meta_pressure.store_event)

    # --- Step 4: Phase B — must not crash ---
    # Before fix: AssertionError: Expected 32 hit tokens, got 0
    # After fix: Phase A pinned the hits → Step 3 evicted req_c instead,
    # Phase B consumes the cached (cpu_hit_blocks, hit_length) tuple.
    gpu_blocks_b = fix.gpu_block_pool.get_new_blocks(2)
    kv_blocks_b = KVCacheBlocks(blocks=(gpu_blocks_b,))
    sched.update_state_after_alloc(req_b, kv_blocks_b, num_external_tokens=hit_tokens)

    # --- Step 5: with the fix, the load is queued correctly ---
    sched_out_b = make_scheduler_output(
        {req_b.request_id: 1},
        new_reqs={req_b.request_id: kv_blocks_b.get_block_ids()},
    )
    meta_b = sched.build_connector_meta(sched_out_b)
    assert meta_b.load_event >= 0, (
        "Phase B should queue a load event using the pinned CPU hit blocks"
    )
    assert len(meta_b.load_gpu_blocks) == 2
    assert len(meta_b.load_cpu_blocks) == 2


# ---------------------------------------------------------------------------
# Test 12: Reset with pending eager stores waits for completion
# ---------------------------------------------------------------------------
def test_reset_pending_eager_stores() -> None:
    """Eager mode: reset() abandons in-flight stores until they complete."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)

    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )

    meta = sched.build_connector_meta(sched_out)
    assert meta.store_event >= 0
    assert len(sched._store_event_to_blocks) > 0

    # GPU blocks should have elevated ref_cnt from touch()
    for bid in meta.store_gpu_blocks:
        assert gpu_pool.blocks[bid].ref_cnt > 0

    # Free the request's own block refs (simulates preemption)
    gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids[0])

    # Reset should keep DMA refs pinned until the worker reports completion.
    assert sched.reset() is False
    assert len(sched._store_event_to_blocks) == 0
    assert len(sched._abandoned_store_event_to_blocks) == 1
    assert len(sched._reqs_to_store) == 0
    assert len(sched._store_event_to_reqs) == 0

    num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
    assert num_used > 1

    simulate_store_completion(sched, meta.store_event)
    assert len(sched._abandoned_store_event_to_blocks) == 0

    # All GPU blocks should now be free (ref_cnt == 0) except null block
    num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
    assert num_used == 1, f"Expected only null block in use, got {num_used}"

    # GPU prefix cache reset should now succeed
    assert gpu_pool.reset_prefix_cache() is True
    assert sched.reset() is True


# ---------------------------------------------------------------------------
# Test 13: Reset with pending lazy stores waits for completion
# ---------------------------------------------------------------------------
def test_reset_pending_lazy_stores() -> None:
    """Lazy mode: reset() abandons in-flight stores until they complete."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=8, lazy=True)
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool

    num_blocks = 2
    req = make_request(num_blocks=num_blocks)

    # Allocate, hash, and free — blocks move to free queue with hashes
    gpu_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)
    gpu_pool.free_blocks(gpu_blocks)

    # Push hashed blocks to LRU head
    fillers = _flush_old_blocks_to_lru_head(gpu_pool, num_filler_blocks=5)

    # Lazy scanner offloads old hashed blocks
    sched_out = make_scheduler_output({})
    meta = sched.build_connector_meta(sched_out)
    assert meta.store_event >= 0
    assert len(sched._store_event_to_blocks) > 0

    gpu_pool.free_blocks(fillers)

    # Reset should keep DMA refs pinned until the worker reports completion.
    assert sched.reset() is False
    assert len(sched._store_event_to_blocks) == 0
    assert len(sched._abandoned_store_event_to_blocks) == 1
    assert sched._cursor is None

    simulate_store_completion(sched, meta.store_event)
    assert len(sched._abandoned_store_event_to_blocks) == 0
    assert sched.reset() is True

    # No CPU cache hits after reset
    req2 = Request(
        request_id="req-after-lazy-reset",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    assert hit_tokens == 0, "CPU cache should be empty after reset"


# ---------------------------------------------------------------------------
# Test 14: Reset with pending loads waits for completion
# ---------------------------------------------------------------------------
def test_reset_pending_loads() -> None:
    """reset() abandons in-flight loads until they complete."""
    fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool

    num_blocks = 2

    # First store blocks to CPU
    req = make_request(num_blocks=num_blocks)
    kv_blocks = _alloc_and_register(fix, req, num_blocks)
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * BLOCK_SIZE},
        new_reqs={req.request_id: block_ids},
    )
    meta = sched.build_connector_meta(sched_out)
    simulate_store_completion(sched, meta.store_event)

    # Start a load — CPU cache hit
    req2 = Request(
        request_id="req-load-reset",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )
    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    assert hit_tokens > 0

    gpu_blocks2 = gpu_pool.get_new_blocks(num_blocks)
    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

    block_ids2 = kv_blocks2.get_block_ids()
    sched_out2 = make_scheduler_output(
        {req2.request_id: 1},
        new_reqs={req2.request_id: block_ids2},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.load_event >= 0
    assert req2.request_id in sched._reqs_to_load

    # Free request block refs (simulates preemption)
    gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids[0])
    gpu_pool.free_blocks(gpu_pool.blocks[bid] for bid in block_ids2[0])

    # Reset should keep load touch refs until the worker reports completion.
    assert sched.reset() is False
    assert len(sched._reqs_to_load) == 0
    assert len(sched._abandoned_reqs_to_load) == 1
    assert len(sched._load_event_to_reqs) == 1

    num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
    assert num_used > 1

    simulate_load_completion(sched, {req2.request_id})
    assert len(sched._abandoned_reqs_to_load) == 0
    assert len(sched._load_event_to_reqs) == 0
    assert sched.reset() is True

    # All GPU blocks free
    num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
    assert num_used == 1, f"Expected only null block in use, got {num_used}"


def _make_cp_vllm_config(
    dcp_world_size: int = 1,
    pcp_world_size: int = 1,
) -> VllmConfig:
    """VllmConfig with context-parallel sizes set for scheduler-only tests."""
    cfg = _make_vllm_config()

    cfg.parallel_config.decode_context_parallel_size = dcp_world_size
    cfg.parallel_config.prefill_context_parallel_size = pcp_world_size
    return cfg


def _make_cp_scheduler(
    *,
    dcp_world_size: int = 1,
    pcp_world_size: int = 1,
    num_cpu_blocks: int = 8,
    num_gpu_blocks: int = 16,
    lazy: bool = False,
) -> SchedulerFixture:
    """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size."""
    cp_world_size = dcp_world_size * pcp_world_size
    virtual_block_size = BLOCK_SIZE * cp_world_size

    kv_cache_config = _make_kv_cache_config(num_gpu_blocks)
    vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size)
    cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks

    sched = SimpleCPUOffloadScheduler(
        vllm_config=vllm_config,
        kv_cache_config=kv_cache_config,
        cpu_capacity_bytes=cpu_capacity_bytes,
        scheduler_block_size=virtual_block_size,
        hash_block_size=virtual_block_size,
        lazy_offload=lazy,
    )

    gpu_block_pool = BlockPool(
        num_gpu_blocks=num_gpu_blocks,
        enable_caching=True,
        hash_block_size=virtual_block_size,
    )
    sched.bind_gpu_block_pool(gpu_block_pool)

    return SchedulerFixture(
        scheduler=sched,
        gpu_block_pool=gpu_block_pool,
        vllm_config=vllm_config,
        kv_cache_config=kv_cache_config,
    )


def _make_cp_request(
    num_blocks: int,
    virtual_block_size: int,
    request_id: str | None = None,
) -> Request:
    """Create a request whose block hashes are computed at the virtual
    (CP-scaled) block size, matching what the real scheduler does.
    """
    global _req_counter
    _req_counter += 1
    if request_id is None:
        request_id = f"req-cp-{_req_counter}"

    num_tokens = num_blocks * virtual_block_size + 1
    start = _req_counter * 10000
    prompt_token_ids = list(range(start, start + num_tokens))
    sampling_params = SamplingParams(max_tokens=1)

    return Request(
        request_id=request_id,
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=get_request_block_hasher(virtual_block_size, sha256),
    )


def _allocate_cp_gpu_blocks(
    gpu_block_pool: BlockPool,
    request: Request,
    num_blocks: int,
    virtual_block_size: int,
    group_id: int = 0,
) -> list:
    """Allocate GPU blocks and cache them using the CP-scaled block size."""
    blocks = gpu_block_pool.get_new_blocks(num_blocks)
    num_full = min(num_blocks, len(request.block_hashes))
    if num_full > 0:
        gpu_block_pool.cache_full_blocks(
            request=request,
            blocks=blocks,
            num_cached_blocks=0,
            num_full_blocks=num_full,
            block_size=virtual_block_size,
            kv_cache_group_id=group_id,
        )
    return blocks


# ---------------------------------------------------------------------------
# Test 15: CP block size scaling is correct
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
    "dcp_world_size, pcp_world_size",
    [
        (2, 1),  # DCP only
        (1, 2),  # PCP only
        (2, 2),  # DCP + PCP
    ],
)
def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None:
    """Verify that the scheduler's block_size and cp_world_size are correctly
    scaled when context parallelism is enabled."""
    fix = _make_cp_scheduler(
        dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
    )
    sched = fix.scheduler

    expected_cp = dcp_world_size * pcp_world_size
    assert sched.cp_world_size == expected_cp
    assert sched.block_size == BLOCK_SIZE * expected_cp
    assert sched.fa_block_size == BLOCK_SIZE * expected_cp


# ---------------------------------------------------------------------------
# Test 16: CP eager store-and-load roundtrip
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
    "dcp_world_size, pcp_world_size",
    [
        (2, 1),
        (1, 2),
    ],
)
def test_cp_eager_store_and_load_roundtrip(
    dcp_world_size: int, pcp_world_size: int
) -> None:
    """With CP enabled, store blocks to CPU and reload them for a new request
    with matching tokens.  Verifies that hash matching and transfer-pair
    construction work with the virtual block size."""
    fix = _make_cp_scheduler(
        dcp_world_size=dcp_world_size,
        pcp_world_size=pcp_world_size,
        num_cpu_blocks=8,
        num_gpu_blocks=16,
        lazy=False,
    )
    sched = fix.scheduler
    cp = dcp_world_size * pcp_world_size
    vbs = BLOCK_SIZE * cp

    num_blocks = 2
    req = _make_cp_request(num_blocks, vbs)

    # Allocate GPU blocks and register hashes
    gpu_blocks = _allocate_cp_gpu_blocks(fix.gpu_block_pool, req, num_blocks, vbs)
    kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
    req.num_computed_tokens = num_blocks * vbs
    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

    block_ids = kv_blocks.get_block_ids()
    sched_out = make_scheduler_output(
        {req.request_id: num_blocks * vbs},
        new_reqs={req.request_id: block_ids},
    )

    meta = sched.build_connector_meta(sched_out)
    assert meta.store_event >= 0, "Expected a store event"
    assert len(meta.store_gpu_blocks) == num_blocks
    assert len(meta.store_cpu_blocks) == num_blocks
    simulate_store_completion(sched, meta.store_event)

    # New request with same tokens — should get a full CPU cache hit.
    req2 = Request(
        request_id="req-cp-load",
        prompt_token_ids=req.prompt_token_ids,
        sampling_params=req.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req._block_hasher,
    )

    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
    assert hit_tokens == num_blocks * vbs
    assert is_async is True

    # Allocate fresh GPU blocks for the load.
    gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

    sched_out2 = make_scheduler_output(
        {req2.request_id: 1},
        new_reqs={req2.request_id: kv_blocks2.get_block_ids()},
    )
    meta2 = sched.build_connector_meta(sched_out2)
    assert meta2.load_event >= 0, "Expected a load event"
    assert len(meta2.load_gpu_blocks) == num_blocks
    assert len(meta2.load_cpu_blocks) == num_blocks


# ---------------------------------------------------------------------------
# Test 18: CP store and load use effective block size
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
    "dcp_world_size, pcp_world_size",
    [
        (2, 1),
        (1, 2),
        (2, 2),
    ],
)
def test_cp_effective_block_size_store_and_load(
    dcp_world_size: int, pcp_world_size: int
) -> None:
    """Verify ready_blocks_g (store) and n_take_g (load) use the effective
    (physical * cp) block size, not the per-rank physical size."""
    fix = _make_cp_scheduler(
        dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
    )
    sched = fix.scheduler
    gpu_pool = fix.gpu_block_pool
    cp = dcp_world_size * pcp_world_size
    vbs = BLOCK_SIZE * cp

    # Store: allocate 2 blocks, confirm only 1. Without the fix,
    # ready_blocks_g = vbs / BLOCK_SIZE = 2, storing both blocks.
    req = _make_cp_request(num_blocks=2, virtual_block_size=vbs)
    gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs)
    kv = KVCacheBlocks(blocks=(gpu_blocks,))
    req.num_computed_tokens = vbs
    sched.update_state_after_alloc(req, kv, num_external_tokens=0)
    m1 = sched.build_connector_meta(
        make_scheduler_output(
            {req.request_id: vbs},
            new_reqs={req.request_id: kv.get_block_ids()},
        )
    )
    assert len(m1.store_gpu_blocks) == 1
    assert len(m1.store_cpu_blocks) == 1
    simulate_store_completion(sched, m1.store_event)

    # Load: store 2 blocks from a second request, accept only 1 as external.
    # Without the fix, n_take_g = vbs / BLOCK_SIZE = 2, loading both.
    req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs)
    kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),))
    req2.num_computed_tokens = 2 * vbs
    sched.update_state_after_alloc(req2, kv2, num_external_tokens=0)
    m2 = sched.build_connector_meta(
        make_scheduler_output(
            {req2.request_id: 2 * vbs},
            new_reqs={req2.request_id: kv2.get_block_ids()},
        )
    )
    simulate_store_completion(sched, m2.store_event)

    req3 = Request(
        request_id="req-cp-partial-load",
        prompt_token_ids=req2.prompt_token_ids,
        sampling_params=req2.sampling_params,
        pooling_params=None,
        mm_features=None,
        block_hasher=req2._block_hasher,
    )
    hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0)
    assert hit == 2 * vbs

    kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(2),))
    sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs)
    m3 = sched.build_connector_meta(
        make_scheduler_output(
            {req3.request_id: vbs},
            new_reqs={req3.request_id: kv3.get_block_ids()},
        )
    )
    assert m3.load_event >= 0
    assert len(m3.load_gpu_blocks) == 1
    assert len(m3.load_cpu_blocks) == 1
    assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]]


# ---------------------------------------------------------------------------
# Test 17: CP lazy target blocks are scaled correctly
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("cp_world_size", [1, 2, 4])
def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None:
    """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1
    because each virtual block covers more tokens."""
    kv_cache_config = _make_kv_cache_config(num_blocks=16)
    max_batched = 64

    target_base = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
        kv_cache_config, max_batched, cp_world_size=1
    )
    target_cp = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
        kv_cache_config, max_batched, cp_world_size=cp_world_size
    )

    if cp_world_size == 1:
        assert target_cp == target_base
    else:
        assert target_cp < target_base, (
            f"cp_world_size={cp_world_size}: target_cp={target_cp} should be "
            f"less than target_base={target_base}"
        )
