vllm.v1.hisparse.runtime

HiSparse resident, host, and hot-cache data plane.

Classes:

• HiSparseCacheHandle –

  Attention-facing handle for resident KV and sparse offload state.

• HiSparseHostPool –

  Keep the single host backing and its exact registration range.

• HiSparseIndexGroup –

  GPU index-resolution state shared by one leader and its followers.

• HiSparsePrefillStagingPlan –
• HiSparseRuntime –

  Per-cache host/hot data plane and GPU replacement state.

• HiSparseSlotInvalidator –

  Drop hot copies of rewritten host slots for many index groups at once.

Functions:

• allocate_hisparse_host_pools –

  Allocate one host mmap shared by local TP ranks.

• allocate_pinned_host_pool –

  Allocate and deterministically register an exact-size host KV region.

• build_hisparse_prefill_staging_plan –

  Build an asynchronous layer-independent host-cache staging remap.

• check_hisparse_host_memory –

  Reject a physical host-pool allocation that cannot fit in RAM.

• get_hisparse_host_pool –

  Return the engine's host pool, which outlives the profiling KV cache.

• hisparse_prefill_staging_remap –

  Renumber blocks without a data-dependent CUDA output allocation.

• initialize_hisparse_runtime_buffers –

  Allocate the request state shared by every HiSparse runtime.

• is_hisparse_group_prefetch_batch –

  Prefetch historical KV for a pure decode batch, including verification.

• release_pinned_state –

  Synchronize and release registered host KV pools.

• update_hisparse_residency –

  Apply scheduled requests' residency changes.

• use_shared_hisparse_host_pool –

  Whether replicated MLA host KV can share one local mmap.

HiSparseCacheHandle

Attention-facing handle for resident KV and sparse offload state.

Methods:

• unbind_cache –

  Detach KV storage and per-binding metadata, retaining the index group.

Source code in vllm/v1/hisparse/runtime.py

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

    def __init__(self, runtime: HiSparseRuntime) -> None:
        self.runtime = runtime
        # Speculator layers write their rows after the target forward.
        self.draft_layer = False
        self.index_group_caches: list[HiSparseCacheHandle] = [self]
        self._reset_binding_state()

    def _reset_binding_state(self) -> None:
        self.view: PagedCacheView | None = None
        self.residency: torch.Tensor | 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.dummy_batch = False
        self.decode_batch = False
        self.prefetch_decode_group = 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_slot_mapping: torch.Tensor | None = None
        self.submit_layer_mirror: Callable[[], None] | None = None

    def unbind_cache(self) -> None:
        """Detach KV storage and per-binding metadata, retaining the index group."""
        self._reset_binding_state()
        self.runtime.unbind_cache()

    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)
        # Verification can prefetch old host rows without publishing tentative
        # writes. Resident writes and host mirroring keep their existing policy.
        self.prefetch_decode_group = self.decode_batch or (
            is_hisparse_group_prefetch_batch(attn_metadata, self.runtime.max_swap_rows)
        )
        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 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,
        num_valid_rows: torch.Tensor,
        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,
            num_valid_rows=num_valid_rows,
        )

unbind_cache()

Detach KV storage and per-binding metadata, retaining the index group.

Source code in vllm/v1/hisparse/runtime.py

def unbind_cache(self) -> None:
    """Detach KV storage and per-binding metadata, retaining the index group."""
    self._reset_binding_state()
    self.runtime.unbind_cache()

HiSparseHostPool

Keep the single host backing and its exact registration range.

Source code in vllm/v1/hisparse/runtime.py

