sglang/python/sglang/srt/mem_cache/hiradix_cache.py

from __future__ import annotations

import atexit
import heapq
import json
import logging
import os
import threading
import time
from dataclasses import dataclass
from queue import Empty
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

import torch

from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
    DecLockRefParams,
    DecLockRefResult,
    EvictParams,
    EvictResult,
    IncLockRefResult,
    InitLoadBackParams,
    InsertParams,
    InsertResult,
    MatchPrefixParams,
    MatchResult,
)
from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.memory_pool import (
    MHATokenToKVPool,
    MLATokenToKVPool,
    NSATokenToKVPool,
)
from sglang.srt.mem_cache.radix_cache import (
    RadixCache,
    RadixKey,
    TreeNode,
    compute_node_hash_values,
    split_node_hash_value,
)
from sglang.srt.mem_cache.utils import convert_to_bigram_key
from sglang.srt.observability.metrics_collector import StorageMetricsCollector
from sglang.srt.utils import bind_to_closest_numa_node_cuda

if TYPE_CHECKING:
    from sglang.srt.mem_cache.cache_init_params import CacheInitParams
    from sglang.srt.server_args import ServerArgs

logger = logging.getLogger(__name__)


@dataclass
class CpHiCacheNodeMetadata:
    logical_len: int
    owned_positions: torch.Tensor
    host_indices: torch.Tensor
    # NEW: one int8 per logical PAGE; identical on all CP ranks (function of
    # the original logical page ids only). Captures the owner pattern of the
    # write-time allocation so load_cp can reproduce it via owner-aware alloc.
    # Without this, the saved owned_positions index a fresh alloc whose
    # per-position owner pattern is arbitrary → load writes to wrong physical
    # slots → attention reads garbage. See _write_cp for the derivation and
    # alloc_pages_with_owners for the consumer.
    page_owners: torch.Tensor
    # NEW: needed at split() time to convert split_len into a page index
    # without plumbing it from the caller.
    page_size: int
    draft_host_indices: Optional[torch.Tensor] = None

    def __post_init__(self):
        if self.logical_len < 0:
            raise ValueError(f"logical_len must be non-negative, got {self.logical_len}")
        if self.page_size <= 0:
            raise ValueError(f"page_size must be positive, got {self.page_size}")
        if self.logical_len % self.page_size != 0:
            raise ValueError(
                f"logical_len ({self.logical_len}) must be a multiple of "
                f"page_size ({self.page_size})"
            )
        self.owned_positions = self.owned_positions.to(
            device="cpu", dtype=torch.int64
        ).detach().clone()
        self.host_indices = self.host_indices.to(
            device="cpu", dtype=torch.int64
        ).detach().clone()
        self.page_owners = self.page_owners.to(
            device="cpu", dtype=torch.int8
        ).detach().clone()
        if self.draft_host_indices is not None:
            self.draft_host_indices = self.draft_host_indices.to(
                device="cpu", dtype=torch.int64
            ).detach().clone()
        expected_num_pages = self.logical_len // self.page_size
        if self.page_owners.numel() != expected_num_pages:
            raise ValueError(
                f"page_owners length ({self.page_owners.numel()}) must equal "
                f"logical_len/page_size ({expected_num_pages})"
            )
        if expected_num_pages > 0 and bool((self.page_owners < 0).any()):
            raise ValueError("page_owners entries must be non-negative")
        if self.owned_positions.numel() != self.host_indices.numel():
            raise ValueError(
                "owned_positions and host_indices must have same length, got "
                f"{self.owned_positions.numel()} and {self.host_indices.numel()}"
            )
        if (
            self.draft_host_indices is not None
            and self.owned_positions.numel() != self.draft_host_indices.numel()
        ):
            raise ValueError(
                "draft_host_indices and owned_positions must have same length, got "
                f"{self.draft_host_indices.numel()} and {self.owned_positions.numel()}"
            )
        if self.owned_positions.numel() > 0:
            if torch.any(self.owned_positions < 0) or torch.any(
                self.owned_positions >= self.logical_len
            ):
                raise ValueError("owned_positions must be in [0, logical_len)")
            if torch.any(self.owned_positions[1:] <= self.owned_positions[:-1]):
                raise ValueError(
                    "owned_positions must be sorted and strictly increasing"
                )

    def split(
        self, split_len: int
    ) -> tuple["CpHiCacheNodeMetadata", "CpHiCacheNodeMetadata"]:
        if split_len < 0 or split_len > self.logical_len:
            raise ValueError(
                f"split_len must be in [0, {self.logical_len}], got {split_len}"
            )
        if split_len % self.page_size != 0:
            raise ValueError(
                f"split_len ({split_len}) must be a multiple of page_size "
                f"({self.page_size})"
            )
        parent_mask = self.owned_positions < split_len
        child_mask = ~parent_mask
        split_pages = split_len // self.page_size
        return (
            CpHiCacheNodeMetadata(
                logical_len=split_len,
                owned_positions=self.owned_positions[parent_mask],
                host_indices=self.host_indices[parent_mask],
                page_owners=self.page_owners[:split_pages],
                page_size=self.page_size,
                draft_host_indices=(
                    self.draft_host_indices[parent_mask]
                    if self.draft_host_indices is not None
                    else None
                ),
            ),
            CpHiCacheNodeMetadata(
                logical_len=self.logical_len - split_len,
                owned_positions=self.owned_positions[child_mask] - split_len,
                host_indices=self.host_indices[child_mask],
                page_owners=self.page_owners[split_pages:],
                page_size=self.page_size,
                draft_host_indices=(
                    self.draft_host_indices[child_mask]
                    if self.draft_host_indices is not None
                    else None
                ),
            ),
        )


