sglang/python/sglang/srt/managers/cache_controller.py

from __future__ import annotations

"""
Copyright 2023-2025 SGLang Team
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import logging
import threading
import time
from dataclasses import dataclass, field
from queue import Empty, Full, Queue
from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Set

import torch

from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig

if TYPE_CHECKING:
    from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
    from sglang.srt.mem_cache.memory_pool_host import HostKVCache

from sglang.srt.distributed import (
    get_tensor_model_parallel_rank,
    get_tensor_model_parallel_world_size,
)
from sglang.srt.layers.dp_attention import (
    get_attention_dp_rank,
    get_attention_tp_rank,
    get_attention_tp_size,
    is_dp_attention_enabled,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool
from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices
from sglang.srt.utils import get_device_module

logger = logging.getLogger(__name__)

device_module = get_device_module()


class LayerLoadingEvent:
    def __init__(self, num_layers: int):
        self._num_layers = num_layers
        self.load_events = [device_module.Event() for _ in range(num_layers)]
        self.start_event = device_module.Event()  # start event on controller stream

    def complete(self, layer_index: int):
        assert 0 <= layer_index < self._num_layers
        self.load_events[layer_index].record()

    def wait(self, layer_index: int):
        device_module.current_stream().wait_event(self.load_events[layer_index])

    @property
    def finish_event(self):
        return self.load_events[-1]


class LayerDoneCounter:
    def __init__(self, num_layers: int):
        self.num_layers = num_layers
        # extra producer and consumer counters for overlap mode
        self.num_counters = 5
        self.events = [LayerLoadingEvent(num_layers) for _ in range(self.num_counters)]
        self.producer_index = -1
        self.consumer_index = -1
        self.consumer_indices: List[int] = []

    def update_producer(self):
        self.producer_index = (self.producer_index + 1) % self.num_counters
        assert self.events[
            self.producer_index
        ].finish_event.query(), (
            "Producer finish event should be ready before being reused."
        )
        return self.producer_index

    def set_consumer(self, index):
        if isinstance(index, (list, tuple, set)):
            self.consumer_indices = [int(i) for i in index if int(i) >= 0]
            self.consumer_index = (
                self.consumer_indices[0] if self.consumer_indices else -1
            )
            return
        self.consumer_index = int(index)
        self.consumer_indices = [self.consumer_index] if self.consumer_index >= 0 else []

    def wait_until(self, threshold: int):
        if not self.consumer_indices:
            return
        for consumer_index in self.consumer_indices:
            self.events[consumer_index].wait(threshold)

    def reset(self):
        self.producer_index = -1
        self.consumer_index = -1
        self.consumer_indices = []


class CacheOperation:

    counter = 0

    def __init__(
        self,
        host_indices: torch.Tensor,
        device_indices: torch.Tensor,
        node_id: int,
        priority: Optional[int] = None,
    ):
        self.host_indices = host_indices
        self.device_indices = device_indices
        self.node_ids = [node_id]
        self.data = None

        self.id = CacheOperation.counter
        CacheOperation.counter += 1
        # default priority is the order of creation
        self.priority = priority if priority is not None else self.id

    @staticmethod
    def merge_ops(ops: List[CacheOperation]) -> CacheOperation:
        assert len(ops) > 0
        if len(ops) == 1:
            return ops[0]

        host_indices = torch.cat([op.host_indices for op in ops])
        device_indices = torch.cat([op.device_indices for op in ops])
        node_ids = []
        priority = min(op.priority for op in ops)
        for op in ops:
            node_ids.extend(op.node_ids)
        merged_op = CacheOperation(host_indices, device_indices, -1, priority)
        merged_op.node_ids = node_ids
        return merged_op

    def __lt__(self, other: CacheOperation):
        return self.priority < other.priority


class HiCacheAck(NamedTuple):
    start_event: device_module.Event
    finish_event: device_module.Event
    node_ids: List[int]


class TransferBuffer:
    """
    Overlapping buffer preparation and transfer operations to improve throughput.
    """

    def __init__(
        self, stop_event, buffer_count: int = 3, max_buffer_size: int = 1024
    ) -> None:
        self.stop_event = stop_event
        self.buffers = Queue(maxsize=buffer_count)
        # todo: adjust the buffer size based on throughput profile of the system
        self.max_buffer_size = max_buffer_size

    def full(self) -> bool:
        return self.buffers.full()

    def empty(self) -> bool:
        return self.buffers.empty()

    def put(self, item, block=True, timeout=1) -> None:
        while not self.stop_event.is_set():
            try:
                self.buffers.put(item, block=block, timeout=timeout)
                break
            except Full:
                if not block:
                    break
                continue
            except Exception as e:
                logger.error(e)

    def get(self, block=True, timeout=1) -> Optional[CacheOperation]:
        try:
            return self.buffers.get(block=block, timeout=timeout)
        except Empty:
            return None
        except Exception as e:
            logger.error(e)

    def clear(self):
        self.buffers.queue.clear()


class StorageOperation:
    counter = 0

    def __init__(
        self,
        host_indices: torch.Tensor,
        token_ids: List[int],
        last_hash: Optional[str] = None,
        hash_value: Optional[List[str]] = None,
        prefix_keys: Optional[List[str]] = None,
    ):
        self.host_indices = host_indices
        self.token_ids = token_ids
        self.last_hash = last_hash
        self.completed_tokens = 0
        self.hash_value = hash_value if hash_value is not None else []
        self.prefix_keys = prefix_keys

        self.id = StorageOperation.counter
        StorageOperation.counter += 1

    def __lt__(self, other: "StorageOperation"):
        return self.id < other.id


class PrefetchOperation(StorageOperation):
    def __init__(
        self,
        request_id: str,
        host_indices: torch.Tensor,
        token_ids: List[int],
        last_hash: Optional[str] = None,
        prefix_keys: Optional[List[str]] = None,
    ):
        self.request_id = request_id

        self._lock = threading.Lock()
        self._terminated_flag = False
        self.start_time = time.monotonic()

        super().__init__(host_indices, token_ids, last_hash, prefix_keys=prefix_keys)

    def increment(self, num_tokens: int):
        with self._lock:
            if self._terminated_flag:
                return False
            self.completed_tokens += num_tokens
            return True

    def mark_terminate(self):
        with self._lock:
            self._terminated_flag = True

    def is_terminated(self) -> bool:
        return self._terminated_flag


@dataclass
class HiCacheWriteResult:
    metadata: object
    required_host_slots: int = 0


@dataclass
class HiCacheWriteFailure:
    required_host_slots: int
    metadata: object = None


@dataclass
class HiCacheWriteReservation:
    metadata: object
    host_indices: torch.Tensor
    physical_device_indices: torch.Tensor
    node_id: int = -1
    priority: Optional[int] = None
    draft_host_indices: Optional[torch.Tensor] = None
    required_host_slots: int = 0


@dataclass
class HiCacheLayerWriteState:
    reservation: HiCacheWriteReservation
    total_layers: int
    start_event: object
    finish_event: object
    layer_events: List[object] = field(default_factory=list)
    completed_target_layers: Set[int] = field(default_factory=set)
    completed_draft_layers: Set[int] = field(default_factory=set)
    host_indices: Optional[torch.Tensor] = None
    physical_device_indices: Optional[torch.Tensor] = None
    draft_host_indices: Optional[torch.Tensor] = None
    ack_appended: bool = False


class HiCacheController:

    def __init__(
        self,
        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
        mem_pool_host: HostKVCache,
        page_size: int,
        tp_group: torch.distributed.ProcessGroup,
        load_cache_event: threading.Event,
        write_policy: str = "write_through_selective",
        io_backend: str = "",
        storage_backend: Optional[str] = None,
        prefetch_threshold: int = 256,
        model_name: Optional[str] = None,
        storage_backend_extra_config: Optional[dict] = None,
        pp_rank: int = 0,
        pp_size: int = 1,
        enable_storage_metrics: bool = False,
        cp_shared_kv_layout: Optional[CpSharedKVLayout] = None,
        draft_mem_pool_host: Optional["HostKVCache"] = None,
        draft_mem_pool_device=None,
    ):
        self.tp_group = tp_group
        self.mem_pool_device_allocator = token_to_kv_pool_allocator
        mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
        from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool

        if isinstance(mem_pool_device, HybridLinearKVPool):
            mem_pool_device = mem_pool_device.full_kv_pool
        self.mem_pool_device = mem_pool_device
        self.cp_shared_kv_layout = cp_shared_kv_layout

        self.has_draft = False
        self.mem_pool_device_draft = None
        self.mem_pool_host_draft = None
        self.uses_cp_hicache = cp_shared_kv_layout is not None
        if self.uses_cp_hicache and not isinstance(
            self.mem_pool_device, NSATokenToKVPool
        ):
            raise ValueError(
                "CP shared KV HiCache host integration requires NSATokenToKVPool."
            )
        self.mem_pool_host = mem_pool_host
        self.write_policy = write_policy
        self.page_size = page_size
        self.io_backend = io_backend
        self.enable_storage = False
        self.storage_backend = None
        self.storage_backend_type = None
        self.pp_rank = pp_rank
        self.pp_size = pp_size
        self.enable_storage_metrics = enable_storage_metrics

        # Default storage page IO functions (may be overridden by attach).
        self.page_get_func = self._generic_page_get
        self.page_set_func = self._generic_page_set

        # Dedicated stop event for storage background threads (prefetch/backup).
        # NOTE: Do NOT reuse `self.stop_event` here since it also guards core HiCache
        # transfer buffers (CPU<->GPU). We want to allow runtime attach/detach of
        # storage without stopping the whole controller.
        self.storage_stop_event = threading.Event()

        self.device = self.mem_pool_device.device
        self.layer_num = self.mem_pool_device.layer_num
        self.layer_done_counter = LayerDoneCounter(self.layer_num)
        self.mem_pool_device.register_layer_transfer_counter(self.layer_done_counter)
        if hasattr(self.mem_pool_device, "register_layer_backup_notifier"):
            self.mem_pool_device.register_layer_backup_notifier(
                lambda layer_id: self.on_layer_end(layer_id, source="target")
            )
        self.draft_mem_pool_host = None
        self.draft_mem_pool_device = None

        if write_policy not in [
            "write_through",
            "write_through_selective",
            "write_back",
            "write_behind",
        ]:
            raise ValueError(f"Invalid write policy: {write_policy}")

        # self.write_queue = PriorityQueue[CacheOperation]()
        self.load_queue: List[CacheOperation] = []
        self.write_queue: List[CacheOperation] = []
        self.draft_load_queue: List[CacheOperation] = []
        self.draft_write_queue: List[CacheOperation] = []
        self.ack_load_queue: List[HiCacheAck] = []
        self.ack_write_queue: List[HiCacheAck] = []
        self.pending_layer_writes: Dict[int, HiCacheLayerWriteState] = {}

        if draft_mem_pool_host is not None or draft_mem_pool_device is not None:
            self.attach_draft_pool(draft_mem_pool_device, draft_mem_pool_host)

        self.stop_event = threading.Event()
        self.write_buffer = TransferBuffer(self.stop_event)
        self.load_buffer = TransferBuffer(
            self.stop_event, buffer_count=10, max_buffer_size=100
        )

        self.write_stream = device_module.Stream()
        self.load_stream = device_module.Stream()

        # If a storage backend is provided at startup, treat it as an implicit attach,
        # so init/runtime share the same lifecycle semantics and code paths.
        if storage_backend is not None:
            try:
                self.attach_storage_backend(
                    storage_backend=storage_backend,
                    prefetch_threshold=prefetch_threshold,
                    model_name=model_name,
                    storage_backend_extra_config=storage_backend_extra_config,
                )
            except ValueError as e:
                # Preserve the historical error shape on init for unknown backends.
                raise ValueError(f"Failed to create storage backend: {e}") from e

    @property
    def has_draft_hicache(self) -> bool:
        return (
            self.draft_mem_pool_host is not None
            and self.draft_mem_pool_device is not None
        )

    def attach_draft_pool(self, draft_mem_pool_device, draft_mem_pool_host) -> None:
        if draft_mem_pool_device is None and draft_mem_pool_host is None:
            return
        if draft_mem_pool_device is None or draft_mem_pool_host is None:
            raise ValueError(
                "draft_mem_pool_device and draft_mem_pool_host must be provided together"
            )
        self.draft_mem_pool_device = draft_mem_pool_device
        self.draft_mem_pool_host = draft_mem_pool_host
        if hasattr(draft_mem_pool_device, "register_layer_transfer_counter"):
            draft_mem_pool_device.register_layer_transfer_counter(
                self.layer_done_counter
            )
        if hasattr(draft_mem_pool_device, "register_layer_backup_notifier"):
            draft_mem_pool_device.register_layer_backup_notifier(
                lambda layer_id: self.on_layer_end(layer_id, source="draft")
            )

    def on_layer_end(self, layer_id: int, source: str = "target") -> None:
        if not self.pending_layer_writes:
            return
        if source not in ("target", "draft"):
            raise ValueError(f"Unknown CP HiCache layer backup source: {source}")
        reservations = [
            state.reservation for state in list(self.pending_layer_writes.values())
        ]
        for reservation in reservations:
            if reservation.node_id not in self.pending_layer_writes:
                continue
            self.submit_write_cp_layer(
                reservation,
                layer_id,
                submit_target=(source == "target"),
                submit_draft=(source == "draft"),
            )

    def _start_storage_threads(self):
        """Start storage prefetch/backup threads and their queues.

        This is used by runtime attach, and also by reset when storage is enabled.
        """
        assert self.enable_storage
        assert not self.storage_stop_event.is_set()

        self.prefetch_thread = threading.Thread(
            target=self.prefetch_thread_func, daemon=True
        )
        self.backup_thread = threading.Thread(
            target=self.backup_thread_func, daemon=True
        )
        self.prefetch_queue = Queue()
        self.backup_queue = Queue()

        self.prefetch_revoke_queue = Queue()
        self.ack_backup_queue = Queue()
        self.host_mem_release_queue = Queue()

        self.prefetch_thread.start()
        self.backup_thread.start()

    def _stop_storage_threads(self):
        """Stop storage prefetch/backup threads and drain internal queues.

        Caller should ensure no in-flight requests.
        """
        # Always request stop. This is safe even when storage is already disabled,
        # and makes detach truly idempotent (previous partial detach may have left
        # threads alive).
        # NOTE: do NOT clear stop_event unless threads have fully stopped; otherwise
        # a still-alive thread may resume and touch released state.
        self.storage_stop_event.set()

        # Best-effort wakeups so threads exit promptly even if blocked on queues.
        try:
            if hasattr(self, "prefetch_queue"):
                self.prefetch_queue.put_nowait(None)
            if hasattr(self, "backup_queue"):
                self.backup_queue.put_nowait(None)
            if hasattr(self, "prefetch_buffer"):
                self.prefetch_buffer.put_nowait(None)
        except Exception:
            pass

        # Best-effort joins (threads are daemon, but join keeps state clean).
        threads = []
        if hasattr(self, "prefetch_thread"):
            threads.append(self.prefetch_thread)
        if hasattr(self, "backup_thread"):
            threads.append(self.backup_thread)
        if hasattr(self, "prefetch_io_aux_thread"):
            threads.append(self.prefetch_io_aux_thread)

        for t in threads:
            try:
                t.join(timeout=10)
            except Exception:
                pass

        alive = [t for t in threads if getattr(t, "is_alive", lambda: False)()]
        if alive:
            logger.error(
                "Failed to stop HiCache storage threads cleanly: %s",
                [getattr(t, "name", repr(t)) for t in alive],
            )
            raise RuntimeError("Failed to stop HiCache storage threads cleanly.")

    def attach_storage_backend(
        self,
        storage_backend: str,
        prefetch_threshold: int = 256,
        model_name: Optional[str] = None,
        storage_backend_extra_config: Optional[dict] = None,
    ):
        """Attach (enable) storage backend at runtime.

        Requirement: no in-flight requests. This call is expected to run on the scheduler
        thread (control path), not concurrently with prefetch/backup.
        """
        if self.uses_cp_hicache:
            raise RuntimeError(
                "CP shared KV HiCache does not support attaching a storage backend at runtime."
            )

        if self.enable_storage:
            raise RuntimeError("Storage backend already attached.")

        # Defensive: a previous partial detach may have flipped `enable_storage` but
        # left background threads alive. Attaching on top of them is unsafe.
        try:
            self._stop_storage_threads()
        except Exception as e:
            raise RuntimeError(
                "Cannot attach storage backend: previous detach did not stop storage threads cleanly."
            ) from e

        # Rollback-safe init: if creation fails, keep controller state consistent
        # for future attach attempts.
        self.storage_backend_type = storage_backend
        from sglang.srt.mem_cache.hicache_storage import get_hash_str

        self.get_hash_str = get_hash_str
        self.storage_config = self._generate_storage_config(
            model_name, storage_backend_extra_config
        )
        # for MLA models, only one rank needs to backup the KV cache
        self.backup_skip = (
            self.storage_config.is_mla_model
            # todo: load balancing
            and self.storage_config.tp_rank != 0
        )

        # Use storage backend factory for dynamic backend creation
        from sglang.srt.mem_cache.storage import StorageBackendFactory

        try:
            self.storage_backend = StorageBackendFactory.create_backend(
                storage_backend, self.storage_config, self.mem_pool_host
            )
            self.storage_backend.register_mem_pool_host(self.mem_pool_host)

            self.enable_storage = True
            # todo: threshold policy for prefetching
            self.prefetch_threshold = max(prefetch_threshold, self.page_size)
            self.prefetch_capacity_limit = max(
                0, int(0.8 * (self.mem_pool_host.size - self.mem_pool_device.size))
            )
            # granularity of batch storage IO operations, in number of pages
            self.storage_batch_size = 128
            # tracking the number of tokens locked in prefetching, updated by the main scheduler thread
            self.prefetch_tokens_occupied = 0

            # create a new communication group for synchronizing storage operations across TP workers
            self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
            if self.tp_world_size > 1:
                from sglang.srt.distributed.parallel_state import (
                    create_custom_parallel_group,
                )

                group_ranks = torch.distributed.get_process_group_ranks(self.tp_group)
                self.prefetch_tp_group = create_custom_parallel_group(
                    group_ranks=group_ranks, backend="gloo"
                )

            # Select the get and set functions
            self.page_get_func = self._generic_page_get
            self.page_set_func = self._generic_page_set

            if (self.storage_backend_type in ["hf3fs", "mooncake", "eic", "nixl"]) or (
                self.storage_backend_type == "dynamic"
                and bool(self.storage_config.extra_config.get("interface_v1", 0))
            ):
                self.page_get_func = self._page_get_zero_copy
                self.page_set_func = self._page_set_zero_copy

            # Ensure stop_event is clear before starting threads.
            self.storage_stop_event.clear()
            self._start_storage_threads()
        except Exception:
            # Best-effort cleanup for partial init.
            try:
                self._stop_storage_threads()
            except Exception:
                pass
            try:
                if hasattr(self, "prefetch_tp_group"):
                    try:
                        torch.distributed.destroy_process_group(self.prefetch_tp_group)
                    except Exception:
                        pass
                    self.prefetch_tp_group = None
            except Exception:
                pass
            try:
                if (
                    hasattr(self, "storage_backend")
                    and self.storage_backend is not None
                ):
                    if hasattr(self.storage_backend, "close"):
                        self.storage_backend.close()
            except Exception:
                pass
            self.storage_backend = None
            self.storage_backend_type = None
            self.enable_storage = False
            self.page_get_func = self._generic_page_get
            self.page_set_func = self._generic_page_set
            raise

    def detach_storage_backend(self):
        """Detach (disable) storage backend at runtime.

        Requirement: no in-flight requests. This will stop storage threads and release
        the backend instance (best-effort close).
        """
        # Idempotent cleanup: even if `enable_storage` is already False,
        # we may still have leftover resources (threads/backend/process group) from a
        # previous partial detach. We attempt cleanup whenever possible.
        try:
            self._stop_storage_threads()
        except Exception as e:
            # Do not proceed tearing down backend/process group if threads are not
            # fully stopped; otherwise still-alive threads may touch released state.
            # Caller can retry detach.
            logger.exception("Stop storage threads failed: %s", e)
            # IMPORTANT: Do not silently succeed. Upper layers rely on exceptions here
            # to avoid flipping `enable_storage` flags while threads are still alive.
            raise RuntimeError("Stop storage threads failed; detach aborted.") from e

        # Best-effort destroy process group created for storage ops.
        try:
            if (
                hasattr(self, "prefetch_tp_group")
                and self.prefetch_tp_group is not None
            ):
                try:
                    torch.distributed.destroy_process_group(self.prefetch_tp_group)
                except Exception:
                    pass
                self.prefetch_tp_group = None
        except Exception:
            pass

        # Best-effort close (some backends rely on GC/destructor).
        try:
            if (
                hasattr(self, "storage_backend")
                and self.storage_backend is not None
                and hasattr(self.storage_backend, "close")
            ):
                self.storage_backend.close()
        except Exception:
            logger.exception("Failed to close storage backend cleanly.")

        self.storage_backend = None
        self.storage_backend_type = None
        self.enable_storage = False
        self.page_get_func = self._generic_page_get
        self.page_set_func = self._generic_page_set
        # Now it's safe to clear the stop event for future re-attach.
        self.storage_stop_event.clear()

    def _generate_storage_config(
        self,
        model_name: Optional[str] = None,
        storage_backend_extra_config: Optional[dict] = None,
    ):
        if storage_backend_extra_config is None:
            storage_backend_extra_config = {}

        if is_dp_attention_enabled():
            self.tp_rank = get_attention_tp_rank()
            self.tp_size = get_attention_tp_size()
            self.dp_rank = get_attention_dp_rank()
        else:
            self.tp_rank = get_tensor_model_parallel_rank()
            self.tp_size = get_tensor_model_parallel_world_size()
            self.dp_rank = 0

        # Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool.
        is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool)
        # Least Common Multiple among heterogeneous tp size
        tp_lcm_size = storage_backend_extra_config.pop("tp_lcm_size", None)
        should_split_heads = False

        if tp_lcm_size:
            assert (
                tp_lcm_size % self.tp_size == 0
            ), "tp_lcm_size must be divisible by tp_size."
            should_split_heads = (
                not is_mla_backend
                and self.mem_pool_host.layout == "page_head"
                and tp_lcm_size > self.tp_size
            )

        return HiCacheStorageConfig(
            tp_rank=self.tp_rank,
            tp_size=self.tp_size,
            pp_rank=self.pp_rank,
            pp_size=self.pp_size,
            is_mla_model=is_mla_backend,
            enable_storage_metrics=self.enable_storage_metrics,
            is_page_first_layout=self.mem_pool_host.layout == "page_first",
            model_name=model_name,
            tp_lcm_size=tp_lcm_size,
            should_split_heads=should_split_heads,
            extra_config=storage_backend_extra_config,
        )

    def reset(self):
        self.stop_event.set()
        self.storage_stop_event.set()

        self.write_queue.clear()
        self.load_queue.clear()
        self.draft_write_queue.clear()
        self.draft_load_queue.clear()
        self.write_buffer.clear()
        self.load_buffer.clear()
        self.ack_write_queue.clear()
        self.ack_load_queue.clear()
        self.pending_layer_writes.clear()
        if self.enable_storage:
            self.prefetch_thread.join()
            self.backup_thread.join()
            self.prefetch_queue.queue.clear()
            self.backup_queue.queue.clear()
            self.prefetch_revoke_queue.queue.clear()
            self.ack_backup_queue.queue.clear()

        self.stop_event.clear()
        self.storage_stop_event.clear()

        if self.enable_storage:
            self.prefetch_thread = threading.Thread(
                target=self.prefetch_thread_func, daemon=True
            )
            self.backup_thread = threading.Thread(
                target=self.backup_thread_func, daemon=True
            )
            self.prefetch_thread.start()
            self.backup_thread.start()

    def clear_draft_host_pool(self) -> None:
        if self.draft_mem_pool_host is not None:
            self.draft_mem_pool_host.clear()

    def write(
        self,
        device_indices: torch.Tensor,
        priority: Optional[int] = None,
        node_id: int = -1,
    ) -> Optional[torch.Tensor | HiCacheWriteResult | HiCacheWriteFailure]:
        """
        Back up KV caches from device memory to host memory.
        """
        logger.info(
            "[CacheCtrl-write] write: node_id=%d len=%d cp=%s",
            node_id,
            len(device_indices),
            self.uses_cp_hicache,
        )
        if self.uses_cp_hicache:
            return self._write_cp(device_indices, priority, node_id)

        host_indices = self.mem_pool_host.alloc(len(device_indices))
        if host_indices is None:
            logger.info(
                "[CacheCtrl-write] write non-CP FAILED (host full): node_id=%d len=%d",
                node_id,
                len(device_indices),
            )
            return None
        self.write_queue.append(
            CacheOperation(host_indices, device_indices, node_id, priority)
        )
        if self.has_draft and not self.uses_cp_hicache:
            self.draft_write_queue.append(
                CacheOperation(host_indices, device_indices, node_id, priority)
            )
        self.start_writing()
        logger.info(
            "[CacheCtrl-write] write non-CP submitted: node_id=%d len=%d",
            node_id,
            len(device_indices),
        )
        return host_indices

    def start_writing(self) -> None:
        if len(self.write_queue) == 0 and len(self.draft_write_queue) == 0:
            return

        op = CacheOperation.merge_ops(self.write_queue) if self.write_queue else None
        draft_op = (
            CacheOperation.merge_ops(self.draft_write_queue)
            if self.draft_write_queue
            else None
        )
        node_ids = op.node_ids if op is not None else draft_op.node_ids
        if op is not None:
            host_indices, device_indices = self.move_indices(op, self.mem_pool_host)
        else:
            host_indices = device_indices = None
        if draft_op is not None:
            draft_host_indices, draft_device_indices = self.move_indices(
                draft_op, self.draft_mem_pool_host
            )
        else:
            draft_host_indices = draft_device_indices = None
        self.write_queue.clear()
        self.draft_write_queue.clear()

        start_event = device_module.Event()
        finish_event = device_module.Event()

        start_event.record()
        with device_module.stream(self.write_stream):
            start_event.wait(self.write_stream)
            if op is not None:
                self.mem_pool_host.backup_from_device_all_layer(
                    self.mem_pool_device, host_indices, device_indices, self.io_backend
                )
            if draft_op is not None:
                self.draft_mem_pool_host.backup_from_device_all_layer(
                    self.draft_mem_pool_device,
                    draft_host_indices,
                    draft_device_indices,
                    self.io_backend,
                )
            finish_event.record()
            # NOTE: We must save the host indices and device indices here,
            # this is because we need to guarantee that these tensors are
            # still alive when the write stream is executing.
            if host_indices is not None and host_indices.is_cuda:
                host_indices.record_stream(self.write_stream)
            if device_indices is not None and device_indices.is_cuda:
                device_indices.record_stream(self.write_stream)
            if draft_host_indices is not None and draft_host_indices.is_cuda:
                draft_host_indices.record_stream(self.write_stream)
            if draft_device_indices is not None and draft_device_indices.is_cuda:
                draft_device_indices.record_stream(self.write_stream)

        self.ack_write_queue.append(HiCacheAck(start_event, finish_event, node_ids))

    def _append_completed_write_ack(self, node_id: int) -> None:
        event = device_module.Event()
        event.record()
        self.ack_write_queue.append(HiCacheAck(event, event, [node_id]))

    def _append_completed_load_ack(self, node_id: int) -> None:
        event = device_module.Event()
        event.record()
        self.ack_load_queue.append(HiCacheAck(event, event, [node_id]))

    def _validate_cp_hicache_page_indices(
        self,
        host_indices: torch.Tensor,
        device_indices: torch.Tensor,
    ) -> None:
        """Validate page-shaped CP HiCache indices without CUDA host sync.

        The production invariant is construction-based: HostKVCache alloc/free
        preserves page-shaped host spans, and CpSharedKVLayout preserves
        per-page contiguity when mapping logical locs to physical locs.  The
        generic validator uses Tensor truth values and would synchronize CUDA
        tensors, so keep the hot path sync-free and validate CPU/fake-test
        tensors only.
        """

        if not host_indices.is_cuda:
            validate_page_aligned_token_indices(
                host_indices, self.page_size, "host_indices"
            )
        if not device_indices.is_cuda:
            validate_page_aligned_token_indices(
                device_indices, self.page_size, "physical_device_indices"
            )

    def reserve_write_cp(
        self,
        device_indices: torch.Tensor,
        priority: Optional[int] = None,
        node_id: int = -1,
    ) -> HiCacheWriteReservation | HiCacheWriteFailure:
        from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata

        layout = self.cp_shared_kv_layout
        owned_mask = layout.owned_by_this_rank(device_indices)
        owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu()
        logical_len = len(device_indices)
        # Capture the global owner pattern (one int8 per logical PAGE; identical
        # on all CP ranks since it's a pure function of the logical page ids).
        # load_cp will replay this pattern via alloc_pages_with_owners so the
        # saved owned_positions correctly index the new allocation.
        page_size = self.page_size
        if logical_len % page_size != 0:
            raise ValueError(
                f"_write_cp expects page-aligned device_indices, got "
                f"logical_len={logical_len} page_size={page_size}"
            )
        page_first_locs = device_indices[::page_size]
        logical_pages = torch.div(page_first_locs, page_size, rounding_mode="floor")
        page_owners = layout.owner_for_logical_pages(logical_pages).to(
            dtype=torch.int8, device="cpu"
        )
        if owned_positions.numel() == 0:
            logger.info(
                "[CacheCtrl-write] reserve_write_cp zero-owned rank: node_id=%d logical_len=%d",
                node_id,
                logical_len,
            )
            empty = torch.empty((0,), dtype=torch.int64)
            return HiCacheWriteReservation(
                metadata=CpHiCacheNodeMetadata(
                    logical_len=logical_len,
                    owned_positions=owned_positions,
                    host_indices=empty,
                    page_owners=page_owners,
                    page_size=page_size,
                    draft_host_indices=(empty.clone() if self.has_draft_hicache else None),
                ),
                host_indices=empty,
                physical_device_indices=empty,
                node_id=node_id,
                priority=priority,
                draft_host_indices=(empty.clone() if self.has_draft_hicache else None),
            )

        owned_logical_indices = device_indices[owned_mask]
        physical_device_indices = layout.logical_locs_to_physical(owned_logical_indices)
        host_indices = self.mem_pool_host.alloc(len(physical_device_indices))
        if host_indices is None:
            logger.info(
                "[CacheCtrl-write] reserve_write_cp FAILED (host full): node_id=%d logical_len=%d owned=%d",
                node_id,
                logical_len,
                owned_positions.numel(),
            )
            return HiCacheWriteFailure(required_host_slots=len(physical_device_indices))

        draft_host_indices = None
        if self.has_draft_hicache:
            draft_host_indices = self.draft_mem_pool_host.alloc(
                len(physical_device_indices)
            )
            if draft_host_indices is None:
                self.mem_pool_host.free(host_indices)
                logger.info(
                    "[CacheCtrl-write] reserve_write_cp FAILED (draft host full): node_id=%d logical_len=%d owned=%d",
                    node_id,
                    logical_len,
                    owned_positions.numel(),
                )
                return HiCacheWriteFailure(
                    required_host_slots=len(physical_device_indices)
                )

        try:
            self._validate_cp_hicache_page_indices(
                host_indices, physical_device_indices
            )
            if draft_host_indices is not None:
                self._validate_cp_hicache_page_indices(
                    draft_host_indices, physical_device_indices
                )
        except Exception:
            self.mem_pool_host.free(host_indices)
            if draft_host_indices is not None:
                self.draft_mem_pool_host.free(draft_host_indices)
            raise

        return HiCacheWriteReservation(
            metadata=CpHiCacheNodeMetadata(
                logical_len=logical_len,
                owned_positions=owned_positions,
                host_indices=host_indices.cpu(),
                page_owners=page_owners,
                page_size=page_size,
                draft_host_indices=(
                    draft_host_indices.cpu() if draft_host_indices is not None else None
                ),
            ),
            host_indices=host_indices,
            physical_device_indices=physical_device_indices,
            node_id=node_id,
            priority=priority,
            draft_host_indices=draft_host_indices,
        )

    def submit_write_cp_all_layer(self, reservation: HiCacheWriteReservation) -> None:
        logger.warning(
            "[CacheCtrl-write] CP HiCache all-layer backup fallback: node_id=%d logical_len=%d owned=%d physical=%d draft=%s",
            reservation.node_id,
            reservation.metadata.logical_len,
            reservation.metadata.owned_positions.numel(),
            len(reservation.physical_device_indices),
            reservation.draft_host_indices is not None,
        )
        if len(reservation.physical_device_indices) == 0:
            self._append_completed_write_ack(reservation.node_id)
            logger.info(
                "[CacheCtrl-write] submit_write_cp_all_layer zero-owned ack: node_id=%d logical_len=%d",
                reservation.node_id,
                reservation.metadata.logical_len,
            )
            return

        self.write_queue.append(
            CacheOperation(
                reservation.host_indices,
                reservation.physical_device_indices,
                reservation.node_id,
                reservation.priority,
            )
        )
        if reservation.draft_host_indices is not None:
            self.draft_write_queue.append(
                CacheOperation(
                    reservation.draft_host_indices,
                    reservation.physical_device_indices,
                    reservation.node_id,
                    reservation.priority,
                )
            )
        self.start_writing()
        logger.info(
            "[CacheCtrl-write] submit_write_cp_all_layer submitted: node_id=%d logical_len=%d owned=%d physical=%d draft=%s",
            reservation.node_id,
            reservation.metadata.logical_len,
            reservation.metadata.owned_positions.numel(),
            len(reservation.physical_device_indices),
            reservation.draft_host_indices is not None,
        )

    def _get_or_create_layer_write_state(
        self, reservation: HiCacheWriteReservation
    ) -> HiCacheLayerWriteState:
        state = self.pending_layer_writes.get(reservation.node_id)
        if state is not None:
            if state.reservation is not reservation:
                raise RuntimeError(
                    f"Conflicting CP HiCache per-layer backup reservation for "
                    f"node_id={reservation.node_id}"
                )
            return state

        draft_layer_num = (
            self.draft_mem_pool_device.layer_num
            if reservation.draft_host_indices is not None
            else 0
        )
        total_layers = max(self.layer_num, draft_layer_num)
        if total_layers <= 0:
            raise RuntimeError(
                f"Invalid CP HiCache per-layer backup layer count: {total_layers}"
            )

        state = HiCacheLayerWriteState(
            reservation=reservation,
            total_layers=total_layers,
            start_event=device_module.Event(),
            finish_event=device_module.Event(),
        )
        if len(reservation.physical_device_indices) > 0:
            op = CacheOperation(
                reservation.host_indices,
                reservation.physical_device_indices,
                reservation.node_id,
                reservation.priority,
            )
            state.host_indices, state.physical_device_indices = self.move_indices(
                op, self.mem_pool_host
            )
            if reservation.draft_host_indices is not None:
                draft_op = CacheOperation(
                    reservation.draft_host_indices,
                    reservation.physical_device_indices,
                    reservation.node_id,
                    reservation.priority,
                )
                state.draft_host_indices, _ = self.move_indices(
                    draft_op, self.draft_mem_pool_host
                )
        self.pending_layer_writes[reservation.node_id] = state
        state.start_event.record()
        return state

    def submit_write_cp_layer(
        self,
        reservation: HiCacheWriteReservation,
        layer_id: int,
        *,
        submit_target: bool = True,
        submit_draft: bool = True,
    ) -> None:
        """Submit one layer of a reserved CP host backup.

        Target and draft D2H are still one logical operation: this method only
        queues a final write ack when all target/draft layers have been submitted.
        """

        state = self._get_or_create_layer_write_state(reservation)
        if layer_id < 0 or layer_id >= state.total_layers:
            raise ValueError(
                f"layer_id={layer_id} is outside CP HiCache backup layer range "
                f"[0, {state.total_layers}) for node_id={reservation.node_id}"
            )
        needs_target = (
            submit_target
            and layer_id < self.layer_num
            and layer_id not in state.completed_target_layers
        )
        needs_draft = (
            submit_draft
            and state.draft_host_indices is not None
            and layer_id < self.draft_mem_pool_device.layer_num
            and layer_id not in state.completed_draft_layers
        )
        if not needs_target and not needs_draft:
            return

        layer_event = device_module.Event()
        layer_event.record()
        state.layer_events.append(layer_event)

        if len(reservation.physical_device_indices) > 0:
            with device_module.stream(self.write_stream):
                state.start_event.wait(self.write_stream)
                layer_event.wait(self.write_stream)
                if needs_target:
                    self.mem_pool_host.backup_from_device_per_layer(
                        self.mem_pool_device,
                        state.host_indices,
                        state.physical_device_indices,
                        layer_id,
                        self.io_backend,
                    )
                if needs_draft:
                    self.draft_mem_pool_host.backup_from_device_per_layer(
                        self.draft_mem_pool_device,
                        state.draft_host_indices,
                        state.physical_device_indices,
                        layer_id,
                        self.io_backend,
                    )

        if needs_target:
            state.completed_target_layers.add(layer_id)
        if needs_draft:
            state.completed_draft_layers.add(layer_id)

        target_done = len(state.completed_target_layers) >= self.layer_num
        draft_done = (
            state.draft_host_indices is None
            or len(state.completed_draft_layers)
            >= self.draft_mem_pool_device.layer_num
        )
        if not target_done or not draft_done:
            return
        if state.ack_appended:
            return

        with device_module.stream(self.write_stream):
            state.finish_event.record()
            if state.host_indices is not None and state.host_indices.is_cuda:
                state.host_indices.record_stream(self.write_stream)
            if (
                state.physical_device_indices is not None
                and state.physical_device_indices.is_cuda
            ):
                state.physical_device_indices.record_stream(self.write_stream)
            if (
                state.draft_host_indices is not None
                and state.draft_host_indices.is_cuda
            ):
                state.draft_host_indices.record_stream(self.write_stream)

        state.ack_appended = True
        self.pending_layer_writes.pop(reservation.node_id, None)
        self.ack_write_queue.append(
            HiCacheAck(state.start_event, state.finish_event, [reservation.node_id])
        )
        logger.info(
            "[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d logical_len=%d layers=%d draft=%s",
            reservation.node_id,
            reservation.metadata.logical_len,
            state.total_layers,
            reservation.draft_host_indices is not None,
        )

    def submit_write_cp_per_layer(
        self,
        reservation: HiCacheWriteReservation,
        *,
        catch_up_all_layers: bool = True,
    ) -> None:
        """Register a CP write reservation for per-layer host backup.

        Current radix insertion calls write_backup after the request KV has already
        been materialized.  For that path, catch up by submitting all layers
        immediately through the per-layer API.  Future early reservations can use
        catch_up_all_layers=False and rely on explicit model layer-end hooks.
        """

        state = self._get_or_create_layer_write_state(reservation)
        if not catch_up_all_layers:
            logger.info(
                "[CacheCtrl-write] submit_write_cp_per_layer registered: node_id=%d logical_len=%d layers=%d draft=%s",
                reservation.node_id,
                reservation.metadata.logical_len,
                state.total_layers,
                reservation.draft_host_indices is not None,
            )
            return

        total_layers = state.total_layers
        for layer_id in range(total_layers):
            self.submit_write_cp_layer(reservation, layer_id)

    def _write_cp(
        self,
        device_indices: torch.Tensor,
        priority: Optional[int] = None,
        node_id: int = -1,
    ) -> HiCacheWriteResult | HiCacheWriteFailure:
        reservation = self.reserve_write_cp(device_indices, priority, node_id)
        if isinstance(reservation, HiCacheWriteFailure):
            return reservation
        self.submit_write_cp_per_layer(reservation)
        return HiCacheWriteResult(metadata=reservation.metadata)

    def set_draft_kv_pool(self, draft_device_pool, draft_host_pool) -> None:
        """Register draft KV pools so L2 ops piggyback draft transfers."""
        self.has_draft = True
        self.mem_pool_device_draft = draft_device_pool
        self.mem_pool_host_draft = draft_host_pool
        self.attach_draft_pool(draft_device_pool, draft_host_pool)
        logger.info(
            "HiCache draft KV registered: %s (host %d slots)",
            type(draft_device_pool).__name__,
            draft_host_pool.size,
        )

    def load(
        self,
        host_indices: torch.Tensor,
        priority: Optional[int] = None,
        node_id: int = -1,
    ) -> Optional[torch.Tensor]:
        """
        Load KV caches from host memory to device memory.
        """
        device_indices = self.mem_pool_device_allocator.alloc(len(host_indices))
        if device_indices is None:
            return None
        self.load_queue.append(
            CacheOperation(host_indices, device_indices, node_id, priority)
        )
        if self.has_draft and not self.uses_cp_hicache:
            self.draft_load_queue.append(
                CacheOperation(host_indices, device_indices, node_id, priority)
            )
        return device_indices

    def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]:
        # Reproduce the original (write-time) CP owner pattern. Each node
        # carries `page_owners` (one int8 per logical page, identical on all
        # CP ranks) so we can ask the allocator for a fresh device range
        # whose per-page owner sequence matches what was backed up. Without
        # this, the saved `owned_positions` (positions WITHIN the original
        # alloc that this rank owned) would index a new alloc with arbitrary
        # owner pattern → each rank loads its host bytes into physical slots
        # whose corresponding logical page is owned by some other rank →
        # attention reads garbage at forward time.
        page_owners: List[int] = []
        for node in nodes_to_load:
            meta = node.cp_hicache
            if meta is None:
                raise RuntimeError(
                    f"load_cp called with node {getattr(node, 'id', '?')} "
                    "that has no cp_hicache metadata"
                )
            # page_owners is CPU int8; .tolist() returns list[int] directly.
            page_owners.extend(meta.page_owners.tolist())

        device_indices = self.mem_pool_device_allocator.alloc_pages_with_owners(
            page_owners
        )
        # Fail closed: returning None lets the caller drop to cache miss
        # (cold prefill). Never proceed with a non-matching owner pattern.
        if device_indices is None:
            return None

        logical_len_expected = sum(node.host_len for node in nodes_to_load)
        if device_indices.numel() != logical_len_expected:
            self.mem_pool_device_allocator.free(device_indices)
            raise RuntimeError(
                "alloc_pages_with_owners returned unexpected length: "
                f"got {device_indices.numel()}, expected {logical_len_expected}"
            )

        host_chunks = []
        draft_host_chunks = []
        physical_chunks = []
        offset = 0
        for node in nodes_to_load:
            node_device_indices = device_indices[offset : offset + node.host_len]
            offset += node.host_len
            if self.has_draft_hicache:
                draft_host_indices = getattr(
                    node.cp_hicache, "draft_host_indices", None
                )
                if draft_host_indices is None:
                    self.mem_pool_device_allocator.free(device_indices)
                    raise RuntimeError(
                        "CP HiCache draft KV restore requested but node metadata "
                        "does not contain draft_host_indices"
                    )
            else:
                draft_host_indices = None

            owned_positions = node.cp_hicache.owned_positions.to(device_indices.device)
            if owned_positions.numel() == 0:
                continue
            selected_logical_locs = node_device_indices[owned_positions]
            # Note: NOT re-validating owner_by_this_rank here. The invariant
            # — every position in `owned_positions` lands on a page owned by
            # this rank — is guaranteed by construction:
            #   (a) at write time, `owned_positions = owned_mask.nonzero(...)`
            #       where `owned_mask = owned_by_this_rank(device_indices)`;
            #   (b) at load time, `alloc_pages_with_owners(page_owners)` returns
            #       pages whose owner sequence matches `page_owners` by
            #       construction (and is debug-asserted inside the allocator).
            # Both sides use the same `page_owners` (write-derived, identical
            # on all CP ranks). A redundant `.all().item()` check here would
            # force a CUDA host-sync — the exact anti-pattern commit 97a9f850c
            # removed from the hot path.
            physical_chunks.append(
                self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs)
            )
            host_chunks.append(node.cp_hicache.host_indices)
            if draft_host_indices is not None:
                draft_host_chunks.append(draft_host_indices)

        if not host_chunks:
            # Keep CP load ACK rows identical across ranks.  A zero-owned rank
            # still queues a zero-length op with the logical node id so a later
            # batched start_loading() merges the same node_ids as owning ranks.
            self.load_queue.append(
                CacheOperation(
                    torch.empty((0,), dtype=torch.int64),
                    torch.empty(
                        (0,), dtype=device_indices.dtype, device=device_indices.device
                    ),
                    node_id,
                )
            )
            return device_indices

        host_indices = torch.cat(host_chunks)
        physical_device_indices = torch.cat(physical_chunks)
        draft_host_indices = None
        if self.has_draft_hicache:
            draft_host_indices = torch.cat(draft_host_chunks)

        try:
            self._validate_cp_hicache_page_indices(
                host_indices, physical_device_indices
            )
            if draft_host_indices is not None:
                self._validate_cp_hicache_page_indices(
                    draft_host_indices, physical_device_indices
                )
        except Exception:
            self.mem_pool_device_allocator.free(device_indices)
            raise

        self.load_queue.append(
            CacheOperation(
                host_indices,
                physical_device_indices,
                node_id,
            )
        )
        if draft_host_indices is not None:
            self.draft_load_queue.append(
                CacheOperation(
                    draft_host_indices,
                    physical_device_indices,
                    node_id,
                )
            )

        return device_indices

    def move_indices(self, op: CacheOperation, mem_pool_host=None):
        mem_pool_host = mem_pool_host or self.mem_pool_host
        host_indices, device_indices = op.host_indices, op.device_indices
        # move indices to GPU if using kernels, to host if using direct indexing
        if self.io_backend == "kernel":
            if not host_indices.is_cuda:
                host_indices = host_indices.to(self.device, non_blocking=True)
            return host_indices, device_indices
        elif self.io_backend == "direct":
            if mem_pool_host.layout == "layer_first":
                device_indices = device_indices.cpu()
                host_indices, idx = host_indices.sort()
                return host_indices, device_indices.index_select(0, idx)
            elif mem_pool_host.layout == "page_first_direct":
                return host_indices, device_indices.cpu()
            else:
                raise ValueError(
                    f"Unsupported layout {mem_pool_host.layout!r} for io backend 'direct'"
                )
        elif self.io_backend == "kernel_ascend":
            return host_indices, device_indices.cpu()
        else:
            raise ValueError(f"Unsupported io backend")

    def start_loading(self) -> int:
        if len(self.load_queue) == 0 and len(self.draft_load_queue) == 0:
            return -1

        producer_id = self.layer_done_counter.update_producer()
        op = CacheOperation.merge_ops(self.load_queue) if self.load_queue else None
        draft_op = (
            CacheOperation.merge_ops(self.draft_load_queue)
            if self.draft_load_queue
            else None
        )
        node_ids = op.node_ids if op is not None else draft_op.node_ids
        if op is not None:
            host_indices, device_indices = self.move_indices(op, self.mem_pool_host)
        else:
            host_indices = device_indices = None
        if draft_op is not None:
            draft_host_indices, draft_device_indices = self.move_indices(
                draft_op, self.draft_mem_pool_host
            )
        else:
            draft_host_indices = draft_device_indices = None
        self.load_queue.clear()
        self.draft_load_queue.clear()
        producer_event = self.layer_done_counter.events[producer_id]
        producer_event.start_event.record()

        with device_module.stream(self.load_stream):
            producer_event.start_event.wait(self.load_stream)
            draft_layer_num = (
                self.draft_mem_pool_device.layer_num if draft_op is not None else 0
            )
            for i in range(max(self.layer_num, draft_layer_num)):
                if draft_op is not None and i < draft_layer_num:
                    if len(draft_host_indices) > 0:
                        self.draft_mem_pool_host.load_to_device_per_layer(
                            self.draft_mem_pool_device,
                            draft_host_indices,
                            draft_device_indices,
                            i,
                            self.io_backend,
                        )
                if op is not None and i < self.layer_num:
                    if len(host_indices) > 0:
                        self.mem_pool_host.load_to_device_per_layer(
                            self.mem_pool_device,
                            host_indices,
                            device_indices,
                            i,
                            self.io_backend,
                        )
                    producer_event.complete(i)
                elif op is None and i < self.layer_num:
                    producer_event.complete(i)
            # NOTE: We must save the host indices and device indices here,
            # this is because we need to guarantee that these tensors are
            # still alive when the load stream is executing.
            if host_indices is not None and host_indices.is_cuda:
                host_indices.record_stream(self.load_stream)
            if device_indices is not None and device_indices.is_cuda:
                device_indices.record_stream(self.load_stream)
            if draft_host_indices is not None and draft_host_indices.is_cuda:
                draft_host_indices.record_stream(self.load_stream)
            if draft_device_indices is not None and draft_device_indices.is_cuda:
                draft_device_indices.record_stream(self.load_stream)

        self.ack_load_queue.append(
            HiCacheAck(
                start_event=producer_event.start_event,
                finish_event=producer_event.finish_event,
                node_ids=node_ids,
            )
        )
        return producer_id

    def evict_device(self, device_indices: torch.Tensor) -> int:
        self.mem_pool_device_allocator.free(device_indices)
        return len(device_indices)

    def evict_host(self, host_indices: torch.Tensor, backup_only: bool = True) -> int:
        if not backup_only:
            raise ValueError("Other eviction policies are not supported yet.")

        self.mem_pool_host.free(host_indices)
        return len(host_indices)

    def evict_cp_host(self, metadata, backup_only: bool = True) -> int:
        if not backup_only:
            raise ValueError("Other eviction policies are not supported yet.")

        draft_host_indices = getattr(metadata, "draft_host_indices", None)
        if self.has_draft_hicache and draft_host_indices is None:
            raise RuntimeError(
                "CP HiCache draft KV host evict requested but node metadata "
                "does not contain draft_host_indices"
            )

        freed = self.evict_host(metadata.host_indices, backup_only=backup_only)
        if self.has_draft_hicache:
            self.draft_mem_pool_host.free(draft_host_indices)
        return freed

    def prefetch(
        self,
        request_id: str,
        host_indices: torch.Tensor,
        new_input_tokens: List[int],
        last_hash: Optional[str] = None,
        prefix_keys: Optional[List[str]] = None,
    ) -> PrefetchOperation:
        """
        Prefetch KV caches from storage backend to host memory.
        """
        operation = PrefetchOperation(
            request_id, host_indices, new_input_tokens, last_hash, prefix_keys
        )
        self.prefetch_queue.put(operation)
        return operation

    def terminate_prefetch(self, operation):
        operation.mark_terminate()
        return operation.completed_tokens, operation.hash_value

    def append_host_mem_release(self, host_indices: torch.Tensor):
        if host_indices.numel() == 0:
            return
        pages = host_indices.split(self.mem_pool_host.page_size)
        for page in pages:
            self.host_mem_release_queue.put(page)

    def _page_get_zero_copy(self, operation, hash_values, host_indices, extra_info=None):
        results = self.storage_backend.batch_get_v1(
            hash_values, host_indices, extra_info
        )
        inc = 0
        for i in range(len(hash_values)):
            if not results[i]:
                logger.warning(
                    f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
                )
                break
            inc += self.page_size
        operation.increment(inc)

    # todo: deprecate
    def _generic_page_get(self, operation, hash_values, host_indices, extra_info=None):
        dummy_page_dst = [
            self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values
        ]
        page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst)
        if page_data is None:
            return
        for i in range(len(hash_values)):
            if page_data[i] is None:
                logger.warning(
                    f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
                )
                break
            # Must set the data before increasing the completed tokens.
            # Otherwise this page may be read before being set.
            self.mem_pool_host.set_from_flat_data_page(
                host_indices[i * self.page_size],
                page_data[i],
            )
            if not operation.increment(self.page_size):
                break  # Operation terminated by controller

    def _page_transfer(self, operation):
        # Transfer batch by batch
        prefix_keys = operation.prefix_keys
        for i in range(0, len(operation.hash_value), self.storage_batch_size):
            batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
            batch_host_indices = operation.host_indices[
                i * self.page_size : (i + len(batch_hashes)) * self.page_size
            ]
            prev_completed_tokens = operation.completed_tokens
            # Get one batch token, and update the completed_tokens if succeed
            from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo

            extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
            self.page_get_func(operation, batch_hashes, batch_host_indices, extra_info)
            # Check termination
            if (
                operation.completed_tokens
                != prev_completed_tokens + len(batch_hashes) * self.page_size
            ):
                operation.mark_terminate()
                break  # Some operations fail or operation terminated by controller

            if prefix_keys and len(prefix_keys) > 0:
                prefix_keys += batch_hashes

    def prefetch_io_aux_func(self):
        """
        Auxiliary function conducting IO operations for prefetching.
        """
        while not self.storage_stop_event.is_set():
            try:
                operation = self.prefetch_buffer.get(block=True, timeout=1)
                if operation is None:
                    continue
                self._page_transfer(operation)
                # operation terminated by controller, release pre-allocated memory
                self.append_host_mem_release(
                    operation.host_indices[operation.completed_tokens :]
                )
            except Empty:
                continue

    def prefetch_rate_limited(self) -> bool:
        """
        Rate limit the prefetching operations to avoid overwhelming the storage backend.
        """
        # cancel prefetch if too much memory is occupied
        if self.prefetch_tokens_occupied >= self.prefetch_capacity_limit:
            return True
        # todo: more sophisticated rate limiting based on storage backend performance
        return False

    def _storage_hit_query(self, operation) -> tuple[list[str], int]:
        last_hash = operation.last_hash
        tokens_to_fetch = operation.token_ids
        prefix_keys = operation.prefix_keys.copy() if operation.prefix_keys else None

        storage_query_count = 0
        hash_value = []

        for start in range(
            0, len(tokens_to_fetch), self.page_size * self.storage_batch_size
        ):
            end = min(
                start + self.page_size * self.storage_batch_size, len(tokens_to_fetch)
            )
            batch_tokens = tokens_to_fetch[start:end]
            batch_hashes = []
            for i in range(0, len(batch_tokens), self.page_size):
                last_hash = self.get_hash_str(
                    batch_tokens[i : i + self.page_size], last_hash
                )
                batch_hashes.append(last_hash)
            from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo

            extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
            hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info)
            hash_value.extend(batch_hashes[:hit_page_num])
            storage_query_count += hit_page_num * self.page_size
            if hit_page_num < len(batch_hashes):
                break
            if prefix_keys and len(prefix_keys) > 0:
                prefix_keys += batch_hashes

        return hash_value, storage_query_count

    def prefetch_thread_func(self):
        """
        Manage prefetching operations from storage backend to host memory.
        """
        self.prefetch_buffer = Queue()
        self.prefetch_io_aux_thread = threading.Thread(
            target=self.prefetch_io_aux_func, daemon=True
        )
        self.prefetch_io_aux_thread.start()
        while (not self.storage_stop_event.is_set()) or not self.prefetch_queue.empty():
            try:
                operation = self.prefetch_queue.get(block=True, timeout=1)
                if operation is None:
                    continue
                hash_value, storage_hit_count = self._storage_hit_query(operation)
                if self.tp_world_size > 1:
                    storage_hit_count_tensor = torch.tensor(
                        storage_hit_count, dtype=torch.int
                    )
                    torch.distributed.all_reduce(
                        storage_hit_count_tensor,
                        op=torch.distributed.ReduceOp.MIN,
                        group=self.prefetch_tp_group,
                    )
                    storage_hit_count = storage_hit_count_tensor.item()

                if storage_hit_count < self.prefetch_threshold:
                    # not to prefetch if not enough benefits
                    self.prefetch_revoke_queue.put(operation.request_id)
                    self.append_host_mem_release(operation.host_indices)
                    logger.info(
                        f"Revoking prefetch for request {operation.request_id} due to insufficient hits ({storage_hit_count})."
                    )
                else:
                    operation.hash_value = hash_value[
                        : (storage_hit_count // self.page_size)
                    ]
                    # free the pre-allocated memory for pages that are not hit
                    self.append_host_mem_release(
                        operation.host_indices[storage_hit_count:]
                    )
                    operation.host_indices = operation.host_indices[:storage_hit_count]
                    logger.info(
                        f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}."
                    )
                    self.prefetch_buffer.put(operation)

            except Empty:
                continue

    def write_storage(
        self,
        host_indices: torch.Tensor,
        token_ids: List[int],
        hash_value: Optional[List[str]] = None,
        prefix_keys: Optional[List[str]] = None,
    ) -> int:
        """
        Write KV caches from host memory to storage backend.
        """
        operation = StorageOperation(
            host_indices, token_ids, hash_value=hash_value, prefix_keys=prefix_keys
        )
        self.backup_queue.put(operation)
        return operation.id

    # todo: deprecate
    def _generic_page_set(self, hash_values, host_indices, extra_info=None) -> bool:
        data = [
            self.mem_pool_host.get_data_page(host_indices[i * self.page_size])
            for i in range(len(hash_values))
        ]
        return self.storage_backend.batch_set(hash_values, data)

    def _page_set_zero_copy(self, hash_values, host_indices, extra_info=None) -> bool:
        return all(
            self.storage_backend.batch_set_v1(hash_values, host_indices, extra_info)
        )

    # Backup batch by batch
    def _page_backup(self, operation):
        # Backup batch by batch
        prefix_keys = operation.prefix_keys
        for i in range(0, len(operation.hash_value), self.storage_batch_size):
            batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
            batch_host_indices = operation.host_indices[
                i * self.page_size : (i + len(batch_hashes)) * self.page_size
            ]
            # Set one batch token, and record if success.
            # todo: allow partial success
            from sglang.srt.mem_cache.hicache_storage import HiCacheStorageExtraInfo

            extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
            success = self.page_set_func(batch_hashes, batch_host_indices, extra_info)
            if not success:
                logger.warning(
                    f"Write page to storage: {len(batch_hashes)} pages failed."
                )
                break

            if prefix_keys and len(prefix_keys) > 0:
                prefix_keys += batch_hashes
            operation.completed_tokens += self.page_size * len(batch_hashes)

    def backup_thread_func(self):
        """
        Manage backup operations from host memory to storage backend.
        """
        while not self.storage_stop_event.is_set():
            try:
                operation = self.backup_queue.get(block=True, timeout=1)
                if operation is None:
                    continue

                if not self.backup_skip:
                    self._page_backup(operation)
                self.ack_backup_queue.put(operation)

            except Empty:
                continue