class HiSparseHostPool:
    """Keep the single host backing and its exact 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
        self.pinned: list[torch.Tensor] = []

    def allocate(self, size: int) -> torch.Tensor:
        if self.backing is not None:
            assert self.backing.numel() == size, "HiSparse host pool size changed."
            return self.backing
        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.pinned, self.shared_region = allocate_hisparse_host_pools(
                self.vllm_config,
                tensor_sizes,
                config.hisparse_host_num_blocks,
                config.hisparse_host_block_stride,
            )
            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)
            self.pinned = [self.registered]
        return self.backing

HiSparseIndexGroup

GPU index-resolution state shared by one leader and its followers.

Source code in vllm/v1/hisparse/runtime.py

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,
        )
        # NPU arange does not support int16 output, so generate int32 then cast.
        self.lru_init = torch.arange(
            region_stride, dtype=torch.int32, device=device
        ).to(torch.int16)
        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] = []
        stats_dtype = get_hisparse_backend().stats_dtype
        self.swap_stats = torch.zeros(2, dtype=stats_dtype, device=device)
        self.swap_stats_host = torch.empty(
            2, dtype=stats_dtype, device="cpu", pin_memory=True
        )
        self.stats_row_bytes = 0

HiSparsePrefillStagingPlan dataclass

Methods:

• ensure_gpu_sources –

  Resolve which staged rows can be served from the resident cache.

Source code in vllm/v1/hisparse/runtime.py

@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_state_rows: torch.Tensor,
        state_indices: torch.Tensor,
        resident_block_size: int,
    ) -> None:
        """Resolve which staged rows can be served from the resident cache.

        ``resident_state_rows`` is a resident group's persistent table by
        request state row and ``state_indices`` holds each staged request's
        state row. Computed once per plan and resident group (every layer in
        the group shares the table); non-null resident pages become
        miss_mask=0 rows gathered device-to-device by ``gather_prefill_cache``.
        """
        source_key = (resident_state_rows.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 state_indices.shape[0] < num_rows:
            return
        resident_block_table = resident_state_rows.index_select(
            0, state_indices[:num_rows].clamp(min=0)
        )
        # 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

ensure_gpu_sources(resident_state_rows, state_indices, resident_block_size)

Resolve which staged rows can be served from the resident cache.

resident_state_rows is a resident group's persistent table by request state row and state_indices holds each staged request's state row. Computed once per plan and resident group (every layer in the group shares the table); non-null resident pages become miss_mask=0 rows gathered device-to-device by gather_prefill_cache.

Source code in vllm/v1/hisparse/runtime.py

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

    ``resident_state_rows`` is a resident group's persistent table by
    request state row and ``state_indices`` holds each staged request's
    state row. Computed once per plan and resident group (every layer in
    the group shares the table); non-null resident pages become
    miss_mask=0 rows gathered device-to-device by ``gather_prefill_cache``.
    """
    source_key = (resident_state_rows.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 state_indices.shape[0] < num_rows:
        return
    resident_block_table = resident_state_rows.index_select(
        0, state_indices[:num_rows].clamp(min=0)
    )
    # 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

HiSparseRuntime

Per-cache host/hot data plane and GPU replacement state.

Methods:

• bind_hot_cache –

  Bind this runtime's strided view into the shared GPU HMA slab.

• gather_prefill_cache –

  Gather this runtime's host cache using a shared staging plan.

• invalidate_sorted_slots –

  Drop hot copies using shared, preprocessed invalidation inputs.

• reset_hot_state –

  Drop all hot-buffer bookkeeping (hits become misses).

• swap_in –

  Resolve once per group and ensure this layer's rows are available.

• unbind_cache –

  Restore the unbound state without unregistering the host pool.

Source code in vllm/v1/hisparse/runtime.py

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:
        self.backend = get_hisparse_backend()
        self.backend.check_available()
        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
        self.max_union_rows = config.max_union_rows
        if self.device.type == "cuda":
            _check_residency_shared_memory(self.device, config)

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

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

    def _reset_binding_state(self) -> None:
        self.cache_ops: Any = None
        self.host_pool: HiSparseHostPool | None = None
        self.request_state_indices: torch.Tensor | None = None
        self.resident_source_index = -1
        self._layer_ready_event: torch.Event | None = None
        self._swap_staged = False
        self._swap_step = 0
        for name in (
            "hot",
            "hot_backing",
            "hot_block_table",
            "_host_cache",
            "registered_host_pool",
        ):
            if hasattr(self, name):
                delattr(self, name)

    def unbind_cache(self) -> None:
        """Restore the unbound state without unregistering the host pool."""
        self._reset_binding_state()

    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
        )
        self.cache_ops = self.backend.create_cache_ops(self)

    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]

        num_blocks = plan.row_ids.shape[1] // plan.block_size
        min_blocks = cdiv(_MIN_STAGED_PREFILL_KV_ROWS, plan.block_size)
        staged = torch.empty(
            (max(num_blocks, min_blocks), plan.block_size, row_width),
            dtype=kv_cache.dtype,
            device=plan.block_table.device,
        )[:num_blocks]
        if plan.gpu_row_ids is not None and resident_cache is not None:
            self.backend.gather_resident_rows(
                resident_cache, staged.view(-1, row_width), plan.gpu_row_ids[0]
            )
        self.cache_ops.hisparse_gather_plan(
            kv_cache.view(-1, row_width),
            staged,
            plan.row_ids,
            plan.dst_rows,
            plan.miss_mask,
        )
        return staged

    def reset_hot_state(self) -> None:
        """Drop all hot-buffer bookkeeping (hits become misses)."""
        group = self.index_group
        compute_stream = torch.accelerator.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 begin_forward(self) -> None:
        self._swap_step = 0

    @property
    def max_swap_rows(self) -> int:
        return self.index_group.shared_topk.physical_topk_indices.shape[0]

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

    def _resolve_residency(
        self,
        *,
        resident: HiSparseCacheHandle | None = None,
        req_id_per_token: torch.Tensor,
        num_valid_rows: 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]

        # Rows past num_valid_rows and rows without request state come out as
        # -1 so the attention kernel masks them.
        self.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,
            self.max_union_rows,
            group.swap_stats,
            physical_topk_indices,
            hot.attention_block_stride,
            req_id_per_token[:num_tokens].contiguous(),
            block_table,
            block_size,
            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,
            num_valid_rows,
        )

    def _swap_rows(self, shared_rows: slice) -> None:
        hot = self.hot
        shared = self.index_group.shared_topk
        self.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,
        num_valid_rows: torch.Tensor,
    ) -> None:
        group = self.index_group
        compute_stream = torch.accelerator.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,
                num_valid_rows=num_valid_rows,
                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.prefetch_decode_group:
                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 = torch.accelerator.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.")
        torch.accelerator.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,
        num_valid_rows: torch.Tensor,
        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,
                num_valid_rows=num_valid_rows,
            )

        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

