"""
Test script to verify CPU offload correctness using distinctive KV patterns.

Strategy:
1. Hook into attention forward pass
2. Overwrite K/V with distinctive patterns based on chunk_idx (e.g., K=chunk_idx, V=-chunk_idx)
3. After offload to CPU, verify CPU cache contains correct patterns
4. On subsequent chunks, verify loaded KV from CPU has correct patterns

This catches bugs like:
- Wrong block being offloaded
- Wrong block being loaded
- Data corruption during transfer
"""

import os
os.environ["NANOVLLM_LOG_LEVEL"] = "INFO"

import torch
from random import randint, seed
from nanovllm import LLM, SamplingParams
from nanovllm.utils.context import get_context


# ============================================================
# Configuration
# ============================================================

MODEL_PATH = os.path.expanduser("~/models/Qwen3-0.6B/")
MAX_MODEL_LEN = 64 * 1024
NUM_GPU_BLOCKS = 4
INPUT_LEN = 32 * 1024  # 32K tokens = 32 chunks (fits in 40 CPU blocks)
BLOCK_SIZE = 1024

# Test state
errors = []
chunk_patterns = {}  # chunk_idx -> (k_pattern, v_pattern)
block_coverage = {}  # chunk_idx -> set of blocks that were actually computed
load_operations = []  # List of (chunk_idx, slot_id, cpu_block_id, k_ok, v_ok) tuples
current_chunk_for_load = [0]  # Mutable container to track current chunk during loads


# ============================================================
# Pattern Helpers
# ============================================================

def get_expected_pattern(chunk_idx: int):
    """Get expected K/V pattern for a chunk."""
    # Use float values that are easy to identify
    k_val = float(chunk_idx + 1)       # 1.0, 2.0, 3.0, ...
    v_val = float(-(chunk_idx + 1))    # -1.0, -2.0, -3.0, ...
    return k_val, v_val


def fill_with_pattern(tensor: torch.Tensor, value: float):
    """Fill tensor with a constant value."""
    tensor.fill_(value)


def check_pattern(tensor: torch.Tensor, expected: float, name: str, tolerance: float = 1e-3):
    """Check if tensor contains expected pattern."""
    actual_mean = tensor.float().mean().item()
    if abs(actual_mean - expected) > tolerance:
        return False, f"{name}: expected mean={expected}, got {actual_mean}"
    return True, None


# ============================================================
# Load Verification Instrumentation
# ============================================================

_original_load_to_slot_layer = None
_offload_engine_ref = None

def make_verified_load_to_slot_layer(original_func, offload_engine):
    """
    Create a wrapper around load_to_slot_layer that verifies each load operation.

    After each H2D transfer, checks that the GPU slot contains the expected
    pattern from the source CPU block.
    """
    def verified_load(slot_idx: int, layer_id: int, cpu_block_id: int):
        # Call original load
        original_func(slot_idx, layer_id, cpu_block_id)

        # Only verify layer 0 to reduce overhead
        if layer_id != 0:
            return

        # IMPORTANT: Synchronize CUDA to ensure async transfer is complete
        # The transfer happens on a per-slot stream, and wait_slot_layer only
        # makes compute_stream wait. We need full sync to read on default stream.
        torch.cuda.synchronize()

        # Get the expected pattern for this CPU block
        # cpu_block_id == chunk_idx in our sequential test
        expected_k, expected_v = get_expected_pattern(cpu_block_id)

        # Read GPU slot data (GPU cache has no layer dimension)
        gpu_k = offload_engine.k_cache_gpu[slot_idx]
        gpu_v = offload_engine.v_cache_gpu[slot_idx]

        actual_k = gpu_k.float().mean().item()
        actual_v = gpu_v.float().mean().item()

        k_ok = abs(actual_k - expected_k) < 1e-3
        v_ok = abs(actual_v - expected_v) < 1e-3

        chunk_idx = current_chunk_for_load[0]
        load_operations.append({
            'chunk_idx': chunk_idx,
            'slot_idx': slot_idx,
            'cpu_block_id': cpu_block_id,
            'expected_k': expected_k,
            'expected_v': expected_v,
            'actual_k': actual_k,
            'actual_v': actual_v,
            'k_ok': k_ok,
            'v_ok': v_ok,
        })

        if not (k_ok and v_ok):
            errors.append(f"Load verification failed: chunk {chunk_idx}, "
                         f"CPU block {cpu_block_id} -> GPU slot {slot_idx}: "
                         f"expected K={expected_k:.1f}/V={expected_v:.1f}, "
                         f"got K={actual_k:.4f}/V={actual_v:.4f}")

    return verified_load


