vLLM/vllm/v1/hisparse/runtime.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NVIDIA HiSparse resident, host, and hot-cache data plane."""

from __future__ import annotations

import math
import mmap
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, TypeAlias

import psutil
import torch

from vllm.config import CUDAGraphMode, VllmConfig
from vllm.distributed import get_tp_group
from vllm.forward_context import get_forward_context, in_piecewise_cudagraph
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.math_utils import round_up
from vllm.utils.torch_utils import current_stream
from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor

logger = init_logger(__name__)

HOST_REGISTER_CHUNK_BYTES = 256 * 2**30

if TYPE_CHECKING:
    from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
    from vllm.v1.kv_cache_interface import KVCacheConfig

# fp8_ds_mla KV row: 512 B quantized NoPE + 16 B scales + 128 B RoPE.
FP8_DS_MLA_ROW_BYTES = 656
HiSparseTopKResult: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor]


def is_hisparse_decode_batch(attn_metadata: Any | None) -> bool:
    return (
        attn_metadata is not None
        and attn_metadata.num_decode_tokens == attn_metadata.num_actual_tokens
        and attn_metadata.max_query_len == 1
        and attn_metadata.num_reqs == attn_metadata.num_actual_tokens
    )


def _get_max_decode_query_len(vllm_config: VllmConfig) -> int:
    speculative_config = getattr(vllm_config, "speculative_config", None)
    if (
        speculative_config is not None
        and speculative_config.num_speculative_tokens is not None
    ):
        return 1 + speculative_config.num_speculative_tokens * (
            2 if speculative_config.parallel_drafting else 1
        )
    return 1


def _get_max_swap_rows(vllm_config: VllmConfig) -> int:
    max_query_len = _get_max_decode_query_len(vllm_config)
    scheduler_config = vllm_config.scheduler_config
    return min(
        scheduler_config.max_num_batched_tokens,
        scheduler_config.max_num_seqs * max_query_len,
    )


@dataclass(frozen=True)
class PagedCacheView:
    cache: torch.Tensor
    attention_cache: torch.Tensor
    block_size: int
    attention_block_stride: int

    @classmethod
    def bind(
        cls,
        raw_tensor: torch.Tensor,
        *,
        dtype: torch.dtype,
        row_width: int,
        byte_offset: int,
        block_stride: int,
        num_blocks: int,
        block_size: int,
    ) -> PagedCacheView:
        itemsize = dtype.itemsize
        assert byte_offset % itemsize == 0 and block_stride % itemsize == 0
        cache = torch.as_strided(
            raw_tensor.view(dtype),
            size=(num_blocks, block_size, row_width),
            stride=(block_stride // itemsize, row_width, 1),
            storage_offset=byte_offset // itemsize,
        )
        row_bytes = row_width * itemsize
        if block_stride % (block_size * row_bytes):
            return cls(cache, cache, block_size, block_size)
        attention_cache = (
            raw_tensor[byte_offset:].view(dtype).view(-1, block_size, row_width)
        )
        return cls(cache, attention_cache, block_size, block_stride // row_bytes)


@dataclass(frozen=True)
class ResolvedHiSparseConfig:
    top_k: int
    device_buffer_size: int
    eager_host_mirror: bool = True

    @classmethod
    def from_vllm_config(
        cls,
        vllm_config: VllmConfig,
        model_top_k: int,
        block_size: int | None = None,
    ) -> ResolvedHiSparseConfig | None:
        config = vllm_config.attention_config.hisparse_config
        if config is None:
            return None

        max_query_len = _get_max_decode_query_len(vllm_config)
        min_device_buffer_size = max_query_len * model_top_k
        # Retain the worst-case union of every speculative step's top-k and
        # leave one additional top-k of LRU slack.
        configured_size = config.device_buffer_size
        if configured_size is None:
            device_buffer_size = (max_query_len + 1) * model_top_k
        else:
            device_buffer_size = configured_size

        if device_buffer_size < min_device_buffer_size:
            raise ValueError(
                "HiSparse device_buffer_size must cover every decode query's "
                "index_topk rows. Got "
                f"device_buffer_size={device_buffer_size}, "
                f"max_decode_query_len={max_query_len}, "
                f"index_topk={model_top_k}; expected at least "
                f"{min_device_buffer_size}."
            )
        max_device_buffer_size = torch.iinfo(torch.int16).max + 1
        if device_buffer_size > max_device_buffer_size:
            raise ValueError(
                "HiSparse device_buffer_size exceeds the int16 slot-index "
                f"limit: got {device_buffer_size}, maximum is "
                f"{max_device_buffer_size}."
            )
        if configured_size is not None and block_size is not None:
            padding = -device_buffer_size % block_size
            if padding:
                logger.warning(
                    "HiSparse device_buffer_size=%d is not aligned to the "
                    "%d-token kernel block size and allocates %d unused "
                    "padding rows per hot group. Use %d to use all allocated "
                    "rows.",
                    device_buffer_size,
                    block_size,
                    padding,
                    device_buffer_size + padding,
                )
        return cls(
            top_k=model_top_k,
            device_buffer_size=device_buffer_size,
            eager_host_mirror=config.eager_host_mirror,
        )


def check_hisparse_host_memory(pool_bytes: int) -> None:
    """Reject a physical host-pool allocation that cannot fit in RAM."""
    mem = psutil.virtual_memory()
    if pool_bytes > mem.available * 0.95:
        raise ValueError(
            f"HiSparse pinned host pool needs ~{pool_bytes / 2**30:.0f} GiB "
            f"but only {mem.available / 2**30:.0f} GiB of RAM is available. "
            "Lower the HiSparseConnector host_pool_gib or leave headroom "
            "for co-tenants."
        )


def allocate_pinned_host_pool(size: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Allocate and deterministically register an exact-size host KV region."""
    page = 4096
    padded_size = round_up(size, page)
    backing = torch.empty(padded_size + page, dtype=torch.int8, device="cpu")
    aligned_offset = (-backing.data_ptr()) % page
    registered = backing[aligned_offset : aligned_offset + padded_size]
    pin_tensor(registered)
    return registered[:size], registered