_resolve_residency(*, resident=None, req_id_per_token, num_valid_rows, block_table, topk_indices, block_size, return_valid_counts=False, shared_rows, attention_indices_out=None, valid_counts_out=None)

Resolve logical top-k positions and compact rows requiring swaps.

Source code in vllm/v1/hisparse/runtime.py

def _resolve_residency(
    self,
    *,
    resident: HiSparseCacheHandle | None = None,
    req_id_per_token: torch.Tensor,
    num_valid_rows: 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]

    # Rows past num_valid_rows and rows without request state come out as
    # -1 so the attention kernel masks them.
    self.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,
        self.max_union_rows,
        group.swap_stats,
        physical_topk_indices,
        hot.attention_block_stride,
        req_id_per_token[:num_tokens].contiguous(),
        block_table,
        block_size,
        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,
        num_valid_rows,
    )

bind_hot_cache(raw_tensor, *, byte_offset, block_stride, num_blocks, block_size, block_table)

Bind this runtime's strided view into the shared GPU HMA slab.

Source code in vllm/v1/hisparse/runtime.py

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

gather_prefill_cache(kv_cache, plan, resident_cache=None)

Gather this runtime's host cache using a shared staging plan.

Source code in vllm/v1/hisparse/runtime.py

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]

    num_blocks = plan.row_ids.shape[1] // plan.block_size
    min_blocks = cdiv(_MIN_STAGED_PREFILL_KV_ROWS, plan.block_size)
    staged = torch.empty(
        (max(num_blocks, min_blocks), plan.block_size, row_width),
        dtype=kv_cache.dtype,
        device=plan.block_table.device,
    )[:num_blocks]
    if plan.gpu_row_ids is not None and resident_cache is not None:
        self.backend.gather_resident_rows(
            resident_cache, staged.view(-1, row_width), plan.gpu_row_ids[0]
        )
    self.cache_ops.hisparse_gather_plan(
        kv_cache.view(-1, row_width),
        staged,
        plan.row_ids,
        plan.dst_rows,
        plan.miss_mask,
    )
    return staged