def install_load_verification(llm):
    """Install verification wrapper on load_to_slot_layer."""
    global _original_load_to_slot_layer, _offload_engine_ref

    oe = llm.model_runner.kvcache_manager.offload_engine
    _offload_engine_ref = oe
    _original_load_to_slot_layer = oe.load_to_slot_layer

    oe.load_to_slot_layer = make_verified_load_to_slot_layer(
        _original_load_to_slot_layer, oe
    )
    print("Installed load verification wrapper on load_to_slot_layer")


def uninstall_load_verification():
    """Restore original load_to_slot_layer."""
    global _original_load_to_slot_layer, _offload_engine_ref

    if _offload_engine_ref is not None and _original_load_to_slot_layer is not None:
        _offload_engine_ref.load_to_slot_layer = _original_load_to_slot_layer
        print("Restored original load_to_slot_layer")

    _original_load_to_slot_layer = None
    _offload_engine_ref = None


# ============================================================
# Attention Hook
# ============================================================

def make_kv_pattern_pre_hook(layer_id: int):
    """
    Create a PRE-forward hook that overwrites K/V with distinctive patterns BEFORE
    they are stored to cache. This is called before attention.forward().

    register_forward_pre_hook receives (module, inputs) and can modify inputs in-place.
    """
    def hook(module, inputs):
        ctx = get_context()
        if not ctx.is_prefill:
            return

        chunk_idx = ctx.current_chunk_idx if hasattr(ctx, 'current_chunk_idx') else 0
        kvcache_manager = ctx.kvcache_manager if hasattr(ctx, 'kvcache_manager') else None

        if kvcache_manager is None:
            return

        # Only process layer 0 for cleaner output
        if layer_id != 0:
            return

        q, k, v = inputs
        k_pattern, v_pattern = get_expected_pattern(chunk_idx)

        # === Overwrite current chunk's K/V with distinctive pattern ===
        # This happens BEFORE forward(), so these values will be stored to cache
        k.fill_(k_pattern)
        v.fill_(v_pattern)

        # Only print for first few and last few chunks to reduce noise
        num_chunks = INPUT_LEN // BLOCK_SIZE
        if chunk_idx < 3 or chunk_idx >= num_chunks - 2:
            print(f"[Chunk {chunk_idx:3d}] Set K={k_pattern:.1f}, V={v_pattern:.1f}")
        elif chunk_idx == 3:
            print(f"... (chunks 3 to {num_chunks - 3} omitted) ...")

    return hook


def make_block_coverage_pre_hook(layer_id: int):
    """
    Create a PRE-forward hook to verify that all previous blocks are included
    in the cpu_block_table for chunked prefill attention.

    This catches bugs where:
    - Some blocks are missing from the computation
    - Sparse policy incorrectly filters out blocks (when not intended)
    - Block table construction has off-by-one errors
    """
    def hook(module, inputs):
        ctx = get_context()
        if not ctx.is_prefill:
            return

        chunk_idx = ctx.current_chunk_idx if hasattr(ctx, 'current_chunk_idx') else 0
        kvcache_manager = ctx.kvcache_manager if hasattr(ctx, 'kvcache_manager') else None

        if kvcache_manager is None:
            return

        # Only process layer 0 for cleaner output
        if layer_id != 0:
            return

        # Update current chunk for load verification tracking
        current_chunk_for_load[0] = chunk_idx

        # No previous blocks for chunk 0
        if chunk_idx == 0:
            return

        # Get the sequence and its block table (same logic as _chunked_prefill_attention)
        seq = ctx.chunked_seq if hasattr(ctx, 'chunked_seq') else None
        if seq is None:
            return

        # Get the CPU block table that will be used for attention
        cpu_block_table = kvcache_manager.get_prefilled_cpu_blocks(seq)

        # Expected blocks: 0 to chunk_idx-1 (all previous chunks)
        expected_blocks = set(range(chunk_idx))
        actual_blocks = set(cpu_block_table) if cpu_block_table else set()

        # Store for later summary
        block_coverage[chunk_idx] = {
            'expected': expected_blocks,
            'actual': actual_blocks,
        }

        # Check for missing blocks
        missing_blocks = expected_blocks - actual_blocks
        extra_blocks = actual_blocks - expected_blocks

        num_chunks = INPUT_LEN // BLOCK_SIZE
        if chunk_idx < 4 or chunk_idx >= num_chunks - 2 or missing_blocks:
            if not missing_blocks and not extra_blocks:
                print(f"  Block coverage chunk {chunk_idx:2d}: {len(actual_blocks)}/{len(expected_blocks)} blocks [OK]")
            else:
                status_parts = []
                if missing_blocks:
                    status_parts.append(f"MISSING {sorted(missing_blocks)}")
                if extra_blocks:
                    status_parts.append(f"EXTRA {sorted(extra_blocks)}")
                print(f"  Block coverage chunk {chunk_idx:2d}: {len(actual_blocks)}/{len(expected_blocks)} blocks [{', '.join(status_parts)}]")
        elif chunk_idx == 4:
            # Indicate that middle chunks are being verified silently
            print(f"  ... (verifying chunks 4-{num_chunks - 3} silently) ...")

        if missing_blocks:
            errors.append(f"Chunk {chunk_idx} missing blocks: {sorted(missing_blocks)}")

    return hook