class HiSparseHostPool:
    """Keep the single host backing and its exact CUDA registration range."""

    def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig) -> None:
        self.vllm_config = vllm_config
        self.kv_cache_config = kv_cache_config
        self.backing: torch.Tensor | None = None
        self.registered: torch.Tensor | None = None
        self.shared_region: SharedOffloadRegion | None = None

    def allocate(self, size: int) -> torch.Tensor:
        assert self.backing is None, "HiSparse host pool is already allocated."
        config = self.kv_cache_config
        if config.hisparse_shared_host_pool:
            tensor_sizes = [
                size
                for _, size in sorted(
                    (tensor.offset + i * tensor.layer_stride, tensor.layer_stride)
                    for tensor in config.kv_cache_tensors
                    if tensor.host_resident
                    for i in range(len(tensor.layers))
                )
            ]
            assert config.hisparse_host_num_blocks is not None
            assert config.hisparse_host_block_stride is not None
            _, _, self.shared_region = allocate_hisparse_host_pools(
                self.vllm_config,
                tensor_sizes,
                config.hisparse_host_num_blocks,
                config.hisparse_host_block_stride,
                use_shared_host_pool=True,
            )
            assert self.shared_region is not None
            self.registered = self.shared_region.base_tensor
            self.backing = self.registered[:size]
        else:
            check_hisparse_host_memory(size)
            self.backing, self.registered = allocate_pinned_host_pool(size)
        return self.backing


def use_shared_hisparse_host_pool(vllm_config: VllmConfig) -> bool:
    """Whether replicated MLA host KV can share one local mmap."""
    parallel = vllm_config.parallel_config
    return (
        current_platform.is_cuda_alike()
        and parallel.tensor_parallel_size > 1
        and parallel.pipeline_parallel_size == 1
        and parallel.prefill_context_parallel_size == 1
        and parallel.decode_context_parallel_size == 1
        and parallel.world_size == parallel.tensor_parallel_size
        and parallel.distributed_executor_backend == "mp"
        and parallel.nnodes_within_dp == 1
    )


def get_hisparse_host_block_stride(
    page_size_bytes: int, *, use_shared_host_pool: bool
) -> int:
    if use_shared_host_pool:
        return round_up(page_size_bytes, SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT)
    return page_size_bytes