invalidate_sorted_slots(sorted_slots, state_indices)

Drop hot copies using shared, preprocessed invalidation inputs.

Source code in vllm/v1/hisparse/runtime.py

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)

reset_hot_state()

Drop all hot-buffer bookkeeping (hits become misses).

Source code in vllm/v1/hisparse/runtime.py

def reset_hot_state(self) -> None:
    """Drop all hot-buffer bookkeeping (hits become misses)."""
    group = self.index_group
    compute_stream = torch.accelerator.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)

swap_in(*, resident, req_id_per_token, block_table, logical_topk_indices, block_size, num_valid_rows, return_valid_counts=False, attention_indices_out=None, valid_counts_out=None)

Resolve once per group and ensure this layer's rows are available.

Source code in vllm/v1/hisparse/runtime.py

def swap_in(
    self,
    *,
    resident: HiSparseCacheHandle,
    req_id_per_token: torch.Tensor,
    block_table: torch.Tensor,
    logical_topk_indices: torch.Tensor,
    block_size: int,
    num_valid_rows: torch.Tensor,
    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,
            num_valid_rows=num_valid_rows,
        )

    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

unbind_cache()

Restore the unbound state without unregistering the host pool.

Source code in vllm/v1/hisparse/runtime.py

def unbind_cache(self) -> None:
    """Restore the unbound state without unregistering the host pool."""
    self._reset_binding_state()

HiSparseSlotInvalidator

Drop hot copies of rewritten host slots for many index groups at once.

Source code in vllm/v1/hisparse/runtime.py

class HiSparseSlotInvalidator:
    """Drop hot copies of rewritten host slots for many index groups at once."""

    def __init__(self, request_state_indices: torch.Tensor) -> None:
        self.request_state_indices = request_state_indices
        self._group_tables: dict[tuple[int, ...], torch.Tensor] = {}

    def __call__(
        self,
        leaders: list[HiSparseRuntime],
        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 or not leaders:
            return
        tables = [leader.index_group.device_global_indices for leader in leaders]
        key = tuple(table.data_ptr() for table in tables)
        group_table = self._group_tables.get(key)
        if group_table is None:
            if any(table.shape != tables[0].shape for table in tables):
                raise ValueError("HiSparse index groups must share one hot layout.")
            group_table = async_tensor_h2d(
                list(key), self.request_state_indices.device, torch.int64
            )
            self._group_tables[key] = group_table
        num_state_rows, region_stride = tables[0].shape
        leaders[0].cache_ops.hisparse_invalidate_written_slots(
            group_table,
            self.request_state_indices,
            req_id_per_token[:num_tokens],
            written_slots[:num_tokens],
            num_state_rows,
            region_stride,
        )

_SharedTopKState

CUDA-graph-safe residency result shared by index-sharing layers.

Source code in vllm/v1/hisparse/runtime.py

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)

_check_residency_shared_memory(device, config)

Reject a config whose residency resolver exceeds per-block shared memory.

Source code in vllm/v1/hisparse/runtime.py