def make_gpu_write_verification_post_hook(layer_id: int):
    """
    Create a POST-forward hook to verify the current chunk's KV was correctly
    written to the GPU ring buffer write_slot.

    This is a more reliable verification than checking load slots, because:
    1. Post-hook runs AFTER forward() writes to GPU cache
    2. write_slot mapping is deterministic: chunk_idx % num_ring_slots
    3. We injected known patterns in pre-hook, now verify they're in GPU cache
    """
    def hook(module, inputs, output):
        ctx = get_context()
        if not ctx.is_prefill:
            return

        chunk_idx = ctx.current_chunk_idx if hasattr(ctx, 'current_chunk_idx') else 0
        kvcache_manager = ctx.kvcache_manager if hasattr(ctx, 'kvcache_manager') else None

        if kvcache_manager is None:
            return

        # Only process layer 0 for cleaner output
        if layer_id != 0:
            return

        oe = kvcache_manager.offload_engine
        num_ring_slots = oe.num_ring_slots
        write_slot = chunk_idx % num_ring_slots

        # Get expected pattern for current chunk
        expected_k, expected_v = get_expected_pattern(chunk_idx)

        # Verify write_slot contains current chunk's data (GPU cache has no layer dimension)
        gpu_k = oe.k_cache_gpu[write_slot]
        gpu_v = oe.v_cache_gpu[write_slot]

        actual_k_mean = gpu_k.float().mean().item()
        actual_v_mean = gpu_v.float().mean().item()

        k_ok, _ = check_pattern(gpu_k, expected_k, f"GPU slot {write_slot}")
        v_ok, _ = check_pattern(gpu_v, expected_v, f"GPU slot {write_slot}")

        num_chunks = INPUT_LEN // BLOCK_SIZE
        # Print for first/last chunks, or if there's an error
        if True or chunk_idx >= num_chunks - 2 or not (k_ok and v_ok):
            if k_ok and v_ok:
                print(f"  GPU write_slot[{write_slot}] chunk {chunk_idx:2d}: K={expected_k:.1f}, V={expected_v:.1f} [OK]")
            else:
                print(f"  GPU write_slot[{write_slot}] chunk {chunk_idx:2d}: expected K={expected_k:.1f}/V={expected_v:.1f}, "
                      f"got K={actual_k_mean:.2f}/V={actual_v_mean:.2f} [FAIL]")
        elif chunk_idx == 4:
            print(f"  ... (GPU write verification for chunks 4-{num_chunks - 3} silently) ...")

        if not (k_ok and v_ok):
            errors.append(f"GPU write_slot {write_slot} at chunk {chunk_idx}: "
                         f"expected K={expected_k}, V={expected_v}, got K={actual_k_mean:.4f}, V={actual_v_mean:.4f}")

    return hook