class HiRadixCache(RadixCache):

    def _create_token_to_kv_pool_host(self, kv_cache, server_args: ServerArgs):
        from sglang.srt.mem_cache.memory_pool_host import (
            MHATokenToKVPoolHost,
            MLATokenToKVPoolHost,
            NSATokenToKVPoolHost,
        )

        if isinstance(kv_cache, MHATokenToKVPool):
            return MHATokenToKVPoolHost(
                kv_cache,
                server_args.hicache_ratio,
                server_args.hicache_size,
                self.page_size,
                server_args.hicache_mem_layout,
                allocator_type=server_args.hicache_storage_backend,
            )
        if isinstance(kv_cache, NSATokenToKVPool):
            return NSATokenToKVPoolHost(
                kv_cache,
                server_args.hicache_ratio,
                server_args.hicache_size,
                self.page_size,
                server_args.hicache_mem_layout,
                allocator_type=server_args.hicache_storage_backend,
            )
        if isinstance(kv_cache, MLATokenToKVPool):
            return MLATokenToKVPoolHost(
                kv_cache,
                server_args.hicache_ratio,
                server_args.hicache_size,
                self.page_size,
                server_args.hicache_mem_layout,
                allocator_type=server_args.hicache_storage_backend,
            )
        raise ValueError("HiRadixCache only supports MHA, MLA, and NSA KV pools")

    def attach_draft_kv_pool(self, draft_token_to_kv_pool) -> None:
        if not self._uses_cp_hicache or draft_token_to_kv_pool is None:
            return
        if (
            getattr(self.cache_controller, "draft_mem_pool_device", None)
            is draft_token_to_kv_pool
        ):
            return
        draft_token_to_kv_pool_host = self._create_token_to_kv_pool_host(
            draft_token_to_kv_pool, self._server_args
        )
        self.cache_controller.attach_draft_pool(
            draft_token_to_kv_pool, draft_token_to_kv_pool_host
        )
        self.draft_token_to_kv_pool_host = draft_token_to_kv_pool_host
        logger.info(
            "[HiCache-draft] attached CP draft KV host pool: pool=%s host_pool=%s page_size=%d",
            draft_token_to_kv_pool.__class__.__name__,
            draft_token_to_kv_pool_host.__class__.__name__,
            self.page_size,
        )

    def __init__(self, params: CacheInitParams, server_args: ServerArgs):
        self._enable_metrics_flag = params.enable_metrics
        self._server_args = server_args

        if not server_args.disable_hicache_numa_detect:
            bind_to_closest_numa_node_cuda()

        self.page_size = params.page_size
        self._uses_cp_hicache = isinstance(
            params.token_to_kv_pool_allocator,
            CPSharedPagedTokenToKVPoolAllocator,
        )
        if self._uses_cp_hicache:
            if server_args.hicache_storage_backend is not None:
                raise ValueError(
                    "CP shared KV HiCache host integration does not support storage backends."
                )
            if not isinstance(
                params.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool
            ):
                raise ValueError(
                    "CP shared KV HiCache host integration requires NSATokenToKVPool."
                )
        self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
        self.token_to_kv_pool_host = self._create_token_to_kv_pool_host(
            self.kv_cache, server_args
        )

        self.tp_group = params.tp_cache_group
        self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
        self.pp_rank = params.pp_rank
        self.pp_size = params.pp_size
        self.enable_storage = server_args.hicache_storage_backend is not None
        self.enable_storage_metrics = self.enable_storage and params.enable_metrics
        self.extra_metric_labels = server_args.extra_metric_labels

        (
            extra_config,
            prefetch_threshold,
            prefetch_timeout_base,
            prefetch_timeout_per_ki_token,
            hicache_storage_pass_prefix_keys,
        ) = self._parse_storage_backend_extra_config(
            server_args.hicache_storage_backend_extra_config
        )
        # TODO: support more timeout check functions
        self.is_prefetch_timeout = self._prefetch_timeout_check_linear_func
        self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy

        self.load_cache_event = threading.Event()
        cp_shared_kv_layout = None
        if self._uses_cp_hicache:
            cp_shared_kv_layout = CpSharedKVLayout(
                page_size=self.page_size,
                cp_size=params.token_to_kv_pool_allocator.cp_size,
                cp_rank=params.token_to_kv_pool_allocator.cp_rank,
            )
        self.cache_controller = HiCacheController(
            params.token_to_kv_pool_allocator,
            self.token_to_kv_pool_host,
            self.page_size,
            self.tp_group,
            load_cache_event=self.load_cache_event,
            write_policy=server_args.hicache_write_policy,
            io_backend=server_args.hicache_io_backend,
            storage_backend=server_args.hicache_storage_backend,
            prefetch_threshold=prefetch_threshold,
            model_name=server_args.served_model_name,
            storage_backend_extra_config=extra_config,
            pp_rank=self.pp_rank,
            pp_size=self.pp_size,
            enable_storage_metrics=self.enable_storage_metrics,
            cp_shared_kv_layout=cp_shared_kv_layout,
        )
        self._apply_storage_runtime_config(
            storage_backend=server_args.hicache_storage_backend,
            prefetch_threshold=prefetch_threshold,
            prefetch_timeout_base=prefetch_timeout_base,
            prefetch_timeout_per_ki_token=prefetch_timeout_per_ki_token,
            hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys,
            enable_storage=self.enable_storage,
            enable_storage_metrics=self.enable_storage_metrics,
            extra_metric_labels=self.extra_metric_labels,
        )

        # record the nodes with ongoing write through
        self.ongoing_write_through = {}
        # record the node segments with ongoing load back
        self.ongoing_load_back = {}
        # record the ongoing prefetch requests
        self.ongoing_prefetch = {}
        self.ongoing_backup = {}
        # track per-request tokens loaded from storage (L3 hits)
        # key: request_id, value: number of tokens actually loaded from storage
        self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
        # todo: dynamically adjust the threshold
        self.write_through_threshold = (
            1 if server_args.hicache_write_policy == "write_through" else 2
        )
        self.load_back_threshold = 10

        # Detach storage backend automatically on process shutdown
        atexit.register(self.shutdown)

        self.evictable_host_leaves = set()

        # Pin budget: max tokens that can be pinned = ratio * host pool capacity.
        pin_ratio = envs.SGLANG_HICACHE_MAX_PINNED_RATIO.get()
        if pin_ratio < 0 or pin_ratio >= 1:
            raise ValueError(
                f"SGLANG_HICACHE_MAX_PINNED_RATIO must be in [0, 1), got {pin_ratio}"
            )
        self._max_pinned_tokens = int(self.token_to_kv_pool_host.size * pin_ratio)
        self.pinned_size_ = 0
        logger.info(
            "Pin budget: %d tokens (ratio=%.3f)", self._max_pinned_tokens, pin_ratio
        )

        super().__init__(params=params)

    def shutdown(self):
        """Best-effort auto-detach of storage backend on process shutdown.

        This keeps startup and runtime behavior consistent: if a backend was attached
        (either via CLI args or via admin API), we attempt to detach it on exit.
        """
        try:
            if self.enable_storage:
                self.detach_storage_backend()
        except Exception:
            logger.exception("Failed to detach storage backend on process shutdown.")

    def _apply_storage_runtime_config(
        self,
        *,
        storage_backend: Optional[str],
        prefetch_threshold: int,
        prefetch_timeout_base: float,
        prefetch_timeout_per_ki_token: float,
        hicache_storage_pass_prefix_keys: bool,
        enable_storage: bool,
        enable_storage_metrics: bool,
        extra_metric_labels: Optional[Dict[str, str]],
    ) -> None:
        prefetch_timeout_per_page = (
            self.page_size / 1024 * prefetch_timeout_per_ki_token
        )

        self.enable_storage = enable_storage
        self.prefetch_threshold = prefetch_threshold
        self.prefetch_timeout_base = prefetch_timeout_base
        self.prefetch_timeout_per_page = prefetch_timeout_per_page
        self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys
        self.enable_storage_metrics = enable_storage_metrics

        if self.enable_storage_metrics:
            labels = {
                "storage_backend": storage_backend,
                "tp_rank": self.cache_controller.tp_rank,
                "dp_rank": self.cache_controller.dp_rank,
                "pp_rank": self.cache_controller.pp_rank,
                "pp_size": self.cache_controller.pp_size,
            }
            if extra_metric_labels:
                labels.update(extra_metric_labels)
            existing_collector = getattr(self, "storage_metrics_collector", None)
            if existing_collector is None:
                self.storage_metrics_collector = StorageMetricsCollector(labels=labels)
            elif set(existing_collector.labels.keys()) == set(labels.keys()):
                existing_collector.labels = labels
            else:
                logger.warning(
                    "Storage metrics labels changed (%s -> %s). Keep existing labels to "
                    "avoid duplicate metric registration.",
                    sorted(existing_collector.labels.keys()),
                    sorted(labels.keys()),
                )

    def attach_storage_backend(
        self,
        storage_backend: str,
        storage_backend_extra_config_json: Optional[str] = None,
        served_model_name: Optional[str] = None,
        hicache_storage_prefetch_policy: Optional[str] = None,
        hicache_write_policy: Optional[str] = None,
    ) -> tuple[bool, str]:
        """Attach (enable) storage backend at runtime.

        This will start storage threads inside `HiCacheController` and enable
        prefetch/backup paths. Caller must ensure there are no running/queued
        requests to avoid races.
        """
        if self._uses_cp_hicache:
            return (
                False,
                "CP shared KV HiCache does not support attaching a storage backend at runtime.",
            )

        # Validate inputs first (no side effects).
        if hicache_storage_prefetch_policy is not None:
            allowed = ["best_effort", "wait_complete", "timeout"]
            if hicache_storage_prefetch_policy not in allowed:
                return (
                    False,
                    f"Invalid hicache_storage_prefetch_policy: {hicache_storage_prefetch_policy!r}. "
                    f"Expected one of {allowed}.",
                )

        if hicache_write_policy is not None:
            allowed = ["write_back", "write_through", "write_through_selective"]
            if hicache_write_policy not in allowed:
                return (
                    False,
                    f"Invalid hicache_write_policy: {hicache_write_policy!r}. "
                    f"Expected one of {allowed}.",
                )

        # If already enabled:
        # - backend unchanged: treat as success, update policies only.
        # - backend changed: treat as failure, do NOT update policies.
        if self.enable_storage:
            current_backend = self.cache_controller.storage_backend_type

            if current_backend == storage_backend:
                if hicache_storage_prefetch_policy is not None:
                    self.prefetch_stop_policy = hicache_storage_prefetch_policy
                    logger.info(
                        f"Set hicache_storage_prefetch_policy to {hicache_storage_prefetch_policy}"
                    )
                if hicache_write_policy is not None:
                    self.cache_controller.write_policy = hicache_write_policy
                    self.write_through_threshold = (
                        1 if hicache_write_policy == "write_through" else 2
                    )
                    logger.info(f"Set hicache_write_policy to {hicache_write_policy}")
                return (
                    True,
                    "HiCache storage backend already enabled with same backend; policies updated.",
                )

            return (
                False,
                f"HiCache storage backend is already enabled with backend '{current_backend}'. "
                f"Cannot attach different backend '{storage_backend}'. Detach first.",
            )

        # Not enabled: update policies before controller attach so storage threads observe new values.
        if hicache_storage_prefetch_policy is not None:
            self.prefetch_stop_policy = hicache_storage_prefetch_policy
            logger.info(
                f"Set hicache_storage_prefetch_policy to {hicache_storage_prefetch_policy}"
            )

        if hicache_write_policy is not None:
            self.cache_controller.write_policy = hicache_write_policy
            self.write_through_threshold = (
                1 if hicache_write_policy == "write_through" else 2
            )
            logger.info(f"Set hicache_write_policy to {hicache_write_policy}")

        logger.info(f"Attaching HiCache storage backend: {storage_backend}")
        try:
            (
                extra_config,
                prefetch_threshold,
                prefetch_timeout_base,
                prefetch_timeout_per_ki_token,
                hicache_storage_pass_prefix_keys,
            ) = self._parse_storage_backend_extra_config(
                storage_backend_extra_config_json
            )
        except Exception as e:
            logger.exception(f"Failed to parse storage_backend_extra_config_json: {e}")
            return (
                False,
                f"Failed to parse storage_backend_extra_config_json '{storage_backend_extra_config_json}': {e}",
            )

        try:
            self.cache_controller.attach_storage_backend(
                storage_backend=storage_backend,
                prefetch_threshold=prefetch_threshold,
                model_name=served_model_name,
                storage_backend_extra_config=extra_config,
            )
        except Exception as e:
            logger.exception(
                f"Failed to attach storage backend '{storage_backend}': {e}"
            )
            return False, f"Failed to attach storage backend '{storage_backend}': {e}"

        self._apply_storage_runtime_config(
            storage_backend=storage_backend,
            prefetch_threshold=prefetch_threshold,
            prefetch_timeout_base=prefetch_timeout_base,
            prefetch_timeout_per_ki_token=prefetch_timeout_per_ki_token,
            hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys,
            enable_storage=True,
            enable_storage_metrics=self._enable_metrics_flag,
            extra_metric_labels=self.extra_metric_labels,
        )
        return True, "Attached HiCache storage backend successfully."

    def detach_storage_backend(self) -> tuple[bool, str]:
        """Detach (disable) storage backend at runtime.

        Caller must ensure there are no running/queued requests to avoid races.
        """
        try:
            # Drain any pending control queues before tearing down storage threads/backend.
            # IMPORTANT: this must happen before we clear `ongoing_*`, otherwise acks/releases
            # cannot be matched to nodes and may leak host pages / locks.
            self._drain_storage_control_queues_local()
            # Idempotent detach: always ask controller to best-effort cleanup, even if
            # `self.enable_storage` is already False (may be leftover state from a
            # previous partial detach).
            self.cache_controller.detach_storage_backend()
        except Exception as e:
            logger.exception("Failed to detach storage backend.")
            # Do NOT crash the server for admin operations. Return failure with detail.
            return False, f"Failed to detach HiCache storage backend: {e}"

        # Best-effort cleanup of any leftover bookkeeping.
        self._drain_storage_control_queues_local()
        # After controller threads are fully stopped, it's safe to force-release any
        # leftover pending ops (e.g., async prefetch/backup that didn't get a revoke/ack).
        self._force_release_pending_storage_ops()

        self.enable_storage = False
        self.enable_storage_metrics = False
        return True, "Detached HiCache storage backend successfully."

    def _force_release_pending_storage_ops(self):
        """Force release any leftover pending prefetch/backup bookkeeping.

        This is a safety net for detach/shutdown paths. It assumes storage threads
        have been stopped already (via controller.detach), so no concurrent access
        to these structures should happen.
        """
        cc = self.cache_controller

        # Force release leftover prefetch ops: free pre-allocated host pages and
        # drop the host protection on the matched prefix node.
        try:
            for req_id, info in list(self.ongoing_prefetch.items()):
                try:
                    last_host_node, token_ids, host_indices, _operation = info
                except Exception:
                    # Unexpected shape; just drop it.
                    self.ongoing_prefetch.pop(req_id, None)
                    continue

                try:
                    if host_indices is not None:
                        cc.mem_pool_host.free(host_indices)
                except Exception:
                    logger.exception(
                        "Failed to free host indices for prefetch %s", req_id
                    )

                try:
                    last_host_node.release_host()
                except Exception:
                    logger.exception(
                        "Failed to release host protection for prefetch %s", req_id
                    )

                try:
                    cc.prefetch_tokens_occupied -= len(token_ids)
                    if cc.prefetch_tokens_occupied < 0:
                        cc.prefetch_tokens_occupied = 0
                except Exception:
                    pass

                self.ongoing_prefetch.pop(req_id, None)
        except Exception:
            logger.exception("Force release pending prefetch ops failed.")

        # Force release leftover backup ops: drop host protection on nodes.
        try:
            for ack_id, node in list(self.ongoing_backup.items()):
                try:
                    node.release_host()
                except Exception:
                    logger.exception(
                        "Failed to release host protection for backup op %s", ack_id
                    )
                self.ongoing_backup.pop(ack_id, None)
        except Exception:
            logger.exception("Force release pending backup ops failed.")

    def _drain_storage_control_queues_local(self):
        """Drain storage control queues without TP synchronization.

        This is intended for shutdown/detach paths where we want to make best-effort
        cleanup even if queue sizes temporarily differ across ranks.
        """
        self._drain_storage_control_queues_impl(
            n_revoke=None,
            n_backup=None,
            n_release=None,
            log_metrics=False,
        )

    def _drain_storage_control_queues_impl(
        self,
        n_revoke: Optional[int],
        n_backup: Optional[int],
        n_release: Optional[int],
        log_metrics: bool,
    ):
        cc = self.cache_controller

        def _drain_queue(q, limit: Optional[int]):
            drained = 0
            while limit is None or drained < limit:
                try:
                    item = q.get_nowait()
                except Empty:
                    break
                drained += 1
                yield item

        def _drain_revoke():
            for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke):
                info = self.ongoing_prefetch.pop(req_id, None)
                if info is not None:
                    last_host_node, token_ids, _, _ = info
                    last_host_node.release_host()
                    cc.prefetch_tokens_occupied -= len(token_ids)
                    if cc.prefetch_tokens_occupied < 0:
                        cc.prefetch_tokens_occupied = 0

        def _drain_backup():
            for operation in _drain_queue(cc.ack_backup_queue, n_backup):
                ack_id = operation.id
                entry = self.ongoing_backup.pop(ack_id, None)
                if entry is not None:
                    entry.release_host()
                if log_metrics and self.enable_storage_metrics:
                    self.storage_metrics_collector.log_backuped_tokens(
                        operation.completed_tokens
                    )

        def _drain_release():
            host_indices_list = []
            for host_indices in _drain_queue(cc.host_mem_release_queue, n_release):
                host_indices_list.append(host_indices)
            if host_indices_list:
                host_indices = torch.cat(host_indices_list, dim=0)
                cc.mem_pool_host.free(host_indices)

        _drain_revoke()
        _drain_backup()
        _drain_release()

    def _parse_storage_backend_extra_config(
        self, storage_backend_extra_config: Optional[str]
    ):
        """
        Parse storage backend extra config JSON and extract specific parameters.

        Args:
            storage_backend_extra_config: JSON string containing extra configuration

        Returns:
            tuple: (extra_config_dict, prefetch_threshold, prefetch_timeout_base, prefetch_timeout_per_ki_token, hicache_storage_pass_prefix_keys)
        """
        # Parse extra config if provided. Extra config can be a JSON string or a json/toml/yaml file path prefixed with "@".
        extra_config = {}
        if storage_backend_extra_config:
            try:
                if storage_backend_extra_config.startswith("@"):
                    # Read config from a json/toml/yaml file
                    path = storage_backend_extra_config[1:]
                    ext = os.path.splitext(path)[1].lower()
                    with open(path, "rb" if ext == ".toml" else "r") as f:
                        if ext == ".json":
                            extra_config = json.load(f)
                        elif ext == ".toml":
                            import tomllib

                            extra_config = tomllib.load(f)
                        elif ext in (".yaml", ".yml"):
                            import yaml

                            extra_config = yaml.safe_load(f)
                        else:
                            raise ValueError(
                                f"Unsupported config file {path} (config format: {ext})"
                            )
                else:
                    # read config from JSON string
                    extra_config = json.loads(storage_backend_extra_config)
            except Exception as e:
                logger.error(f"Invalid backend extra config JSON: {e}")
                raise e

        prefetch_threshold = extra_config.pop("prefetch_threshold", 256)  # tokens
        prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", 1)  # seconds
        prefetch_timeout_per_ki_token = extra_config.pop(
            "prefetch_timeout_per_ki_token", 0.25
        )  # seconds per 1024 tokens
        hicache_storage_pass_prefix_keys = extra_config.pop(
            "hicache_storage_pass_prefix_keys", False
        )

        if not isinstance(prefetch_threshold, int):
            raise ValueError(
                f"prefetch_threshold must be int, got {type(prefetch_threshold).__name__}"
            )
        if not isinstance(prefetch_timeout_base, (int, float)):
            raise ValueError(
                f"prefetch_timeout_base must be number, got {type(prefetch_timeout_base).__name__}"
            )
        if not isinstance(prefetch_timeout_per_ki_token, (int, float)):
            raise ValueError(
                f"prefetch_timeout_per_ki_token must be number, got {type(prefetch_timeout_per_ki_token).__name__}"
            )
        if not isinstance(hicache_storage_pass_prefix_keys, bool):
            raise ValueError(
                "hicache_storage_pass_prefix_keys must be bool, got "
                f"{type(hicache_storage_pass_prefix_keys).__name__}"
            )

        return (
            extra_config,
            prefetch_threshold,
            float(prefetch_timeout_base),
            float(prefetch_timeout_per_ki_token),
            hicache_storage_pass_prefix_keys,
        )

    def reset(self):
        TreeNode.counter = 0
        self.cache_controller.reset()
        self.token_to_kv_pool_host.clear()
        # Clear per-request tracking dicts
        self.prefetch_loaded_tokens_by_reqid.clear()
        self.evictable_host_leaves.clear()
        self.pinned_size_ = 0
        super().reset()

    def get_height(self, node: TreeNode):
        height = 0
        while node != self.root_node:
            node = node.parent
            height += 1
        return height

    def clear_storage_backend(self) -> bool:
        if self.enable_storage:
            try:
                # Check if the storage backend has a clear method (for nixl backends)
                if hasattr(self.cache_controller.storage_backend, "clear"):
                    self.cache_controller.storage_backend.clear()
                    logger.info(
                        "Hierarchical cache storage backend cleared successfully!"
                    )
                    return True
                else:
                    logger.warning(
                        f"Storage backend {type(self.cache_controller.storage_backend).__name__} does not support clear operation."
                    )
                    return False
            except Exception as e:
                logger.error(f"Failed to clear hierarchical cache storage backend: {e}")
                return False
        else:
            logger.warning("Hierarchical cache storage backend is not enabled.")
            return False

    def _node_backuped(self, node: TreeNode) -> bool:
        if self._uses_cp_hicache:
            return node.host_len > 0 and node.cp_hicache is not None
        return node.host_value is not None

    def _node_host_len(self, node: TreeNode) -> int:
        if self._uses_cp_hicache:
            return node.host_len
        return len(node.host_value)

    def _node_host_evict_indices(self, node: TreeNode) -> torch.Tensor:
        if self._uses_cp_hicache:
            return node.cp_hicache.host_indices
        return node.host_value

    def write_backup(self, node: TreeNode, write_back=False):
        logger.info(
            "[HiCache-write] write_backup triggered: node_id=%d len=%d cp=%s write_back=%s",
            node.id,
            len(node.value) if node.value is not None else 0,
            self._uses_cp_hicache,
            write_back,
        )
        if self._uses_cp_hicache:
            result = self.cache_controller.write(
                device_indices=node.value,
                node_id=node.id,
            )
            metadata = getattr(result, "metadata", None)
            required_host_slots = 0
            if metadata is None:
                required_host_slots = result.required_host_slots
            self._evict_host_for_physical_slots(
                required_host_slots,
                synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1,
            )
            if metadata is None:
                logger.info(
                    "[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d",
                    node.id,
                    required_host_slots,
                )
                result = self.cache_controller.write(
                    device_indices=node.value,
                    node_id=node.id,
                )
                metadata = getattr(result, "metadata", None)
            if metadata is None:
                logger.info(
                    "[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d",
                    node.id,
                    len(node.value) if node.value is not None else 0,
                )
                return 0

            node.host_len = len(node.value)
            node.cp_hicache = metadata
            node.host_value = None
            self.ongoing_write_through[node.id] = node
            if not write_back:
                self.inc_node_lock_ref(node)
            logger.info(
                "[HiCache-write] write_backup CP SUCCESS: node_id=%d host_len=%d owned_positions=%d ongoing_writes=%d",
                node.id,
                node.host_len,
                node.cp_hicache.owned_positions.numel(),
                len(self.ongoing_write_through),
            )
            return len(node.value)

        host_indices = self.cache_controller.write(
            device_indices=node.value,
            node_id=node.id,
        )
        if host_indices is None:
            logger.info(
                "[HiCache-write] write_backup non-CP retry after host eviction: node_id=%d len=%d",
                node.id,
                len(node.value) if node.value is not None else 0,
            )
            self.evict_host(len(node.value))
            host_indices = self.cache_controller.write(
                device_indices=node.value,
                node_id=node.id,
            )
        if host_indices is not None:
            node.host_value = host_indices
            assert len(node.host_value) > 0
            self.ongoing_write_through[node.id] = node
            if not write_back:
                # Only lock the specific node being written through, not the
                # entire path to root.  Ancestors cannot be evicted while this
                # node holds a device value (_collect_leaves_device skips
                # parents with non-evicted children), so path locking is
                # unnecessary and would over-reduce evictable_size_.
                self.inc_node_lock_ref(node)
            logger.info(
                "[HiCache-write] write_backup non-CP SUCCESS: node_id=%d len=%d ongoing_writes=%d",
                node.id,
                len(host_indices),
                len(self.ongoing_write_through),
            )
        else:
            logger.info(
                "[HiCache-write] write_backup non-CP FAILED (host full): node_id=%d len=%d",
                node.id,
                len(node.value) if node.value is not None else 0,
            )
            return 0

        return len(host_indices)

    def write_backup_storage(self, node: TreeNode):
        prefix_keys = (
            node.get_prefix_hash_values(node.parent)
            if self.hicache_storage_pass_prefix_keys
            else None
        )

        operation_id = self.cache_controller.write_storage(
            node.host_value, node.key, node.hash_value, prefix_keys
        )
        self.ongoing_backup[operation_id] = node
        node.protect_host()

    def _inc_hit_count(self, node: TreeNode, chunked=False):
        # skip the hit count update for chunked requests
        if self.cache_controller.write_policy == "write_back" or chunked:
            return
        node.hit_count += 1

        if not self._node_backuped(node):
            if node.hit_count >= self.write_through_threshold:
                logger.info(
                    "[HiCache-write] _inc_hit_count triggering write_through: node_id=%d hit_count=%d threshold=%d len=%d",
                    node.id,
                    node.hit_count,
                    self.write_through_threshold,
                    len(node.value) if node.value is not None else 0,
                )
                # write to host if the node is not backuped
                self.write_backup(node)

    def writing_check(self, write_back=False):
        if write_back:
            # blocking till all write back complete
            while len(self.ongoing_write_through) > 0:
                for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
                    finish_event.synchronize()
                    for ack_id in ack_list:
                        backuped_node = self.ongoing_write_through.pop(ack_id)
                        if self.enable_storage:
                            self.write_backup_storage(backuped_node)
                self.cache_controller.ack_write_queue.clear()
                assert len(self.ongoing_write_through) == 0
            return

        # NOTE: all ranks has the same ongoing_write_through, can skip sync if empty
        if len(self.ongoing_write_through) == 0:
            return

        ack_queue_len = len(self.cache_controller.ack_write_queue)
        finish_count = 0
        for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
            if not finish_event.query():
                break
            finish_count += 1
        queue_size = torch.tensor(finish_count, dtype=torch.int, device="cpu")
        if self.tp_world_size > 1:
            # synchronize TP workers to make the same update to radix cache
            torch.distributed.all_reduce(
                queue_size,
                op=torch.distributed.ReduceOp.MIN,
                group=self.tp_group,
            )

        finish_count = int(queue_size.item())
        logger.info(
            "[HiCache-write] writing_check: ongoing=%d ack_queue=%d local_finished=%d sync_finished=%d tp_size=%d",
            len(self.ongoing_write_through),
            ack_queue_len,
            finish_count if self.tp_world_size <= 1 else queue_size.item(),
            finish_count,
            getattr(self, "tp_world_size", 1),
        )
        released_nodes = []
        while finish_count > 0:
            _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0)
            finish_event.synchronize()
            for ack_id in ack_list:
                backuped_node = self.ongoing_write_through.pop(ack_id)
                released_nodes.append(ack_id)
                self.dec_node_lock_ref(backuped_node)
                if self.enable_storage:
                    self.write_backup_storage(backuped_node)
            finish_count -= 1
        if released_nodes:
            logger.info(
                "[HiCache-write] writing_check released %d write locks: node_ids=%s remaining_ongoing=%d",
                len(released_nodes),
                released_nodes,
                len(self.ongoing_write_through),
            )

    def loading_check(self):
        finish_count = 0
        for _, finish_event, ack_list in self.cache_controller.ack_load_queue:
            if not finish_event.query():
                # the KV cache loading is still ongoing
                break
            finish_count += 1
            # no need to sync across TP workers as batch forwarding is synced
            for ack_id in ack_list:
                end_node = self.ongoing_load_back.pop(ack_id)
                self.dec_lock_ref(end_node)

        # ACK until all events are processed
        del self.cache_controller.ack_load_queue[:finish_count]

    def evictable_size(self):
        return self.evictable_size_

    def _is_pinned(self, node: TreeNode) -> bool:
        """Check if a node has an active (non-expired) pin."""
        return node.pin_expiry > 0 and time.monotonic() <= node.pin_expiry

    def _clear_pin(self, node: TreeNode):
        """Clear expired pin state and release host_ref_counter hold."""
        if node.pin_expiry > 0:
            self.pinned_size_ = max(0, self.pinned_size_ - len(node.key))
            node.host_ref_counter = max(0, node.host_ref_counter - 1)
        node.pin_expiry = 0.0
        node.pin_ttl = 0

    def pin_prefix(
        self, token_ids: List[int], ttl_seconds: int = 300
    ) -> Tuple[int, Optional[str]]:
        """Pin nodes along a prefix path. Returns (nodes_pinned, reject_reason)."""
        if self.disable or not token_ids:
            return (0, None)

        key, _ = self.maybe_bigram_convert(self._to_radix_key(token_ids))
        if self.page_size != 1:
            page_aligned_len = len(key) // self.page_size * self.page_size
            key = key[:page_aligned_len]
        if len(key) == 0:
            return (0, None)

        expiry = time.monotonic() + ttl_seconds
        nodes_pinned = 0
        budget_exceeded = False
        node = self.root_node
        child_key = self.get_child_key_fn(key)

        while len(key) > 0 and child_key in node.children:
            child = node.children[child_key]
            prefix_len = self.key_match_fn(child.key, key)

            # First pin on this node: check budget, then acquire hold
            if child.pin_expiry == 0:
                if self.pinned_size_ + len(child.key) > self._max_pinned_tokens:
                    budget_exceeded = True
                    break
                child.host_ref_counter += 1
                self.pinned_size_ += len(child.key)

                # Eagerly back up to host so eviction finds pinned nodes
                # already backuped and never enters the write_back drain
                # path, which would leak lock_ref on in-flight
                # write-through entries. No-op under write_back policy.
                self._inc_hit_count(child)

            # Extend expiry and store TTL for refresh-on-hit
            child.pin_expiry = max(child.pin_expiry, expiry)
            child.pin_ttl = max(child.pin_ttl, ttl_seconds)
            nodes_pinned += 1

            if prefix_len < len(child.key):
                break

            node = child
            key = key[prefix_len:]
            if len(key):
                child_key = self.get_child_key_fn(key)

        logger.info(
            "[PIN] pin_prefix: nodes_pinned=%d, ttl=%ds", nodes_pinned, ttl_seconds
        )
        if budget_exceeded:
            msg = f"Pin budget exhausted ({self.pinned_size_}/{self._max_pinned_tokens} tokens pinned)"
            if nodes_pinned == 0:
                return (0, msg)
            return (nodes_pinned, f"prefix partially pinned; {msg}")
        return (nodes_pinned, None)

    def _to_radix_key(self, token_ids: List[int]) -> RadixKey:
        """Convert raw token_ids to a RadixKey for tree walking.

        Must use list (not tuple) to match scheduler's RadixKey format,
        since _key_match_paged compares slices directly and list != tuple.
        """
        return RadixKey(token_ids=list(token_ids))

    def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
        if self.disable:
            return IncLockRefResult(delta=0)

        delta = 0
        while node != self.root_node:
            if node.lock_ref == 0:
                self.evictable_size_ -= len(node.key)
                self.protected_size_ += len(node.key)
                delta -= len(node.key)
            node.lock_ref += 1
            self._update_leaf_status(node)
            self._update_host_leaf_status(node)
            node = node.parent
        return IncLockRefResult(delta=delta)

    def dec_lock_ref(
        self, node: TreeNode, params: Optional[DecLockRefParams] = None
    ) -> DecLockRefResult:
        if self.disable:
            return DecLockRefResult(delta=0)

        delta = 0
        while node != self.root_node:
            if node.lock_ref == 1:
                self.evictable_size_ += len(node.key)
                self.protected_size_ -= len(node.key)
                delta += len(node.key)
            node.lock_ref -= 1
            self._update_leaf_status(node)
            self._update_host_leaf_status(node)
            if node.parent is None:
                assert (
                    node is self.root_node
                ), f"This request holds the node from another tree"
            node = node.parent
        return DecLockRefResult(delta=delta)

    def _update_host_leaf_status(self, node: TreeNode):
        if not node.evicted or node.lock_ref > 0:
            if node in self.evictable_host_leaves:
                self.evictable_host_leaves.remove(node)
            return

        for child in node.children.values():
            if child.evicted:
                if node in self.evictable_host_leaves:
                    self.evictable_host_leaves.remove(node)
                return

        if node not in self.evictable_host_leaves:
            self.evictable_host_leaves.add(node)

    def evict(self, params: EvictParams) -> EvictResult:
        start_time = time.perf_counter()
        num_tokens = params.num_tokens
        leaves = list(self.evictable_leaves)
        eviction_heap = [
            (self.eviction_strategy.get_priority(node), node) for node in leaves
        ]
        heapq.heapify(eviction_heap)

        logger.info(
            "[HiCache-evict] evict START: num_tokens=%d heap_size=%d evictable_size=%d available_size=%d",
            num_tokens,
            len(eviction_heap),
            self.evictable_size_,
            self.token_to_kv_pool_allocator.available_size(),
        )
        num_evicted = 0
        num_locked_skipped = 0
        write_back_nodes = []
        while num_evicted < num_tokens and len(eviction_heap):
            _priority, x = heapq.heappop(eviction_heap)

            if x.lock_ref > 0:
                num_locked_skipped += 1
                continue

            if self._is_pinned(x):
                # Still active: demote to host if possible
                if self._node_backuped(x):
                    num_evicted += self._evict_backuped(x)
                    continue
                written = self.write_backup(x, write_back=True)
                if written > 0:
                    num_evicted += written
                    write_back_nodes.append(x)
                    continue  # backup succeeded, pin holds on host
                # Host full -- drop pin so GPU can be freed
                self._clear_pin(x)
                logger.warning(
                    "[PIN] evict: can't backup node %d to host, releasing pin",
                    x.id,
                )
            elif x.pin_expiry > 0:
                # Expired pin: clear and fall through to normal eviction
                self._clear_pin(x)

            if not self._node_backuped(x):
                if self.cache_controller.write_policy == "write_back":
                    # write to host if the node is not backuped
                    num_evicted += self.write_backup(x, write_back=True)
                    write_back_nodes.append(x)
                else:
                    num_evicted += self._evict_regular(x)
            else:
                num_evicted += self._evict_backuped(x)

            for child in x.parent.children.values():
                if child in write_back_nodes:
                    continue
                if not child.evicted:
                    break
            else:
                # all children are evicted or no children
                new_priority = self.eviction_strategy.get_priority(x.parent)
                heapq.heappush(eviction_heap, (new_priority, x.parent))

        if self.cache_controller.write_policy == "write_back":
            self.writing_check(write_back=True)
            for node in write_back_nodes:
                assert self._node_backuped(node)
                self._evict_backuped(node)

        self.update_eviction_metrics(num_evicted, start_time)
        logger.info(
            "[HiCache-evict] evict END: num_tokens=%d num_evicted=%d num_locked_skipped=%d evictable_size_after=%d available_size_after=%d",
            num_tokens,
            num_evicted,
            num_locked_skipped,
            self.evictable_size_,
            self.token_to_kv_pool_allocator.available_size(),
        )
        return EvictResult(num_tokens_evicted=num_evicted)

    def _evict_backuped(self, node: TreeNode):
        # GPU -> CPU demotion: no BlockRemoved since block is still reachable via load_back
        num_evicted = self.cache_controller.evict_device(node.value)
        assert num_evicted > 0
        self.evictable_size_ -= num_evicted
        logger.info(
            "[HiCache-evict] _evict_backuped: node_id=%d num_evicted=%d lock_ref=%d backed=%s",
            node.id,
            num_evicted,
            node.lock_ref,
            self._node_backuped(node),
        )
        node.value = None
        self._update_leaf_status(node)
        self._update_host_leaf_status(node)
        # update leaf status for the parent because the node is evicted
        self._update_leaf_status(node.parent)
        self._update_host_leaf_status(node.parent)
        return num_evicted

    def _evict_regular(self, node: TreeNode):
        # evict a node not initiated write to host -- emit BlockRemoved
        num_evicted = len(node.value)
        logger.info(
            "[HiCache-evict] _evict_regular: node_id=%d num_evicted=%d",
            node.id,
            num_evicted,
        )
        self._record_remove_event(node)
        self.cache_controller.mem_pool_device_allocator.free(node.value)
        self._delete_leaf(node)
        return num_evicted

    def _remove_host_leaf(self, node: TreeNode) -> TreeNode:
        parent = node.parent
        key = self.get_child_key_fn(node.key)
        v = parent.children.pop(key, None)
        assert v == node, f"parent does not have child key, {key}"
        if node in self.evictable_host_leaves:
            self.evictable_host_leaves.remove(node)
        self._update_host_leaf_status(parent)
        return parent

    def _evict_host_for_physical_slots(
        self, required_host_slots: int, synchronize_across_ranks: bool = False
    ) -> int:
        synchronize_across_ranks = (
            synchronize_across_ranks
            and getattr(self, "tp_group", None) is not None
            and getattr(self, "tp_world_size", 1) > 1
        )
        if required_host_slots <= 0 and not synchronize_across_ranks:
            return 0

        leaves = list(self.evictable_host_leaves)
        logger.info(
            "[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d sync=%s leaves=%d",
            required_host_slots,
            synchronize_across_ranks,
            len(leaves),
        )
        eviction_heap = [
            (self.eviction_strategy.get_priority(node), node) for node in leaves
        ]
        heapq.heapify(eviction_heap)

        num_evicted = 0

        def all_ranks_done() -> bool:
            local_done = int(num_evicted >= required_host_slots)
            if not synchronize_across_ranks:
                return bool(local_done)
            done = torch.tensor(local_done, dtype=torch.int, device="cpu")
            torch.distributed.all_reduce(
                done, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
            )
            return bool(done.item())

        while len(eviction_heap) and not all_ranks_done():
            _priority, x = heapq.heappop(eviction_heap)
            if x == self.root_node:
                break
            if not x.evicted:
                continue
            parent_key = self.get_child_key_fn(x.key)
            if x.parent is None or x.parent.children.get(parent_key) is not x:
                continue

            if not self._node_backuped(x):
                if len(x.children) == 0:
                    parent = self._remove_host_leaf(x)
                    if (
                        len(parent.children) == 0
                        and parent.evicted
                        and self._node_backuped(parent)
                    ):
                        new_priority = self.eviction_strategy.get_priority(parent)
                        heapq.heappush(eviction_heap, (new_priority, parent))
                continue

            if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
                self._clear_pin(x)

            if x.host_ref_counter > 0:
                continue

            host_indices = self._node_host_evict_indices(x)
            physical_count = len(host_indices) if host_indices is not None else 0

            self._record_remove_event(x)
            if physical_count > 0:
                if self._uses_cp_hicache and hasattr(
                    self.cache_controller, "evict_cp_host"
                ):
                    num_evicted += self.cache_controller.evict_cp_host(x.cp_hicache)
                else:
                    num_evicted += self.cache_controller.evict_host(host_indices)
            x.host_len = 0
            x.cp_hicache = None
            x.host_value = None

            parent = self._remove_host_leaf(x)
            if len(parent.children) == 0 and parent.evicted:
                new_priority = self.eviction_strategy.get_priority(parent)
                heapq.heappush(eviction_heap, (new_priority, parent))

        return num_evicted

    def evict_host(self, num_tokens: int):
        if self._uses_cp_hicache:
            return self._evict_host_for_physical_slots(num_tokens)

        leaves = list(self.evictable_host_leaves)
        eviction_heap = [
            (self.eviction_strategy.get_priority(node), node) for node in leaves
        ]
        heapq.heapify(eviction_heap)

        num_evicted = 0
        while num_evicted < num_tokens and len(eviction_heap):
            _priority, x = heapq.heappop(eviction_heap)
            if x == self.root_node:
                break
            # only evict the host value of evicted nodes
            if not x.evicted:
                continue

            # Expire stale pins before checking host_ref_counter
            if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
                self._clear_pin(x)

            # node is protected from eviction as it has ongoing prefetch, backup, or pin
            if x.host_ref_counter > 0:
                continue

            # Block deleted entirely (GPU already evicted, now CPU freed) --
            # emit BlockRemoved so the router removes this block from its index.
            self._record_remove_event(x)
            num_evicted += self.cache_controller.evict_host(
                self._node_host_evict_indices(x)
            )

            parent = self._remove_host_leaf(x)

            if len(parent.children) == 0 and parent.evicted:
                new_priority = self.eviction_strategy.get_priority(parent)
                heapq.heappush(eviction_heap, (new_priority, parent))

    def load_back(
        self, node: TreeNode, mem_quota: Optional[int] = None
    ) -> Optional[torch.Tensor]:

        if self._uses_cp_hicache:
            start_time = time.perf_counter()
            last_hit_node = node
            nodes_to_load = []
            while node.evicted:
                assert self._node_backuped(
                    node
                ), "No backup available on evicted nodes, should not happen"
                nodes_to_load.insert(0, node)
                node = node.parent
            else:
                ancester_node = node

            # protect the ancestor nodes from eviction
            result = self.inc_lock_ref(ancester_node)
            delta = result.delta

            # load it all or not at all
            host_hit_len = sum(self._node_host_len(n) for n in nodes_to_load)
            logger.info(
                "[HiCache-load] load_back CP: node_id=%d nodes_to_load=%d host_hit_len=%d threshold=%d",
                last_hit_node.id,
                len(nodes_to_load),
                host_hit_len,
                self.load_back_threshold,
            )
            if host_hit_len < self.load_back_threshold or (
                host_hit_len > mem_quota + delta if mem_quota is not None else False
            ):
                # skip loading back if the total size is too small or exceeding the memory quota
                self.dec_lock_ref(ancester_node)
                logger.info(
                    "[HiCache-load] load_back CP SKIP: host_hit_len=%d below threshold=%d or over quota",
                    host_hit_len,
                    self.load_back_threshold,
                )
                return None

            device_indices = self.cache_controller.load_cp(
                nodes_to_load, node_id=last_hit_node.id
            )
            if device_indices is None:
                logger.info(
                    "[HiCache-load] load_back CP retry with lane-aware eviction: "
                    "node_id=%d tokens_needed=%d",
                    last_hit_node.id,
                    host_hit_len,
                )
                # Lane-aware eviction: alloc_pages_with_owners failed because
                # at least one owner lane is short of free pages. Targeted
                # eviction frees pages from THOSE specific lanes, leaving
                # other lanes untouched. Mirrors cold prefill's retry path
                # in common.py:alloc_paged_token_slots_extend.
                _retry_page_owners: List[int] = []
                for _node in nodes_to_load:
                    # page_owners is CPU int8; tolist() returns list[int].
                    _retry_page_owners.extend(_node.cp_hicache.page_owners.tolist())
                _evict_for_compute_owner_lanes(
                    tree_cache=self,
                    allocator=self.token_to_kv_pool_allocator,
                    page_compute_owners=_retry_page_owners,
                )
                device_indices = self.cache_controller.load_cp(
                    nodes_to_load, node_id=last_hit_node.id
                )
            self.dec_lock_ref(ancester_node)
            if device_indices is None:
                # no sufficient GPU memory to load back KV caches
                logger.warning(
                    "load_back: FAILED to load %d tokens for node %d "
                    "even after eviction (evictable_size=%d)",
                    host_hit_len,
                    last_hit_node.id,
                    self.evictable_size_,
                )
                return None

            self.ongoing_load_back[last_hit_node.id] = last_hit_node
            offset = 0
            for loaded_node in nodes_to_load:
                host_len = self._node_host_len(loaded_node)
                loaded_node.value = device_indices[offset : offset + host_len].clone()
                offset += host_len
            self.evictable_size_ += len(device_indices)
            self.inc_lock_ref(last_hit_node)

            if self.metrics_collector is not None:
                self.metrics_collector.observe_load_back_duration(
                    time.perf_counter() - start_time
                )
                self.metrics_collector.increment_load_back_num_tokens(
                    len(device_indices)
                )

            logger.info(
                "[HiCache-load] load_back CP SUCCESS: node_id=%d loaded_tokens=%d",
                last_hit_node.id,
                len(device_indices),
            )
            return device_indices

        start_time = time.perf_counter()
        last_hit_node = node
        nodes_to_load = []
        while node.evicted:
            assert (
                node.backuped
            ), "No backup available on evicted nodes, should not happen"
            nodes_to_load.insert(0, node)
            node = node.parent
        else:
            ancester_node = node

        # protect the ancestor nodes from eviction
        result = self.inc_lock_ref(ancester_node)
        delta = result.delta

        # load it all or not at all
        host_indices = torch.cat([n.host_value for n in nodes_to_load])
        logger.info(
            "[HiCache-load] load_back non-CP: node_id=%d nodes_to_load=%d host_hit_len=%d threshold=%d",
            last_hit_node.id,
            len(nodes_to_load),
            len(host_indices),
            self.load_back_threshold,
        )
        if len(host_indices) < self.load_back_threshold or (
            len(host_indices) > mem_quota + delta if mem_quota is not None else False
        ):
            # skip loading back if the total size is too small or exceeding the memory quota
            self.dec_lock_ref(ancester_node)
            logger.info(
                "[HiCache-load] load_back non-CP SKIP: host_hit_len=%d below threshold=%d or over quota",
                len(host_indices),
                self.load_back_threshold,
            )
            return None

        device_indices = self.cache_controller.load(
            host_indices=host_indices, node_id=last_hit_node.id
        )
        if device_indices is None:
            logger.info(
                "[HiCache-load] load_back non-CP retry with eviction: node_id=%d tokens_needed=%d",
                last_hit_node.id,
                len(host_indices),
            )
            self.evict(EvictParams(num_tokens=len(host_indices)))
            device_indices = self.cache_controller.load(
                host_indices=host_indices, node_id=last_hit_node.id
            )
        self.dec_lock_ref(ancester_node)
        if device_indices is None:
            # no sufficient GPU memory to load back KV caches
            logger.warning(
                "load_back: FAILED to load %d tokens for node %d "
                "even after eviction (evictable_size=%d)",
                len(host_indices),
                last_hit_node.id,
                self.evictable_size_,
            )
            return None

        self.ongoing_load_back[last_hit_node.id] = last_hit_node
        offset = 0
        for node in nodes_to_load:
            node.value = device_indices[offset : offset + len(node.host_value)].clone()
            offset += len(node.host_value)
        self.evictable_size_ += len(device_indices)
        self.inc_lock_ref(last_hit_node)

        if self.metrics_collector is not None:
            self.metrics_collector.observe_load_back_duration(
                time.perf_counter() - start_time
            )
            self.metrics_collector.increment_load_back_num_tokens(len(device_indices))

        logger.info(
            "[HiCache-load] load_back non-CP SUCCESS: node_id=%d loaded_tokens=%d",
            last_hit_node.id,
            len(device_indices),
        )

        return device_indices

    def init_load_back(
        self,
        params: InitLoadBackParams,
    ):
        last_node = params.last_host_node
        mem_quota = params.mem_quota
        if last_node.evicted:
            loading_values = self.load_back(last_node, mem_quota)
            if loading_values is not None:
                logger.info(
                    f"loading back {len(loading_values)} tokens for node {last_node.id}"
                )
                return loading_values, last_node

            while last_node.evicted:
                last_node = last_node.parent

        return (
            torch.empty((0,), dtype=torch.int64, device=self.device),
            last_node,
        )

    def ready_to_load_host_cache(self) -> int:
        """
        Notify the cache controller to start the KV cache loading.
        Return the consumer index for the schedule batch manager to track.
        """
        return self.cache_controller.start_loading()

    def flush_write_through_acks(self) -> None:
        ongoing_before = len(self.ongoing_write_through)
        self.writing_check()
        ongoing_after = len(self.ongoing_write_through)
        if ongoing_before > 0:
            logger.info(
                "[HiCache-write] flush_write_through_acks: before=%d after=%d",
                ongoing_before,
                ongoing_after,
            )

    def check_hicache_events(self):
        self.writing_check()
        self.loading_check()
        if self.enable_storage:
            self.drain_storage_control_queues()
        if self.enable_storage_metrics:
            self.storage_metrics_collector.log_storage_metrics(
                self.cache_controller.storage_backend.get_stats()
            )

    def drain_storage_control_queues(self):
        """
        Combine prefetch revoke, backup ack, and host mem release checks
        to minimize TP synchronization and Python overhead.
        """
        cc = self.cache_controller

        qsizes = torch.tensor(
            [
                cc.prefetch_revoke_queue.qsize(),
                cc.ack_backup_queue.qsize(),
                cc.host_mem_release_queue.qsize(),
            ],
            dtype=torch.int,
        )
        if self.tp_world_size > 1:
            torch.distributed.all_reduce(
                qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
            )

        n_revoke, n_backup, n_release = map(int, qsizes.tolist())
        self._drain_storage_control_queues_impl(
            n_revoke=n_revoke,
            n_backup=n_backup,
            n_release=n_release,
            log_metrics=True,
        )

    # Timeout is linearly increasing with the number of pages
    def _prefetch_timeout_check_linear_func(self, operation: PrefetchOperation):
        # If hash_value has not been computed in timeout_base seconds, terminate it.
        return (
            time.monotonic() - operation.start_time
            > self.prefetch_timeout_base
            + len(operation.hash_value) * self.prefetch_timeout_per_page
        )

    def can_terminate_prefetch(self, operation: PrefetchOperation):
        can_terminate = True

        if self.prefetch_stop_policy == "best_effort":
            return can_terminate

        if len(operation.hash_value) == 0:
            completed = False
        else:
            completed = (
                operation.completed_tokens == len(operation.hash_value) * self.page_size
            )

        if self.prefetch_stop_policy == "wait_complete":
            can_terminate = completed
        elif self.prefetch_stop_policy == "timeout":
            can_terminate = completed or self.is_prefetch_timeout(operation)
        else:
            # unknown prefetch stop policy, just return True
            return True

        operation_terminated = operation.is_terminated()
        if self.tp_world_size > 1:
            states = torch.tensor(
                [1 - int(can_terminate), int(operation_terminated)],
                dtype=torch.int,
            )
            torch.distributed.all_reduce(
                states,
                op=torch.distributed.ReduceOp.MAX,
                group=self.tp_group,
            )
            can_terminate = states[0].item() == 0
            operation_terminated = states[1].item() == 1
        # the operation should be terminated if it is already terminated on any TP worker
        # or it meets the termination condition on all TP workers
        can_terminate = can_terminate or operation_terminated
        return can_terminate

    def check_prefetch_progress(self, req_id: str) -> bool:
        if req_id not in self.ongoing_prefetch:
            # there is no ongoing prefetch for this request or it has been revoked
            return True

        # todo: more policies for prefetch progress such as timeout
        # the current policy is to prefetch with best effort and terminate when queuing is over
        last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[
            req_id
        ]

        if operation.host_indices is None:
            # prefetch has not been issued due to insufficient host memory
            return True

        if not self.can_terminate_prefetch(operation):
            return False

        completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
            operation
        )
        logger.info(f"Prefetch {req_id} completed with {completed_tokens} tokens")

        min_completed_tokens = completed_tokens
        if self.tp_world_size > 1:
            # synchrnoize TP workers to make the same update to hiradix cache
            completed_tokens_tensor = torch.tensor(
                min_completed_tokens, dtype=torch.int
            )
            torch.distributed.all_reduce(
                completed_tokens_tensor,
                op=torch.distributed.ReduceOp.MIN,
                group=self.tp_group,
            )
            min_completed_tokens = completed_tokens_tensor.item()
        fetched_token_ids = token_ids[:min_completed_tokens]
        written_indices = host_indices[:min_completed_tokens]
        matched_length = self._insert_helper_host(
            last_host_node,
            RadixKey(
                token_ids=fetched_token_ids, extra_key=last_host_node.key.extra_key
            ),
            written_indices,
            hash_value[: min_completed_tokens // self.page_size],
        )

        self.cache_controller.mem_pool_host.free(host_indices[:matched_length])
        self.cache_controller.append_host_mem_release(
            host_indices[min_completed_tokens:completed_tokens]
        )
        last_host_node.release_host()
        del self.ongoing_prefetch[req_id]
        self.cache_controller.prefetch_tokens_occupied -= len(token_ids)

        # Track tokens actually loaded from storage for this request (L3 hits)
        loaded_from_storage = min_completed_tokens - matched_length
        self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage

        if self.enable_storage_metrics:
            self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage)

        return True

    def terminate_prefetch(self, req_id: str):
        if req_id not in self.ongoing_prefetch:
            return

        _, _, _, operation = self.ongoing_prefetch[req_id]
        if operation.host_indices is None:
            return
        operation.mark_terminate()

    def pop_prefetch_loaded_tokens(self, req_id: str) -> int:
        """
        Pop and return the number of tokens loaded from storage for a request.
        Returns 0 if no prefetch was done or was revoked.
        This should be called after check_prefetch_progress() returns True.
        """
        return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0)

    def match_prefix(self, params: MatchPrefixParams):
        key = params.key
        empty_value = torch.empty((0,), dtype=torch.int64, device=self.device)
        key, _ = self.maybe_bigram_convert(key)
        if self.disable or len(key) == 0:
            return MatchResult(
                device_indices=empty_value,
                last_device_node=self.root_node,
                last_host_node=self.root_node,
                host_hit_length=0,
            )

        page_aligned_len = len(key)
        if self.page_size != 1:
            page_aligned_len = len(key) // self.page_size * self.page_size
            key = key[:page_aligned_len]

        if len(key) == 0:
            return MatchResult(
                device_indices=empty_value,
                last_device_node=self.root_node,
                last_host_node=self.root_node,
                host_hit_length=0,
            )

        value, last_node = self._match_prefix_helper(self.root_node, key)
        if value:
            value = torch.cat(value)
        else:
            value = empty_value

        host_hit_length = 0
        last_host_node = last_node
        if self._uses_cp_hicache:
            while last_node.evicted:
                host_hit_length += self._node_host_len(last_node)
                last_node = last_node.parent
            while (
                last_host_node != self.root_node
                and not self._node_backuped(last_host_node)
            ):
                last_host_node = last_host_node.parent
            if not self._node_backuped(last_host_node):
                last_host_node = self.root_node
        else:
            while last_node.evicted:
                host_hit_length += len(last_node.host_value)
                last_node = last_node.parent
            while last_host_node != self.root_node and not last_host_node.backuped:
                last_host_node = last_host_node.parent
            if not last_host_node.backuped:
                last_host_node = self.root_node

        return MatchResult(
            device_indices=value,
            last_device_node=last_node,
            last_host_node=last_host_node,
            host_hit_length=host_hit_length,
        )

    def prefetch_from_storage(
        self,
        req_id: str,
        last_host_node: TreeNode,
        new_input_tokens: List[int],
        last_hash: Optional[str] = None,
        prefix_keys: Optional[List[str]] = None,
    ):
        new_input_tokens = (
            convert_to_bigram_key(new_input_tokens)
            if self.is_eagle
            else new_input_tokens
        )
        # align the number of fetching tokens to the page size
        prefetch_length = len(new_input_tokens) - (
            len(new_input_tokens) % self.page_size
        )
        new_input_tokens = new_input_tokens[:prefetch_length]
        if (
            not self.enable_storage
            or prefetch_length < self.prefetch_threshold
            or self.cache_controller.prefetch_rate_limited()
        ):
            return

        last_host_node.protect_host()
        host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
        if host_indices is None:
            self.evict_host(prefetch_length)
            host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
        if host_indices is None:
            last_host_node.release_host()
            # no sufficient host memory for prefetch
            return
        operation = self.cache_controller.prefetch(
            req_id, host_indices, new_input_tokens, last_hash, prefix_keys
        )
        self.ongoing_prefetch[req_id] = (
            last_host_node,
            new_input_tokens,
            host_indices,
            operation,
        )
        self.cache_controller.prefetch_tokens_occupied += len(new_input_tokens)

    def _insert_helper_host(
        self, node: TreeNode, key: RadixKey, host_value, hash_value
    ):
        node.last_access_time = time.monotonic()
        if len(key) == 0:
            return 0

        child_key = self.get_child_key_fn(key)

        matched_length = 0
        while len(key) > 0 and child_key in node.children.keys():
            node = node.children[child_key]
            node.last_access_time = time.monotonic()
            # Refresh pin TTL on host insert hit
            if self._is_pinned(node):
                node.pin_expiry = time.monotonic() + node.pin_ttl
            prefix_len = self.key_match_fn(node.key, key)
            key = key[prefix_len:]
            host_value = host_value[prefix_len:]
            hash_value = hash_value[prefix_len // self.page_size :]
            matched_length += prefix_len

            if prefix_len < len(node.key):
                new_node = self._split_node(node.key, node, prefix_len)
                node = new_node

            if len(key):
                child_key = self.get_child_key_fn(key)

        if len(key):
            new_node = TreeNode(priority=node.priority)
            new_node.parent = node
            new_node.key = key
            new_node.value = None
            new_node.host_value = host_value.clone()
            new_node.hash_value = hash_value
            node.children[child_key] = new_node
            self._update_host_leaf_status(new_node)
            self._update_leaf_status(node)
            self._update_host_leaf_status(node)

        return matched_length

    def _match_prefix_helper(self, node: TreeNode, key: RadixKey):
        node.last_access_time = time.monotonic()
        child_key = self.get_child_key_fn(key)
        value = []

        while len(key) > 0 and child_key in node.children.keys():
            child = node.children[child_key]
            child.last_access_time = time.monotonic()
            # Refresh pin TTL on cache hit
            if self._is_pinned(child):
                child.pin_expiry = time.monotonic() + child.pin_ttl
            prefix_len = self.key_match_fn(child.key, key)
            if prefix_len < len(child.key):
                new_node = self._split_node(child.key, child, prefix_len)
                if not new_node.evicted:
                    value.append(new_node.value)
                node = new_node
                break
            else:
                if not child.evicted:
                    value.append(child.value)
                node = child
                key = key[prefix_len:]

                if len(key):
                    child_key = self.get_child_key_fn(key)

        return value, node

    def _split_node(self, key: RadixKey, child: TreeNode, split_len: int):
        # child node split into new_node -> child
        new_node = TreeNode(priority=child.priority)
        new_node.children = {self.get_child_key_fn(key[split_len:]): child}
        new_node.parent = child.parent
        new_node.lock_ref = child.lock_ref
        new_node.pin_expiry = child.pin_expiry
        new_node.pin_ttl = child.pin_ttl
        # If child is pinned, new parent inherits a host_ref_counter hold
        if child.pin_expiry > 0:
            new_node.host_ref_counter += 1
        new_node.key = child.key[:split_len]
        new_node.hit_count = child.hit_count

        # split value and host value if exists
        if child.evicted:
            new_node.value = None
        else:
            new_node.value = child.value[:split_len].clone()
            child.value = child.value[split_len:].clone()
        if self._uses_cp_hicache:
            if self._node_backuped(child):
                new_node.cp_hicache, child.cp_hicache = child.cp_hicache.split(
                    split_len
                )
                new_node.host_len = split_len
                child.host_len = child.host_len - split_len
        elif child.backuped:
            new_node.host_value = child.host_value[:split_len].clone()
            child.host_value = child.host_value[split_len:].clone()

        new_node.hash_value, child.hash_value = split_node_hash_value(
            child.hash_value, split_len, self.page_size
        )
        child.parent = new_node
        child.key = child.key[split_len:]
        new_node.parent.children[self.get_child_key_fn(key)] = new_node

        return new_node

    def insert(self, params: InsertParams) -> InsertResult:
        key = params.key
        value = params.value
        chunked = params.chunked
        priority = params.priority

        if priority is None:
            priority = 0
        key, value = self.maybe_bigram_convert(key, value)

        if len(key) == 0:
            return InsertResult(prefix_len=0)

        if self.is_eagle and value is not None:
            # Make sure the value len equal to the EAGLE bigram key len
            value = value[: len(key)]

        node = self.root_node
        child_key = self.get_child_key_fn(key)
        total_prefix_length = 0

        while len(key) > 0 and child_key in node.children.keys():
            node = node.children[child_key]
            node.last_access_time = time.monotonic()
            node.priority = max(node.priority, priority)
            prefix_len = self.key_match_fn(node.key, key)

            if prefix_len == len(node.key):
                if node.evicted:
                    # change the reference if the node is evicted
                    # this often happens in the case of KV cache recomputation
                    node.value = value[:prefix_len].clone()
                    self.evictable_size_ += len(node.value)
                    self._update_leaf_status(node)
                    self._update_host_leaf_status(node)
                    # update parent status as a new leaf is added into device
                    self._update_leaf_status(node.parent)
                    self._update_host_leaf_status(node.parent)
                else:
                    self._inc_hit_count(node, chunked)
                    total_prefix_length += prefix_len
            else:
                # partial match, split the node
                new_node = self._split_node(node.key, node, prefix_len)
                # shared-prefix node should also reflect max priority
                new_node.priority = max(new_node.priority, priority)
                if new_node.evicted:
                    new_node.value = value[:prefix_len].clone()
                    self.evictable_size_ += len(new_node.value)
                    self._update_leaf_status(new_node)
                    self._update_host_leaf_status(new_node)
                    # update parent status as a new leaf is added into device
                    self._update_leaf_status(new_node.parent)
                    self._update_host_leaf_status(new_node.parent)
                else:
                    self._inc_hit_count(new_node, chunked)
                    total_prefix_length += prefix_len
                node = new_node

            key = key[prefix_len:]
            value = value[prefix_len:]

            if len(key):
                child_key = self.get_child_key_fn(key)

        if len(key):
            new_node = TreeNode(priority=priority)
            new_node.parent = node
            new_node.key = key
            new_node.value = value.clone()
            node.children[child_key] = new_node
            self.evictable_size_ += len(value)
            self._update_leaf_status(node)
            self._update_leaf_status(new_node)

            # Compute hash_value if storage or kv events are enabled
            if self.enable_storage or self.enable_kv_cache_events:
                new_node.hash_value = compute_node_hash_values(new_node, self.page_size)

            # Emit BlockStored so the router indexes this block.
            self._record_store_event(new_node)

            if self.cache_controller.write_policy != "write_back":
                self._inc_hit_count(new_node, chunked)
        return InsertResult(prefix_len=total_prefix_length)

    def release_aborted_request(self, rid: str):
        # Clean up storage hit tracking for aborted request
        self.prefetch_loaded_tokens_by_reqid.pop(rid, None)

        if rid not in self.ongoing_prefetch:
            return

        last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[rid]
        if operation.host_indices is None:
            return

        completed_tokens, _ = self.cache_controller.terminate_prefetch(operation)
        if self.tp_world_size > 1:
            torch.distributed.barrier(group=self.tp_group)
        last_host_node.release_host()
        del self.ongoing_prefetch[rid]
        self.cache_controller.append_host_mem_release(host_indices[:completed_tokens])
        self.cache_controller.prefetch_tokens_occupied -= len(token_ids)