def _check_residency_shared_memory(
    device: torch.device, config: ResolvedHiSparseConfig
) -> None:
    """Reject a config whose residency resolver exceeds per-block shared memory."""
    # Minimum resolver layout: union hash keys and values (plus an empty slot),
    # per-chunk offsets, five counters, done bits, and the int16 compacted LRU.
    # The kernel uses remaining shared memory to lower the hash-table load.
    hot_size = config.device_buffer_size
    num_chunks = cdiv(hot_size, 32)
    required = 4 * (2 * (config.max_union_rows + 1) + 3 * num_chunks + 7) + 2 * hot_size
    device_index = (
        device.index
        if device.index is not None
        else torch.accelerator.current_device_index()
    )
    available = get_max_shared_memory_bytes(device_index)
    if required > available:
        raise ValueError(
            "HiSparse residency resolution needs "
            f"{required} B of shared memory per block for "
            f"device_buffer_size={config.device_buffer_size} and "
            f"{config.max_union_rows} rows per request "
            "(decode query length x index_topk), but the device allows "
            f"{available} B. Lower device_buffer_size or num_speculative_tokens."
        )

allocate_hisparse_host_pools(vllm_config, tensor_sizes, num_blocks, host_block_stride)

Allocate one host mmap shared by local TP ranks.

Source code in vllm/v1/hisparse/runtime.py

def allocate_hisparse_host_pools(
    vllm_config: VllmConfig,
    tensor_sizes: list[int],
    num_blocks: int,
    host_block_stride: int,
) -> tuple[list[torch.Tensor], SharedOffloadRegion]:
    """Allocate one host mmap shared by local TP ranks."""
    tp_group = get_tp_group()
    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=tp_group.barrier,
        creator_memory_check=check_hisparse_host_memory,
        populate_shard=(tp_group.rank_in_group, tp_group.world_size),
        num_segments=REGION_SEGMENTS,
    )
    pinned: list[torch.Tensor] = []
    try:
        for start, end in host_registration_ranges(
            tensor_sizes, num_blocks, num_blocks * host_block_stride
        ):
            tensor = region.base_tensor[start:end]
            pin_tensor(tensor)
            pinned.append(tensor)
    except Exception:
        release_pinned_state(pinned, region)
        raise
    return pinned, region

allocate_pinned_host_pool(size)

Allocate and deterministically register an exact-size host KV region.

Source code in vllm/v1/hisparse/runtime.py

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

build_hisparse_prefill_staging_plan(block_table, seq_lens, block_size, staging_block_capacity)

Build an asynchronous layer-independent host-cache staging remap.

Source code in vllm/v1/hisparse/runtime.py

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

check_hisparse_host_memory(pool_bytes)

Reject a physical host-pool allocation that cannot fit in RAM.

Source code in vllm/v1/hisparse/runtime.py

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."
        )

get_hisparse_host_pool(vllm_config, kv_cache_config)

Return the engine's host pool, which outlives the profiling KV cache.

Source code in vllm/v1/hisparse/runtime.py

def get_hisparse_host_pool(
    vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
) -> HiSparseHostPool:
    """Return the engine's host pool, which outlives the profiling KV cache."""
    instance_id = vllm_config.instance_id
    if instance_id not in _host_pools:
        _host_pools[instance_id] = HiSparseHostPool(vllm_config, kv_cache_config)
    return _host_pools[instance_id]

hisparse_prefill_staging_remap(block_table, block_size)

Renumber blocks without a data-dependent CUDA output allocation.

Source code in vllm/v1/hisparse/runtime.py

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

initialize_hisparse_runtime_buffers(cache_handles, *, max_num_reqs)

Allocate the request state shared by every HiSparse runtime.

Source code in vllm/v1/hisparse/runtime.py

def initialize_hisparse_runtime_buffers(
    cache_handles: list[HiSparseCacheHandle],
    *,
    max_num_reqs: int,
) -> None:
    """Allocate the request state shared by every HiSparse runtime."""
    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
    )
    for cache_handle in cache_handles:
        cache_handle.runtime.request_state_indices = request_state_indices