def make_kv_verification_post_hook(layer_id: int):
    """
    Create a POST-forward hook to verify CPU cache contains correct patterns
    from previously offloaded blocks.
    """
    def hook(module, inputs, output):
        ctx = get_context()
        if not ctx.is_prefill:
            return

        chunk_idx = ctx.current_chunk_idx if hasattr(ctx, 'current_chunk_idx') else 0
        kvcache_manager = ctx.kvcache_manager if hasattr(ctx, 'kvcache_manager') else None

        if kvcache_manager is None:
            return

        # Only process layer 0 for cleaner output
        if layer_id != 0:
            return

        # === Verify previously offloaded blocks in CPU cache ===
        if chunk_idx >= 1:
            oe = kvcache_manager.offload_engine
            num_ok = 0
            num_fail = 0

            # Check all previously offloaded blocks
            for prev_chunk in range(chunk_idx):
                # CPU block ID = prev_chunk (in simple sequential case)
                cpu_block_id = prev_chunk

                # Get expected pattern for this block
                expected_k, expected_v = get_expected_pattern(prev_chunk)

                # Read from CPU cache (layer 0)
                cpu_k = oe.k_cache_cpu[layer_id, cpu_block_id]
                cpu_v = oe.v_cache_cpu[layer_id, cpu_block_id]

                # Verify patterns
                k_ok, k_err = check_pattern(cpu_k, expected_k, f"CPU K block {cpu_block_id}")
                v_ok, v_err = check_pattern(cpu_v, expected_v, f"CPU V block {cpu_block_id}")

                if k_ok and v_ok:
                    num_ok += 1
                else:
                    num_fail += 1
                    if k_err:
                        errors.append(k_err)
                    if v_err:
                        errors.append(v_err)

            # Only print summary for each chunk verification
            num_chunks = INPUT_LEN // BLOCK_SIZE
            if chunk_idx < 4 or chunk_idx >= num_chunks - 2 or num_fail > 0:
                status = "OK" if num_fail == 0 else f"FAIL({num_fail})"
                print(f"  CPU verify chunk {chunk_idx:2d}: {num_ok} blocks OK [{status}]")
            elif chunk_idx == 4:
                print(f"  ... (CPU cache verification for chunks 4-{num_chunks - 3} silently) ...")

    return hook


def make_post_chunk_verification_hook(layer_id: int):
    """
    Post-forward hook to verify GPU ring buffer state after attention.
    """
    def hook(module, inputs, output):
        ctx = get_context()
        if not ctx.is_prefill or layer_id != 0:
            return

        chunk_idx = ctx.current_chunk_idx if hasattr(ctx, 'current_chunk_idx') else 0
        kvcache_manager = ctx.kvcache_manager if hasattr(ctx, 'kvcache_manager') else None

        if kvcache_manager is None:
            return

        oe = kvcache_manager.offload_engine

        # After attention, the current chunk's KV should be in the GPU ring buffer
        # Ring slot = chunk_idx % num_ring_slots
        ring_slot = chunk_idx % oe.num_ring_slots

        expected_k, expected_v = get_expected_pattern(chunk_idx)

        # Check GPU ring buffer (GPU cache has no layer dimension)
        gpu_k = oe.k_cache_gpu[ring_slot]
        gpu_v = oe.v_cache_gpu[ring_slot]

        k_ok, k_err = check_pattern(gpu_k, expected_k, f"GPU K slot {ring_slot}")
        v_ok, v_err = check_pattern(gpu_v, expected_v, f"GPU V slot {ring_slot}")

        if k_ok and v_ok:
            print(f"  [OK] GPU slot {ring_slot} (chunk {chunk_idx}): K={expected_k}, V={expected_v}")
        else:
            if k_err:
                print(f"  [FAIL] {k_err}")
                errors.append(k_err)
            if v_err:
                print(f"  [FAIL] {v_err}")
                errors.append(v_err)

    return hook


def register_hooks(llm):
    """Register pre and post forward hooks."""
    hooks = []
    model = llm.model_runner.model

    for layer_idx, decoder_layer in enumerate(model.model.layers):
        attn_module = decoder_layer.self_attn.attn

        # PRE-forward hook 1: Verify all previous blocks are in cpu_block_table
        coverage_hook = attn_module.register_forward_pre_hook(make_block_coverage_pre_hook(layer_idx))
        hooks.append(coverage_hook)

        # PRE-forward hook 2: Inject K/V patterns before they're stored to cache
        pattern_hook = attn_module.register_forward_pre_hook(make_kv_pattern_pre_hook(layer_idx))
        hooks.append(pattern_hook)

        # POST-forward hook 1: Verify GPU write_slot contains current chunk's data
        gpu_verify_hook = attn_module.register_forward_hook(make_gpu_write_verification_post_hook(layer_idx))
        hooks.append(gpu_verify_hook)

        # POST-forward hook 2: Verify CPU cache contains correct patterns after offload
        cpu_verify_hook = attn_module.register_forward_hook(make_kv_verification_post_hook(layer_idx))
        hooks.append(cpu_verify_hook)

    return hooks