def _hisparse_registration_ranges(
    tensor_sizes: list[int],
    num_blocks: int,
    host_block_stride: int,
    max_chunk_bytes: int = HOST_REGISTER_CHUNK_BYTES,
) -> tuple[tuple[int, int], ...]:
    """Split a shared pool without bisecting any tensor's host block.

    CUDA batch copies reject descriptors spanning adjacent registrations.
    """
    total_size = num_blocks * host_block_stride
    if total_size % mmap.PAGESIZE:
        raise ValueError("HiSparse host pool must be OS-page aligned.")
    if max_chunk_bytes < mmap.PAGESIZE:
        raise ValueError("HiSparse registration chunks must span at least one page.")

    tensors: list[tuple[int, int, int]] = []
    tensor_start = 0
    for tensor_size in tensor_sizes:
        if tensor_size % num_blocks:
            raise ValueError(
                f"HiSparse host tensor size {tensor_size} is not divisible by "
                f"{num_blocks} blocks."
            )
        tensor_end = tensor_start + tensor_size
        tensors.append((tensor_start, tensor_end, tensor_size // num_blocks))
        tensor_start = tensor_end
    if tensor_start > total_size:
        raise ValueError("HiSparse host tensors exceed the shared host pool.")

    def last_safe_boundary(lower: int, upper: int) -> int | None:
        best = None
        for start, end, block_bytes in tensors:
            first_block = max(0, (lower - start) // block_bytes + 1)
            last_block = min(num_blocks, (min(upper, end) - start) // block_bytes)
            if first_block > last_block:
                continue
            divisor = math.gcd(block_bytes, mmap.PAGESIZE)
            if start % divisor:
                continue
            period = mmap.PAGESIZE // divisor
            if period == 1:
                residue = 0
            else:
                # Find block boundaries that are also OS-page boundaries.
                residue = (
                    (-start // divisor) * pow(block_bytes // divisor, -1, period)
                ) % period
            block = last_block - (last_block - residue) % period
            if block >= first_block:
                candidate = start + block * block_bytes
                best = candidate if best is None else max(best, candidate)

        if upper >= tensor_start:
            candidate = upper // mmap.PAGESIZE * mmap.PAGESIZE
            if candidate >= tensor_start and candidate > lower:
                best = candidate if best is None else max(best, candidate)
        return best

    ranges: list[tuple[int, int]] = []
    start = 0
    while total_size - start > max_chunk_bytes:
        end = last_safe_boundary(start, start + max_chunk_bytes)
        if end is None:
            raise ValueError(
                "A HiSparse host block is too large for the registration chunk limit."
            )
        ranges.append((start, end))
        start = end
    ranges.append((start, total_size))
    return tuple(ranges)


def allocate_hisparse_host_pools(
    vllm_config: VllmConfig,
    tensor_sizes: list[int],
    num_blocks: int,
    host_block_stride: int,
    *,
    use_shared_host_pool: bool,
) -> tuple[list[torch.Tensor], list[torch.Tensor], SharedOffloadRegion | None]:
    """Allocate private host tensors or one mmap shared by local TP ranks."""
    if not use_shared_host_pool:
        check_hisparse_host_memory(sum(tensor_sizes))
        private_pools = [allocate_pinned_host_pool(size) for size in tensor_sizes]
        return (
            [pool for pool, _ in private_pools],
            [registered for _, registered in private_pools],
            None,
        )

    region = SharedOffloadRegion(
        engine_id=(
            f"hisparse_{vllm_config.instance_id}_"
            f"dp{vllm_config.parallel_config.data_parallel_index}"
        ),
        num_chunks=1,
        rank=0,
        kv_bytes_per_chunk=num_blocks * host_block_stride,
        cpu_page_size=sum(tensor_sizes),
        barrier=get_tp_group().barrier,
        creator_memory_check=check_hisparse_host_memory,
        populate_only_on_creator=True,
    )
    try:
        for start, end in _hisparse_registration_ranges(
            tensor_sizes, num_blocks, host_block_stride
        ):
            tensor = region.base_tensor[start:end]
            pin_tensor(tensor)
            region.pinned_addresses.append(tensor.data_ptr())
            region.is_pinned = True
        pools = [
            region.create_next_canonical_view(size).view(-1) for size in tensor_sizes
        ]
    except Exception:
        region.cleanup()
        raise
    return pools, [], region


def release_pinned_state(
    runtimes: list[HiSparseRuntime],
    pinned_host_pools: list[torch.Tensor],
    shared_host_region: SharedOffloadRegion | None = None,
) -> None:
    """Synchronize and release registered host KV pools."""
    if pinned_host_pools or shared_host_region is not None:
        try:
            torch.accelerator.synchronize()
        except RuntimeError as e:
            logger.warning(
                "HiSparse: CUDA context unusable at teardown (%s); leaving "
                "%d host-pool tensors pinned for kernel exit reclaim.",
                e,
                len(pinned_host_pools) + int(shared_host_region is not None),
            )
            return

        cudart = torch.cuda.cudart()
        release_start = time.perf_counter()
        freed_bytes = 0
        while pinned_host_pools:
            tensor = pinned_host_pools[-1]
            err = cudart.cudaHostUnregister(tensor.data_ptr())
            if err.value != 0:
                logger.warning(
                    "HiSparse: cudaHostUnregister failed (code=%d); leaving "
                    "%d host-pool tensors pinned for kernel exit reclaim.",
                    err.value,
                    len(pinned_host_pools),
                )
                cudart.cudaGetLastError()
                break
            freed_bytes += tensor.nbytes
            pinned_host_pools.pop()
        if freed_bytes:
            logger.info(
                "HiSparse: unpinned %.1f GiB of host pool in %.1fs.",
                freed_bytes / 2**30,
                time.perf_counter() - release_start,
            )

    for runtime in runtimes:
        del runtime._host_cache
        del runtime.registered_host_pool
        del runtime.hot_backing
    if shared_host_region is not None:
        shared_host_region.cleanup()


def hisparse_prefill_staging_remap(
    block_table: torch.Tensor, block_size: int
) -> tuple[torch.Tensor, torch.Tensor]:
    """Renumber blocks without a data-dependent CUDA output allocation."""
    block_ids = torch.cat(
        (
            torch.zeros(1, dtype=torch.int32, device=block_table.device),
            block_table.flatten().clamp(min=0).to(torch.int32),
        )
    )
    sorted_ids, permutation = torch.sort(block_ids)
    is_new = torch.cat(
        (
            torch.ones(1, dtype=torch.bool, device=block_table.device),
            sorted_ids[1:] != sorted_ids[:-1],
        )
    )
    compact_sorted = torch.cumsum(is_new, dim=0, dtype=torch.int32) - 1
    inverse = torch.empty_like(compact_sorted)
    inverse.scatter_(0, permutation, compact_sorted)

    unique_ids = torch.full_like(sorted_ids, -1)
    unique_ids.scatter_(0, compact_sorted, sorted_ids)
    valid = (
        torch.arange(sorted_ids.numel(), device=block_table.device)
        <= compact_sorted[-1]
    )
    unique_ids = torch.where(valid, unique_ids, -1)
    new_bt = inverse[1:].view_as(block_table)
    offsets = torch.arange(block_size, dtype=torch.int32, device=block_table.device)
    row_ids = (
        torch.where(
            valid[:, None],
            unique_ids[:, None] * block_size + offsets,
            -1,
        )
    ).view(1, -1)
    return new_bt, row_ids


@dataclass
class HiSparsePrefillStagingPlan:
    block_table: torch.Tensor
    row_ids: torch.Tensor
    dst_rows: torch.Tensor
    miss_mask: torch.Tensor
    block_size: int
    # Host rows with a GPU-resident copy (adopted shadow pages): the flat
    # resident-cache row to read instead of DMAing from host, -1 for misses.
    gpu_row_ids: torch.Tensor | None = None
    gpu_source_key: tuple[int, int] | None = None

    def ensure_gpu_sources(
        self,
        resident_block_table: torch.Tensor,
        resident_block_size: int,
    ) -> None:
        """Resolve which staged rows can be served from the resident cache.

        Computed once per plan (the resident block table is shared by every
        layer in the group); non-null resident pages become miss_mask=0 rows
        gathered device-to-device by ``gather_prefill_cache``.
        """
        source_key = (resident_block_table.data_ptr(), resident_block_size)
        if self.gpu_source_key == source_key:
            return
        block_size = self.block_size
        if resident_block_size <= 0 or block_size % resident_block_size != 0:
            return
        device = self.row_ids.device
        num_unique = self.row_ids.shape[1] // block_size
        host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size
        new_bt = self.block_table.to(torch.int64)
        num_rows, num_cols = new_bt.shape
        if num_rows == 0 or resident_block_table.shape[0] < num_rows:
            return
        # One representative (row, col) per unique host block: any request
        # referencing the block holds an equivalent (refcounted) resident view.
        flat_pos = torch.arange(num_rows * num_cols, device=device)
        rep = torch.full(
            (num_unique,), num_rows * num_cols, device=device, dtype=torch.int64
        )
        rep.scatter_reduce_(
            0, new_bt.reshape(-1), flat_pos, reduce="amin", include_self=True
        )
        rep_row = (rep // num_cols).clamp(max=resident_block_table.shape[0] - 1)
        rep_col = rep % num_cols
        pages_per_block = block_size // resident_block_size
        res_pos = rep_col[:, None] * pages_per_block + torch.arange(
            pages_per_block, device=device
        )
        res_cols = resident_block_table.shape[1]
        res_blocks = torch.where(
            res_pos < res_cols,
            torch.gather(
                resident_block_table.to(torch.int64)[rep_row],
                1,
                res_pos.clamp(max=max(res_cols - 1, 0)),
            ),
            torch.zeros_like(res_pos),
        )
        offsets = torch.arange(block_size, device=device)
        per_off_block = res_blocks[:, offsets // resident_block_size]
        gpu_rows = (
            per_off_block * resident_block_size
            + (offsets % resident_block_size)[None, :]
        )
        # Block id 0 is the null block in both id spaces.
        valid_rows = self.row_ids.view(num_unique, block_size) >= 0
        hit = (per_off_block > 0) & (host_ids[:, None] > 0) & valid_rows
        self.gpu_row_ids = torch.where(hit, gpu_rows, -1).reshape(1, -1).to(torch.int32)
        self.miss_mask = ((self.gpu_row_ids < 0) & valid_rows.view(1, -1)).int()
        self.gpu_source_key = source_key


def build_hisparse_prefill_staging_plan(
    block_table: torch.Tensor,
    seq_lens: torch.Tensor,
    block_size: int,
    staging_block_capacity: int,
) -> HiSparsePrefillStagingPlan:
    """Build an asynchronous layer-independent host-cache staging remap."""
    device = block_table.device
    used = (seq_lens.to(torch.int64) + block_size - 1) // block_size
    packed_offsets = torch.arange(staging_block_capacity, device=device)
    row_ends = torch.cumsum(used, dim=0)
    rows = torch.searchsorted(row_ends, packed_offsets, right=True)
    valid = rows < block_table.shape[0]
    rows = rows.clamp(max=block_table.shape[0] - 1)
    row_starts = torch.cat((torch.zeros_like(row_ends[:1]), row_ends[:-1]))
    columns = packed_offsets - row_starts[rows]
    valid &= columns < block_table.shape[1]
    flat_positions = rows * block_table.shape[1] + columns.clamp(
        max=block_table.shape[1] - 1
    )
    referenced = torch.where(valid, block_table.flatten()[flat_positions], 0)
    referenced_bt, row_ids = hisparse_prefill_staging_remap(referenced, block_size)
    scratch_index = block_table.numel()
    scatter_positions = torch.where(valid, flat_positions, scratch_index)
    remapped = torch.where(valid, referenced_bt.flatten(), 0)
    new_bt_storage = torch.zeros(scratch_index + 1, dtype=torch.int32, device=device)
    new_bt_storage.scatter_(0, scatter_positions, remapped)
    new_bt = new_bt_storage[:-1].view_as(block_table)
    valid_rows = row_ids >= 0
    dst_rows = torch.arange(row_ids.shape[1], dtype=torch.int32, device=device).view(
        1, -1
    )
    return HiSparsePrefillStagingPlan(
        block_table=new_bt,
        row_ids=row_ids,
        dst_rows=dst_rows,
        miss_mask=valid_rows.to(torch.int32),
        block_size=block_size,
    )


def _has_hisparse_ops() -> bool:
    if not hasattr(torch.ops, "_C_cache_ops"):
        return False
    return (
        hasattr(torch.ops._C_cache_ops, "hisparse_resolve_residency")
        and hasattr(torch.ops._C_cache_ops, "hisparse_invalidate_written_slots")
        and hasattr(torch.ops._C_cache_ops, "hisparse_gather_plan")
        and hasattr(torch.ops._C_cache_ops, "hisparse_gather_compact")
    )


class _SharedTopKState:
    """CUDA-graph-safe residency result shared by index-sharing layers."""

    def __init__(self, device: torch.device, max_rows: int, top_k: int) -> None:
        self.device_topk_rows = torch.full(
            (max_rows, top_k), -1, dtype=torch.int32, device=device
        )
        self.physical_topk_indices = torch.empty_like(self.device_topk_rows)
        self.swap_host_physical_rows = torch.empty(
            (max_rows, top_k), dtype=torch.int32, device=device
        )
        self.swap_device_physical_rows = torch.empty_like(self.swap_host_physical_rows)
        self.swap_counts = torch.empty(max_rows, dtype=torch.int32, device=device)
        self.valid_topk_counts = torch.empty(max_rows, dtype=torch.int32, device=device)


def _create_shared_topk_state(
    device: torch.device, max_rows: int, top_k: int
) -> _SharedTopKState:
    return _SharedTopKState(device, max_rows, top_k)


def _create_copy_stream(device: torch.device) -> torch.Stream:
    return torch.Stream(device=device)


class HiSparseIndexGroup:
    """GPU index-resolution state shared by one leader and its followers."""

    leader: HiSparseRuntime

    def __init__(
        self,
        device: torch.device,
        max_num_reqs: int,
        region_stride: int,
        max_swap_rows: int,
        top_k: int,
        copy_stream: torch.Stream | None = None,
        logical_topk_ready: torch.Event | None = None,
    ) -> None:
        self.device_global_indices = torch.full(
            (max_num_reqs, region_stride),
            -1,
            dtype=torch.int32,
            device=device,
        )
        self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device)
        self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous()
        self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k)
        self.copy_stream = (
            copy_stream if copy_stream is not None else _create_copy_stream(device)
        )
        self.logical_topk_ready = logical_topk_ready
        self.followers: list[HiSparseRuntime] = []
        self.swap_stats = torch.zeros(2, dtype=torch.uint64, device=device)
        self.swap_stats_host = torch.empty(
            2, dtype=torch.uint64, device="cpu", pin_memory=True
        )
        self.stats_row_bytes = 0


class HiSparseRuntime:
    """Per-cache host/hot data plane and GPU replacement state."""

    hot: PagedCacheView
    hot_backing: torch.Tensor
    hot_block_table: torch.Tensor
    _host_cache: torch.Tensor
    registered_host_pool: torch.Tensor

    def __init__(
        self,
        config: ResolvedHiSparseConfig,
        max_num_reqs: int,
        row_width: int,
        kv_dtype: torch.dtype,
        device: torch.device | str,
        max_swap_rows: int | None = None,
        index_group: HiSparseIndexGroup | None = None,
        copy_stream: torch.Stream | None = None,
        logical_topk_ready: torch.Event | None = None,
    ) -> None:
        if not _has_hisparse_ops():
            raise RuntimeError(
                "HiSparse requires its compiled _C_cache_ops CUDA kernels "
                "(host-resident decode has no Python fallback). Rebuild vLLM "
                "from source so "
                "csrc/libtorch_stable/hisparse_kernels.cu is included."
            )
        self.max_num_reqs = max_num_reqs
        self.row_width = row_width
        self.kv_dtype = kv_dtype
        self.device = torch.device(device)
        # Logical slots per request. Physical rows come from its ephemeral HMA
        # block table and need not be contiguous in the shared slab.
        self.region_stride = config.device_buffer_size

        row_bytes = row_width * kv_dtype.itemsize
        if row_bytes % 16 != 0:
            raise ValueError(
                f"HiSparse requires 16-byte aligned KV rows, got {row_bytes}B."
            )

        self._layer_ready_event: torch.Event | None = None
        self._swap_staged = False
        self._swap_step = 0
        self.is_group_leader = index_group is None
        if index_group is None:
            index_group = HiSparseIndexGroup(
                self.device,
                max_num_reqs,
                self.region_stride,
                max_swap_rows or max_num_reqs,
                config.top_k,
                copy_stream,
                logical_topk_ready,
            )
            index_group.leader = self
        else:
            index_group.followers.append(self)
        self.index_group = index_group
        index_group.stats_row_bytes += row_bytes

        self.eager_host_mirror = config.eager_host_mirror
        self.resident_source_index = -1
        self.request_state_indices: torch.Tensor | None = None
        self.shared_host_region: SharedOffloadRegion | None = None

    @property
    def host_cache(self) -> torch.Tensor:
        return self._host_cache

    def bind_hot_cache(
        self,
        raw_tensor: torch.Tensor,
        *,
        byte_offset: int,
        block_stride: int,
        num_blocks: int,
        block_size: int,
        block_table: torch.Tensor,
    ) -> None:
        """Bind this runtime's strided view into the shared GPU HMA slab."""
        self.hot = PagedCacheView.bind(
            raw_tensor,
            dtype=self.kv_dtype,
            row_width=self.row_width,
            byte_offset=byte_offset,
            block_stride=block_stride,
            num_blocks=num_blocks,
            block_size=block_size,
        )
        self.hot_backing = raw_tensor
        self.hot_block_table = block_table

    def bind_source_cache(
        self,
        kv_cache: torch.Tensor,
        *,
        registered_host_pool: torch.Tensor | None = None,
    ) -> None:
        if kv_cache.dtype != self.kv_dtype or kv_cache.shape[-1] != self.row_width:
            raise ValueError(
                "HiSparse runtime bound to a KV cache with mismatched "
                f"layout: expected ({self.row_width}, {self.kv_dtype}), got "
                f"({kv_cache.shape[-1]}, {kv_cache.dtype})."
            )
        # Host-resident pool: the cache itself is the only full-size store;
        # every allocated slot is written before the indexer can select it.
        if kv_cache.device.type != "cpu":
            raise ValueError(
                "HiSparse requires a host-resident KV pool; got a KV cache "
                f"on {kv_cache.device}."
            )
        # The pool is pinned via cudaHostRegister (exact-size, deterministic
        # unpin at shutdown); torch's is_pinned() only recognizes its own
        # caching-host-allocator memory, so also accept ranges the model
        # allocator explicitly registered.
        if not (kv_cache.is_pinned() or registered_host_pool is not None):
            raise ValueError("HiSparse host-resident KV pool must be pinned memory.")

        self._host_cache = kv_cache.view(-1, kv_cache.shape[-1])
        self.registered_host_pool = (
            registered_host_pool if registered_host_pool is not None else kv_cache
        )

    def gather_prefill_cache(
        self,
        kv_cache: torch.Tensor,
        plan: HiSparsePrefillStagingPlan,
        resident_cache: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Gather this runtime's host cache using a shared staging plan."""
        if kv_cache.shape[1] != plan.block_size:
            raise ValueError(
                f"HiSparse staging block size {plan.block_size} does not match "
                f"the KV cache block size {kv_cache.shape[1]}."
            )
        row_width = kv_cache.shape[-1]

        staged = torch.empty(
            (
                plan.row_ids.shape[1] // plan.block_size,
                plan.block_size,
                row_width,
            ),
            dtype=kv_cache.dtype,
            device=plan.block_table.device,
        )
        if plan.gpu_row_ids is not None and resident_cache is not None:
            gpu_rows = plan.gpu_row_ids[0].to(torch.long)
            src = gpu_rows.clamp_min(0)
            resident_block_size = resident_cache.shape[1]
            rows = resident_cache[src // resident_block_size, src % resident_block_size]
            if rows.dtype != staged.dtype:
                rows = rows.contiguous().view(staged.dtype)
            staged.view(-1, row_width).copy_(rows)
        torch.ops._C_cache_ops.hisparse_gather_plan(
            kv_cache.view(-1, row_width),
            staged,
            plan.row_ids,
            plan.dst_rows,
            plan.miss_mask,
            None,
            None,
            0,
        )
        return staged

    def reset_hot_state(self) -> None:
        """Drop all hot-buffer bookkeeping (hits become misses)."""
        group = self.index_group
        compute_stream = current_stream()
        compute_stream.wait_stream(group.copy_stream)
        group.device_global_indices.fill_(-1)
        group.lru_slots.copy_(group.lru_init.expand_as(group.lru_slots))
        group.swap_stats.zero_()
        group.copy_stream.wait_stream(compute_stream)

    def invalidate_sorted_slots(
        self,
        sorted_slots: torch.Tensor,
        state_indices: torch.Tensor,
    ) -> None:
        """Drop hot copies using shared, preprocessed invalidation inputs."""
        device_global_indices = self.index_group.device_global_indices
        active_indices = device_global_indices.index_select(0, state_indices)
        positions = torch.searchsorted(sorted_slots, active_indices)
        positions.clamp_(max=sorted_slots.numel() - 1)
        active_indices.masked_fill_(
            sorted_slots[positions] == active_indices,
            -1,
        )
        device_global_indices.index_copy_(0, state_indices, active_indices)

    def invalidate_written_slots(
        self,
        written_slots: torch.Tensor,
        req_id_per_token: torch.Tensor,
    ) -> None:
        num_tokens = min(written_slots.numel(), req_id_per_token.numel())
        if num_tokens == 0:
            return
        assert self.request_state_indices is not None
        torch.ops._C_cache_ops.hisparse_invalidate_written_slots(
            self.index_group.device_global_indices,
            self.request_state_indices,
            req_id_per_token[:num_tokens],
            written_slots[:num_tokens],
        )

    def begin_forward(self) -> None:
        self._swap_step = 0

    def _step_rows(self, num_tokens: int) -> slice:
        start = self._swap_step * num_tokens
        stop = start + num_tokens
        if stop > self.index_group.shared_topk.physical_topk_indices.shape[0]:
            raise ValueError(
                "HiSparse swap rows exceed the configured speculative decode "
                f"capacity: stop={stop}, capacity="
                f"{self.index_group.shared_topk.physical_topk_indices.shape[0]}."
            )
        self._swap_step += 1
        return slice(start, stop)

    def _resolve_residency(
        self,
        *,
        resident: HiSparseCacheHandle | None = None,
        req_id_per_token: torch.Tensor,
        block_table: torch.Tensor,
        topk_indices: torch.Tensor,
        block_size: int,
        return_valid_counts: bool = False,
        shared_rows: slice,
        attention_indices_out: torch.Tensor | None = None,
        valid_counts_out: torch.Tensor | None = None,
    ) -> None:
        """Resolve logical top-k positions and compact rows requiring swaps."""
        num_tokens = topk_indices.shape[0]
        hot = self.hot
        host_cache = self.host_cache
        assert hot.block_size == block_size
        request_state_indices = self.request_state_indices
        assert request_state_indices is not None
        group = self.index_group
        shared = group.shared_topk

        device_topk_rows = shared.device_topk_rows[shared_rows]
        valid_topk_counts = None
        if return_valid_counts:
            valid_topk_counts = (
                shared.valid_topk_counts[shared_rows]
                if valid_counts_out is None
                else valid_counts_out
            )
        physical_topk_indices = (
            shared.physical_topk_indices[shared_rows]
            if attention_indices_out is None
            else attention_indices_out
        )
        swap_host_physical_rows = shared.swap_host_physical_rows[shared_rows]
        swap_device_physical_rows = shared.swap_device_physical_rows[shared_rows]
        swap_counts = shared.swap_counts[shared_rows]

        # Padded rows are skipped by the kernel (request_state_indices) and must
        # come out as -1 so the attention kernel masks them.
        torch.ops._C_cache_ops.hisparse_resolve_residency(
            host_cache,
            hot.cache,
            self.hot_block_table,
            topk_indices,
            device_topk_rows,
            group.device_global_indices,
            group.lru_slots,
            request_state_indices,
            self.region_stride,
            None,
            group.swap_stats,
            physical_topk_indices,
            hot.attention_block_stride,
            req_id_per_token[:num_tokens].contiguous(),
            block_table,
            block_size,
            None,
            valid_topk_counts,
            swap_host_physical_rows,
            swap_device_physical_rows,
            swap_counts,
            resident.block_table if resident is not None else None,
            resident.view.block_size
            if resident is not None and resident.view is not None
            else 0,
            0,
        )

    def _swap_rows(self, shared_rows: slice) -> None:
        hot = self.hot
        shared = self.index_group.shared_topk
        torch.ops._C_cache_ops.hisparse_gather_compact(
            self.host_cache,
            hot.cache,
            shared.swap_host_physical_rows[shared_rows],
            shared.swap_device_physical_rows[shared_rows],
            shared.swap_counts[shared_rows],
        )

    def _resolve_and_stage_group(
        self,
        *,
        resident: HiSparseCacheHandle,
        req_id_per_token: torch.Tensor,
        block_table: torch.Tensor,
        logical_topk_indices: torch.Tensor,
        block_size: int,
        return_valid_counts: bool,
        shared_rows: slice,
        attention_indices_out: torch.Tensor | None,
        valid_counts_out: torch.Tensor | None,
    ) -> None:
        group = self.index_group
        compute_stream = current_stream()
        if group.logical_topk_ready is not None and not in_piecewise_cudagraph():
            group.copy_stream.wait_event(group.logical_topk_ready)
        else:
            group.copy_stream.wait_stream(compute_stream)
        with group.copy_stream:
            self._resolve_residency(
                resident=resident,
                req_id_per_token=req_id_per_token,
                block_table=block_table,
                topk_indices=logical_topk_indices,
                block_size=block_size,
                return_valid_counts=return_valid_counts,
                shared_rows=shared_rows,
                attention_indices_out=attention_indices_out,
                valid_counts_out=valid_counts_out,
            )
            runtimes = [self]
            if resident.decode_batch:
                runtimes.extend(group.followers)
            for runtime in runtimes:
                runtime._swap_rows(shared_rows)
                if runtime._layer_ready_event is None:
                    runtime._layer_ready_event = torch.Event()
                runtime._layer_ready_event.record(group.copy_stream)
                runtime._swap_staged = True

    def _stage_this_layer(self, shared_rows: slice) -> None:
        group = self.index_group
        compute_stream = current_stream()
        group.copy_stream.wait_stream(compute_stream)
        with group.copy_stream:
            self._swap_rows(shared_rows)
            if self._layer_ready_event is None:
                self._layer_ready_event = torch.Event()
            self._layer_ready_event.record(group.copy_stream)
            self._swap_staged = True

    def wait_for_staged_swap(self) -> None:
        if not self._swap_staged or self._layer_ready_event is None:
            raise RuntimeError("HiSparse layer swap was not prefetched.")
        current_stream().wait_event(self._layer_ready_event)
        self._swap_staged = False

    def swap_in(
        self,
        *,
        resident: HiSparseCacheHandle,
        req_id_per_token: torch.Tensor,
        block_table: torch.Tensor,
        logical_topk_indices: torch.Tensor,
        block_size: int,
        return_valid_counts: bool = False,
        attention_indices_out: torch.Tensor | None = None,
        valid_counts_out: torch.Tensor | None = None,
    ) -> HiSparseTopKResult:
        """Resolve once per group and ensure this layer's rows are available."""
        num_tokens = logical_topk_indices.shape[0]
        shared_rows = self._step_rows(num_tokens)
        group = self.index_group
        hot = self.hot
        assert hot.block_size == block_size
        assert hot.attention_block_stride == group.leader.hot.attention_block_stride

        if self.is_group_leader:
            self._resolve_and_stage_group(
                resident=resident,
                req_id_per_token=req_id_per_token,
                block_table=block_table,
                logical_topk_indices=logical_topk_indices,
                block_size=block_size,
                return_valid_counts=return_valid_counts,
                shared_rows=shared_rows,
                attention_indices_out=attention_indices_out,
                valid_counts_out=valid_counts_out,
            )

        if not self._swap_staged:
            self._stage_this_layer(shared_rows)
        self.wait_for_staged_swap()

        physical_topk_indices = group.shared_topk.physical_topk_indices[shared_rows]
        if return_valid_counts:
            return (
                physical_topk_indices,
                group.shared_topk.valid_topk_counts[shared_rows],
            )
        return physical_topk_indices


class HiSparseCacheHandle:
    """Attention-facing handle for resident KV and sparse offload state."""

    def __init__(self, runtime: HiSparseRuntime) -> None:
        self.view: PagedCacheView | None = None
        self.block_table: torch.Tensor | None = None
        self.source_block_table: torch.Tensor | None = None
        self.slot_mapping: torch.Tensor | None = None
        self.runtime = runtime
        self.dummy_batch = False
        self.decode_batch = False
        self.all_context_pages_resident = True
        self.num_actual_tokens = 0
        self.num_decode_tokens = 0
        self.req_id_per_token: torch.Tensor | None = None
        self.host_mirror_required = False
        self.mirror_from_resident = False
        self.mirror_slot_mapping: torch.Tensor | None = None
        self.mirror_staging_cache: torch.Tensor | None = None
        self.mirror_staging_slots: torch.Tensor | None = None
        self.submit_layer_mirror: Callable[[], None] | None = None
        self.index_group_caches: list[HiSparseCacheHandle] = [self]

    def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:
        assert self.runtime.is_group_leader
        for cache in self.index_group_caches:
            cache._prepare_for_batch(attn_metadata)

    def _prepare_for_batch(self, attn_metadata: Any | None) -> None:
        self.dummy_batch = attn_metadata is None
        self.runtime.begin_forward()
        self.num_actual_tokens = (
            attn_metadata.num_actual_tokens if attn_metadata is not None else 0
        )
        self.num_decode_tokens = (
            attn_metadata.num_decode_tokens if attn_metadata is not None else 0
        )
        self.req_id_per_token = (
            attn_metadata.req_id_per_token if attn_metadata is not None else None
        )
        self.decode_batch = is_hisparse_decode_batch(attn_metadata)
        self.host_mirror_required = attn_metadata is not None and (
            not self.decode_batch or self.runtime.eager_host_mirror
        )

    def write_target(
        self, num_input_rows: int, num_slot_rows: int
    ) -> tuple[torch.Tensor, torch.Tensor, int]:
        assert self.view is not None and self.slot_mapping is not None
        num_rows = min(num_input_rows, num_slot_rows, self.num_actual_tokens)
        if self.decode_batch:
            num_rows = min(num_rows, self.runtime.max_num_reqs)
        return self.view.cache, self.slot_mapping[:num_rows], num_rows

    def mirror_write_target(
        self, num_rows: int
    ) -> tuple[torch.Tensor, torch.Tensor] | None:
        if (
            not self.host_mirror_required
            or self.decode_batch
            or self.mirror_from_resident
        ):
            return None
        cache = self.mirror_staging_cache
        slots = self.mirror_staging_slots
        if cache is None or slots is None:
            raise RuntimeError("HiSparse prefill mirror staging is not bound.")
        if num_rows > slots.numel():
            raise RuntimeError(
                "HiSparse prefill mirror exceeds staging capacity: "
                f"{num_rows} > {slots.numel()}."
            )
        return cache, slots[:num_rows]

    def finish_kv_update(self) -> None:
        if self.dummy_batch:
            return
        if (
            self.submit_layer_mirror is not None
            and not self.decode_batch
            and get_forward_context().cudagraph_runtime_mode != CUDAGraphMode.FULL
        ):
            self.submit_layer_mirror()

    def bind_cache(
        self,
        raw_tensor: torch.Tensor,
        *,
        byte_offset: int,
        block_stride: int,
        num_blocks: int,
        block_size: int,
        block_table: torch.Tensor,
        slot_mapping: torch.Tensor,
    ) -> None:
        self.view = PagedCacheView.bind(
            raw_tensor,
            dtype=self.runtime.kv_dtype,
            row_width=self.runtime.row_width,
            byte_offset=byte_offset,
            block_stride=block_stride,
            num_blocks=num_blocks,
            block_size=block_size,
        )
        self.block_table = block_table
        self.slot_mapping = slot_mapping

    def swap_in(
        self,
        req_id_per_token: torch.Tensor,
        block_table: torch.Tensor,
        logical_topk_indices: torch.Tensor,
        *,
        block_size: int,
        return_valid_counts: bool = False,
        attention_indices_out: torch.Tensor | None = None,
        valid_counts_out: torch.Tensor | None = None,
    ) -> HiSparseTopKResult:
        return self.runtime.swap_in(
            resident=self,
            req_id_per_token=req_id_per_token[: logical_topk_indices.shape[0]],
            block_table=block_table,
            logical_topk_indices=logical_topk_indices,
            block_size=block_size,
            return_valid_counts=return_valid_counts,
            attention_indices_out=attention_indices_out,
            valid_counts_out=valid_counts_out,
        )


def initialize_hisparse_runtime_buffers(
    cache_handles: list[HiSparseCacheHandle],
    *,
    max_num_reqs: int,
    max_num_batched_tokens: int,
) -> None:
    """Allocate shared request state and per-layer prefill staging buffers."""
    assert cache_handles
    resident = cache_handles[0].view
    assert resident is not None
    device = resident.cache.device
    request_state_indices = torch.full(
        (max_num_reqs,), -1, dtype=torch.int32, device=device
    )
    staging_blocks = (
        max_num_batched_tokens + resident.block_size - 1
    ) // resident.block_size
    mirror_staging_caches = torch.empty(
        (
            len(cache_handles),
            staging_blocks,
            resident.block_size,
            resident.cache.shape[-1],
        ),
        dtype=resident.cache.dtype,
        device=device,
    )
    mirror_staging_slots = torch.arange(
        max_num_batched_tokens, dtype=torch.int64, device=device
    )
    for layer_index, cache_handle in enumerate(cache_handles):
        cache_handle.runtime.request_state_indices = request_state_indices
        cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index]
        cache_handle.mirror_staging_slots = mirror_staging_slots


def create_hisparse_cache_handle(
    vllm_config: VllmConfig,
    model_top_k: int,
    *,
    is_index_group_leader: bool,
    row_width: int,
    kv_dtype: torch.dtype,
    index_group: HiSparseMLAIndexGroup | None = None,
    device: torch.device | str | None = None,
) -> HiSparseCacheHandle | None:
    config = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k)
    if config is None:
        return None

    max_num_reqs = vllm_config.scheduler_config.max_num_seqs
    max_swap_rows = _get_max_swap_rows(vllm_config)
    if device is None:
        device = torch.device(
            current_platform.device_type, torch.accelerator.current_device_index()
        )

    hisparse_group = (
        None
        if is_index_group_leader or index_group is None
        else index_group.hisparse_group
    )
    runtime = HiSparseRuntime(
        config=config,
        max_num_reqs=max_num_reqs,
        max_swap_rows=max_swap_rows,
        row_width=row_width,
        kv_dtype=kv_dtype,
        device=device,
        index_group=hisparse_group,
        copy_stream=index_group.side_stream if index_group is not None else None,
        logical_topk_ready=(
            index_group.logical_topk_ready
            if index_group is not None and index_group.has_indexer
            else None
        ),
    )
    if is_index_group_leader and index_group is not None:
        index_group.hisparse_group = runtime.index_group
    logger.info_once(
        "Enabled experimental HiSparse HMA hot cache: top_k=%d, "
        "device_buffer_size=%d (%d LRU rows), max_num_seqs=%d.",
        config.top_k,
        config.device_buffer_size,
        config.device_buffer_size,
        max_num_reqs,
    )
    handle = HiSparseCacheHandle(runtime)
    speculative_config = vllm_config.speculative_config
    handle.mirror_from_resident = bool(
        vllm_config.scheduler_config.async_scheduling
        and speculative_config is not None
        and speculative_config.uses_draft_kv_cache()
    )
    return handle