is_hisparse_group_prefetch_batch(attn_metadata, max_rows)

Prefetch historical KV for a pure decode batch, including verification.

Source code in vllm/v1/hisparse/runtime.py

def is_hisparse_group_prefetch_batch(attn_metadata: Any | None, max_rows: int) -> bool:
    """Prefetch historical KV for a pure decode batch, including verification."""
    return (
        attn_metadata is not None
        and 0 < attn_metadata.num_actual_tokens <= max_rows
        and attn_metadata.num_decode_tokens == attn_metadata.num_actual_tokens
        and 0 < attn_metadata.num_reqs <= attn_metadata.num_actual_tokens
        and 0 < attn_metadata.max_query_len <= max_rows
    )

release_pinned_state(pinned_host_pools, shared_host_region=None)

Synchronize and release registered host KV pools.

Source code in vllm/v1/hisparse/runtime.py

def release_pinned_state(
    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:
        backend = get_hisparse_backend()
        try:
            torch.accelerator.synchronize()
        except RuntimeError as e:
            logger.warning(
                "HiSparse: device context unusable at teardown (%s); leaving "
                "%d host-pool tensors pinned for kernel exit reclaim.",
                e,
                len(pinned_host_pools),
            )
            return

        release_start = time.perf_counter()
        freed_bytes = 0
        while pinned_host_pools:
            tensor = pinned_host_pools[-1]
            try:
                backend.unpin_tensor(tensor)
            except RuntimeError as error:
                logger.warning(
                    "HiSparse: host unregister failed (%s); leaving "
                    "%d host-pool tensors pinned for kernel exit reclaim.",
                    error,
                    len(pinned_host_pools),
                )
                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,
            )

    if shared_host_region is not None:
        shared_host_region.cleanup()

update_hisparse_residency(residency, updates, request_ids, request_state_indices)

Apply scheduled requests' residency changes.

residency is [resident groups, state rows, pages]: the GPU block of each page, or block 0 when the page is read from the host.

Source code in vllm/v1/hisparse/runtime.py

def update_hisparse_residency(
    residency: torch.Tensor,
    updates: Mapping[str, SparseKVResidencyUpdate],
    request_ids: Sequence[str],
    request_state_indices: torch.Tensor,
) -> None:
    """Apply scheduled requests' residency changes.

    ``residency`` is ``[resident groups, state rows, pages]``: the GPU block of
    each page, or block 0 when the page is read from the host.
    """
    if not updates:
        return
    batch_rows = {request_id: row for row, request_id in enumerate(request_ids)}
    num_groups = residency.shape[0]
    rows: list[int] = []
    pages: list[int] = []
    block_ids: list[list[int]] = [[] for _ in range(num_groups)]
    for request_id, update in updates.items():
        rows.append(batch_rows[request_id])
        pages.extend(update.pages)
        for group_block_ids, update_block_ids in zip(
            block_ids, update.block_ids, strict=True
        ):
            group_block_ids.extend(update_block_ids)
    host_values = np.empty((len(pages), 2 + num_groups), dtype=np.int32)
    host_values[:, 0] = np.repeat(
        rows, [len(update.pages) for update in updates.values()]
    )
    host_values[:, 1] = pages
    host_values[:, 2:] = np.asarray(block_ids, dtype=np.int32).T
    values = async_tensor_h2d(
        host_values, device=request_state_indices.device, dtype=torch.int32
    )
    state_rows = request_state_indices[values[:, 0]]
    residency[:, state_rows, values[:, 1]] = values[:, 2:].T

use_shared_hisparse_host_pool(vllm_config)

Whether replicated MLA host KV can share one local mmap.

Source code in vllm/v1/hisparse/runtime.py

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 (
        get_hisparse_backend().supports_shared_host_pool
        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
    )