# ============================================================
# Final Verification
# ============================================================

def verify_final_cpu_state(llm, num_chunks: int):
    """Verify all CPU blocks have correct patterns after prefill completes."""
    print("\n" + "=" * 60)
    print("Final CPU Cache Verification")
    print("=" * 60)

    kvcache_manager = llm.model_runner.kvcache_manager
    oe = kvcache_manager.offload_engine

    num_ok = 0
    num_fail = 0
    fail_details = []

    # After prefill, all chunks should be in CPU
    for chunk_idx in range(num_chunks):
        cpu_block_id = chunk_idx
        expected_k, expected_v = get_expected_pattern(chunk_idx)

        # Check layer 0
        cpu_k = oe.k_cache_cpu[0, cpu_block_id]
        cpu_v = oe.v_cache_cpu[0, cpu_block_id]

        k_ok, k_err = check_pattern(cpu_k, expected_k, f"Final CPU K block {cpu_block_id}")
        v_ok, v_err = check_pattern(cpu_v, expected_v, f"Final CPU V block {cpu_block_id}")

        if k_ok and v_ok:
            num_ok += 1
            # Only print first few and last few
            if chunk_idx < 3 or chunk_idx >= num_chunks - 2:
                actual_k_mean = cpu_k.float().mean().item()
                actual_v_mean = cpu_v.float().mean().item()
                print(f"  Block {cpu_block_id:3d}: K={expected_k:.1f} ({actual_k_mean:.4f}), "
                      f"V={expected_v:.1f} ({actual_v_mean:.4f}) [OK]")
            elif chunk_idx == 3:
                print(f"  ... (blocks 3 to {num_chunks - 3} verified OK) ...")
        else:
            num_fail += 1
            actual_k_mean = cpu_k.float().mean().item()
            actual_v_mean = cpu_v.float().mean().item()
            print(f"  Block {cpu_block_id:3d}: K={expected_k:.1f} ({actual_k_mean:.4f}), "
                  f"V={expected_v:.1f} ({actual_v_mean:.4f}) [FAIL]")
            if k_err:
                errors.append(k_err)
            if v_err:
                errors.append(v_err)

    print(f"\nTotal: {num_ok} OK, {num_fail} FAIL out of {num_chunks} blocks")


def verify_block_coverage_summary(num_chunks: int):
    """Verify that all chunks had complete block coverage during prefill."""
    print("\n" + "=" * 60)
    print("Block Coverage Verification Summary")
    print("=" * 60)

    num_ok = 0
    num_fail = 0
    total_blocks_expected = 0
    total_blocks_computed = 0

    for chunk_idx in range(1, num_chunks):  # Start from 1 (chunk 0 has no previous)
        if chunk_idx not in block_coverage:
            print(f"  Chunk {chunk_idx}: NO COVERAGE DATA [FAIL]")
            errors.append(f"Chunk {chunk_idx} has no block coverage data")
            num_fail += 1
            continue

        coverage = block_coverage[chunk_idx]
        expected = coverage['expected']
        actual = coverage['actual']
        missing = expected - actual

        total_blocks_expected += len(expected)
        total_blocks_computed += len(actual)

        if not missing:
            num_ok += 1
        else:
            num_fail += 1

    # Print summary
    if num_fail == 0:
        print(f"  All {num_ok} chunks had complete block coverage [OK]")
        print(f"  Total blocks computed: {total_blocks_computed} (expected: {total_blocks_expected})")
    else:
        print(f"  {num_ok} chunks OK, {num_fail} chunks with missing blocks [FAIL]")
        print(f"  Total blocks computed: {total_blocks_computed} (expected: {total_blocks_expected})")

    # Verify the total is correct: sum of 0+1+2+...+(n-1) = n*(n-1)/2
    expected_total = num_chunks * (num_chunks - 1) // 2
    if total_blocks_expected == expected_total:
        print(f"  Expected total blocks matches formula: {expected_total} [OK]")
    else:
        print(f"  Expected total mismatch: got {total_blocks_expected}, formula gives {expected_total} [FAIL]")
        errors.append(f"Block coverage total mismatch")


def verify_load_operations_summary(num_chunks: int):
    """Verify all H2D load operations transferred correct data."""
    print("\n" + "=" * 60)
    print("H2D Load Operations Verification Summary")
    print("=" * 60)

    if not load_operations:
        print("  WARNING: No load operations recorded!")
        print("  (This may indicate load verification was not installed)")
        return

    num_ok = 0
    num_fail = 0
    loads_per_chunk = {}

    for op in load_operations:
        chunk_idx = op['chunk_idx']
        if chunk_idx not in loads_per_chunk:
            loads_per_chunk[chunk_idx] = []
        loads_per_chunk[chunk_idx].append(op)

        if op['k_ok'] and op['v_ok']:
            num_ok += 1
        else:
            num_fail += 1

    # Print per-chunk summary for first/last chunks
    for chunk_idx in sorted(loads_per_chunk.keys()):
        ops = loads_per_chunk[chunk_idx]
        chunk_ok = sum(1 for op in ops if op['k_ok'] and op['v_ok'])
        chunk_fail = len(ops) - chunk_ok

        if chunk_idx < 4 or chunk_idx >= num_chunks - 2 or chunk_fail > 0:
            # Show loaded block IDs in order
            block_ids = [op['cpu_block_id'] for op in ops]
            if chunk_fail == 0:
                print(f"  Chunk {chunk_idx:2d}: loaded {len(ops)} blocks {block_ids} [OK]")
            else:
                print(f"  Chunk {chunk_idx:2d}: loaded {len(ops)} blocks, {chunk_fail} FAILED [FAIL]")
                for op in ops:
                    if not (op['k_ok'] and op['v_ok']):
                        print(f"    CPU block {op['cpu_block_id']} -> slot {op['slot_idx']}: "
                              f"expected K={op['expected_k']:.1f}/V={op['expected_v']:.1f}, "
                              f"got K={op['actual_k']:.4f}/V={op['actual_v']:.4f}")
        elif chunk_idx == 4:
            print(f"  ... (chunks 4-{num_chunks - 3} load verification running silently) ...")

    # Print overall summary
    print(f"\n  Total load operations: {len(load_operations)}")
    print(f"  Successful: {num_ok}, Failed: {num_fail}")

    if num_fail == 0:
        print(f"  All H2D transfers verified correct [OK]")
    else:
        print(f"  {num_fail} H2D transfers had incorrect data [FAIL]")


# ============================================================
# Main Test Script
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Test: CPU Offload Correctness with Distinctive KV Patterns")
    print("=" * 60)
    print(f"Input: {INPUT_LEN} tokens, {INPUT_LEN // BLOCK_SIZE} chunks")
    print(f"GPU blocks: {NUM_GPU_BLOCKS}, Block size: {BLOCK_SIZE}")
    print(f"Pattern: K=chunk_idx+1, V=-(chunk_idx+1)")
    print()

    # 1. Initialize LLM
    print("Initializing LLM...")
    llm = LLM(
        MODEL_PATH,
        enforce_eager=True,
        max_model_len=MAX_MODEL_LEN,
        max_num_batched_tokens=MAX_MODEL_LEN,
        enable_cpu_offload=True,
        kvcache_block_size=BLOCK_SIZE,
        num_gpu_blocks=NUM_GPU_BLOCKS,
        dtype="float16",
    )

    # 2. Register hooks
    hooks = register_hooks(llm)
    print(f"Registered {len(hooks)} hooks")

    # 3. Install load verification (instrument load_to_slot_layer)
    install_load_verification(llm)

    # 4. Generate prompt
    seed(42)
    prompt_token_ids = [[randint(0, 10000) for _ in range(INPUT_LEN)]]
    num_chunks = INPUT_LEN // BLOCK_SIZE

    # 5. Run prefill
    print("\n" + "=" * 60)
    print("Running Prefill with KV Pattern Injection...")
    print("=" * 60)

    sampling_params = SamplingParams(temperature=0.6, ignore_eos=True, max_tokens=1)
    outputs = llm.generate(prompt_token_ids, sampling_params, use_tqdm=False)

    # 6. Remove hooks and uninstall load verification
    for hook in hooks:
        hook.remove()
    uninstall_load_verification()

    # 7. Final verification
    verify_final_cpu_state(llm, num_chunks)

    # 8. Block coverage summary
    verify_block_coverage_summary(num_chunks)

    # 9. H2D load operations summary
    verify_load_operations_summary(num_chunks)

    # 10. Report results
    print("\n" + "=" * 60)
    if errors:
        print(f"test_offload_correctness: FAILED ({len(errors)} errors)")
        for err in errors[:10]:  # Show first 10 errors
            print(f"  - {err}")
        exit(1)
    else:
        print("test_offload_correctness: PASSED")
    print("=" * 60)