chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity
Files
4e9b4d05c09ae3aa4f52ab56e3339b74892e9bf8
sglang/python/sglang/srt/mem_cache/hiradix_cache.py
leavelet 4e9b4d05c0 CP HiCache: kill the O(victims*candidates) GPU-sync storm in the owner-lane eviction planner 
A CP4 prefill server hung 79s in _plan_cp_load_back_owner_lane_evictions (the [HiCache-load]
slow-scan) and was killed by the detokenizer health-check. Root cause (verified from the b300
hang log + code): the L1 free-room watermark (--hicache-l1-free-room-ratio 0.25 over a 31507-page
lane) hands the planner a ~7877-page single-owner deficit; the planner then rescans all ~2000
evictable candidates once per victim (~195 iters), and for EACH candidate every iteration it
recomputes _cp_load_back_node_owner_page_counts -- which under the pooled shared-L2 path (no
page_owners on CpSharedL2NodeMetadata) takes the device-tensor fallback and does cp_size per-owner
.item() device->host syncs. That is ~195 * 2000 * 4 ~= 1.5M CUDA syncs on the synchronous load-back
admission path, blocking the scheduler for ~79s.

Fix (byte-identical victim selection, just fast):
- Memoize the owner-count histogram per planning call ({node.id: counts}); the counts are invariant
  while node.value is fixed (the plan does not mutate values), so node.id is a safe key for the plan's
  duration. Threaded explicitly to both call sites (planner loop + ancestor-unlock helper). Turns
  O(victims*candidates) recomputes into O(distinct nodes).
- Replace the cp_size per-owner sum().item() loop with one bincount().tolist() device->host sync.

Net: ~79s -> ~1s; the eviction plan (victims, planned_freed) is unchanged. bincount == the per-owner
loop proven over 8000 random vectors incl. the (-1)%cp==cp-1 zero-loc edge. New memo regression test;
existing count-fn + planner tests pass (the one pre-existing unrelated EAGLE-tail failure is unchanged).
(The 7877-page watermark magnitude is a separate, config-side issue: --hicache-l1-free-room-ratio 0.25
reserves ~25% of a ~2M-token lane -- ~30x more than a 64K chunk needs; lower it.)

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-24 01:13:59 +00:00

6371 lines
279 KiB
Python
Raw Blame History

from __future__ import annotations
import atexit
import collections
import heapq
import json
import logging
import math
import os
import threading
import time
from dataclasses import dataclass, field
from queue import Empty
from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
import torch
from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import (
HiCacheController,
HiCacheWriteFailure,
HiCacheWriteReservation,
PrefetchOperation,
)
from sglang.srt.mem_cache.allocator import (
CPSharedPagedTokenToKVPoolAllocator,
compute_owner_lane_free_room_deficits,
)
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
InitLoadBackParams,
InsertParams,
InsertResult,
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import (
CpSharedKVLayout,
pad_token_locs_to_page_boundary,
)
from sglang.srt.mem_cache.hicache_storage import get_hash_str
from sglang.srt.mem_cache.memory_pool import (
MHATokenToKVPool,
MLATokenToKVPool,
NSATokenToKVPool,
)
from sglang.srt.mem_cache.radix_cache import (
RadixCache,
RadixKey,
TreeNode,
ceil_to_page_len,
compute_node_hash_values,
floor_to_page_len,
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__)
def _estimate_nsa_index_size_per_token(kv_cache) -> int:
if not isinstance(kv_cache, NSATokenToKVPool):
raise ValueError(
f"NSA index size is only defined for NSATokenToKVPool; "
f"got {type(kv_cache).__name__}"
)
indexer_size_per_token = (
kv_cache.index_head_dim
+ kv_cache.index_head_dim // kv_cache.quant_block_size * 4
)
index_layer_count = len(
getattr(
kv_cache,
"index_active_layer_ids",
range(kv_cache.layer_num),
)
)
return (
indexer_size_per_token
* index_layer_count
* NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize
)
def _estimate_hicache_kv_size_per_token(kv_cache) -> int:
dtype_size = kv_cache.store_dtype.itemsize
if isinstance(kv_cache, MHATokenToKVPool):
return kv_cache.head_dim * kv_cache.head_num * kv_cache.layer_num * dtype_size * 2
if isinstance(kv_cache, MLATokenToKVPool):
kv_cache_dim = getattr(
kv_cache,
"kv_cache_dim",
kv_cache.kv_lora_rank + kv_cache.qk_rope_head_dim,
)
return kv_cache_dim * dtype_size * kv_cache.layer_num
raise ValueError(f"Unsupported HiCache KV pool type: {type(kv_cache).__name__}")
def _estimate_hicache_size_per_token(kv_cache) -> int:
size_per_token = _estimate_hicache_kv_size_per_token(kv_cache)
if isinstance(kv_cache, NSATokenToKVPool):
size_per_token += _estimate_nsa_index_size_per_token(kv_cache)
return size_per_token
def _get_cp_shared_l2_rank_and_group(token_to_kv_pool_allocator):
cp_rank = getattr(token_to_kv_pool_allocator, "cp_rank", None)
cp_size = getattr(token_to_kv_pool_allocator, "cp_size", None)
if cp_rank is None or cp_size is None:
raise ValueError(
"[CP_SHARED_L2_FAILFAST][missing_cp_layout] "
"CP shared physical L2 HiCache requires an allocator with "
"explicit cp_rank/cp_size; refusing rank-local host allocation."
)
cp_group = None
cp_cpu_group = None
if int(cp_size) > 1:
try:
from sglang.srt.layers.dp_attention import get_attention_cp_group
cp_group = get_attention_cp_group()
except Exception as exc:
raise ValueError(
"[CP_SHARED_L2_FAILFAST][missing_cp_group] "
"CP shared physical L2 HiCache requires an attention CP "
"group so rank 0 can broadcast one shared slab handle."
) from exc
if cp_group is None:
raise ValueError(
"[CP_SHARED_L2_FAILFAST][missing_cp_group] "
"CP shared physical L2 HiCache requires a non-null attention "
"CP group for multi-rank CP."
)
cp_cpu_group = getattr(cp_group, "cpu_group", cp_group)
if cp_cpu_group is None:
raise ValueError(
"[CP_SHARED_L2_FAILFAST][missing_cp_group] "
"CP shared physical L2 HiCache requires a usable CP CPU "
"group for multi-rank CP."
)
return int(cp_rank), int(cp_size), cp_cpu_group
def _make_cp_shared_l2_host_tensor_allocator(
*,
server_args,
token_to_kv_pool_allocator,
payload_kind: str,
slab_name_suffix: str,
slabs_by_payload=None,
validate_production: bool = True,
):
from sglang.srt.mem_cache.memory_pool_host import (
SharedHostTensorAllocator,
SharedHostTensorGroupAllocator,
)
cp_rank, _cp_size, cp_cpu_group = _get_cp_shared_l2_rank_and_group(
token_to_kv_pool_allocator
)
slab_name = f"sglang-cp-shared-l2-{payload_kind}-{slab_name_suffix}"
slabs = tuple((slabs_by_payload or {}).get(payload_kind, ()))
common_kwargs = dict(
directory=server_args.cp_shared_l2_hugetlbfs_dir,
name=slab_name,
creator_rank=0,
validate_production=validate_production,
rank=cp_rank,
source_rank=0,
cp_cpu_group=cp_cpu_group,
)
if len(slabs) > 1:
return SharedHostTensorGroupAllocator(slabs=slabs, **common_kwargs)
return SharedHostTensorAllocator(**common_kwargs)
def _cp_shared_l2_host_slab_suffix(
kv_cache, *, page_size: int, draft: bool = False
) -> str:
return (
f"{'draft' if draft else 'target'}-"
f"layers{getattr(kv_cache, 'start_layer', 0)}-"
f"{getattr(kv_cache, 'end_layer', 0)}-"
f"ps{int(page_size)}"
)
def _compute_shared_hicache_token_capacities(
*,
total_host_bytes: int,
target_size_per_token: int,
draft_size_per_token: int,
page_size: int,
) -> Tuple[int, int]:
if total_host_bytes <= 0:
raise ValueError(f"total_host_bytes must be positive, got {total_host_bytes}")
if target_size_per_token <= 0:
raise ValueError(
f"target_size_per_token must be positive, got {target_size_per_token}"
)
if draft_size_per_token <= 0:
raise ValueError(
f"draft_size_per_token must be positive, got {draft_size_per_token}"
)
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
target_pages = total_host_bytes // (
(target_size_per_token + draft_size_per_token) * page_size
)
if target_pages <= 0:
raise ValueError(
"HiCache shared target+draft budget cannot hold a single page: "
f"total_host_bytes={total_host_bytes}, "
f"target_size_per_token={target_size_per_token}, "
f"draft_size_per_token={draft_size_per_token}, page_size={page_size}"
)
target_tokens = int(target_pages * page_size)
remaining_bytes = total_host_bytes - target_tokens * target_size_per_token
draft_pages = remaining_bytes // (draft_size_per_token * page_size)
draft_tokens = int(draft_pages * page_size)
if draft_tokens < target_tokens:
raise ValueError(
"HiCache shared budget computation produced draft capacity below target "
f"capacity: target_tokens={target_tokens}, draft_tokens={draft_tokens}"
)
return target_tokens, draft_tokens
def _cp_shared_l2_bytes_per_page_by_payload(
*,
kv_cache,
page_size: int,
draft_token_to_kv_pool=None,
) -> Dict[str, int]:
from sglang.srt.mem_cache.cp_shared_l2_pool import (
PAYLOAD_DRAFT_KV,
PAYLOAD_INDEX_K,
PAYLOAD_TARGET_KV,
)
page_size = int(page_size)
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
bytes_per_page = {
PAYLOAD_TARGET_KV: _estimate_hicache_kv_size_per_token(kv_cache) * page_size
}
if draft_token_to_kv_pool is not None:
bytes_per_page[PAYLOAD_DRAFT_KV] = (
_estimate_hicache_kv_size_per_token(draft_token_to_kv_pool) * page_size
)
if isinstance(kv_cache, NSATokenToKVPool):
bytes_per_page[PAYLOAD_INDEX_K] = (
_estimate_nsa_index_size_per_token(kv_cache) * page_size
)
return bytes_per_page
def _cp_shared_l2_slab_pages_by_payload(
*,
server_args,
page_size: int,
bytes_per_page_by_payload: Dict[str, int],
) -> Dict[str, int]:
from sglang.srt.mem_cache.memory_pool_host import (
cp_hicache_max_single_register_bytes,
)
ceiling = cp_hicache_max_single_register_bytes()
slab_size_gb = int(getattr(server_args, "cp_shared_l2_slab_size_gb", 0) or 0)
if slab_size_gb <= 0:
# AUTO: split each payload into slabs no larger than ONE cudaHostRegister
# call can pin. A single >~1 TiB registration OOMs, and a registration
# cannot be chunked (a transfer memcpy cannot cross a registration
# boundary), so large host caches must be physically multi-slab. Payloads
# that already fit in one slab stay single (build_cp_shared_l2_slabs_by_payload
# collapses slab_pages >= total_pages to one slab).
slab_size_bytes = ceiling
else:
slab_size_bytes = slab_size_gb * 1000 * 1000 * 1000
if slab_size_bytes > ceiling:
logger.warning(
"cp_shared_l2_slab_size_gb=%d GB exceeds the safe single-"
"cudaHostRegister ceiling (%.0f GiB); capping to the ceiling.",
slab_size_gb,
ceiling / (1024**3),
)
slab_size_bytes = ceiling
slab_pages: Dict[str, int] = {}
for payload_kind, bytes_per_page in bytes_per_page_by_payload.items():
bytes_per_page = int(bytes_per_page)
if bytes_per_page <= 0:
raise ValueError(f"bytes_per_page for {payload_kind!r} must be positive")
if slab_size_bytes < bytes_per_page:
suggested_gb = math.ceil(bytes_per_page / (1024**3))
raise ValueError(
"CP shared-L2 single-registration slab budget is smaller than one "
f"logical page for payload {payload_kind!r}: slab_bytes={slab_size_bytes} "
f"bytes_per_page={bytes_per_page}. One logical page must fit in one "
f"cudaHostRegister; raise SGLANG_CP_HICACHE_MAX_SLAB_GB to at least "
f"{suggested_gb}."
)
slab_pages[payload_kind] = slab_size_bytes // bytes_per_page
return slab_pages
def _cp_shared_l2_pages_for_kv_cache_host(
*,
kv_cache,
server_args,
page_size: int,
host_size: int,
host_token_capacity: Optional[int],
) -> int:
if host_token_capacity is not None:
tokens = int(host_token_capacity)
return max(1, (tokens + int(page_size) - 1) // int(page_size))
size_per_token = _estimate_hicache_size_per_token(kv_cache)
if int(host_size) > 0:
tokens = int(int(host_size) * 1e9 // size_per_token)
else:
tokens = int(kv_cache.size * server_args.hicache_ratio)
return max(1, tokens // int(page_size) + 1)
def _cp_shared_l2_pages_per_payload(
*,
token_to_kv_pool_host,
page_size: int,
draft_token_to_kv_pool=None,
draft_host_token_capacity: Optional[int] = None,
) -> Dict[str, int]:
from sglang.srt.mem_cache.cp_shared_l2_pool import (
PAYLOAD_DRAFT_KV,
PAYLOAD_INDEX_K,
PAYLOAD_TARGET_KV,
)
target_pages = max(
1, int(getattr(token_to_kv_pool_host, "size", 0)) // int(page_size)
)
pages_per_payload = {PAYLOAD_TARGET_KV: target_pages}
if draft_token_to_kv_pool is not None:
draft_tokens = (
int(draft_host_token_capacity)
if draft_host_token_capacity is not None
else int(getattr(token_to_kv_pool_host, "size", 0))
)
pages_per_payload[PAYLOAD_DRAFT_KV] = max(1, draft_tokens // int(page_size))
if int(getattr(token_to_kv_pool_host, "index_active_layer_num", 0) or 0) > 0:
pages_per_payload[PAYLOAD_INDEX_K] = max(
1,
int(getattr(token_to_kv_pool_host, "indexer_page_num", target_pages)),
)
return pages_per_payload
def _page_aligned_room_from_ratio(capacity: int, ratio: float, page_size: int) -> int:
if page_size <= 0:
raise ValueError(f"page_size must be positive, got {page_size}")
if capacity <= 0 or ratio <= 0:
return 0
raw_room = int(math.ceil(float(capacity) * float(ratio)))
return ((raw_room + page_size - 1) // page_size) * page_size
def _free_room_deficit(
*,
required: int,
available: int,
capacity: int,
page_size: int,
target_ratio: float,
trigger_ratio: float,
) -> int:
target_room = _page_aligned_room_from_ratio(capacity, target_ratio, page_size)
trigger_room = _page_aligned_room_from_ratio(capacity, trigger_ratio, page_size)
if int(available) >= int(required) + trigger_room:
return 0
return max(max(0, int(required) - int(available)), target_room)
@dataclass
class CpHiCacheNodeMetadata:
# Legacy name kept for existing call sites. Under the page-aligned cache
# contract this is the radix/scheduler-visible valid token length.
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
# Physical host/device span. When omitted, legacy callers keep the old
# page-aligned invariant by using logical_len as the physical length.
padded_len: Optional[int] = None
draft_host_indices: Optional[torch.Tensor] = None
_owner_page_counts_cache: Dict[int, Tuple[int, ...]] = field(
default_factory=dict, init=False, repr=False
)
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.padded_len is None:
self.padded_len = self.logical_len
self.padded_len = int(self.padded_len)
if self.padded_len < 0:
raise ValueError(f"padded_len must be non-negative, got {self.padded_len}")
if self.padded_len < self.logical_len:
raise ValueError(
f"padded_len ({self.padded_len}) must be >= logical_len "
f"({self.logical_len})"
)
if self.padded_len % self.page_size != 0:
raise ValueError(
f"padded_len ({self.padded_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.padded_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"padded_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.padded_len
):
raise ValueError(
"owned_positions must be in [0, padded_len) "
"(legacy [0, logical_len) for unpadded metadata)"
)
if torch.any(self.owned_positions[1:] <= self.owned_positions[:-1]):
raise ValueError(
"owned_positions must be sorted and strictly increasing"
)
def owner_page_counts(self, cp_size: int) -> Tuple[int, ...]:
cp_size = int(cp_size)
if cp_size <= 0:
raise ValueError(f"cp_size must be positive, got {cp_size}")
cached = self._owner_page_counts_cache.get(cp_size)
if cached is not None:
return cached
counts = [0 for _ in range(cp_size)]
for owner in self.page_owners.tolist():
owner = int(owner)
if owner < 0:
continue
if owner >= cp_size:
raise RuntimeError(
"CP HiCache metadata owner exceeds CP size: "
f"owner={owner} cp_size={cp_size}"
)
counts[owner] += 1
result = tuple(counts)
self._owner_page_counts_cache[cp_size] = result
return result
@property
def valid_len(self) -> int:
return self.logical_len
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,
padded_len=split_len,
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,
padded_len=self.padded_len - split_len,
draft_host_indices=(
self.draft_host_indices[child_mask]
if self.draft_host_indices is not None
else None
),
),
)
@dataclass
class PendingHiCacheBackup:
node: TreeNode
metadata: CpHiCacheNodeMetadata
logical_len: int
submitted: bool = False
local_done: bool = False
locked: bool = True
@dataclass
class PreparedCpHiCacheBackup:
node_id: int
reservation: HiCacheWriteReservation
metadata: CpHiCacheNodeMetadata
logical_len: int
attached: bool = False
@dataclass(frozen=True)
class CpWriteBackupCandidate:
req: object
kv_indices: torch.Tensor
node_id: int
@dataclass(frozen=True)
class CpHiCacheCapacitySnapshot:
target_capacity: Tuple[int, ...]
draft_capacity: Optional[Tuple[int, ...]]
committed_target: Tuple[int, ...]
committed_draft: Tuple[int, ...]
pending_target: Tuple[int, ...]
pending_draft: Tuple[int, ...]
@property
def target_used(self) -> Tuple[int, ...]:
return tuple(
committed + pending
for committed, pending in zip(self.committed_target, self.pending_target)
)
@property
def draft_used(self) -> Tuple[int, ...]:
return tuple(
committed + pending
for committed, pending in zip(self.committed_draft, self.pending_draft)
)
class _CpRunningHostCounts:
"""Running per-owner host TOKEN counts for the CP-HiCache capacity snapshot.
Equals, by construction, the full radix-walk result -- committed nodes and
pending backups, summed per CP owner lane, gated on draft presence -- but is
maintained incrementally at each state transition so the snapshot is
O(cp_size) rather than O(total-cached-pages) per write-admission. Created
EMPTY (production starts with an empty tree) and then maintained by the
accounting hooks; ``_cp_walk_capacity_counts`` is the reference the env-gated
drift verifier compares against, NOT a lazy-build source.
"""
__slots__ = (
"cp_size",
"committed_target",
"committed_draft",
"pending_target",
"pending_draft",
)
def __init__(self, cp_size: int):
self.cp_size = int(cp_size)
self.committed_target = [0] * self.cp_size
self.committed_draft = [0] * self.cp_size
self.pending_target = [0] * self.cp_size
self.pending_draft = [0] * self.cp_size
def set_from_tuples(
self, committed_target, committed_draft, pending_target, pending_draft
) -> None:
self.committed_target = list(committed_target)
self.committed_draft = list(committed_draft)
self.pending_target = list(pending_target)
self.pending_draft = list(pending_draft)
@staticmethod
def _token_counts(metadata, cp_size: int) -> List[int]:
# owner_page_counts caches on the metadata object (and survives commit --
# pending.metadata IS node.cp_hicache -- and split, which builds fresh
# objects with fresh caches); tokens == pages * page_size.
page_counts = metadata.owner_page_counts(cp_size)
if len(page_counts) != cp_size:
raise RuntimeError(
"CP HiCache running-count cp_size mismatch: "
f"metadata={len(page_counts)} expected={cp_size}"
)
page_size = int(metadata.page_size)
return [int(c) * page_size for c in page_counts]
def _add(self, target, draft, metadata) -> None:
counts = self._token_counts(metadata, self.cp_size)
for i, c in enumerate(counts):
target[i] += c
if getattr(metadata, "draft_host_indices", None) is not None:
for i, c in enumerate(counts):
draft[i] += c
def _sub(self, target, draft, metadata) -> None:
counts = self._token_counts(metadata, self.cp_size)
has_draft = getattr(metadata, "draft_host_indices", None) is not None
# Underflow is a hard error, not a floor: the full walk can never go
# negative, so a negative here means a missed add or a double sub.
for i, c in enumerate(counts):
new_t = target[i] - c
if new_t < 0:
raise RuntimeError(
f"CP HiCache running token underflow at lane {i}: "
f"{target[i]} - {c} < 0"
)
target[i] = new_t
if has_draft:
new_d = draft[i] - c
if new_d < 0:
raise RuntimeError(
f"CP HiCache running draft underflow at lane {i}: "
f"{draft[i]} - {c} < 0"
)
draft[i] = new_d
def enter_committed(self, metadata) -> None:
self._add(self.committed_target, self.committed_draft, metadata)
def leave_committed(self, metadata) -> None:
self._sub(self.committed_target, self.committed_draft, metadata)
def enter_pending(self, metadata) -> None:
self._add(self.pending_target, self.pending_draft, metadata)
def leave_pending(self, metadata) -> None:
self._sub(self.pending_target, self.pending_draft, metadata)
def reset(self) -> None:
n = self.cp_size
self.committed_target = [0] * n
self.committed_draft = [0] * n
self.pending_target = [0] * n
self.pending_draft = [0] * n
def as_tuples(self):
return (
tuple(self.committed_target),
tuple(self.committed_draft),
tuple(self.pending_target),
tuple(self.pending_draft),
)
def sanity_check(self) -> None:
"""Cheap always-on invariants (O(cp_size)): non-negativity and
draft<=target lane-wise (draft counts are a subset of target counts for
the same metadata). Does NOT bound per-lane sums against pool .size --
availability floors at zero (`_cp_sub_counts_floor_zero`) so used may
legitimately exceed capacity. The env-gated full-walk verifier is the
real drift net."""
for name, vec in (
("committed_target", self.committed_target),
("committed_draft", self.committed_draft),
("pending_target", self.pending_target),
("pending_draft", self.pending_draft),
):
for i, v in enumerate(vec):
if v < 0:
raise RuntimeError(
f"CP HiCache running {name} lane {i} negative: {v}"
)
for i in range(self.cp_size):
if (
self.committed_draft[i] > self.committed_target[i]
or self.pending_draft[i] > self.pending_target[i]
):
raise RuntimeError(
f"CP HiCache running draft>target lane {i}: "
f"cd={self.committed_draft[i]} ct={self.committed_target[i]} "
f"pd={self.pending_draft[i]} pt={self.pending_target[i]}"
)
@dataclass(frozen=True)
class CpHiCacheEvictionPlan:
victims: Tuple[TreeNode, ...]
planned_freed: Tuple[int, ...]
remaining_deficit: Tuple[int, ...]
@dataclass(frozen=True)
class CpLoadBackPlan:
page_owners: List[int]
required_by_owner: List[int]
available_by_owner: List[int]
# Exact capacity deficit. This is the only deficit allowed to block the
# synchronous load-back admission path.
deficit_by_owner: List[int]
# Advisory free-room deficit. This is logged for observability but must not
# force synchronous load-back eviction when exact capacity already fits.
free_room_deficit_by_owner: List[int]
host_hit_len: int
@dataclass(frozen=True)
class CpLoadBackEvictionPlan:
victims: Tuple[TreeNode, ...]
planned_freed_by_owner: Tuple[int, ...]
remaining_deficit_by_owner: Tuple[int, ...]
@dataclass(frozen=True)
class CpWriteAdmission:
node_id: int
phase: str
required_by_owner: Tuple[int, ...]
target_available_by_owner: Tuple[int, ...]
draft_available_by_owner: Tuple[int, ...]
deficit_by_owner: Tuple[int, ...]
eviction_plan: CpHiCacheEvictionPlan
class HiCachePendingBackupSplit(Exception):
def __init__(self, node: TreeNode):
self.node = node
super().__init__(
f"Cannot split node_id={getattr(node, 'id', None)} while CP HiCache backup is pending"
)
class HiRadixCache(RadixCache):
def _create_token_to_kv_pool_host(
self,
kv_cache,
server_args: ServerArgs,
*,
host_size: Optional[int] = None,
host_token_capacity: Optional[int] = None,
shared_l2_payload_kind: Optional[str] = None,
slabs_by_payload=None,
):
from sglang.srt.mem_cache.cp_shared_l2_pool import (
PAYLOAD_DRAFT_KV,
PAYLOAD_INDEX_K,
PAYLOAD_TARGET_KV,
)
from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
)
host_size = server_args.hicache_size if host_size is None else host_size
host_tensor_allocator = None
index_host_tensor_allocator = None
if getattr(server_args, "enable_cp_shared_physical_l2_hicache", False):
payload_kind = shared_l2_payload_kind or PAYLOAD_TARGET_KV
suffix = _cp_shared_l2_host_slab_suffix(
kv_cache,
page_size=self.page_size,
draft=(payload_kind == PAYLOAD_DRAFT_KV),
)
host_tensor_allocator = _make_cp_shared_l2_host_tensor_allocator(
server_args=server_args,
token_to_kv_pool_allocator=self._token_to_kv_pool_allocator,
payload_kind=payload_kind,
slab_name_suffix=suffix,
slabs_by_payload=slabs_by_payload,
)
if isinstance(kv_cache, NSATokenToKVPool):
index_host_tensor_allocator = _make_cp_shared_l2_host_tensor_allocator(
server_args=server_args,
token_to_kv_pool_allocator=self._token_to_kv_pool_allocator,
payload_kind=PAYLOAD_INDEX_K,
slab_name_suffix=suffix,
slabs_by_payload=slabs_by_payload,
)
logger.info(
"[CP_SHARED_L2] constructing shared physical host L2 pool: "
"hicache_size_gb=%s is_node_local_total=true payload=%s index_payload=%s slabs=%s numa_policy=%s",
server_args.hicache_size,
payload_kind,
index_host_tensor_allocator is not None,
{key: len(value) for key, value in (slabs_by_payload or {}).items()},
getattr(server_args, "cp_shared_l2_numa_policy", "interleave_2m"),
)
if isinstance(kv_cache, MHATokenToKVPool):
return MHATokenToKVPoolHost(
kv_cache,
server_args.hicache_ratio,
host_size,
self.page_size,
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
host_token_capacity=host_token_capacity,
host_tensor_allocator=host_tensor_allocator,
)
if isinstance(kv_cache, NSATokenToKVPool):
return NSATokenToKVPoolHost(
kv_cache,
server_args.hicache_ratio,
host_size,
self.page_size,
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
host_token_capacity=host_token_capacity,
host_tensor_allocator=host_tensor_allocator,
index_host_tensor_allocator=index_host_tensor_allocator,
)
if isinstance(kv_cache, MLATokenToKVPool):
return MLATokenToKVPoolHost(
kv_cache,
server_args.hicache_ratio,
host_size,
self.page_size,
server_args.hicache_mem_layout,
allocator_type=server_args.hicache_storage_backend,
host_token_capacity=host_token_capacity,
host_tensor_allocator=host_tensor_allocator,
)
raise ValueError("HiRadixCache only supports MHA, MLA, and NSA KV pools")
def attach_draft_kv_pool(
self, draft_token_to_kv_pool, host_token_capacity: Optional[int] = None
) -> 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,
host_size=0 if host_token_capacity is not None else None,
host_token_capacity=host_token_capacity,
shared_l2_payload_kind="draft_kv",
slabs_by_payload=getattr(self, "_cp_shared_l2_slabs_by_payload", None),
)
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_allocator = params.token_to_kv_pool_allocator
draft_token_to_kv_pool = getattr(params, "draft_token_to_kv_pool", None)
target_host_token_capacity = None
draft_host_token_capacity = None
if (
self._uses_cp_hicache
and draft_token_to_kv_pool is not None
and server_args.hicache_size > 0
):
target_size_per_token = _estimate_hicache_size_per_token(self.kv_cache)
draft_size_per_token = _estimate_hicache_size_per_token(
draft_token_to_kv_pool
)
target_host_token_capacity, draft_host_token_capacity = (
_compute_shared_hicache_token_capacities(
total_host_bytes=int(server_args.hicache_size * 1e9),
target_size_per_token=target_size_per_token,
draft_size_per_token=draft_size_per_token,
page_size=self.page_size,
)
)
logger.info(
"[HiCache-draft] sharing --hicache-size across target+draft: "
"budget_gb=%d target_size_per_token=%d draft_size_per_token=%d "
"target_tokens=%d draft_tokens=%d",
server_args.hicache_size,
target_size_per_token,
draft_size_per_token,
target_host_token_capacity,
draft_host_token_capacity,
)
cp_shared_l2_slabs_by_payload = None
cp_shared_l2_pages_per_payload_planned = None
if getattr(server_args, "enable_cp_shared_physical_l2_hicache", False):
from sglang.srt.mem_cache.cp_shared_l2_pool import (
PAYLOAD_DRAFT_KV,
PAYLOAD_INDEX_K,
PAYLOAD_TARGET_KV,
build_cp_shared_l2_slabs_by_payload,
)
target_host_size = (
0 if target_host_token_capacity is not None else server_args.hicache_size
)
target_pages = _cp_shared_l2_pages_for_kv_cache_host(
kv_cache=self.kv_cache,
server_args=server_args,
page_size=self.page_size,
host_size=target_host_size,
host_token_capacity=target_host_token_capacity,
)
pages_per_payload = {PAYLOAD_TARGET_KV: target_pages}
bytes_per_page_by_payload = _cp_shared_l2_bytes_per_page_by_payload(
kv_cache=self.kv_cache,
page_size=self.page_size,
draft_token_to_kv_pool=draft_token_to_kv_pool,
)
if draft_token_to_kv_pool is not None:
draft_pages = _cp_shared_l2_pages_for_kv_cache_host(
kv_cache=draft_token_to_kv_pool,
server_args=server_args,
page_size=self.page_size,
host_size=(
0
if draft_host_token_capacity is not None
else server_args.hicache_size
),
host_token_capacity=draft_host_token_capacity,
)
pages_per_payload[PAYLOAD_DRAFT_KV] = draft_pages
if isinstance(self.kv_cache, NSATokenToKVPool):
pages_per_payload[PAYLOAD_INDEX_K] = target_pages + 1
cp_shared_l2_pages_per_payload_planned = dict(pages_per_payload)
slab_pages_by_payload = _cp_shared_l2_slab_pages_by_payload(
server_args=server_args,
page_size=self.page_size,
bytes_per_page_by_payload=bytes_per_page_by_payload,
)
cp_shared_l2_slabs_by_payload = build_cp_shared_l2_slabs_by_payload(
pages_per_payload,
slab_pages_by_payload=slab_pages_by_payload,
numa_policy=getattr(
server_args, "cp_shared_l2_numa_policy", "interleave_2m"
),
)
logger.info(
"[CP_SHARED_L2] planned shared physical L2 slabs: policy=%s counts=%s pages=%s",
getattr(server_args, "cp_shared_l2_numa_policy", "interleave_2m"),
{
key: len(value)
for key, value in cp_shared_l2_slabs_by_payload.items()
},
pages_per_payload,
)
self._cp_shared_l2_slabs_by_payload = cp_shared_l2_slabs_by_payload
self.token_to_kv_pool_host = self._create_token_to_kv_pool_host(
self.kv_cache,
server_args,
host_size=0 if target_host_token_capacity is not None else None,
host_token_capacity=target_host_token_capacity,
slabs_by_payload=cp_shared_l2_slabs_by_payload,
)
# PREREQ-1 (CP8DP2 scoping): the B1 commit/evict consensus collectives
# (writing_check MIN, placement_digest, drain/evict MINs, the prefetch group,
# the flush barrier) MUST reduce over the CP group -- the ranks that share the
# replicated event stream + the MAP_SHARED slab. tp_cache_group equals the CP
# group ONLY at dp_size=1 (CP==TP); under CP8DP2 (DP attention) tp_cache_group
# is the size-1 attn-TP group, so every `tp_world_size > 1` collective would
# silently no-op per rank -> divergent placement -> shared-slab corruption.
# Scope to the CP cpu group (a no-op handle change at dp_size=1).
_cp_hicache_cp_size = (
int(getattr(params.token_to_kv_pool_allocator, "cp_size", 1) or 1)
if self._uses_cp_hicache
else 1
)
if self._uses_cp_hicache and _cp_hicache_cp_size > 1:
from sglang.srt.layers.dp_attention import get_attention_cp_group
_cp_group = get_attention_cp_group()
self.tp_group = getattr(_cp_group, "cpu_group", _cp_group)
else:
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
# CP L3 durable floor (disk tier): own content-addressed path, NOT enable_storage. Built on the
# pooled shared-L2 (server_args validation already enforced the prerequisites + mutual exclusion).
self.enable_cp_l3 = bool(getattr(server_args, "enable_cp_l3", False)) and self._uses_cp_hicache
self.cp_l3_store = None # constructed after the host pool + allocator (slab accessors need them)
self.ongoing_l3_spill = {} # op_id -> pinned (protect_host'd) resident TreeNode; submitted..durable-acked
# PROACTIVE spill: FIFO of (object_key, node) enqueued at the replicated commit frontier
# (_commit_pending_backup) + at radix splits, drained by the maintainer in commit order -> L3 stays
# caught up with L2. Capped: if the disk stalls, the bg write thread blocks -> ongoing pins at
# max_inflight -> budget hits 0 -> the maintainer stops popping; without a cap the deque would then grow
# unbounded with every new commit (OOM). The cap is a replicated constant (rank-uniform drop) sized well
# above the max committed-object backlog, so it only fires under a genuine disk stall (fail-soft: the
# dropped objects recompute on a miss). drop count is for a rate-limited warn.
self.cp_l3_spill_pending = collections.deque()
self.cp_l3_spill_pending_cap = 1 << 18 # 262144 entries (~tens of MB of (str,node) refs)
self._cp_l3_spill_dropped = 0
# backpressure: cap concurrent background spills (bounds the bg queue + staging RAM, rank-uniform)
self.cp_l3_spill_max_inflight = 32
# L3 RELOAD (3.2): disk->L2 + radix re-insert, holding the triggering request (mirrors the storage
# prefetch hold, which is dead under CP). All cross-rank agreement folds into the existing per-tick
# _drain_l3_control_queues MIN -- no new collective. ongoing keyed by reload op_id (like spill).
self.cp_l3_reload_candidates = [] # [(rid, last_host_node, suffix_key, suffix_hashes)] marked at entry
self.ongoing_l3_reload = {} # op_id -> dict(reload_key,last_host_node,suffix_key,suffix_hashes,n_pages,waiters)
self.cp_l3_reload_inflight_keys = {} # reload_key -> op_id (same-suffix dedup: a 2nd req piggybacks)
self.cp_l3_reload_done = {} # rid -> bool; read (no collective) by the scheduler waiting-queue skip
self.cp_l3_reload_max_inflight = 32 # cap concurrent reloads (bounds reserved-but-pending L2 pages)
# Bound the request hold (rank-uniform tick countdown): the default SGLANG_REQ_WAITING_TIMEOUT is -1
# (off), so a lost/hung reload op (e.g. a dead disk) could otherwise wedge a held request forever.
self.cp_l3_reload_ttl_ticks = 4096
# CP-HiCache metrics (L2 pooled allocator + L3 store usage; the mainline storage metrics are dead under
# CP). Created in _maybe_init_cp_l3 when enable_metrics; pushed periodically from check_hicache_events.
self.cp_hicache_metrics = None
self._cp_metrics_last_emit = 0.0
(
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
cp_shared_l2_page_allocator = 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,
)
if getattr(server_args, "enable_cp_shared_physical_l2_hicache", False):
from sglang.srt.mem_cache.cp_shared_l2_pool import (
CpSharedL2PageAllocator,
)
if cp_shared_l2_pages_per_payload_planned is not None:
pages_per_payload = dict(cp_shared_l2_pages_per_payload_planned)
else:
pages_per_payload = _cp_shared_l2_pages_per_payload(
token_to_kv_pool_host=self.token_to_kv_pool_host,
page_size=self.page_size,
draft_token_to_kv_pool=draft_token_to_kv_pool,
draft_host_token_capacity=draft_host_token_capacity,
)
# B1 (Tier-2): the allocator carries only placement state; the
# option-A ctor args (expected_ranks/expected_layers/required_payloads
# feeding the dead commit quorum) are gone.
cp_shared_l2_page_allocator = CpSharedL2PageAllocator(
pages_per_payload=pages_per_payload,
slabs_by_payload=cp_shared_l2_slabs_by_payload,
)
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,
cp_shared_l2_page_allocator=cp_shared_l2_page_allocator,
)
if draft_token_to_kv_pool is not None:
self.attach_draft_kv_pool(
draft_token_to_kv_pool,
host_token_capacity=draft_host_token_capacity,
)
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 CP host reservations/backups that are not request-visible yet.
self.pending_host_backups: Dict[int, PendingHiCacheBackup] = {}
# CP-HiCache running host-capacity counts are lazily built from the tree
# on first access (see the _cp_counts property) and then maintained
# incrementally at every state transition, so _cp_host_capacity_snapshot
# is O(cp_size) rather than a full O(total-cached-pages) radix walk per
# write-admission.
# record the node segments with ongoing load back
self.ongoing_load_back = {}
# record the ongoing prefetch requests
self.ongoing_prefetch = {}
self.ongoing_backup = {}
# Temporary aggregate accounting for CP HiCache collectives. These
# calls are on latency-sensitive scheduler paths; keep logging coarse
# enough to avoid per-request spam while still exposing hot collectives.
self._cp_hicache_collective_stats = {}
# 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
self.hicache_host_free_room_ratio = server_args.hicache_host_free_room_ratio
self.hicache_host_free_room_trigger_ratio = (
server_args.hicache_host_free_room_trigger_ratio
)
self.hicache_l1_free_room_ratio = server_args.hicache_l1_free_room_ratio
self.hicache_l1_free_room_trigger_ratio = (
server_args.hicache_l1_free_room_trigger_ratio
)
# 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)
if self._uses_cp_hicache:
# CP shared-KV eviction must be COLLECTIVE-FREE: every rank must pick
# identical victims from rank-replicated state, else partial backups /
# all-reduce-shape deadlock. The base LRU/SLRU strategies tiebreak on
# per-rank wall-clock last_access_time and diverge. Override with the
# replicated-clock SLRU value (is_protected, last_access_time[logical],
# node.id) — a strict total order over replicated state. Reaches every
# eviction key site that uses eviction_strategy.get_priority (device
# owner-lane score, host physical-slot heap, generic evict); the host
# write-admission key _cp_host_evict_key is switched to it explicitly.
from sglang.srt.mem_cache.evict_policy import CpReplicatedSLRUStrategy
self.eviction_strategy = CpReplicatedSLRUStrategy()
self._maybe_init_cp_l3(server_args, params)
def _maybe_init_cp_l3(self, server_args, params) -> None:
"""Construct the CP L3 disk store from the live host pool slabs (Model B: content-addressed,
rank-local I/O + the shared LMDB). Default-off; only when --enable-cp-l3."""
if not self.enable_cp_l3:
return
from sglang.srt.mem_cache.cp_l3_config import CpL3Config
from sglang.srt.mem_cache.cp_l3_slab_accessor import (
CpL3SlabLayout,
CpSharedL2SlabAccessor,
)
from sglang.srt.mem_cache.cp_l3_store import CpL3Store
from sglang.srt.mem_cache.memory_pool_host import NSATokenToKVPoolHost
cp_allocator = params.token_to_kv_pool_allocator
cp_rank, cp_size = int(cp_allocator.cp_rank), int(cp_allocator.cp_size)
cfg = CpL3Config.from_file(server_args.cp_l3_config)
def _build_accessor(allocator, layout, payload_kind):
# One accessor over N physical slabs (N==1 for a single slab); the
# layout supplies n_layers + slice_bytes (slab-count invariant), the
# spans supply each slab's global page range + its own mmap.
return CpSharedL2SlabAccessor(
self._cp_l3_slab_spans(allocator, layout.page_num, payload_kind),
n_layers=layout.n_layers,
slice_bytes=layout.slice_bytes,
)
host = self.token_to_kv_pool_host
accessors = {
"target_kv": _build_accessor(
host.allocator,
CpL3SlabLayout.for_mla(
layer_num=host.layer_num, page_num=host.page_num,
page_size=self.page_size, kv_cache_dim=host.kv_cache_dim,
itemsize=host.dtype.itemsize,
),
"target_kv",
)
}
draft = getattr(self, "draft_token_to_kv_pool_host", None)
if draft is not None:
accessors["draft_kv"] = _build_accessor(
draft.allocator,
CpL3SlabLayout.for_mla(
layer_num=draft.layer_num, page_num=draft.page_num,
page_size=self.page_size, kv_cache_dim=draft.kv_cache_dim,
itemsize=draft.dtype.itemsize,
),
"draft_kv",
)
if isinstance(host, NSATokenToKVPoolHost):
accessors["index_k"] = _build_accessor(
host.index_host_tensor_allocator,
CpL3SlabLayout.for_index(
n_active_layers=host.index_active_layer_num,
indexer_page_num=host.indexer_page_num,
indexer_page_stride_size=host.indexer_page_stride_size,
),
"index_k",
)
self.cp_l3_store = CpL3Store.from_config(
cfg, cp_rank=cp_rank, cp_size=cp_size, accessors=accessors
)
cold_start = envs.SGLANG_CP_L3_COLD_START.get()
self.cp_l3_store.connect(cfg, cold_start=cold_start)
logger.info(
"[CP-L3] enabled rank=%d/%d payloads=%s disk_dir=%s cold_start=%s",
cp_rank, cp_size, list(accessors), self.cp_l3_store.disk_dir, cold_start,
)
# (CP-HiCache metrics are created lazily in _cp_maybe_collect_metrics, gated on L2 pooling -- so they
# cover PURE-L2-without-L3 too, not just this L3 path.)
@staticmethod
def _cp_l3_slab_spans(allocator, total_pages: int, payload_kind: str):
"""Per-slab page ranges + mmaps for an L3 accessor.
Handles both the single-slab allocator (one span over the whole payload)
and the multi-slab group allocator (one span per physical slab, carrying
each slab's global_base_page/num_pages from its slab_info). L3 addresses
host pages by *global* page id, so a multi-slab payload must dispatch each
page to its owning slab -- which these spans + CpSharedL2SlabAccessor do.
"""
from sglang.srt.mem_cache.cp_l3_slab_accessor import CpL3SlabSpan
from sglang.srt.mem_cache.memory_pool_host import (
SharedHostTensorGroupAllocator,
)
if isinstance(allocator, SharedHostTensorGroupAllocator):
spans = []
for slab_info, sub in zip(allocator.slabs, allocator.allocators):
mapping = getattr(sub, "mapping", None)
if mapping is None or getattr(mapping, "mmap", None) is None:
raise RuntimeError(
f"[CP_L3_FAILFAST] {payload_kind}: group slab "
f"{int(slab_info.slab_id)} has no live mapping for L3 access."
)
spans.append(
CpL3SlabSpan(
int(slab_info.global_base_page),
int(slab_info.num_pages),
mapping.mmap,
)
)
return spans
mapping = getattr(allocator, "mapping", None)
if mapping is None or getattr(mapping, "mmap", None) is None:
raise RuntimeError(
f"[CP_L3_FAILFAST] {payload_kind}: shared-L2 slab has no live "
f"mapping for L3 access."
)
return [CpL3SlabSpan(0, int(total_pages), mapping.mmap)]
def _drain_l3_control_queues(self) -> None:
"""Per-tick CP-cpu-group MIN over the three L3 ack-queue sizes [gather, durable, reload], then drain
MIN of each: gather-ack -> release the eviction pin (phase 1), durable-ack -> mark l3_durable under an
ok-AND (phase 2), reload-ack -> 3.2. Idle early-return: has_inflight() is a replicated quantity (submits
are lockstep on the replicated plan), so an idle tick issues no collective."""
store = self.cp_l3_store
if store is None:
return
marked = self.cp_l3_reload_candidates
if not store.has_inflight() and not marked:
return # idle (replicated quantity) -> no collective
# Re-validate marked reload candidates against the REPLICATED radix -> rank-uniform valid set + order;
# each contributes its LOCAL exists_prefix count (skew-prone value, reconciled by the MIN below).
self.cp_l3_reload_candidates = []
payloads = store.payloads
valid_cands = [] # (rid, last_host_node, suffix_key, suffix_hashes)
cand_counts = []
for (rid, lhn, skey, shashes) in marked:
if rid in self.cp_l3_reload_done or not self._cp_l3_reload_attach_ok(lhn, skey):
continue # already in flight, or lhn evicted / suffix no longer a clean miss -> recompute
cand_counts.append(int(store.exists_prefix(shashes, payloads)))
valid_cands.append((rid, lhn, skey, shashes))
# (1) ONE MIN reduce: the three ack qsizes + the per-candidate counts (rank-uniform vector shape).
vec = torch.tensor(
[store.ack_gather_qsize(), store.ack_durable_qsize(), store.ack_reload_qsize()] + cand_counts,
dtype=torch.int,
)
if self.tp_world_size > 1:
self._cp_hicache_all_reduce(
vec, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_queue_min"
)
vals = [int(v) for v in vec.tolist()]
n_gather, n_durable, n_reload = vals[0], vals[1], vals[2]
agreed_counts = vals[3:]
# Phase 1 -- GATHER done: release the eviction pin (MIN-gated -> rank-uniform). The staged copy is
# slab-independent, so the object may be evicted now; the disk write still completes off it.
for op_id in store.drain_gather_acks(n_gather):
node = self.ongoing_l3_spill.get(op_id)
if node is not None and int(getattr(node, "host_ref_counter", 0) or 0) > 0:
node.release_host()
# (2) ONE ok-AND (MIN over ok bits) covering BOTH spill-durable and reload acks: an op is marked
# durable / admitted only if EVERY rank succeeded -> rank-uniform (a per-rank disk failure -> uniform
# fail-soft, never a divergent l3_durable or divergent insert that would crash placement_digest). Acks
# are FIFO in replicated submit order, so the ok vector aligns per-op. (Spill ok=False is NOT re-enqueued
# -- it retries on natural re-commit; 3.3 removes the slab-full cause; recompute is correctness-safe.)
durable = store.drain_durable_acks(n_durable)
reload_acks = store.drain_reload_acks(n_reload)
ok_bits = [1 if ok else 0 for (_i, ok) in durable] + [1 if ok else 0 for (_i, ok) in reload_acks]
if ok_bits:
ok_t = torch.tensor(ok_bits, dtype=torch.int)
if self.tp_world_size > 1:
self._cp_hicache_all_reduce(
ok_t, op=torch.distributed.ReduceOp.MIN, group=self.tp_group, tag="l3_ok_and"
)
agreed_ok = [int(a) for a in ok_t.tolist()]
else:
agreed_ok = []
nd = len(durable)
for (op_id, _ok), a in zip(durable, agreed_ok[:nd]): # spill durable
node = self.ongoing_l3_spill.pop(op_id, None)
if node is not None and a == 1:
node.l3_durable = True
for (op_id, _ok), a in zip(reload_acks, agreed_ok[nd:]): # reload admission
self._cp_l3_admit_reload(op_id, a == 1)
# Reserve NEW reloads from the rank-uniform agreed counts (collective-free deterministic reserve).
self._cp_l3_reserve_reloads(valid_cands, agreed_counts)
# Bound the hold: release waiters of any reload op that hasn't acked within ttl ticks (rank-uniform).
self._cp_l3_expire_stale_reloads()
def cp_l3_reload_lookup(self, req, last_host_node: TreeNode, suffix_tokens) -> None:
"""Entry hook (scheduler _prefetch_kvcache, rank-uniform: requests are broadcast). Mark a radix-miss
request as an L3 reload candidate. Cheap + LOCAL: compute the suffix's FULL-page content hashes (the
same SHA chain spill wrote, radix_cache.compute_node_hash_values) and queue the candidate. NO
exists_prefix / collective / reserve here -- the rank-uniform reserve decision (a MIN over the per-rank
exists_prefix counts) happens in _drain_l3_control_queues, which runs before the waiting-queue batch
forms. A non-L3 miss simply falls out there (agreed count 0) and the request proceeds to recompute.
EAGLE/bigram (verified correct): the radix key is bigram-converted over the WHOLE request
(maybe_bigram_convert), so slicing raw fill_ids at the bigram-unit matched_len then re-converting
reproduces EXACTLY the suffix node's bigrams -- the identity convert_to_bigram_key(f[M:]) == bigrams(f)[M:]
holds because the boundary bigram (f[M-1],f[M]) belongs to the matched prefix (index M-1), not the suffix.
So get_hash_str over the bigram pages, chained from the prefix's last hash, reproduces the spilled
node.hash_value (this is the proven storage-prefetch recipe). The content hash is over token_ids only
(compute_node_hash_values ignores extra_key); the RadixKey still carries extra_key + is_bigram so the
attach/insert child-key routing matches the radix. page_size is in bigram units under eagle, so the page
chunking is unchanged."""
if not self.enable_cp_l3 or self.cp_l3_store is None or last_host_node is None:
return
if self.is_eagle:
suffix_tokens = convert_to_bigram_key(list(suffix_tokens))
page = self.page_size
n_full = len(suffix_tokens) // page # only full pages are content pages (spill wrote full pages)
if n_full <= 0:
return
parent_hash = (
last_host_node.get_last_hash_value()
if last_host_node is not self.root_node
else None
)
suffix_hashes: List[str] = []
suffix_key_tokens = [] # bigram tuples under eagle, raw ints otherwise
for i in range(n_full):
pg = list(suffix_tokens[i * page : (i + 1) * page])
h = get_hash_str(pg, prior_hash=parent_hash) # get_hash_str hashes bigram tuples too (compute_node_hash_values)
suffix_hashes.append(h)
parent_hash = h
suffix_key_tokens.extend(pg)
suffix_key = RadixKey(
token_ids=suffix_key_tokens,
extra_key=getattr(req, "extra_key", None),
is_bigram=self.is_eagle,
)
self.cp_l3_reload_candidates.append(
(req.rid, last_host_node, suffix_key, suffix_hashes)
)
def check_cp_l3_reload_progress(self, rid: str) -> bool:
"""Scheduler waiting-queue skip (mirrors check_prefetch_progress). True = the request may proceed
(not an L3 reload candidate, or its reload is admitted); False = an L3 reload is in flight -> hold.
Reads the rank-uniformly-set done flag -- no collective. On a True (admitted) it consumes the flag;
the request then re-matches (init_next_round_input) and hits the inserted L2 node."""
done = self.cp_l3_reload_done.get(rid)
if done is None:
return True # not held for L3 reload
if done:
self.cp_l3_reload_done.pop(rid, None)
return True
return False
def cp_l3_release_request(self, rid: str) -> None:
"""Abort/cleanup hook (scheduler): drop a request from all L3 reload state so an aborted held request
leaves no dangling done-flag / waiter. The reservation is freed when its op acks (a then-empty-waiters
op aborts in _cp_l3_admit_reload). Rank-uniform: aborts are broadcast."""
self.cp_l3_reload_done.pop(rid, None)
for info in self.ongoing_l3_reload.values():
if rid in info["waiters"]:
info["waiters"].remove(rid)
def _cp_l3_expire_stale_reloads(self) -> None:
"""Bound the request hold (rank-uniform tick countdown, called per tick): if a reload op hasn't acked
within ttl ticks (a lost/hung op -- e.g. a dead disk; the default waiting-timeout is off), RELEASE its
waiters to recompute. The op + its reservation stay until the (eventual) ack frees them -- we must NOT
free the reserved pages while the bg thread might still scatter into them. A real disk hang wedges all
of L3 anyway; this just keeps the held request from waiting forever."""
for info in self.ongoing_l3_reload.values():
info["ttl"] -= 1
if info["ttl"] <= 0 and info["waiters"]:
for rid in info["waiters"]:
self.cp_l3_reload_done.pop(rid, None) # release -> recompute
info["waiters"] = []
def _cp_l3_reload_attach_ok(self, lhn: TreeNode, suffix_key) -> bool:
"""Rank-uniform (replicated radix): lhn is still attached to the tree AND the suffix's first page is a
clean miss under it (no existing child) -> a single new child can attach. Else the reload is dropped
(recompute). Checked at reserve AND at admission (the tree may change in between)."""
if lhn is None or len(suffix_key) == 0:
return False
if lhn is not self.root_node:
parent = getattr(lhn, "parent", None)
if parent is None:
return False
try:
pkey = self.get_child_key_fn(lhn.key)
except Exception:
return False
if getattr(parent, "children", {}).get(pkey) is not lhn:
return False # lhn detached from the tree since the candidate was marked
try:
child_key = self.get_child_key_fn(suffix_key)
except Exception:
return False
return child_key not in getattr(lhn, "children", {})
def _cp_l3_reserve_reloads(self, valid_cands, agreed_counts) -> None:
"""Collective-free reserve (rank-uniform inputs => identical mutation on every rank): for each
candidate whose AGREED reloadable page count C >= the load-back floor, reserve a NEW (uncommitted) L2
object (owner-aware, evict-to-fit) + submit the disk->slab reload + hold the request. Same-suffix dedup
(a 2nd request piggybacks on the in-flight reload). Bounded by cp_l3_reload_max_inflight."""
store = self.cp_l3_store
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
if allocator is None:
return
min_pages = max(1, self.load_back_threshold // self.page_size)
reserved_this_tick = {}
for (rid, lhn, skey, shashes), c in zip(valid_cands, agreed_counts):
c = int(c)
if c < min_pages:
continue # below the load-back floor (or 0 = not in L3) -> recompute
reload_key = f"l3rl:{shashes[c - 1]}" # rank-uniform + dedups identical (content,length) reloads
existing = self.cp_l3_reload_inflight_keys.get(reload_key)
if existing is None:
existing = reserved_this_tick.get(reload_key)
if existing is not None:
self.ongoing_l3_reload[existing]["waiters"].append(rid) # piggyback on the in-flight reload
self.cp_l3_reload_done[rid] = False
continue
if len(self.ongoing_l3_reload) >= self.cp_l3_reload_max_inflight:
continue # cap reached (len is replicated) -> recompute
ranges = self._cp_l3_reserve_reload_object(reload_key, store.payloads, c, allocator)
if ranges is None:
continue # could not fit even after evict-to-fit -> recompute
owned_pages = self._cp_l3_build_reload_owned_pages(ranges, shashes[:c], store)
op_id = store.submit_reload(reload_key, owned_pages)
# Slice the candidate key to the C L3-durable pages (matching suffix_hashes[:c] and the C-page
# reservation) so the reloaded node is WELL-FORMED:
# len(key)//page == host_len//page == object.num_pages == len(hash_value) == C.
# The un-durable suffix [C:n_full] is simply not present (a correct cache miss / recompute);
# carrying the full key while only C pages are backed would make the node claim a prefix it
# cannot serve and crash the radix split (split_pages from the long key vs the C-page object).
self.ongoing_l3_reload[op_id] = {
"reload_key": reload_key, "last_host_node": lhn,
"suffix_key": skey[: c * self.page_size],
"suffix_hashes": list(shashes[:c]), "n_pages": c, "waiters": [rid],
"ttl": self.cp_l3_reload_ttl_ticks,
}
self.cp_l3_reload_inflight_keys[reload_key] = op_id
reserved_this_tick[reload_key] = op_id
self.cp_l3_reload_done[rid] = False # hold the request until admission
def _cp_l3_reserve_reload_object(self, reload_key, payloads, n_pages, allocator):
"""Reserve n_pages per payload for a reload as a NEW uncommitted object, evicting cold committed
shared-L2 leaves (replicated SLRU order -> collective-free, mirrors _reserve_write_cp_shared_l2_evict_to_fit)
to open contiguous room. Returns the ranges dict, or None if it cannot fit (abort + recompute)."""
victims = None
vi = 0
try:
for pk in payloads:
while True:
try:
allocator.reserve(reload_key, pk, n_pages)
break
except ValueError as e:
if "insufficient contiguous" not in str(e):
raise # a real programming/validation error -> fail loud
if victims is None:
victims = sorted(
(
nd
for nd in getattr(self, "evictable_host_leaves", set())
if self._cp_host_leaf_is_plannable_victim(nd)
and self._is_cp_shared_l2_metadata(getattr(nd, "cp_hicache", None))
),
key=self._cp_host_evict_key,
)
if vi >= len(victims):
allocator.abort(reload_key) # free any payload already reserved under this key
return None
victim = victims[vi]
vi += 1
if not self._cp_host_leaf_is_plannable_victim(victim):
raise RuntimeError(
"CP shared-L2 reload evict-to-fit plan diverged before reservation: "
f"reload_key={reload_key} victim_id={getattr(victim, 'id', None)}"
)
self._record_remove_event(victim)
self._cp_account_leave_node_evict(victim, victim.cp_hicache)
self.cache_controller.evict_cp_host(victim.cp_hicache)
victim.host_len = 0
victim.cp_hicache = None
victim.host_value = None
self._remove_host_leaf(victim)
except Exception:
allocator.abort(reload_key)
raise
return allocator.object_ranges(reload_key)
def _cp_l3_build_reload_owned_pages(self, ranges, page_hashes, store):
"""This rank's owned (dest_slab_page, content_hash) per payload (owner = object-local i % cp_size,
identical to spill's _cp_l3_build_owned_pages -> each rank reloads exactly the pages it spilled)."""
cp_size, cp_rank = store.cp_size, store.cp_rank
n = len(page_hashes)
owned = {}
for pk, rng in ranges.items():
if n > rng.num_pages:
# Symmetric with the spill builder _cp_l3_build_owned_pages: base + i
# past the object would address the neighbour object's pages in the same
# slab. Safe by construction (n == c == ranges), fail loud on any desync.
raise RuntimeError(
f"[CP_L3_FAILFAST] reload page count {n} exceeds object pages "
f"{rng.num_pages} for payload {pk!r} (node key/host_len desync); "
f"base_page={rng.base_page} slab_id={rng.slab_id}"
)
base = rng.base_page
owned[pk] = [(base + i, page_hashes[i]) for i in range(n) if i % cp_size == cp_rank]
return owned
def _cp_l3_admit_reload(self, op_id: int, ok: bool) -> None:
"""Reload op done (rank-uniform ok via the ok-AND). ok + still-clean-attach -> commit the reserved L2
object + insert ONE CP-aware radix node + mark waiters done (they re-match -> L2 hit -> load_back). Else
-> abort the reservation + release waiters to recompute. Rank-uniform: same op, same ok, replicated radix."""
info = self.ongoing_l3_reload.pop(op_id, None)
if info is None:
return
self.cp_l3_reload_inflight_keys.pop(info["reload_key"], None)
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
inserted = False
if not info["waiters"]:
# every waiter was released (aborted, or the ttl bound fired) -> nobody needs this reload now;
# free the reservation rather than insert an unrequested node (frees L2 for active work).
if allocator is not None:
allocator.abort(info["reload_key"])
return
if (
ok
and allocator is not None
and self._cp_l3_reload_attach_ok(info["last_host_node"], info["suffix_key"])
):
allocator.mark_object_committed(info["reload_key"])
node = self._cp_l3_insert_reloaded_node(info)
inserted = node is not None
if inserted:
# 3.3 GC: bump the L3 last_access of THIS rank's owned reloaded pages (owner = i%cp_size), so
# the reloaded prefix stays warm in L3 and isn't GC'd as cold. Touch ALL reloaded pages (not
# just the tail) so the prefix ages uniformly. last_access = the inserted node's replicated clock.
store = self.cp_l3_store
la = int(node.last_access_time)
shashes = info["suffix_hashes"]
touches = [
(shashes[i], pk, la)
for pk in store.payloads
for i in range(len(shashes))
if i % store.cp_size == store.cp_rank
]
if touches:
store.submit_touch(touches)
else:
allocator.abort(info["reload_key"])
elif allocator is not None:
allocator.abort(info["reload_key"]) # failed reload / stale attach -> free the reservation
for rid in info["waiters"]:
if inserted:
self.cp_l3_reload_done[rid] = True # admitted -> re-match hits the inserted node
else:
self.cp_l3_reload_done.pop(rid, None) # released -> proceed (recompute)
def _cp_l3_insert_reloaded_node(self, info) -> Optional[TreeNode]:
"""Insert ONE CP-aware radix node for the reloaded prefix under last_host_node (clean attach
re-validated here against the live tree). cp_hicache = a committed CpSharedL2NodeMetadata over the
reserved ranges; value=None (device-evicted, L2-resident) so the normal load_back serves L2->L1; the
node is L3-durable (just read from L3). Returns the new node (its last_access_time seeds the GC touch),
or None if the attach raced to non-clean. Rank-uniform (replicated radix + rank-uniform inputs)."""
from sglang.srt.mem_cache.cp_shared_l2_pool import CpSharedL2NodeMetadata
lhn = info["last_host_node"]
skey = info["suffix_key"]
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
ranges = allocator.object_ranges(info["reload_key"]) if allocator is not None else {}
if not ranges:
return None
child_key = self.get_child_key_fn(skey)
if child_key in getattr(lhn, "children", {}):
return None # raced to non-clean since the attach_ok check
# Fail-soft anchor guard: re-derive the first suffix hash from lhn's LIVE last hash; if it no longer
# matches the candidate's recorded hash the anchor moved (lhn's chain changed since the mark) -> abort
# (recompute). Traced unreachable today (a retained device-evicted node is provably backuped + its hash
# stable), but L3 reload is the first CP consumer of this match boundary, so guard explicitly.
parent_basis = lhn.get_last_hash_value() if lhn is not self.root_node else None
first_page = list(skey.token_ids[: self.page_size])
if not first_page or get_hash_str(first_page, prior_hash=parent_basis) != info["suffix_hashes"][0]:
return None
n_pages = info["n_pages"]
meta = CpSharedL2NodeMetadata(
logical_len=n_pages * self.page_size,
padded_len=n_pages * self.page_size,
page_size=self.page_size,
object_ranges=ranges,
required_payloads=tuple(ranges.keys()),
object_key=info["reload_key"],
)
new_node = TreeNode(priority=lhn.priority)
new_node.parent = lhn
new_node.key = skey.compacted()
new_node.value = None # device-evicted; load_back fills L1 from L2
new_node.host_value = None
new_node.host_len = n_pages * self.page_size
new_node.cp_hicache = meta
new_node.hash_value = list(info["suffix_hashes"])
new_node.l3_durable = True # just reloaded from L3 -> the spill maintainer skips it (O(1))
new_node.last_access_time = TreeNode.next_access_time()
lhn.children[child_key] = new_node
self._cp_account_enter_committed(new_node.cp_hicache)
self._update_host_leaf_status(new_node)
self._update_host_leaf_status(lhn)
return new_node
def _cp_l3_enqueue_spill(self, object_key: str, node: TreeNode) -> None:
"""Enqueue a committed object for proactive spill (rank-uniform: called only from the replicated
commit frontier + radix split). Drops (fail-soft) past the cap -- only reachable under a disk stall
(the deque can't drain), so the drop bounds RAM; the object recomputes on a future miss. len() is O(1)
and rank-uniform (the deque fills + drains identically on every rank), so the cap decision can't diverge."""
if len(self.cp_l3_spill_pending) < self.cp_l3_spill_pending_cap:
self.cp_l3_spill_pending.append((object_key, node))
return
self._cp_l3_spill_dropped += 1
if self._cp_l3_spill_dropped == 1 or self._cp_l3_spill_dropped % 1024 == 0:
logger.warning(
"[CP-L3] spill backlog at cap %d (disk falling behind); dropped %d spill candidates "
"(fail-soft: they recompute on miss). 3.3 disk-eviction reduces this.",
self.cp_l3_spill_pending_cap, self._cp_l3_spill_dropped,
)
def _cp_l3_spill_maintainer(self) -> None:
"""Continuous PROACTIVE spill (3.1): drain the FIFO of just-committed objects (filled at the replicated
commit frontier, _commit_pending_backup) -- NOT a coldest-scan. Each still-resident, not-yet-durable,
unpinned object is submitted as a background COPY to L3; protect_host pins the LIVE node only until the
RAM gather completes (released at the gather-ack), the slow disk write runs off the slab-independent
staged copy. Because we spill HOT just-committed objects (disjoint from eviction's COLD victims), the
cold tail is ALREADY L3-durable when evicted -> eviction just drops it, never contends for the pin.
Rank-synced by construction: the FIFO is filled in the rank-uniform MIN-frontier order, so all ranks
submit the SAME objects in lockstep -> object-granular acks + the MIN-drain stay aligned. Bounded per
tick by the inflight budget; budget-limited entries stay queued (drained next tick)."""
store = self.cp_l3_store
if store is None or not self.cp_l3_spill_pending:
return
budget = self.cp_l3_spill_max_inflight - len(self.ongoing_l3_spill)
if budget <= 0:
return
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
if allocator is None:
return
pending = self.cp_l3_spill_pending
submitted = 0
while pending and submitted < budget:
object_key, node = pending.popleft()
# Every non-submittable state below is TERMINAL -> drop the entry (popped). Only the budget limit
# (loop guard) leaves entries queued. Under the current CP+L3 config host_ref_counter is set ONLY by
# our L3 pin (radix_cache:248) -- storage prefetch/backup is mutually-exclusive (server_args) and
# pin_prefix is a CP no-op (pin_expiry stays 0) -- so >0 here means an in-flight spill of this same
# node -> it will become durable; drop. (Revisit this drop rule if either exclusion is relaxed.)
if getattr(node, "l3_durable", False):
continue # already durable in L3
meta = getattr(node, "cp_hicache", None)
if meta is None or not self._is_cp_shared_l2_metadata(meta):
continue # evicted / rolled back since commit
if int(getattr(node, "host_ref_counter", 0) or 0) > 0:
continue # already being spilled (our pin) -> the in-flight op makes it durable
if not object_key or not allocator.is_committed(object_key):
continue # rolled back / no longer a committed pool object
page_keys = getattr(node, "hash_value", None)
if not page_keys:
continue # no hash chain (shouldn't happen under enable_cp_l3 hash-from-init)
required = tuple(getattr(meta, "required_payloads", ()))
if store.exists_prefix(page_keys, required) >= len(page_keys):
node.l3_durable = True # already fully durable (e.g. a reloaded-from-L3 node) -> memoize O(1)
continue
owned_pages = self._cp_l3_build_owned_pages(meta, page_keys, required, store)
last_access = int(getattr(node, "last_access_time", 0) or 0) # rank-replicated logical clock
node.protect_host()
op_id = store.submit_spill(object_key, owned_pages, last_access)
self.ongoing_l3_spill[op_id] = node
submitted += 1
def _cp_l3_build_owned_pages(self, meta, page_keys, required, store):
"""Per payload: this rank's owned (slab_page, content_hash) pairs. Owner = object-LOCAL page index
i % cp_size (verified vs _shared_l2_current_page_owners = owner_for_logical_pages(arange(1,n+1))).
Logical pages only (len(page_keys)); padding has no content hash. slab_page = payload range base_page+i."""
cp_size, cp_rank = store.cp_size, store.cp_rank
n = len(page_keys)
owned = {}
for pk in required:
if pk not in store.payloads:
continue
rng = meta.object_ranges.get(pk)
if rng is None:
continue
if n > rng.num_pages:
# base + i for i in range(n) would address PAST this object into the
# neighbour object's pages in the same slab -> silent L3 corruption
# (wrong content under a valid hash). After the reload-key slice this
# cannot happen; fail loud on any future key/host_len desync.
raise RuntimeError(
f"[CP_L3_FAILFAST] spill page count {n} exceeds object pages "
f"{rng.num_pages} for payload {pk!r} (node key/host_len desync); "
f"base_page={rng.base_page} slab_id={rng.slab_id}"
)
base = rng.base_page
owned[pk] = [
(base + i, page_keys[i]) for i in range(n) if i % cp_size == cp_rank
]
return owned
def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None:
start = time.perf_counter()
torch.distributed.all_reduce(tensor, op=op, group=group)
elapsed_ms = (time.perf_counter() - start) * 1000
if not hasattr(self, "_cp_hicache_collective_stats"):
self._cp_hicache_collective_stats = {}
stats = self._cp_hicache_collective_stats.setdefault(
tag,
{
"count": 0,
"total_ms": 0.0,
"max_ms": 0.0,
},
)
stats["count"] += 1
stats["total_ms"] += elapsed_ms
stats["max_ms"] = max(stats["max_ms"], elapsed_ms)
count = stats["count"]
if count == 1 or count % 128 == 0 or elapsed_ms >= 10.0:
logger.debug(
"[HiCache-collective] tag=%s count=%d total_ms=%.3f avg_ms=%.3f max_ms=%.3f last_ms=%.3f",
tag,
count,
stats["total_ms"],
stats["total_ms"] / count,
stats["max_ms"],
elapsed_ms,
)
def _cp_hicache_cp_size(self, page_owners: Optional[torch.Tensor] = None) -> int:
layout = getattr(
getattr(self, "cache_controller", None), "cp_shared_kv_layout", None
)
cp_size = getattr(layout, "cp_size", None)
if cp_size is not None:
return int(cp_size)
tp_world_size = int(getattr(self, "tp_world_size", 1) or 1)
if tp_world_size > 1:
return tp_world_size
if page_owners is not None and page_owners.numel() > 0:
non_negative = page_owners.to(device="cpu", dtype=torch.int64)
non_negative = non_negative[non_negative >= 0]
if non_negative.numel() > 0:
return int(non_negative.max().item()) + 1
return 1
def _cp_hicache_cp_rank(self) -> int:
layout = getattr(
getattr(self, "cache_controller", None), "cp_shared_kv_layout", None
)
return int(getattr(layout, "cp_rank", getattr(self, "_tp_group_rank", 0)) or 0)
def _cp_zero_counts(self, cp_size: Optional[int] = None) -> Tuple[int, ...]:
cp_size = self._cp_hicache_cp_size() if cp_size is None else int(cp_size)
return tuple(0 for _ in range(cp_size))
@staticmethod
def _cp_add_counts(lhs: Tuple[int, ...], rhs: Tuple[int, ...]) -> Tuple[int, ...]:
if len(lhs) != len(rhs):
raise RuntimeError(
f"CP HiCache count length mismatch: {len(lhs)} vs {len(rhs)}"
)
return tuple(int(a) + int(b) for a, b in zip(lhs, rhs))
@staticmethod
def _cp_sub_counts_floor_zero(
lhs: Tuple[int, ...], rhs: Tuple[int, ...]
) -> Tuple[int, ...]:
if len(lhs) != len(rhs):
raise RuntimeError(
f"CP HiCache count length mismatch: {len(lhs)} vs {len(rhs)}"
)
return tuple(max(0, int(a) - int(b)) for a, b in zip(lhs, rhs))
def _cp_owner_token_counts(
self,
page_owners: torch.Tensor,
*,
page_size: int,
cp_size: Optional[int] = None,
) -> Tuple[int, ...]:
cp_size = (
self._cp_hicache_cp_size(page_owners)
if cp_size is None
else int(cp_size)
)
counts = [0 for _ in range(cp_size)]
owners = page_owners.to(device="cpu", dtype=torch.int64)
for owner in owners.tolist():
owner = int(owner)
if owner < 0:
continue
if owner >= cp_size:
raise RuntimeError(
"CP HiCache metadata owner exceeds CP size: "
f"owner={owner} cp_size={cp_size}"
)
counts[owner] += int(page_size)
return tuple(counts)
def _cp_metadata_token_counts(
self, metadata: CpHiCacheNodeMetadata
) -> Tuple[int, ...]:
return self._cp_owner_token_counts(
metadata.page_owners,
page_size=metadata.page_size,
)
def _cp_assert_metadata_counts(
self, metadata: CpHiCacheNodeMetadata, *, context: str
) -> Tuple[int, ...]:
counts = self._cp_metadata_token_counts(metadata)
cp_rank = self._cp_hicache_cp_rank()
if cp_rank >= len(counts):
raise RuntimeError(
"CP HiCache owner lane mismatch: "
f"context={context} cp_rank={cp_rank} cp_size={len(counts)}"
)
expected_local = counts[cp_rank]
actual_local = int(metadata.owned_positions.numel())
if actual_local != expected_local:
raise RuntimeError(
"CP HiCache owner lane mismatch: "
f"context={context} cp_rank={cp_rank} expected_local={expected_local} "
f"owned_positions={actual_local}"
)
if int(metadata.host_indices.numel()) != expected_local:
raise RuntimeError(
"CP HiCache host index owner lane mismatch: "
f"context={context} cp_rank={cp_rank} expected_local={expected_local} "
f"host_indices={metadata.host_indices.numel()}"
)
draft_host_indices = getattr(metadata, "draft_host_indices", None)
if draft_host_indices is not None and int(draft_host_indices.numel()) != expected_local:
raise RuntimeError(
"CP HiCache draft host index owner lane mismatch: "
f"context={context} cp_rank={cp_rank} expected_local={expected_local} "
f"draft_host_indices={draft_host_indices.numel()}"
)
return counts
def _cp_walk_radix_nodes(self) -> List[TreeNode]:
root = getattr(self, "root_node", None)
if root is None:
return []
nodes: List[TreeNode] = []
stack = list(getattr(root, "children", {}).values())
while stack:
node = stack.pop()
nodes.append(node)
stack.extend(getattr(node, "children", {}).values())
return nodes
# ----- B1: O(cp_size) running host-capacity counts -------------------
# `_cp_counts` (an _CpRunningHostCounts) holds the per-owner committed /
# pending host TOKEN counts that `_cp_host_capacity_snapshot` used to derive
# by walking the whole radix tree (O(total cached pages)) on every
# write-admission. It is created EMPTY on first access (production starts
# with an empty tree) and then maintained incrementally at each state
# transition (enter/leave pending, commit, evict, split, reset), so the
# snapshot is O(cp_size). `_cp_walk_capacity_counts` recomputes the same
# result from scratch and backs the env-gated drift verifier (it is the
# correctness reference, not a build source).
def _cp_walk_capacity_counts(
self,
) -> Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]:
"""Authoritative (committed_target, committed_draft, pending_target,
pending_draft) per-owner token counts from the full radix walk + pending
dict. Runs the per-node owner-lane integrity assertion. This is the
correctness reference the env-gated drift verifier compares the
incrementally-maintained running counts against -- NOT a build source for
them (the running counts start empty and are maintained by the hooks)."""
cp_size = self._cp_hicache_cp_size()
ct = self._cp_zero_counts(cp_size)
cd = self._cp_zero_counts(cp_size)
pt = self._cp_zero_counts(cp_size)
pd = self._cp_zero_counts(cp_size)
pending = getattr(self, "pending_host_backups", {})
pending_node_ids = set(pending.keys())
for node in self._cp_walk_radix_nodes():
metadata = getattr(node, "cp_hicache", None)
if metadata is None or getattr(node, "host_len", 0) <= 0:
continue
if self._is_cp_shared_l2_metadata(metadata):
continue
if getattr(node, "id", None) in pending_node_ids:
continue
counts = self._cp_assert_metadata_counts(
metadata,
context=f"capacity_committed node_id={getattr(node, 'id', '?')}",
)
ct = self._cp_add_counts(ct, counts)
if getattr(metadata, "draft_host_indices", None) is not None:
cd = self._cp_add_counts(cd, counts)
for node_id, pending_backup in pending.items():
if self._is_cp_shared_l2_metadata(pending_backup.metadata):
continue
counts = self._cp_assert_metadata_counts(
pending_backup.metadata,
context=f"capacity_pending node_id={node_id}",
)
pt = self._cp_add_counts(pt, counts)
if getattr(pending_backup.metadata, "draft_host_indices", None) is not None:
pd = self._cp_add_counts(pd, counts)
return ct, cd, pt, pd
@property
def _cp_counts(self) -> "_CpRunningHostCounts":
"""Running host-capacity counts, created empty on first access and then
maintained incrementally at every state transition.
Production starts with an empty radix tree, so zero-initialisation is
correct: every committed/pending entry is created only via commit /
split (committed) or attach / write_backup (pending), each of which
flows through the accounting hooks below. Only reached when CP HiCache
is enabled (the wrappers and the snapshot guard on _uses_cp_hicache).
Tests that construct nodes directly register them via the same hooks."""
counts = self.__dict__.get("_cp_running_host_counts")
if counts is None:
counts = _CpRunningHostCounts(self._cp_hicache_cp_size())
self.__dict__["_cp_running_host_counts"] = counts
return counts
# CP accounting is a no-op unless CP HiCache is enabled; an absent
# _uses_cp_hicache (a cache built without __init__, e.g. a unit fixture or a
# non-CP path) counts as disabled.
@staticmethod
def _is_cp_shared_l2_metadata(metadata) -> bool:
return hasattr(metadata, "object_ranges") and not hasattr(metadata, "page_owners")
@staticmethod
def _cp_reservation_owned_count(reservation) -> int:
owned_positions = getattr(reservation, "owned_positions", None)
if owned_positions is not None:
return int(owned_positions.numel())
metadata = getattr(reservation, "metadata", None)
metadata_owned_positions = getattr(metadata, "owned_positions", None)
if metadata_owned_positions is not None:
return int(metadata_owned_positions.numel())
return int(getattr(reservation, "host_indices").numel())
def _cp_account_enter_pending(self, metadata) -> None:
if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata):
self._cp_counts.enter_pending(metadata)
def _cp_account_leave_pending(self, metadata) -> None:
if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata):
self._cp_counts.leave_pending(metadata)
def _cp_account_enter_committed(self, metadata) -> None:
if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata):
self._cp_counts.enter_committed(metadata)
def _cp_account_leave_committed(self, metadata) -> None:
if getattr(self, "_uses_cp_hicache", False) and not self._is_cp_shared_l2_metadata(metadata):
self._cp_counts.leave_committed(metadata)
def _cp_account_leave_node_evict(self, node, metadata) -> None:
"""Leave-accounting for an evicted node; metadata captured BEFORE the
caller nulls node.cp_hicache. Picks the bucket by pending membership so
the running total self-corrects; a pending eviction is provably
impossible (pending nodes are locked and skipped by both evictors and
the planner) so it also fails loud."""
if metadata is None or not getattr(self, "_uses_cp_hicache", False):
return
if self._is_cp_shared_l2_metadata(metadata):
return
if getattr(node, "id", None) in getattr(self, "pending_host_backups", {}):
self._cp_counts.leave_pending(metadata)
raise RuntimeError(
"CP HiCache invariant violated: evicting a PENDING node "
f"id={getattr(node, 'id', None)} — should be impossible"
)
self._cp_counts.leave_committed(metadata)
def _cp_reset_running_counts(self) -> None:
if getattr(self, "_uses_cp_hicache", False):
self._cp_counts.reset()
def _cp_host_capacity_snapshot(self) -> CpHiCacheCapacitySnapshot:
cp_size = self._cp_hicache_cp_size()
target_capacity = tuple(
int(getattr(getattr(self, "token_to_kv_pool_host", None), "size", 0))
for _ in range(cp_size)
)
controller = getattr(self, "cache_controller", None)
draft_pool = getattr(controller, "draft_mem_pool_host", None)
has_draft = bool(getattr(controller, "has_draft_hicache", False))
draft_capacity = (
tuple(int(getattr(draft_pool, "size", 0)) for _ in range(cp_size))
if has_draft
else None
)
if not self._uses_cp_hicache:
zero = self._cp_zero_counts(cp_size)
committed_target = committed_draft = pending_target = pending_draft = zero
else:
counts = self._cp_counts
counts.sanity_check()
committed_target, committed_draft, pending_target, pending_draft = (
counts.as_tuples()
)
if envs.SGLANG_CP_HICACHE_VERIFY_SNAPSHOT.get():
walk = self._cp_walk_capacity_counts()
running = (
committed_target,
committed_draft,
pending_target,
pending_draft,
)
if running != walk:
raise RuntimeError(
"CP HiCache running-count DRIFT vs full walk:\n"
f" running={running}\n"
f" walk={walk}"
)
return CpHiCacheCapacitySnapshot(
target_capacity=target_capacity,
draft_capacity=draft_capacity,
committed_target=committed_target,
committed_draft=committed_draft,
pending_target=pending_target,
pending_draft=pending_draft,
)
def _cp_required_host_tokens_by_rank(
self, device_indices: torch.Tensor
) -> Tuple[int, ...]:
layout = getattr(
getattr(self, "cache_controller", None), "cp_shared_kv_layout", None
)
if layout is None:
return self._cp_zero_counts()
logical_len = int(device_indices.numel())
if logical_len == 0:
return self._cp_zero_counts(layout.cp_size)
padded_device_indices = pad_token_locs_to_page_boundary(
device_indices,
self.page_size,
name="CP HiCache capacity device_indices",
)
logical_pages = torch.div(
padded_device_indices[:: self.page_size],
self.page_size,
rounding_mode="floor",
)
page_owners = layout.owner_for_logical_pages(logical_pages).to(
dtype=torch.int8, device="cpu"
)
return self._cp_owner_token_counts(
page_owners,
page_size=self.page_size,
cp_size=layout.cp_size,
)
def _cp_host_available_tokens_by_rank(
self, snapshot: CpHiCacheCapacitySnapshot, *, draft: bool = False
) -> Tuple[int, ...]:
if draft:
if snapshot.draft_capacity is None:
return tuple(2**62 for _ in snapshot.target_capacity)
return self._cp_sub_counts_floor_zero(
snapshot.draft_capacity, snapshot.draft_used
)
return self._cp_sub_counts_floor_zero(
snapshot.target_capacity, snapshot.target_used
)
def _cp_host_evictable_tokens_by_rank(self) -> Tuple[int, ...]:
"""Per-rank host (L2) tokens reclaimable by evicting all currently
plannable host leaves -- the host analog of tree_cache.evictable_size().
Computed from the REPLICATED radix tree (evictable_host_leaves + per-node
owner counts), so it is identical on every CP rank (collective-free).
"""
cp_size = self._cp_hicache_cp_size()
totals = self._cp_zero_counts(cp_size)
for node in getattr(self, "evictable_host_leaves", set()):
if not self._cp_host_leaf_is_plannable_victim(node):
continue
metadata = getattr(node, "cp_hicache", None)
if metadata is None:
continue
counts = self._cp_assert_metadata_counts(
metadata, context="host_evictable_tokens"
)
totals = self._cp_add_counts(totals, counts)
return totals
def cp_host_backup_admission_budget(self) -> Optional[int]:
"""Reserve-at-admission (Phase 2a) host headroom: min over CP ranks of
per-rank (available + evictable) host tokens.
A scalar, conservative budget -- the tightest lane's effective L2 capacity
-- mirroring the device-side ``rem_total_tokens = available_size() +
evictable_size()``. The scheduler debits each admitted backup's worst-lane
footprint and stops admitting (backpressure) when this hits zero, so a
complete backup never fails reactively at prepare. Returns None when not
under CP shared-KV HiCache. Derived purely from the replicated radix
snapshot -> identical on every rank (collective-free).
"""
if not self._uses_cp_hicache:
return None
snapshot = self._cp_host_capacity_snapshot()
available = self._cp_host_available_tokens_by_rank(snapshot)
evictable = self._cp_host_evictable_tokens_by_rank()
cp_size = min(len(available), len(evictable))
if cp_size == 0:
return None
return min(int(available[r]) + int(evictable[r]) for r in range(cp_size))
def _cp_host_evict_key(self, node: TreeNode):
# Replicated-clock SLRU total order (is_protected, last_access_time[logical],
# node.id) — replaces the old FIFO-by-node.id key. Same strategy the device
# owner-lane / physical-slot eviction paths use, so all host/device victim
# selection is consistent and collective-free.
return self.eviction_strategy.get_priority(node)
def _cp_host_leaf_is_plannable_victim(self, node: TreeNode) -> bool:
if node == getattr(self, "root_node", None):
return False
if not getattr(node, "evicted", False):
return False
parent = getattr(node, "parent", None)
if parent is None:
return False
try:
parent_key = self.get_child_key_fn(node.key)
except Exception:
return False
if getattr(parent, "children", {}).get(parent_key) is not node:
return False
pin_expiry = float(getattr(node, "pin_expiry", 0.0) or 0.0)
pin_active = pin_expiry > 0 and time.monotonic() <= pin_expiry
if pin_active:
return False
effective_host_refs = int(getattr(node, "host_ref_counter", 0) or 0)
if pin_expiry > 0:
# Current eviction clears an expired pin before checking host refs.
# Mirror that decision without mutating the tree in observer mode.
effective_host_refs = max(0, effective_host_refs - 1)
if effective_host_refs > 0:
return False
if self._node_host_write_pending(node):
return False
try:
return self._node_backuped(node)
except RuntimeError:
return False
def _plan_cp_host_evictions(
self, required_tokens_by_rank: Tuple[int, ...]
) -> CpHiCacheEvictionPlan:
deficits = [max(0, int(v)) for v in required_tokens_by_rank]
cp_size = len(deficits)
planned_freed = [0 for _ in range(cp_size)]
victims: List[TreeNode] = []
if all(v <= 0 for v in deficits):
return CpHiCacheEvictionPlan((), tuple(planned_freed), tuple(deficits))
leaves = sorted(
list(getattr(self, "evictable_host_leaves", set())),
key=self._cp_host_evict_key,
)
for node in leaves:
if not self._cp_host_leaf_is_plannable_victim(node):
continue
metadata = getattr(node, "cp_hicache", None)
if metadata is None:
continue
counts = self._cp_assert_metadata_counts(
metadata, context=f"eviction_plan node_id={getattr(node, 'id', '?')}"
)
if len(counts) != cp_size:
raise RuntimeError(
"CP HiCache eviction plan CP size mismatch: "
f"required={cp_size} metadata={len(counts)}"
)
victims.append(node)
for rank, count in enumerate(counts):
planned_freed[rank] += int(count)
deficits[rank] = max(0, deficits[rank] - int(count))
if all(v <= 0 for v in deficits):
break
return CpHiCacheEvictionPlan(
victims=tuple(victims),
planned_freed=tuple(planned_freed),
remaining_deficit=tuple(deficits),
)
def _cp_shared_l2_current_page_owners(self, num_pages: int) -> List[int]:
cp_size = self._cp_hicache_cp_size()
if num_pages < 0:
raise ValueError(f"num_pages must be non-negative, got {num_pages}")
return [int(page_index % cp_size) for page_index in range(num_pages)]
def _build_cp_load_back_plan(
self, nodes_to_load: List[TreeNode], *, node_id: int
) -> CpLoadBackPlan:
allocator = self.token_to_kv_pool_allocator
compute_owner_lane_stats = getattr(allocator, "compute_owner_lane_stats", None)
if compute_owner_lane_stats is None:
raise RuntimeError(
"CP HiCache load-back requires an allocator with "
f"compute_owner_lane_stats: node_id={node_id}"
)
page_owners: List[int] = []
saw_shared_l2_metadata = False
saw_legacy_cp_metadata = False
host_hit_len = 0
physical_hit_len = 0
for node in nodes_to_load:
metadata = getattr(node, "cp_hicache", None)
if metadata is None:
raise RuntimeError(
"CP HiCache load-back missing cp_hicache metadata: "
f"load_node_id={getattr(node, 'id', '?')} root_node_id={node_id}"
)
shared_l2_metadata = self._is_cp_shared_l2_metadata(metadata)
saw_shared_l2_metadata = saw_shared_l2_metadata or shared_l2_metadata
saw_legacy_cp_metadata = saw_legacy_cp_metadata or not shared_l2_metadata
if saw_shared_l2_metadata and saw_legacy_cp_metadata:
raise RuntimeError(
"CP HiCache load-back cannot mix shared physical L2 metadata "
"with legacy page_owners metadata in one load plan: "
f"root_node_id={node_id}"
)
node_host_len = int(self._node_host_len(node))
node_padded_len = int(getattr(metadata, "padded_len", node_host_len))
if node_padded_len % self.page_size != 0:
raise RuntimeError(
"CP HiCache load-back requires page-aligned physical lengths: "
f"load_node_id={getattr(node, 'id', '?')} "
f"padded_len={node_padded_len} page_size={self.page_size}"
)
if int(getattr(metadata, "valid_len", node_host_len)) != node_host_len:
raise RuntimeError(
"CP HiCache load-back metadata length mismatch: "
f"load_node_id={getattr(node, 'id', '?')} "
f"host_len={node_host_len} metadata_len={metadata.valid_len}"
)
if node_padded_len < node_host_len:
raise RuntimeError(
"CP HiCache load-back physical length is shorter than "
"valid host length: "
f"load_node_id={getattr(node, 'id', '?')} "
f"host_len={node_host_len} padded_len={node_padded_len}"
)
if not shared_l2_metadata:
page_owners.extend(int(owner) for owner in metadata.page_owners.tolist())
host_hit_len += node_host_len
physical_hit_len += node_padded_len
expected_pages = physical_hit_len // self.page_size
if saw_shared_l2_metadata:
page_owners = self._cp_shared_l2_current_page_owners(expected_pages)
else:
if len(page_owners) != expected_pages:
raise RuntimeError(
"CP HiCache load-back page owner count mismatch: "
f"node_id={node_id} physical_hit_len={physical_hit_len} "
f"page_size={self.page_size} page_owners={len(page_owners)}"
)
required, available, deficits, free_room_deficits = (
self._cp_load_back_owner_lane_stats(page_owners)
)
return CpLoadBackPlan(
page_owners=list(page_owners),
required_by_owner=[int(v) for v in required],
available_by_owner=[int(v) for v in available],
deficit_by_owner=[int(v) for v in deficits],
free_room_deficit_by_owner=[int(v) for v in free_room_deficits],
host_hit_len=host_hit_len,
)
def _cp_load_back_owner_lane_stats(
self, page_owners: List[int]
) -> Tuple[List[int], List[int], List[int], List[int]]:
allocator = self.token_to_kv_pool_allocator
target_ratio = float(getattr(self, "hicache_l1_free_room_ratio", 0.0) or 0.0)
trigger_ratio = float(
getattr(self, "hicache_l1_free_room_trigger_ratio", 0.0) or 0.0
)
required, available, exact_deficits = allocator.compute_owner_lane_stats(
page_owners
)
lane_capacity_pages = getattr(allocator, "compute_owner_lane_capacity_pages", None)
if lane_capacity_pages is None:
return required, available, exact_deficits, exact_deficits
return (
required,
available,
exact_deficits,
compute_owner_lane_free_room_deficits(
required=required,
available=available,
capacities=lane_capacity_pages(),
target_ratio=target_ratio,
trigger_ratio=trigger_ratio,
),
)
def _refresh_cp_load_back_plan(self, plan: CpLoadBackPlan) -> CpLoadBackPlan:
required, available, deficits, free_room_deficits = (
self._cp_load_back_owner_lane_stats(plan.page_owners)
)
return CpLoadBackPlan(
page_owners=plan.page_owners,
required_by_owner=[int(v) for v in required],
available_by_owner=[int(v) for v in available],
deficit_by_owner=[int(v) for v in deficits],
free_room_deficit_by_owner=[int(v) for v in free_room_deficits],
host_hit_len=plan.host_hit_len,
)
def _cp_device_node_is_load_back_victim_base(self, node: TreeNode) -> bool:
if node == getattr(self, "root_node", None):
return False
if getattr(node, "lock_ref", 0) > 0:
return False
if getattr(node, "value", None) is None:
return False
parent = getattr(node, "parent", None)
if parent is None:
return False
try:
parent_key = self.get_child_key_fn(node.key)
except Exception:
return False
if getattr(parent, "children", {}).get(parent_key) is not node:
return False
if self._is_pinned(node):
return False
return True
def _cp_device_node_is_load_back_victim_after_plan(
self, node: TreeNode, planned_evicted_nodes: set
) -> bool:
if not self._cp_device_node_is_load_back_victim_base(node):
return False
for child in getattr(node, "children", {}).values():
if child in planned_evicted_nodes:
continue
if not getattr(child, "evicted", False):
return False
return True
def _cp_device_leaf_is_load_back_victim(self, node: TreeNode) -> bool:
return self._cp_device_node_is_load_back_victim_after_plan(node, set())
def _cp_load_back_node_owner_page_counts(
self, node: TreeNode, cp_size: int, *, memo: Optional[dict] = None
) -> Tuple[int, ...]:
"""Per-owner device page-count histogram for a node.
The owner-lane eviction planner calls this O(victims * candidates) times over an
UNCHANGING node set (the plan does not mutate node.value), so the per-node counts
are invariant for the whole planning call. Pass a planner-scoped ``memo``
({node.id: counts}) to compute each node's counts exactly once -- node.id is a safe
key for the duration of one plan. Without ``memo`` the behaviour is unchanged.
The histogram itself uses a single ``bincount().tolist()`` device->host sync instead
of ``cp_size`` per-owner ``.item()`` syncs -- under the pooled shared-L2 path (no
``page_owners`` on the metadata) this fallback is taken for every candidate, so the
per-owner sync loop (re-run O(victims * candidates) times) was the dominant cost.
"""
if memo is not None:
cached = memo.get(node.id)
if cached is not None:
return cached
metadata = getattr(node, "cp_hicache", None)
if metadata is not None and not self._is_cp_shared_l2_metadata(metadata):
counts = metadata.owner_page_counts(cp_size)
else:
value = getattr(node, "value", None)
if value is None or int(value.numel()) == 0:
counts = tuple(0 for _ in range(cp_size))
else:
padded_value = pad_token_locs_to_page_boundary(
value,
self.page_size,
name=(
"CP HiCache load-back eviction device values "
f"node_id={getattr(node, 'id', '?')}"
),
)
first_locs = padded_value[:: self.page_size]
logical_pages = torch.div(
first_locs, self.page_size, rounding_mode="floor"
)
owners = torch.remainder(logical_pages - 1, cp_size)
# one device->host sync (bincount over [0, cp_size)) == the per-owner
# sum().item() loop, byte-identical counts, cp_size-fewer syncs.
counts = tuple(torch.bincount(owners, minlength=cp_size).tolist())
if memo is not None:
memo[node.id] = counts
return counts
def _cp_load_back_ancestor_unlock_contribution(
self,
node: TreeNode,
deficits: List[int],
planned_evicted_nodes: set,
cp_size: int,
*,
memo: Optional[dict] = None,
) -> int:
ancestor = getattr(node, "parent", None)
while ancestor is not None and ancestor != getattr(self, "root_node", None):
if ancestor in planned_evicted_nodes:
return 0
if getattr(ancestor, "value", None) is None:
ancestor = getattr(ancestor, "parent", None)
continue
if not self._cp_device_node_is_load_back_victim_base(ancestor):
return 0
counts = self._cp_load_back_node_owner_page_counts(
ancestor, cp_size, memo=memo
)
contribution = sum(
min(int(count), int(deficit))
for count, deficit in zip(counts, deficits)
)
if contribution > 0:
return int(contribution)
ancestor = getattr(ancestor, "parent", None)
return 0
def _plan_cp_load_back_owner_lane_evictions(
self, plan: CpLoadBackPlan
) -> CpLoadBackEvictionPlan:
plan_start_time = time.perf_counter()
last_progress_time = plan_start_time
deficits = [max(0, int(v)) for v in plan.deficit_by_owner]
cp_size = len(deficits)
planned_freed = [0 for _ in range(cp_size)]
victims: List[TreeNode] = []
planned_evicted_nodes = set()
candidate_nodes = set(getattr(self, "evictable_leaves", set()))
initial_candidate_count = len(candidate_nodes)
iteration = 0
# Per-plan memo {node.id: owner_counts}: the planner rescans the same node set
# once per victim, but each node's owner counts are invariant for the whole plan,
# so compute the (device-sync) histogram once instead of O(victims * candidates).
owner_counts_memo: dict = {}
while any(v > 0 for v in deficits):
iteration += 1
best_node = None
best_counts = None
best_score = None
scanned_candidates = 0
for node in list(candidate_nodes):
scanned_candidates += 1
if (scanned_candidates & 1023) == 0:
now = time.perf_counter()
if now - last_progress_time >= 5.0:
logger.warning(
"[HiCache-load] slow CP owner-lane eviction planning scan: "
"iteration=%d scanned=%d candidates=%d victims=%d "
"deficits=%s planned_freed=%s elapsed_ms=%.3f",
iteration,
scanned_candidates,
len(candidate_nodes),
len(victims),
deficits,
planned_freed,
(now - plan_start_time) * 1000.0,
)
last_progress_time = now
if node in planned_evicted_nodes:
continue
if not self._cp_device_node_is_load_back_victim_after_plan(
node, planned_evicted_nodes
):
continue
counts = self._cp_load_back_node_owner_page_counts(
node, cp_size, memo=owner_counts_memo
)
contribution = sum(
min(int(count), int(deficit))
for count, deficit in zip(counts, deficits)
)
unlock_contribution = 0
if contribution <= 0:
unlock_contribution = (
self._cp_load_back_ancestor_unlock_contribution(
node,
deficits,
planned_evicted_nodes,
cp_size,
memo=owner_counts_memo,
)
)
if contribution <= 0 and unlock_contribution <= 0:
continue
# get_priority returns the CpReplicatedSLRUStrategy tuple
# (is_protected, last_access_time[logical], node.id) under CP — a
# rank-replicated value; the trailing node.id keeps the full score a
# strict total order even if the strategy is ever swapped.
score = (
-int(contribution),
-int(unlock_contribution),
self.eviction_strategy.get_priority(node),
int(getattr(node, "id", 0) or 0),
)
if best_score is None or score < best_score:
best_score = score
best_node = node
best_counts = counts
if best_node is None or best_counts is None:
break
victims.append(best_node)
planned_evicted_nodes.add(best_node)
candidate_nodes.discard(best_node)
for owner, count in enumerate(best_counts):
planned_freed[owner] += int(count)
deficits[owner] = max(0, deficits[owner] - int(count))
ancestor = getattr(best_node, "parent", None)
while ancestor is not None and ancestor != getattr(self, "root_node", None):
if ancestor in planned_evicted_nodes:
break
if self._cp_device_node_is_load_back_victim_after_plan(
ancestor, planned_evicted_nodes
):
candidate_nodes.add(ancestor)
break
if getattr(ancestor, "value", None) is not None:
break
ancestor = getattr(ancestor, "parent", None)
plan_elapsed_ms = (time.perf_counter() - plan_start_time) * 1000.0
if plan_elapsed_ms >= 1000.0 or any(v > 0 for v in deficits):
logger.warning(
"[HiCache-load] CP owner-lane eviction planning done: "
"elapsed_ms=%.3f initial_candidates=%d remaining_candidates=%d "
"iterations=%d victims=%d original_deficit_by_owner=%s "
"planned_freed_by_owner=%s remaining_deficit_by_owner=%s",
plan_elapsed_ms,
initial_candidate_count,
len(candidate_nodes),
iteration,
len(victims),
plan.deficit_by_owner,
planned_freed,
deficits,
)
return CpLoadBackEvictionPlan(
victims=tuple(victims),
planned_freed_by_owner=tuple(planned_freed),
remaining_deficit_by_owner=tuple(deficits),
)
def _evict_cp_load_back_owner_lanes(
self, plan: CpLoadBackPlan, *, node_id: int
) -> CpLoadBackPlan:
evict_start_time = time.perf_counter()
eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan)
plan_elapsed_ms = (time.perf_counter() - evict_start_time) * 1000.0
if plan_elapsed_ms >= 1000.0:
logger.warning(
"[HiCache-load] CP owner-lane eviction plan ready: node_id=%d "
"elapsed_ms=%.3f victims=%d planned_freed_by_owner=%s "
"remaining_deficit_by_owner=%s",
node_id,
plan_elapsed_ms,
len(eviction_plan.victims),
eviction_plan.planned_freed_by_owner,
eviction_plan.remaining_deficit_by_owner,
)
if any(v > 0 for v in eviction_plan.remaining_deficit_by_owner):
logger.warning(
"[CP_HICACHE_FALLBACK][cp_load_back_owner_lane_eviction_insufficient] "
"node_id=%d required_by_owner=%s available_by_owner=%s "
"deficit_by_owner=%s planned_freed_by_owner=%s "
"remaining_deficit_by_owner=%s evictable_size=%d protected_size=%d "
"victims=%s allocator_state=%s",
node_id,
plan.required_by_owner,
plan.available_by_owner,
plan.deficit_by_owner,
eviction_plan.planned_freed_by_owner,
eviction_plan.remaining_deficit_by_owner,
int(getattr(self, "evictable_size_", 0)),
int(getattr(self, "protected_size_", 0)),
[getattr(node, "id", None) for node in eviction_plan.victims],
self.token_to_kv_pool_allocator.allocator_state_str(),
)
return self._refresh_cp_load_back_plan(plan)
if len(eviction_plan.victims) == 0:
return plan
num_evicted = 0
write_back_nodes: List[TreeNode] = []
for victim in eviction_plan.victims:
victim_id = getattr(victim, "id", None)
if not self._cp_device_leaf_is_load_back_victim(victim):
raise RuntimeError(
"CP HiCache load-back owner-lane eviction plan diverged "
f"before reservation: node_id={node_id} victim_id={victim_id}"
)
if victim.pin_expiry > 0 and time.monotonic() > victim.pin_expiry:
self._clear_pin(victim)
if not self._node_backuped(victim):
if self.cache_controller.write_policy == "write_back":
logger.warning(
"[CP_HICACHE_FALLBACK][cp_load_back_owner_lane_write_back] "
"node_id=%d victim_id=%s reason=unbacked_write_back_victim",
node_id,
victim_id,
)
written = self.write_backup(victim, write_back=True)
if written > 0:
write_back_nodes.append(victim)
continue
num_evicted += self._evict_regular(victim)
else:
num_evicted += self._evict_backuped(victim)
for child in victim.parent.children.values():
if child in write_back_nodes:
continue
if not child.evicted:
break
else:
self._update_leaf_status(victim.parent)
if write_back_nodes:
self.writing_check(write_back=True)
for victim in write_back_nodes:
if self._node_backuped(victim):
num_evicted += self._evict_backuped(victim)
refreshed = self._refresh_cp_load_back_plan(plan)
evict_elapsed_ms = (time.perf_counter() - evict_start_time) * 1000.0
log_fn = (
logger.warning
if evict_elapsed_ms >= 1000.0
or any(v > 0 for v in refreshed.deficit_by_owner)
else logger.debug
)
log_fn(
"[HiCache-load] owner-lane device eviction before CP load-back done: "
"node_id=%d elapsed_ms=%.3f victims=%s num_evicted=%d "
"required_by_owner=%s before_available_by_owner=%s "
"before_deficit_by_owner=%s after_available_by_owner=%s "
"after_deficit_by_owner=%s planned_freed_by_owner=%s",
node_id,
evict_elapsed_ms,
[getattr(node, "id", None) for node in eviction_plan.victims],
num_evicted,
plan.required_by_owner,
plan.available_by_owner,
plan.deficit_by_owner,
refreshed.available_by_owner,
refreshed.deficit_by_owner,
eviction_plan.planned_freed_by_owner,
)
return refreshed
def _evict_cp_owner_lane_deficit_nodes(
self, params: EvictParams
) -> EvictResult:
deficits = [max(0, int(v)) for v in (params.owner_lane_deficits or [])]
if not deficits or all(v <= 0 for v in deficits):
return EvictResult()
plan = CpLoadBackPlan(
page_owners=[],
required_by_owner=[],
available_by_owner=[],
deficit_by_owner=deficits,
free_room_deficit_by_owner=deficits,
host_hit_len=0,
)
eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan)
if len(eviction_plan.victims) == 0:
logger.debug(
"[HiCache-evict] owner-lane evict found no contributing victims: "
"num_tokens=%d deficits=%s evictable_size=%d available_size=%d",
params.num_tokens,
deficits,
self.evictable_size_,
self.token_to_kv_pool_allocator.available_size(),
)
return EvictResult()
num_evicted = 0
num_locked_skipped = 0
write_back_nodes: List[TreeNode] = []
for victim in eviction_plan.victims:
if victim.lock_ref > 0:
num_locked_skipped += 1
continue
if victim.pin_expiry > 0 and time.monotonic() > victim.pin_expiry:
self._clear_pin(victim)
if not self._cp_device_leaf_is_load_back_victim(victim):
logger.debug(
"[HiCache-evict] owner-lane evict victim no longer evictable: "
"victim_id=%s deficits=%s",
getattr(victim, "id", None),
deficits,
)
continue
if self._is_pinned(victim):
if self._node_backuped(victim):
num_evicted += self._evict_backuped(victim)
else:
written = self.write_backup(victim, write_back=True)
if written > 0:
write_back_nodes.append(victim)
else:
self._clear_pin(victim)
continue
if not self._node_backuped(victim):
if self.cache_controller.write_policy == "write_back":
written = self.write_backup(victim, write_back=True)
if written > 0:
write_back_nodes.append(victim)
continue
num_evicted += self._evict_regular(victim)
else:
num_evicted += self._evict_backuped(victim)
for child in victim.parent.children.values():
if child in write_back_nodes:
continue
if not child.evicted:
break
else:
self._update_leaf_status(victim.parent)
if write_back_nodes:
self.writing_check(write_back=True)
for victim in write_back_nodes:
if self._node_backuped(victim):
num_evicted += self._evict_backuped(victim)
logger.debug(
"[HiCache-evict] owner-lane evict END: requested_tokens=%d "
"deficits=%s victims=%s planned_freed_by_owner=%s "
"remaining_deficit_by_owner=%s num_evicted=%d "
"num_locked_skipped=%d evictable_size_after=%d "
"available_size_after=%d",
params.num_tokens,
deficits,
[getattr(node, "id", None) for node in eviction_plan.victims],
eviction_plan.planned_freed_by_owner,
eviction_plan.remaining_deficit_by_owner,
num_evicted,
num_locked_skipped,
self.evictable_size_,
self.token_to_kv_pool_allocator.available_size(),
)
return EvictResult(num_tokens_evicted=num_evicted)
def _cp_build_write_admission_from_required(
self, required: Tuple[int, ...], *, node_id: int, phase: str
) -> CpWriteAdmission:
required = tuple(int(v) for v in required)
snapshot = self._cp_host_capacity_snapshot()
target_available = self._cp_host_available_tokens_by_rank(snapshot)
draft_available = self._cp_host_available_tokens_by_rank(snapshot, draft=True)
target_ratio = float(getattr(self, "hicache_host_free_room_ratio", 0.0) or 0.0)
trigger_ratio = float(
getattr(self, "hicache_host_free_room_trigger_ratio", 0.0) or 0.0
)
page_size = int(getattr(self, "page_size", 1) or 1)
target_deficit = tuple(
_free_room_deficit(
required=req,
available=avail,
capacity=capacity,
page_size=page_size,
target_ratio=target_ratio,
trigger_ratio=trigger_ratio,
)
for req, avail, capacity in zip(
required, target_available, snapshot.target_capacity
)
)
draft_capacity = snapshot.draft_capacity
draft_deficit = tuple(
(
_free_room_deficit(
required=req,
available=avail,
capacity=capacity,
page_size=page_size,
target_ratio=target_ratio,
trigger_ratio=trigger_ratio,
)
if draft_capacity is not None
else 0
)
for req, avail, capacity in zip(
required,
draft_available,
draft_capacity if draft_capacity is not None else snapshot.target_capacity,
)
)
deficit = tuple(max(a, b) for a, b in zip(target_deficit, draft_deficit))
eviction_plan = self._plan_cp_host_evictions(deficit)
return CpWriteAdmission(
node_id=node_id,
phase=phase,
required_by_owner=required,
target_available_by_owner=target_available,
draft_available_by_owner=draft_available,
deficit_by_owner=deficit,
eviction_plan=eviction_plan,
)
def _cp_build_write_admission(
self, device_indices: torch.Tensor, *, node_id: int, phase: str
) -> CpWriteAdmission:
required = self._cp_required_host_tokens_by_rank(device_indices)
return self._cp_build_write_admission_from_required(
required,
node_id=node_id,
phase=phase,
)
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.")
try:
if getattr(self, "cp_l3_store", None) is not None:
self.cp_l3_store.close()
except Exception:
logger.exception("Failed to close CP L3 store 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
TreeNode.last_access_counter = 0
self.cache_controller.reset()
if self.cp_l3_store is not None:
# flush_cache is idle-gated -> no in-flight L3 ops; clear() stops/drains/resets/restarts.
self.cp_l3_store.clear()
self.cp_l3_spill_pending.clear() # drop queued (now-flushed) spill candidates
self.ongoing_l3_spill.clear()
self.cp_l3_reload_candidates.clear()
self.ongoing_l3_reload.clear()
self.cp_l3_reload_inflight_keys.clear()
self.cp_l3_reload_done.clear()
self.token_to_kv_pool_host.clear()
self.cache_controller.clear_draft_host_pool()
# Clear per-request tracking dicts
self.prefetch_loaded_tokens_by_reqid.clear()
if hasattr(self, "pending_host_backups"):
self.pending_host_backups.clear()
self._cp_reset_running_counts()
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_has_required_draft_hicache(self, node: TreeNode) -> bool:
if not self._uses_cp_hicache:
return True
controller = getattr(self, "cache_controller", None)
layout = getattr(controller, "cp_shared_kv_layout", None)
cp_rank = getattr(layout, "cp_rank", "?")
metadata = getattr(node, "cp_hicache", None)
if metadata is None:
raise RuntimeError(
"CP HiCache invariant violation: "
f"node_id={getattr(node, 'id', '?')} "
f"host_len={getattr(node, 'host_len', '?')} "
f"cp_rank={cp_rank} has host_len without cp_hicache metadata"
)
if controller is None or not getattr(controller, "has_draft_hicache", False):
return True
if self._is_cp_shared_l2_metadata(metadata):
from sglang.srt.mem_cache.cp_shared_l2_pool import PAYLOAD_DRAFT_KV
required_payloads = set(getattr(metadata, "required_payloads", ()))
object_ranges = getattr(metadata, "object_ranges", {})
if (
PAYLOAD_DRAFT_KV in required_payloads
and PAYLOAD_DRAFT_KV in object_ranges
):
return True
raise RuntimeError(
"CP HiCache invariant violation: draft HiCache is attached but "
"shared physical L2 metadata is missing draft_kv payload: "
f"node_id={getattr(node, 'id', '?')} "
f"host_len={getattr(node, 'host_len', '?')} cp_rank={cp_rank}"
)
if getattr(metadata, "draft_host_indices", None) is None:
raise RuntimeError(
"CP HiCache invariant violation: draft HiCache is attached but "
f"node_id={getattr(node, 'id', '?')} "
f"host_len={getattr(node, 'host_len', '?')} "
f"cp_rank={cp_rank} is missing draft_host_indices"
)
return True
def _node_host_write_pending(self, node: TreeNode) -> bool:
node_id = getattr(node, "id", None)
return node_id in getattr(self, "ongoing_write_through", {}) or node_id in getattr(
self, "pending_host_backups", {}
)
def _node_host_write_ready(self, node: TreeNode) -> bool:
return not self._node_host_write_pending(node)
def _node_backuped(self, node: TreeNode) -> bool:
if self._uses_cp_hicache:
if node.host_len <= 0:
return False
return (
self._node_has_required_draft_hicache(node)
and self._node_host_write_ready(node)
)
return node.host_value is not None
def _commit_pending_backup(self, node_id: int) -> TreeNode:
pending = self.pending_host_backups[node_id]
node = pending.node
# Leave pending BEFORE the pop, then enter committed AFTER the node owns
# the metadata, so a first-touch lazy rebuild of the running counts
# always sees a tree consistent with the single transition in flight
# (same metadata object; counts net-transfer pending -> committed).
self._cp_account_leave_pending(pending.metadata)
self.pending_host_backups.pop(node_id)
node.host_len = pending.logical_len
node.cp_hicache = pending.metadata
node.host_value = None
self._cp_account_enter_committed(pending.metadata)
if pending.locked:
self.dec_node_lock_ref(node)
# B1: this commit is gated by the writing_check ReduceOp.MIN frontier, which
# is the all-ranks-done consensus. Mark the shared-L2 object committed here
# (replacing the per-(rank,layer,payload) gather quorum) so it transitions
# rank-uniformly. Guarded on object_ranges so it fires only for objects that
# were actually reserved into the pool (non-shared / skipped nodes no-op).
# B1 is validated for the write_through policy (the prod default); under
# write_back this commit is reached via the blocking drain rather than the
# MIN frontier -- still rank-uniform (deterministic registration), but the
# MIN-frontier reasoning above is the write_through path.
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
if allocator is not None:
object_key = self.cache_controller._cp_shared_l2_object_key(node_id)
if allocator.object_ranges(object_key):
allocator.mark_object_committed(object_key)
if self.enable_cp_l3:
# 3.1 proactive spill: enqueue at the replicated commit frontier (writing_check's MIN
# gates this transition rank-uniformly), so all ranks queue the SAME objects in the SAME
# order. The maintainer drains FIFO + validates each entry (a since-evicted node is dropped).
self._cp_l3_enqueue_spill(object_key, node)
return node
def _remove_undrained_write_acks(self, node_id: int) -> bool:
"""Scrub this node's final write acks out of the controller queue.
Every rollback that clears the radix-side registration
(ongoing_write_through / pending_host_backups) MUST also scrub the
node's acks: an orphaned ack re-opens the _node_host_write_pending
guard for a fresh registration while the stale ack is still queued,
producing two acks for one registration (one inc_lock_ref, two
completions) and crashing writing_check on the second.
"""
retained_acks = []
removed = False
for ack in self.cache_controller.ack_write_queue:
if node_id not in ack.node_ids:
retained_acks.append(ack)
continue
ack.finish_event.synchronize()
remaining_node_ids = [nid for nid in ack.node_ids if nid != node_id]
if remaining_node_ids:
retained_acks.append(ack._replace(node_ids=remaining_node_ids))
removed = True
if removed:
self.cache_controller.ack_write_queue = retained_acks
return removed
def _rollback_pending_backup(self, node_id: int) -> TreeNode:
pending = self.pending_host_backups[node_id]
node = pending.node
# Leave pending BEFORE removing the entry (so a first-touch lazy rebuild
# of the running counts still sees this pending backup); unconditional
# because the entry existed regardless of whether node.cp_hicache
# currently aliases pending.metadata (pending/committed are exclusive).
self._cp_account_leave_pending(pending.metadata)
self.pending_host_backups.pop(node_id)
# Cancel the half-submitted per-layer state (so later layer hooks
# cannot write into the host slots evicted below) and scrub any
# already-appended final ack before releasing the registration.
self.cache_controller.cancel_layer_write_state(node_id)
self._remove_undrained_write_acks(node_id)
self.cache_controller.evict_cp_host(pending.metadata)
if node.cp_hicache is pending.metadata:
node.cp_hicache = None
node.host_len = 0
if pending.locked:
self.dec_node_lock_ref(node)
return node
def _evict_cp_host_for_write_admission(
self, admission: CpWriteAdmission, *, node_id: int, phase: str
) -> bool:
eviction_plan = admission.eviction_plan
# Only the HARD deficit -- what the backup actually needs in order to fit
# -- may block admission. admission.deficit_by_owner (hence
# eviction_plan.remaining_deficit) also folds in the PROACTIVE free-room
# watermark (hicache_host_free_room_ratio): once a lane dips below the
# trigger watermark that eviction target balloons to ~20% of capacity.
# Refusing a backup that physically fits, just because a lane is below
# that advisory target, turned a ~1k-slot write into a 20%-of-capacity
# reclaim demand; when locked residual made the target unreachable it
# skipped the backup AND re-fired every tick on every rank (a 177MB log
# storm with the host pool <42% used, 2026-06-19). Evict toward the
# target best-effort, but admit whenever the write itself fits.
# draft_available is 2**62 when there is no draft pool, so the draft term
# is a no-op in that case.
hard_deficit = [
max(max(0, int(req) - int(tgt)), max(0, int(req) - int(drf)))
for req, tgt, drf in zip(
admission.required_by_owner,
admission.target_available_by_owner,
admission.draft_available_by_owner,
)
]
remaining_hard = [
max(0, hard - int(freed))
for hard, freed in zip(hard_deficit, eviction_plan.planned_freed)
]
if any(v > 0 for v in remaining_hard):
self._warn_cp_hicache_fallback(
"cp_host_reservation_plan_insufficient",
"insufficient_evictable_host_slots",
rate_limit_s=10.0,
node_id=node_id,
phase=phase,
hard_deficit=hard_deficit,
required=admission.required_by_owner,
target_avail=admission.target_available_by_owner,
draft_avail=admission.draft_available_by_owner,
planned_freed=eviction_plan.planned_freed,
remaining_hard=remaining_hard,
victims=len(eviction_plan.victims),
)
return False
if len(eviction_plan.victims) == 0:
return True
local_freed = 0
victim_ids = []
for victim in eviction_plan.victims:
victim_id = getattr(victim, "id", None)
if not self._cp_host_leaf_is_plannable_victim(victim):
raise RuntimeError(
"CP HiCache deterministic eviction plan diverged before "
f"reservation: node_id={node_id} phase={phase} victim_id={victim_id}"
)
victim_ids.append(victim_id)
self._record_remove_event(victim)
self._cp_account_leave_node_evict(victim, victim.cp_hicache)
local_freed += self.cache_controller.evict_cp_host(victim.cp_hicache)
victim.host_len = 0
victim.cp_hicache = None
victim.host_value = None
self._remove_host_leaf(victim)
logger.debug(
"[HiCache-evict] deterministic CP host eviction before write: "
"node_id=%d phase=%s victims=%s local_freed=%d planned_freed=%s",
node_id,
phase,
victim_ids,
local_freed,
eviction_plan.planned_freed,
)
return True
def _reserve_write_cp_shared_l2_evict_to_fit(
self, device_indices: torch.Tensor, node_id: int, last_result
):
"""B1 (2b.2) deterministic collective-free evict-to-fit for the shared L2 pool.
Snapshot the plannable shared-L2 host leaves (excluding the node being backed
up) in the replicated SLRU victim order (``_cp_host_evict_key``), then release
them one at a time -- each release is a replicated free-list mutation -- and
retry the pooled reserve after each. Fragmentation is handled by construction:
the retry succeeds as soon as a contiguous run reopens (free coalesces in
``_return_range``). Identical victim set + order + deterministic reserve =>
every rank evicts the same victims and the same retry succeeds at the same
point (design Thm 1), so NO collective is needed -- this is the replicated
replacement for the #5 ``_evict_host_for_physical_slots(synchronize_across_ranks=True)``
all_reduce. If the evictable set is exhausted and the reserve still cannot
fit, the backup is skipped (transient: live/pinned/referenced objects hold the
pool, or the single backup exceeds pool capacity) -- a loud, rate-limited
warn, never a hang.
The victim snapshot is one-level by design: a parent promoted to a new host
leaf when its child is removed is NOT re-pushed mid-loop (unlike the
heap-based _evict_host_for_physical_slots), so deep cascades skip-and-retry
on the next tick's snapshot rather than this round -- missed identically on
every rank, so determinism (Thm 1) is unaffected.
"""
victims = sorted(
(
node
for node in getattr(self, "evictable_host_leaves", set())
if getattr(node, "id", None) != node_id
and self._cp_host_leaf_is_plannable_victim(node)
and self._is_cp_shared_l2_metadata(getattr(node, "cp_hicache", None))
),
key=self._cp_host_evict_key,
)
evicted_ids = []
for victim in victims:
# Re-check plannability immediately before mutating (fail-loud on any
# divergence from the snapshot -- the determinism guard mirrored from
# _evict_cp_host_for_write_admission).
if not self._cp_host_leaf_is_plannable_victim(victim):
raise RuntimeError(
"CP shared-L2 evict-to-fit plan diverged before reservation: "
f"node_id={node_id} victim_id={getattr(victim, 'id', None)}"
)
# Teardown identical to _evict_cp_host_for_write_admission: evict_cp_host
# releases the shared pool range (-> allocator.release) and drops the
# replicated host leaf. All inputs are in R => all ranks evict in lockstep.
self._record_remove_event(victim)
self._cp_account_leave_node_evict(victim, victim.cp_hicache)
self.cache_controller.evict_cp_host(victim.cp_hicache)
evicted_ids.append(getattr(victim, "id", None))
victim.host_len = 0
victim.cp_hicache = None
victim.host_value = None
self._remove_host_leaf(victim)
result = self.cache_controller.reserve_write_cp(
device_indices=device_indices,
node_id=node_id,
)
if not isinstance(result, HiCacheWriteFailure):
logger.debug(
"[HiCache-evict] CP shared-L2 evict-to-fit succeeded: "
"node_id=%d evicted=%s",
node_id,
evicted_ids,
)
return result
if getattr(result, "reason", "host_capacity") != "shared_l2_capacity":
# A non-capacity failure (e.g. undrained_ack) cannot be helped by
# more eviction.
return result
last_result = result
# Exhausted every evictable shared-L2 leaf and the reserve still does not fit:
# the pool is held by non-evictable (pinned / in-flight / referenced) objects,
# or this single backup exceeds pool capacity. Skip this round (the backup
# retries next tick); loud + rate-limited so a real shortfall stays visible.
self._warn_cp_hicache_fallback(
"cp_shared_l2_evict_exhausted",
"shared_l2_capacity_after_evict",
rate_limit_s=10.0,
node_id=node_id,
required_host_slots=int(getattr(last_result, "required_host_slots", 0)),
evicted=len(evicted_ids),
)
return last_result
def _reserve_write_cp_indices_no_collective(
self,
device_indices: torch.Tensor,
node_id: int,
*,
admission_checked: bool = False,
):
if getattr(
self.cache_controller, "cp_shared_l2_page_allocator", None
) is not None:
result = self.cache_controller.reserve_write_cp(
device_indices=device_indices,
node_id=node_id,
)
if not isinstance(result, HiCacheWriteFailure):
return result
if getattr(result, "reason", "host_capacity") != "shared_l2_capacity":
return result
# B1 (2b.2): a shared-L2 capacity miss raises identically on every CP
# rank (replicated free list), so run the deterministic collective-free
# evict-to-fit -- release the coldest plannable shared-L2 host leaves in
# the replicated SLRU order and retry the pooled reserve. Every rank
# frees the same victims in the same order and the same retry succeeds at
# the same point (placement determinism, design Thm 1) -- NO collective,
# NO _evict_host_for_physical_slots(synchronize_across_ranks=True) (#5).
return self._reserve_write_cp_shared_l2_evict_to_fit(
device_indices, node_id, result
)
if not admission_checked:
admission = self._cp_build_write_admission(
device_indices, node_id=node_id, phase="initial"
)
if any(v > 0 for v in admission.deficit_by_owner):
if not self._evict_cp_host_for_write_admission(
admission, node_id=node_id, phase="initial"
):
return HiCacheWriteFailure(
required_host_slots=max(
admission.eviction_plan.remaining_deficit
)
)
result = self.cache_controller.reserve_write_cp(
device_indices=device_indices,
node_id=node_id,
)
if not isinstance(result, HiCacheWriteFailure):
return result
if getattr(result, "reason", "host_capacity") != "host_capacity":
# Not a capacity failure (e.g. undrained_ack: the node's previous
# backup ack is still queued). Eviction cannot help, and entering
# the retry path would trip the predicted-no-deficit fail-fast
# below. Skip this backup round.
return result
retry_admission = self._cp_build_write_admission(
device_indices, node_id=node_id, phase="retry_after_local_failure"
)
if not any(v > 0 for v in retry_admission.deficit_by_owner) and not any(
v > 0 for v in retry_admission.eviction_plan.remaining_deficit
):
raise RuntimeError(
"CP HiCache host reservation failed although deterministic "
"owner-lane admission predicted no host deficit: "
f"node_id={node_id} required_host_slots={result.required_host_slots}"
)
if not self._evict_cp_host_for_write_admission(
retry_admission, node_id=node_id, phase="retry_after_local_failure"
):
return HiCacheWriteFailure(
required_host_slots=max(
retry_admission.eviction_plan.remaining_deficit,
default=int(result.required_host_slots),
)
)
logger.debug(
"[HiCache-write] write_backup CP retry after deterministic host eviction: "
"node_id=%d deficit_by_owner=%s",
node_id,
retry_admission.deficit_by_owner,
)
result = self.cache_controller.reserve_write_cp(
device_indices=device_indices,
node_id=node_id,
)
if not isinstance(result, HiCacheWriteFailure):
return result
logger.warning(
"[HiCache-write] write_backup CP FAILED after deterministic retry: "
"node_id=%d len=%d needed_slots=%d",
node_id,
len(device_indices) if device_indices is not None else 0,
int(result.required_host_slots),
)
return result
def _reserve_write_cp(self, node: TreeNode):
return self._reserve_write_cp_indices_no_collective(node.value, node.id)
def _attach_prepared_cp_backup(
self, node: TreeNode, prepared: PreparedCpHiCacheBackup
) -> None:
if prepared.attached:
raise RuntimeError(
f"CP HiCache prepared backup node_id={prepared.node_id} was attached twice"
)
if len(node.value) != prepared.logical_len:
raise RuntimeError(
f"Prepared CP HiCache backup length mismatch for node_id={prepared.node_id}: "
f"node_len={len(node.value)} prepared_len={prepared.logical_len}"
)
node.id = prepared.node_id
self.ongoing_write_through[node.id] = node
self.inc_node_lock_ref(node)
self.pending_host_backups[node.id] = PendingHiCacheBackup(
node=node,
metadata=prepared.metadata,
logical_len=prepared.logical_len,
submitted=True,
locked=True,
)
self._cp_account_enter_pending(prepared.metadata)
prepared.attached = True
logger.debug(
"[HiCache-write] attached prepared CP backup: node_id=%d logical_len=%d owned_positions=%d pending_backups=%d",
node.id,
prepared.logical_len,
self._cp_reservation_owned_count(prepared.reservation),
len(self.pending_host_backups),
)
def _rollback_prepared_cp_backup(
self, prepared: Optional[PreparedCpHiCacheBackup], reason: str
) -> None:
if prepared is None or prepared.attached:
return
if prepared.node_id in self.cache_controller.pending_layer_writes:
raise RuntimeError(
f"Cannot rollback prepared CP HiCache backup node_id={prepared.node_id} "
f"while per-layer writes are still pending; reason={reason}"
)
self._remove_undrained_write_acks(prepared.node_id)
logger.debug(
"[HiCache-write] rollback unattached prepared CP backup: node_id=%d logical_len=%d reason=%s",
prepared.node_id,
prepared.logical_len,
reason,
)
self.cache_controller.evict_cp_host(prepared.metadata)
def _warn_cp_hicache_fallback(
self, fallback_name: str, reason: str, *, rate_limit_s: float = 0.0, **kwargs
):
# E1 probe: count every fallback/drop by name BEFORE rate-limiting, so the
# true drop rate is captured even when the warning log is suppressed.
fallback_counts = getattr(self, "_cp_fallback_counts", None)
if fallback_counts is None:
fallback_counts = {}
self._cp_fallback_counts = fallback_counts
fallback_counts[fallback_name] = fallback_counts.get(fallback_name, 0) + 1
details = " ".join(f"{key}={value}" for key, value in kwargs.items())
if details:
details = " " + details
if rate_limit_s > 0.0:
# Host-reservation shortfalls are emitted per backup-attempt per rank;
# under genuine host pressure that is thousands of identical lines a
# minute (a 177MB log storm, 2026-06-19). Collapse repeats of the same
# fallback_name within the window so a real shortfall stays loud but
# bounded -- the first line of each window carries the count suppressed
# since the previous emission.
warn_state = getattr(self, "_cp_warn_state", None)
if warn_state is None:
warn_state = {}
self._cp_warn_state = warn_state
now = time.monotonic()
entry = warn_state.get(fallback_name)
if entry is not None and (now - entry[0]) < rate_limit_s:
entry[1] += 1
return
suppressed = entry[1] if entry is not None else 0
warn_state[fallback_name] = [now, 0]
if suppressed:
details = (
f"{details} suppressed_since_last={suppressed} "
f"window_s={int(rate_limit_s)}"
)
logger.warning(
"[CP_HICACHE_FALLBACK][%s] reason=%s%s",
fallback_name,
reason,
details,
)
def _cp_maybe_collect_metrics(self) -> None:
"""Push L2 (pooled allocator) + L3 (store, if present) usage to Prometheus, rate-limited (~1s, per-rank
wall-clock; metrics need no rank-uniformity). LAZILY creates the collector on first emit, gated on L2
pooling being on (the allocator exists) -- so it covers PURE-L2 (no L3) as well as L2+L3. Off the data
path; cheap (occupancy snapshot + queue sizes + counters)."""
if not self._enable_metrics_flag:
return
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
if allocator is None:
return # L2 pooling not on -> nothing to report
now = time.monotonic()
if now - self._cp_metrics_last_emit < 1.0:
return
self._cp_metrics_last_emit = now
if self.cp_hicache_metrics is None:
from sglang.srt.observability.metrics_collector import CpHiCacheMetricsCollector
cc = self.cache_controller
# cp_rank/cp_size are the identity that actually distinguishes the
# per-rank L2 lanes / L3 slabs (what these metrics measure); tp/dp/pp
# give the full parallel coordinate. cp_* come from the CP layout
# (_cp_hicache_cp_*), the rest from the controller (always set).
labels = {
"cp_rank": self._cp_hicache_cp_rank(),
"cp_size": self._cp_hicache_cp_size(),
"tp_rank": cc.tp_rank, "dp_rank": cc.dp_rank,
"pp_rank": cc.pp_rank, "pp_size": cc.pp_size,
}
if self.extra_metric_labels:
labels.update(self.extra_metric_labels)
self.cp_hicache_metrics = CpHiCacheMetricsCollector(labels)
l2_occ = allocator.occupancy_by_payload() if hasattr(allocator, "occupancy_by_payload") else None
l2_stats = allocator.stats() if hasattr(allocator, "stats") else None
l3_stats = self.cp_l3_store.stats() if self.cp_l3_store is not None else None
self.cp_hicache_metrics.collect(l2_occupancy=l2_occ, l2_stats=l2_stats, l3_stats=l3_stats)
def _cp_maybe_log_lane_stats(self) -> None:
"""E1 lane-imbalance probe (read-only, measurement only).
Every ``SGLANG_CP_HICACHE_LANE_STATS_S`` seconds, log the per-owner-lane L2
(host) occupancy vector, the cross-lane imbalance (max/min ratio + spread),
a hot-prefix length histogram (does the hot working set skew owners?), and
the backup-drop counters accumulated since the last emit. The host capacity
snapshot is derived from the REPLICATED radix tree, so every CP rank
computes the identical 8-lane vector; emit only on cp_rank 0 to avoid a
cp_size-fold of duplicate lines. Tests the load-bearing Phase-2 premise:
is per-rank L2 saturation (lane imbalance) real, or just global pressure?
"""
interval = float(envs.SGLANG_CP_HICACHE_LANE_STATS_S.get() or 0.0)
if interval <= 0.0 or not self._uses_cp_hicache:
return
if self._cp_hicache_cp_rank() != 0:
return
now = time.monotonic()
last = getattr(self, "_cp_lane_stats_last_emit", 0.0)
if last and (now - last) < interval:
return
self._cp_lane_stats_last_emit = now
# Real cp_size (owner round-robin period for the hot-prefix skew below) --
# decoupled from the occ-vector length, which is per-payload when pooling.
cp_size = self._cp_hicache_cp_size()
controller = getattr(self, "cache_controller", None)
allocator = (
getattr(controller, "cp_shared_l2_page_allocator", None)
if controller is not None
else None
)
if allocator is not None and hasattr(allocator, "occupancy_by_payload"):
# l2 pooling ON: report the shared pool's per-payload-slab fill. The
# per-rank snapshot (_cp_counts -> target_used) is 0 here because the
# accounting skips shared-L2 metadata (_cp_account_*: `not
# _is_cp_shared_l2_metadata`), so the per-rank gauge is meaningless --
# repoint occ/used/cap at the actual pool.
by_payload = allocator.occupancy_by_payload()
occ_labels = list(by_payload.keys())
used = [by_payload[p][0] for p in occ_labels]
capacity = [by_payload[p][1] for p in occ_labels]
occ_by = "payload"
else:
snapshot = self._cp_host_capacity_snapshot()
capacity = list(snapshot.target_capacity)
used = list(snapshot.target_used)
occ_labels = [f"lane{i}" for i in range(len(capacity))]
occ_by = "rank"
occ = [
(used[i] / capacity[i]) if capacity[i] > 0 else 0.0
for i in range(len(capacity))
]
max_occ = max(occ) if occ else 0.0
min_occ = min(occ) if occ else 0.0
spread = max_occ - min_occ
imbalance = (max_occ / min_occ) if min_occ > 1e-9 else float("inf")
# Hot-prefix length histogram: bucket reused (hit_count>=2) backed-up nodes
# at the owner round-robin period (cp_size * page_size tokens). A prefix
# shorter than one period cannot span all lanes -> it skews ownership; a
# working set dominated by such short hot prefixes is what makes per-rank
# pools imbalance. Long prefixes round-robin evenly and do not skew.
skew_token_threshold = cp_size * self.page_size
hot_short = 0
hot_long = 0
for node in self._cp_walk_radix_nodes():
if int(getattr(node, "host_len", 0) or 0) <= 0:
continue
if int(getattr(node, "hit_count", 0) or 0) < 2:
continue
if int(getattr(node, "host_len", 0)) < skew_token_threshold:
hot_short += 1
else:
hot_long += 1
fallback_counts = getattr(self, "_cp_fallback_counts", {})
last_counts = getattr(self, "_cp_fallback_counts_last", {})
drops_since_last = {
name: fallback_counts[name] - last_counts.get(name, 0)
for name in fallback_counts
if fallback_counts[name] - last_counts.get(name, 0) > 0
}
self._cp_fallback_counts_last = dict(fallback_counts)
# Shared physical L2 pool report -- the REAL occupancy when l2 pooling is on.
# occ/used_tokens above are derived from _cp_counts, which the accounting
# explicitly skips for shared-L2 metadata (_cp_account_*: `not
# _is_cp_shared_l2_metadata`), so they read 0 even when the pool is full. The
# shared pool tracks its own pages_used/objects_committed/objects_evicted +
# load_hits/misses in the replicated CpSharedL2PageAllocator stats.
shared_l2_report = None
if allocator is not None and controller is not None:
report_fn = getattr(controller, "cp_shared_l2_cache_report", None)
if report_fn is not None:
try:
shared_l2_report = report_fn()
except Exception:
shared_l2_report = None
logger.info(
"[CP_HICACHE_LANE_STATS] occ_by=%s occ=%s occ_labels=%s cap=%s "
"used=%s max=%.3f "
"min=%.3f spread=%.3f imbalance=%s hot_prefix_short=%d hot_prefix_long=%d "
"skew_token_threshold=%d drops_since_last=%s drops_total=%s shared_l2_pool=%s",
occ_by,
[round(x, 3) for x in occ],
occ_labels,
list(capacity),
list(used),
max_occ,
min_occ,
spread,
("inf" if imbalance == float("inf") else round(imbalance, 2)),
hot_short,
hot_long,
skew_token_threshold,
drops_since_last,
dict(fallback_counts),
shared_l2_report,
)
def _probe_existing_radix_prefix_len_no_split(self, key: RadixKey) -> int:
"""Return the already-present radix prefix length without mutating the tree.
`prepare_write_backup_for_req` runs before forward and must reserve host
slots only for KV that will become a *new* radix node at insertion time.
The request's `cache_protected_len` is based on scheduler prefix-match
semantics, which can be shorter than the prefix that a final insertion
will find (for example EAGLE/bigram plus page alignment). Probing the
full insertion key here avoids backing up duplicate pages that insertion
will later reject and roll back.
"""
if self.disable or len(key) == 0:
return 0
key, _ = self.maybe_bigram_convert(key)
if len(key) == 0:
return 0
node = getattr(self, "root_node", None)
if node is None:
return 0
child_key = self.get_child_key_fn(key)
total_prefix_len = 0
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
prefix_len = self.key_match_fn(child.key, key)
prefix_len = self._cp_floor_exact_valid_tail_extension_len(
child, prefix_len, len(key), floor_exact_key=True
)
if prefix_len <= 0:
break
if prefix_len < len(child.key):
prefix_len = self._cp_floor_backed_partial_split_len(
child, prefix_len
)
if prefix_len <= 0:
break
total_prefix_len += prefix_len
if prefix_len < len(child.key):
break
node = child
key = key[prefix_len:]
if len(key):
child_key = self.get_child_key_fn(key)
return total_prefix_len
def _cp_floor_backed_partial_split_len(
self, child: TreeNode, prefix_len: int
) -> int:
if (
not self._uses_cp_hicache
or self.page_size <= 1
or prefix_len <= 0
or prefix_len % self.page_size == 0
):
return prefix_len
return prefix_len // self.page_size * self.page_size
def _cp_floor_exact_valid_tail_extension_len(
self,
child: TreeNode,
prefix_len: int,
request_len: int,
*,
floor_exact_key: bool = False,
) -> int:
"""Floor exact CP valid-tail hits to the previous physical page.
CP HiCache may keep radix keys at scheduler-visible valid lengths while
the underlying target/draft pools and host reservations are page-owned.
Even when the incoming radix key exactly equals a non-page-aligned child
key, the scheduler will still compute the current suffix token
(`max_prefix_len = input_len - 1`). Exposing the sub-page tail as a
protected prefix lets later backup/current-reuse paths start inside a
page. Sacrifice that tail and let the new request re-own it.
`floor_exact_key` is enabled for scheduler-visible prefix matching and
prepared-backup probing. Internal cache insertion refreshes can keep
exact sub-page tails for the current request so they do not immediately
invalidate their own just-inserted prefix.
"""
if (
not self._uses_cp_hicache
or self.page_size <= 1
or prefix_len <= 0
or prefix_len != len(child.key)
or (prefix_len >= request_len and not floor_exact_key)
or prefix_len % self.page_size == 0
):
return prefix_len
return floor_to_page_len(prefix_len, self.page_size)
def _cp_subtree_has_unprunable_state(self, node: TreeNode) -> bool:
stack = [node]
while stack:
current = stack.pop()
if (
current.lock_ref > 0
or current.host_ref_counter > 0
or self._node_host_write_pending(current)
):
return True
stack.extend(current.children.values())
return False
def _cp_drain_write_acks_before_pending_split(self) -> None:
"""Do not drain CP write acks from the radix pending-split path.
This helper is reached from radix insert, where CP ranks can legitimately
diverge: one rank may be checking a node-pending-backup split while
another is checking stale-tail pruning or has already returned to the
scheduler event loop. `writing_check()` may issue the final
write-visibility TP collective when a local ack is ready. Calling it
here can therefore interleave a scalar write-ack collective with the
scheduler's inflight-poll collective on other ranks and abort Gloo with
a message-size mismatch.
Keep this path local and conservative: report the split as still blocked
and let the globally ordered scheduler/cache-event path drain write acks.
"""
return
def _cp_node_split_still_pending(self, node: TreeNode) -> bool:
if not self._uses_cp_hicache or not self._node_host_write_pending(node):
return False
self._cp_drain_write_acks_before_pending_split()
return self._node_host_write_pending(node)
def _cp_stale_tail_prune_still_blocked(self, stale_tail: TreeNode) -> bool:
if not self._uses_cp_hicache or not self._cp_subtree_has_unprunable_state(
stale_tail
):
return False
self._cp_drain_write_acks_before_pending_split()
return self._cp_subtree_has_unprunable_state(stale_tail)
def _cp_defer_insert_for_pending_backup_split(
self, node: TreeNode, prefix_len: int, *, reason: str
) -> InsertResult:
self._warn_cp_hicache_fallback(
"insert_deferred_pending_backup_split",
reason,
node_id=getattr(node, "id", None),
prefix_len=prefix_len,
ongoing=len(getattr(self, "ongoing_write_through", {})),
pending=len(getattr(self, "pending_host_backups", {})),
)
return InsertResult(
prefix_len=prefix_len,
pending_backup_deferred_node=node,
)
def _cp_prune_stale_tail_after_page_floor(
self, parent: TreeNode, stale_tail: TreeNode
) -> None:
"""Drop a sub-page CP tail after flooring an extending request to a page.
CP HiCache owns KV at page granularity. If a cached key ends inside a
page and a later request extends past that valid tail, the old sub-page
child must not remain as an independent radix child beside the new
page-boundary suffix. Keeping both creates overlapping children for the
same physical page. We sacrifice the old tail page and let the new
request re-own that page from the boundary.
"""
if not self._uses_cp_hicache:
return
if stale_tail.parent is not parent:
raise RuntimeError(
"CP HiCache stale-tail prune saw a detached split child: "
f"parent_id={getattr(parent, 'id', None)} "
f"tail_id={getattr(stale_tail, 'id', None)}"
)
if self._cp_stale_tail_prune_still_blocked(stale_tail):
raise HiCachePendingBackupSplit(stale_tail)
self._cp_delete_stale_tail_subtree(stale_tail)
def _cp_delete_stale_tail_subtree(self, node: TreeNode) -> None:
for child in list(node.children.values()):
self._cp_delete_stale_tail_subtree(child)
parent = node.parent
if parent is None:
raise RuntimeError(
"CP HiCache stale-tail prune attempted to delete a root node"
)
device_resident_len = self._node_device_resident_len(node)
if hasattr(self, "evictable_leaves") and node in self.evictable_leaves:
self.evictable_leaves.remove(node)
if hasattr(self, "evictable_size_"):
self.evictable_size_ -= device_resident_len
if hasattr(self, "evictable_host_leaves") and node in self.evictable_host_leaves:
self.evictable_host_leaves.remove(node)
if node.value is not None:
self._record_remove_event(node)
cache_controller = getattr(self, "cache_controller", None)
if cache_controller is not None and hasattr(cache_controller, "evict_device"):
cache_controller.evict_device(node.value)
elif cache_controller is not None and hasattr(
cache_controller, "mem_pool_device_allocator"
):
cache_controller.mem_pool_device_allocator.free(node.value)
node.value = None
if node.cp_hicache is not None:
self._cp_account_leave_node_evict(node, node.cp_hicache)
cache_controller = getattr(self, "cache_controller", None)
if cache_controller is not None and hasattr(cache_controller, "evict_cp_host"):
cache_controller.evict_cp_host(node.cp_hicache)
node.cp_hicache = None
node.host_len = 0
elif node.host_value is not None:
cache_controller = getattr(self, "cache_controller", None)
if cache_controller is not None and hasattr(cache_controller, "evict_host"):
cache_controller.evict_host(node.host_value)
node.host_value = None
node.host_len = 0
key = self.get_child_key_fn(node.key)
removed = parent.children.pop(key, None)
if removed is not node:
raise RuntimeError(
"CP HiCache stale-tail prune removed unexpected child: "
f"key={key} node_id={getattr(node, 'id', None)}"
)
if hasattr(self, "evictable_leaves"):
self._update_leaf_status(parent)
if hasattr(self, "evictable_host_leaves"):
self._update_host_leaf_status(parent)
def _build_prepare_write_backup_candidate_for_req(
self, req
) -> Optional[CpWriteBackupCandidate]:
if getattr(req, "cp_hicache_prepared_backup", None) is not None:
prepared = getattr(req, "cp_hicache_prepared_backup", None)
self._warn_cp_hicache_fallback(
"prepare_write_backup_skipped",
"already_prepared",
rid=getattr(req, "rid", "<unknown>"),
node_id=getattr(prepared, "node_id", None),
)
return
token_ids = req.fill_ids
keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
backup_end = len(keys)
if getattr(req, "cp_backup_is_intermediate_chunk", False) and self.page_size > 1:
# Only a genuine INTERMEDIATE chunk floors off its tail page: a later
# chunk of the same request will extend it, so backing up that still
# incomplete page now would create a scheduler-visible stale tail the
# next chunk has to prune while this request still owns/locks it. The
# FINAL chunk's tail page is complete and must be backed up in full,
# to match the radix node that the final (non-chunked) insertion
# builds -- otherwise prepared logical_len != len(value) and the
# exact-equality attach predicate drops the overlap backup.
#
# NOTE: we must NOT key this off `req.is_chunked`. It is a per-tick
# counter whose decrement lags one iteration under overlap scheduling
# and so reads stale (>0) on the final chunk. The scheduler sets
# `cp_backup_is_intermediate_chunk` from the live `chunked_req`
# identity in `_prepare_hicache_write_backups_before_forward`, which
# runs in the same run_batch immediately before this.
backup_end = floor_to_page_len(backup_end, self.page_size)
if backup_end <= 0:
logger.debug(
"[HiCache-write] skip chunked CP backup with no completed page: rid=%s key_len=%d",
getattr(req, "rid", "<unknown>"),
len(keys),
)
return
keys_for_probe = keys[:backup_end]
else:
keys_for_probe = keys
existing_prefix_len = self._probe_existing_radix_prefix_len_no_split(
RadixKey(
keys_for_probe,
getattr(req, "extra_key", None),
is_bigram=self.is_eagle,
)
)
raw_start = min(max(req.cache_protected_len, existing_prefix_len), backup_end)
if backup_end <= raw_start:
return
start = floor_to_page_len(raw_start, self.page_size)
if backup_end <= start:
return
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, start:backup_end
].to(dtype=torch.int64, copy=True)
if len(kv_indices) == 0:
self._warn_cp_hicache_fallback(
"prepare_write_backup_skipped",
"empty_kv_indices",
rid=getattr(req, "rid", "<unknown>"),
start=start,
key_len=backup_end,
)
return
node_id = TreeNode.counter
TreeNode.counter += 1
return CpWriteBackupCandidate(
req=req,
kv_indices=kv_indices,
node_id=node_id,
)
def _submit_prepare_write_backup_candidate(
self,
candidate: CpWriteBackupCandidate,
*,
admission_checked: bool = False,
) -> None:
req = candidate.req
kv_indices = candidate.kv_indices
node_id = candidate.node_id
if admission_checked:
result = self._reserve_write_cp_indices_no_collective(
kv_indices,
node_id,
admission_checked=True,
)
else:
result = self._reserve_write_cp_indices_no_collective(
kv_indices,
node_id,
)
if isinstance(result, HiCacheWriteFailure):
self._warn_cp_hicache_fallback(
"prepare_write_backup_reservation_failed",
"host_reservation_failed",
rate_limit_s=10.0,
node_id=node_id,
rid=getattr(req, "rid", "<unknown>"),
logical_len=len(kv_indices),
required_host_slots=result.required_host_slots,
)
return
try:
self.cache_controller.submit_write_cp_per_layer(
result, catch_up_all_layers=False
)
except Exception:
self.cache_controller.evict_cp_host(result.metadata)
raise
req.cp_hicache_prepared_backup = PreparedCpHiCacheBackup(
node_id=node_id,
reservation=result,
metadata=result.metadata,
logical_len=len(kv_indices),
)
logger.debug(
"[HiCache-write] prepared CP per-layer backup before forward: node_id=%d rid=%s logical_len=%d owned_positions=%d",
node_id,
getattr(req, "rid", "<unknown>"),
len(kv_indices),
self._cp_reservation_owned_count(result),
)
def prepare_write_backup_for_req(self, req) -> None:
if self.disable or not self._uses_cp_hicache:
return
if self.cache_controller.write_policy == "write_back":
self._warn_cp_hicache_fallback(
"prepare_write_backup_skipped",
"write_back_policy",
rid=getattr(req, "rid", "<unknown>"),
)
return
candidate = self._build_prepare_write_backup_candidate_for_req(req)
if candidate is None:
return
self._submit_prepare_write_backup_candidate(candidate)
def prepare_write_backups_for_reqs(self, reqs) -> None:
if self.disable or not self._uses_cp_hicache:
return
reqs = list(reqs)
if len(reqs) == 0:
return
if self.cache_controller.write_policy == "write_back":
self._warn_cp_hicache_fallback(
"prepare_write_backups_skipped",
"write_back_policy",
batch_size=len(reqs),
)
return
candidates: List[CpWriteBackupCandidate] = []
for req in reqs:
candidate = self._build_prepare_write_backup_candidate_for_req(req)
if candidate is not None:
candidates.append(candidate)
if len(candidates) == 0:
return
admission_checked = False
shared_l2_write = (
getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
is not None
)
if len(candidates) > 1 and not shared_l2_write:
required = self._cp_zero_counts()
for candidate in candidates:
required = self._cp_add_counts(
required,
self._cp_required_host_tokens_by_rank(candidate.kv_indices),
)
admission = self._cp_build_write_admission_from_required(
required,
node_id=candidates[0].node_id,
phase="batch_prepare",
)
if any(v > 0 for v in admission.deficit_by_owner):
if self._evict_cp_host_for_write_admission(
admission,
node_id=candidates[0].node_id,
phase="batch_prepare",
):
admission_checked = True
else:
self._warn_cp_hicache_fallback(
"prepare_write_backups_batch_admission_failed",
"host_reservation_failed",
rate_limit_s=10.0,
batch_size=len(candidates),
required_by_owner=required,
remaining_deficit=admission.eviction_plan.remaining_deficit,
)
else:
admission_checked = True
for candidate in candidates:
self._submit_prepare_write_backup_candidate(
candidate,
admission_checked=admission_checked,
)
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:
metadata = node.cp_hicache
if self._is_cp_shared_l2_metadata(metadata):
return torch.empty((0,), dtype=torch.int64)
return metadata.host_indices
return node.host_value
def write_backup(self, node: TreeNode, write_back=False):
logger.debug(
"[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:
if not write_back:
self._warn_cp_hicache_fallback(
"post_forward_catch_up_backup",
"write_backup_called_without_prepared_cp_backup",
node_id=node.id,
logical_len=len(node.value) if node.value is not None else 0,
)
result = self._reserve_write_cp(node)
if isinstance(result, HiCacheWriteFailure):
self._warn_cp_hicache_fallback(
"post_forward_catch_up_backup_failed",
"host_reservation_failed",
rate_limit_s=10.0,
node_id=node.id,
logical_len=len(node.value) if node.value is not None else 0,
required_host_slots=result.required_host_slots,
)
return 0
self.ongoing_write_through[node.id] = node
if not write_back:
self.inc_node_lock_ref(node)
self.pending_host_backups[node.id] = PendingHiCacheBackup(
node=node,
metadata=result.metadata,
logical_len=len(node.value),
submitted=True,
locked=not write_back,
)
# Enter pending AFTER the dict insert and BEFORE the submit try:
# if submit raises, the except calls _rollback_pending_backup which
# does the matching leave_pending (net zero, no double-sub).
self._cp_account_enter_pending(result.metadata)
try:
self.cache_controller.submit_write_cp_per_layer(result)
except Exception:
self.ongoing_write_through.pop(node.id, None)
self._rollback_pending_backup(node.id)
raise
logger.debug(
"[HiCache-write] write_backup CP SUBMITTED: node_id=%d logical_len=%d owned_positions=%d ongoing_writes=%d pending_backups=%d",
node.id,
len(node.value),
self._cp_reservation_owned_count(result),
len(self.ongoing_write_through),
len(self.pending_host_backups),
)
return len(node.value)
host_indices = self.cache_controller.write(
device_indices=node.value,
node_id=node.id,
)
if host_indices is None:
logger.debug(
"[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.debug(
"[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.debug(
"[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 self._uses_cp_hicache and self._node_host_write_pending(node):
return
if not self._node_backuped(node):
if node.hit_count >= self.write_through_threshold:
logger.debug(
"[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):
completed_write_ack_ids = getattr(self, "_completed_write_ack_ids", None)
if completed_write_ack_ids is None:
completed_write_ack_ids = set()
self._completed_write_ack_ids = completed_write_ack_ids
def complete_write_ack_node(ack_id: int) -> Optional[TreeNode]:
if ack_id in self.pending_host_backups:
backuped_node = self._commit_pending_backup(ack_id)
self.ongoing_write_through.pop(ack_id, None)
completed_write_ack_ids.add(ack_id)
return backuped_node
backuped_node = self.ongoing_write_through.pop(ack_id, None)
if backuped_node is None:
if ack_id not in completed_write_ack_ids:
raise KeyError(ack_id)
# A duplicated ready ack can pass the pre-scan while the first
# occurrence still owns radix state, then become already-consumed
# by the time this loop reaches the second occurrence. Keep
# genuinely unattached acks deferred in the pre-scan, but make
# duplicate completion idempotent instead of crashing scheduler.
logger.warning(
"[HiCache-write] writing_check skipped already-consumed write ack: "
"ack_id=%d pending=%s ongoing=%s",
ack_id,
list(self.pending_host_backups.keys()),
list(self.ongoing_write_through.keys()),
)
return None
self.dec_node_lock_ref(backuped_node)
completed_write_ack_ids.add(ack_id)
return backuped_node
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 = complete_write_ack_node(ack_id)
if backuped_node is not None and 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)
# With per-layer CP backup, ongoing_write_through is populated before
# the final write ack is appended. Polling a TP all-reduce while the
# ack queue is empty is pure scheduler overhead: there is no candidate
# radix-state transition to make yet. The ack itself is appended at a
# deterministic layer-end / post-forward point across CP ranks; once it
# exists, the MIN below still gates host visibility on all ranks having
# completed the corresponding local transfer.
if ack_queue_len == 0:
return
def ack_registered_for_radix_state(ack_id: int) -> bool:
return (
ack_id in self.pending_host_backups
or ack_id in self.ongoing_write_through
or ack_id in completed_write_ack_ids
)
finish_count = 0
for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
if not finish_event.query():
break
if not all(ack_registered_for_radix_state(ack_id) for ack_id in ack_list):
# Prepared per-layer CP backups can finish and enqueue their final
# ack before the corresponding radix node is attached. A later
# catch-up/write-back node may make ongoing_write_through non-empty
# and enter writing_check while that prepared ack is still
# unattached. Keep the ack queued until insert either attaches the
# prepared backup or rolls it back.
logger.debug(
"[HiCache-write] writing_check defers unattached ack ids: "
"ack_ids=%s ongoing=%s pending=%s",
ack_list,
list(self.ongoing_write_through.keys()),
list(self.pending_host_backups.keys()),
)
break
finish_count += 1
local_finish_count = finish_count
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
self._cp_hicache_all_reduce(
queue_size,
op=torch.distributed.ReduceOp.MIN,
group=self.tp_group,
tag="writing_check_min",
)
finish_count = int(queue_size.item())
logger.debug(
"[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,
local_finish_count,
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 = complete_write_ack_node(ack_id)
if backuped_node is None:
continue
released_nodes.append(ack_id)
if self.enable_storage:
self.write_backup_storage(backuped_node)
finish_count -= 1
if completed_write_ack_ids:
remaining_ack_ids = set()
for ack in self.cache_controller.ack_write_queue:
remaining_ack_ids.update(ack.node_ids)
completed_write_ack_ids.intersection_update(remaining_ack_ids)
if released_nodes:
logger.debug(
"[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 _node_device_resident_len(self, node: TreeNode) -> int:
"""Allocator-visible device residency for a radix node.
CP HiCache radix keys are valid-token lengths, but CP device KV
ownership is page-granular. Residency accounting therefore uses the
physical padded page span whenever CP HiCache is active, while normal
HiCache keeps historical token-count accounting.
"""
value = getattr(node, "value", None)
if value is None:
return 0
if not getattr(self, "_uses_cp_hicache", False):
return len(value)
metadata = getattr(node, "cp_hicache", None)
padded_len = getattr(metadata, "padded_len", None)
if padded_len is not None:
return int(padded_len)
return ceil_to_page_len(len(value), getattr(self, "page_size", 1))
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._uses_cp_hicache:
# Forbidden under CP shared-KV: pin_expiry is a per-rank wall-clock TTL
# and the eviction victim-eligibility check (time.monotonic() <=
# pin_expiry) would then diverge across CP ranks → non-replicated
# victim set → collective-free eviction breaks. SLRU scan-resistance
# already protects the hot working set; re-enable only with a
# replicated logical TTL.
return (0, "pin_prefix is not supported under CP shared-KV HiCache")
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:
resident_len = self._node_device_resident_len(node)
if node.lock_ref == 0:
self.evictable_size_ -= resident_len
self.protected_size_ += resident_len
delta -= resident_len
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:
resident_len = self._node_device_resident_len(node)
if node.lock_ref == 1:
self.evictable_size_ += resident_len
self.protected_size_ -= resident_len
delta += resident_len
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 inc_node_lock_ref(self, node: TreeNode):
if self.disable:
return
if node == self.root_node:
return
resident_len = self._node_device_resident_len(node)
if node.lock_ref == 0:
self.evictable_size_ -= resident_len
self.protected_size_ += resident_len
node.lock_ref += 1
self._update_leaf_status(node)
if hasattr(self, "evictable_host_leaves"):
self._update_host_leaf_status(node)
def dec_node_lock_ref(self, node: TreeNode):
if self.disable:
return
if node == self.root_node:
return
resident_len = self._node_device_resident_len(node)
if node.lock_ref == 1:
self.evictable_size_ += resident_len
self.protected_size_ -= resident_len
node.lock_ref -= 1
self._update_leaf_status(node)
if hasattr(self, "evictable_host_leaves"):
self._update_host_leaf_status(node)
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
if params.owner_lane_deficits is not None and any(
int(v) > 0 for v in params.owner_lane_deficits
):
result = self._evict_cp_owner_lane_deficit_nodes(params)
self.update_eviction_metrics(result.num_tokens_evicted, start_time)
return result
leaves = list(self.evictable_leaves)
eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves
]
heapq.heapify(eviction_heap)
logger.debug(
"[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.debug(
"[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
device_resident_len = self._node_device_resident_len(node)
freed_len = self.cache_controller.evict_device(node.value)
assert freed_len > 0
self.evictable_size_ -= device_resident_len
logger.debug(
"[HiCache-evict] _evict_backuped: node_id=%d num_evicted=%d physical_tokens=%d lock_ref=%d backed=%s",
node.id,
freed_len,
device_resident_len,
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 device_resident_len
def _evict_regular(self, node: TreeNode):
# evict a node not initiated write to host -- emit BlockRemoved
num_evicted = self._node_device_resident_len(node)
logger.debug(
"[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 _delete_leaf(self, node: TreeNode):
key = self.get_child_key_fn(node.key)
v = node.parent.children.pop(key, None)
assert v == node, f"parent does not have child key, {key}"
self.evictable_size_ -= self._node_device_resident_len(node)
if node in self.evictable_leaves:
self.evictable_leaves.remove(node)
self._update_leaf_status(node.parent)
if hasattr(self, "evictable_host_leaves"):
self._update_host_leaf_status(node.parent)
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) -> int:
# B1 (2b.2): the #5 synchronize_across_ranks all_reduce (`while len(heap) and
# not all_ranks_done()` gated on a ReduceOp.MIN host_evict_done_min) is
# DELETED. It was never reached (no caller passed synchronize_across_ranks=True)
# and was the headline CP-deadlock shape (a collective gated on per-rank
# len(heap)). The shared-L2 capacity path now uses the deterministic
# replicated _reserve_write_cp_shared_l2_evict_to_fit instead.
if required_host_slots <= 0:
return 0
leaves = list(self.evictable_host_leaves)
logger.debug(
"[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d leaves=%d",
required_host_slots,
len(leaves),
)
eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves
]
heapq.heapify(eviction_heap)
num_evicted = 0
while len(eviction_heap) and num_evicted < required_host_slots:
_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
is_shared_l2_metadata = self._is_cp_shared_l2_metadata(x.cp_hicache)
self._record_remove_event(x)
if (
self._uses_cp_hicache
and is_shared_l2_metadata
and hasattr(self.cache_controller, "evict_cp_host")
):
num_evicted += self.cache_controller.evict_cp_host(x.cp_hicache)
elif 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)
# Account before nulling (helper no-ops when not CP / metadata None);
# placed outside the physical_count gate so any committed teardown
# here is subtracted from the running counts.
self._cp_account_leave_node_evict(x, x.cp_hicache)
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()
stage_start_time = start_time
stage_durations_ms: List[Tuple[str, float]] = []
def record_stage(stage: str) -> None:
nonlocal stage_start_time
now = time.perf_counter()
stage_durations_ms.append((stage, (now - stage_start_time) * 1000.0))
stage_start_time = now
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
if len(nodes_to_load) == 0:
self.dec_lock_ref(ancester_node)
return None
# load it all or not at all. The scalar length remains only a
# coarse upper bound; final CP admission is the owner-lane vector
# check below.
try:
load_back_plan = self._build_cp_load_back_plan(
nodes_to_load, node_id=last_hit_node.id
)
record_stage("build_plan")
except Exception:
self.dec_lock_ref(ancester_node)
raise
host_hit_len = load_back_plan.host_hit_len
logger.debug(
"[HiCache-load] load_back CP: node_id=%d nodes_to_load=%d "
"host_hit_len=%d threshold=%d required_by_owner=%s "
"available_by_owner=%s deficit_by_owner=%s "
"free_room_deficit_by_owner=%s",
last_hit_node.id,
len(nodes_to_load),
host_hit_len,
self.load_back_threshold,
load_back_plan.required_by_owner,
load_back_plan.available_by_owner,
load_back_plan.deficit_by_owner,
load_back_plan.free_room_deficit_by_owner,
)
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.debug(
"[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
if any(v > 0 for v in load_back_plan.deficit_by_owner):
load_back_plan = self._evict_cp_load_back_owner_lanes(
load_back_plan, node_id=last_hit_node.id
)
record_stage("owner_lane_evict")
if any(v > 0 for v in load_back_plan.free_room_deficit_by_owner):
# L1 free-room is an advisory signal for future extend batches,
# not a reason to put synchronous eviction on the CP load-back
# hot path when exact owner-lane admission already succeeds.
# Exact deficits are handled above; proactive/free-room refill
# should happen outside this latency-sensitive L2->L1 path.
record_stage("owner_lane_free_room_advisory")
if any(v > 0 for v in load_back_plan.deficit_by_owner):
self.dec_lock_ref(ancester_node)
logger.warning(
"[CP_HICACHE_FALLBACK][cp_load_back_owner_lane_capacity_failed] "
"node_id=%d host_hit_len=%d nodes_to_load=%d "
"required_by_owner=%s available_by_owner=%s "
"deficit_by_owner=%s free_room_deficit_by_owner=%s "
"evictable_size=%d protected_size=%d "
"ongoing_load_back_count=%d allocator_state=%s",
last_hit_node.id,
host_hit_len,
len(nodes_to_load),
load_back_plan.required_by_owner,
load_back_plan.available_by_owner,
load_back_plan.deficit_by_owner,
load_back_plan.free_room_deficit_by_owner,
int(getattr(self, "evictable_size_", 0)),
int(getattr(self, "protected_size_", 0)),
len(getattr(self, "ongoing_load_back", {})),
self.token_to_kv_pool_allocator.allocator_state_str(),
)
return None
device_indices = self.cache_controller.load_cp(
nodes_to_load, node_id=last_hit_node.id
)
record_stage("load_cp_plan")
if device_indices is None:
failed_plan = self._refresh_cp_load_back_plan(load_back_plan)
logger.warning(
"[CP_HICACHE_FALLBACK][cp_load_back_preflight_mismatch] "
"node_id=%d host_hit_len=%d required_by_owner=%s "
"available_by_owner=%s deficit_by_owner=%s "
"free_room_deficit_by_owner=%s "
"allocator_state=%s",
last_hit_node.id,
host_hit_len,
failed_plan.required_by_owner,
failed_plan.available_by_owner,
failed_plan.deficit_by_owner,
failed_plan.free_room_deficit_by_owner,
self.token_to_kv_pool_allocator.allocator_state_str(),
)
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
physical_loaded_len = 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
metadata = getattr(loaded_node, "cp_hicache", None)
physical_loaded_len += int(
getattr(metadata, "padded_len", host_len)
)
record_stage("assign_tree")
self.evictable_size_ += physical_loaded_len
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.debug(
"[HiCache-load] load_back CP SUCCESS: node_id=%d loaded_tokens=%d physical_tokens=%d",
last_hit_node.id,
len(device_indices),
physical_loaded_len,
)
total_duration = time.perf_counter() - start_time
if total_duration >= 1.0:
logger.warning(
"[HiCache-load] slow CP load_back planning: node_id=%d "
"duration_ms=%.3f host_hit_len=%d required_by_owner=%s "
"available_by_owner=%s exact_deficit_by_owner=%s "
"free_room_deficit_by_owner=%s stages_ms=%s",
last_hit_node.id,
total_duration * 1000.0,
host_hit_len,
load_back_plan.required_by_owner,
load_back_plan.available_by_owner,
load_back_plan.deficit_by_owner,
load_back_plan.free_room_deficit_by_owner,
[(stage, round(ms, 3)) for stage, ms in stage_durations_ms],
)
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.debug(
"[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.debug(
"[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.debug(
"[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.debug(
"[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.debug(
"loading back %d tokens for node %d",
len(loading_values),
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.debug(
"[HiCache-write] flush_write_through_acks: before=%d after=%d",
ongoing_before,
ongoing_after,
)
def _cp_assert_placement_replicated(self) -> None:
"""B1 proof-obligation RUNTIME gate (Theorem 1). When
SGLANG_CP_HICACHE_PLACEMENT_ASSERT is on, MIN/MAX-reduce the pooled-L2
allocator's placement_digest across the CP TP group; any cross-rank
divergence (a non-replicated reserve/release/commit) makes MIN != MAX ->
fail loud. Entry is gated ONLY on the global flag + tp_world_size + the
flag-implied allocator (all rank-uniform), never on per-rank state, so the
all_reduce can never be entered divergently (opus SF2). Off in prod.
"""
if not envs.SGLANG_CP_HICACHE_PLACEMENT_ASSERT.get():
return
if self.tp_world_size <= 1:
return
allocator = getattr(self.cache_controller, "cp_shared_l2_page_allocator", None)
if allocator is None:
return
if not getattr(self, "_placement_assert_warned", False):
self._placement_assert_warned = True
logger.warning(
"[CP-HiCache] SGLANG_CP_HICACHE_PLACEMENT_ASSERT is ON: every scheduler tick MIN/MAX-reduces "
"placement_digest across the CP group (2 blocking gloo collectives on the scheduler thread "
"BEFORE forward) + a SHA256 over the whole free-list. This is a DEBUG/bring-up invariant check "
"-- turn it OFF in production (unset the env) for the per-tick critical-path cost to vanish."
)
local = int(allocator.placement_digest(), 16) % (1 << 61)
lo = torch.tensor(local, dtype=torch.int64, device="cpu")
hi = torch.tensor(local, dtype=torch.int64, device="cpu")
self._cp_hicache_all_reduce(
lo, op=torch.distributed.ReduceOp.MIN, group=self.tp_group,
tag="placement_digest_min",
)
self._cp_hicache_all_reduce(
hi, op=torch.distributed.ReduceOp.MAX, group=self.tp_group,
tag="placement_digest_max",
)
if int(lo.item()) != local or int(hi.item()) != local:
raise RuntimeError(
"CP shared-L2 placement digest DIVERGED across ranks (B1 Theorem-1 "
f"violation): local={local} min={int(lo.item())} max={int(hi.item())}"
)
def check_hicache_events(self):
self.writing_check()
# B1 runtime invariant: placement is mutated by reserve (prepare) / release
# (evict) / mark_object_committed (the writing_check just above), so assert
# cross-rank replication here, right after the MIN commit frontier.
self._cp_assert_placement_replicated()
self.loading_check()
if self.enable_storage:
self.drain_storage_control_queues()
if self.enable_cp_l3:
self._drain_l3_control_queues()
self._cp_l3_spill_maintainer()
if self.enable_storage_metrics:
self.storage_metrics_collector.log_storage_metrics(
self.cache_controller.storage_backend.get_stats()
)
self._cp_maybe_log_lane_stats()
self._cp_maybe_collect_metrics()
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:
self._cp_hicache_all_reduce(
qsizes,
op=torch.distributed.ReduceOp.MIN,
group=self.tp_group,
tag="storage_queue_min",
)
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,
)
self._cp_hicache_all_reduce(
states,
op=torch.distributed.ReduceOp.MAX,
group=self.tp_group,
tag="prefetch_terminate_max",
)
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
)
self._cp_hicache_all_reduce(
completed_tokens_tensor,
op=torch.distributed.ReduceOp.MIN,
group=self.tp_group,
tag="prefetch_completed_min",
)
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 and not self._uses_cp_hicache:
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,
)
deferred_node = None
try:
value, last_node = self._match_prefix_helper(
self.root_node,
key,
floor_exact_key=getattr(params, "cp_floor_exact", True),
)
except HiCachePendingBackupSplit as exc:
value = []
last_node = exc.node.parent if exc.node.parent is not None else self.root_node
deferred_node = exc.node
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 != self.root_node and 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,
pending_backup_deferred_node=deferred_node,
)
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
):
# NOTE: dead under CP today (storage forbidden); converted for whole-repo
# correctness so a future CP-aware L3 prefetch insert stays rank-replicated.
# (The prefetch-completion insert is MIN-synchronized, so the bump order is
# replicated; if it ever ran at a non-replicated point this must be excluded.)
node.last_access_time = TreeNode.next_access_time()
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 = TreeNode.next_access_time()
# 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.compacted()
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, *, floor_exact_key: bool = True
):
node.last_access_time = TreeNode.next_access_time()
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 = TreeNode.next_access_time()
# Refresh pin TTL on cache hit
if self._is_pinned(child):
child.pin_expiry = time.monotonic() + child.pin_ttl
raw_prefix_len = self.key_match_fn(child.key, key)
prefix_len = self._cp_floor_exact_valid_tail_extension_len(
child,
raw_prefix_len,
len(key),
floor_exact_key=floor_exact_key,
)
stop_after_page_floor = prefix_len != raw_prefix_len
prune_stale_tail_after_split = stop_after_page_floor
if prefix_len <= 0:
break
if prefix_len < len(child.key):
if self._cp_node_split_still_pending(child):
raise HiCachePendingBackupSplit(child)
if (
self._uses_cp_hicache
and self.page_size > 1
and prefix_len % self.page_size != 0
):
unfloored_prefix_len = prefix_len
prefix_len = self._cp_floor_backed_partial_split_len(
child, prefix_len
)
prune_stale_tail_after_split = (
prune_stale_tail_after_split
or prefix_len != unfloored_prefix_len
)
if prefix_len == 0:
break
if (
prune_stale_tail_after_split
and self._cp_stale_tail_prune_still_blocked(child)
):
raise HiCachePendingBackupSplit(child)
new_node = self._split_node(child.key, child, prefix_len)
if prune_stale_tail_after_split:
self._cp_prune_stale_tail_after_page_floor(new_node, child)
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):
if (
self._uses_cp_hicache
and self._node_host_write_pending(child)
):
raise HiCachePendingBackupSplit(child)
# 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].compacted()
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):
_old_cp_meta = child.cp_hicache
# Invariant: a backed node's host backup spans its full key, so the
# radix split point never exceeds host_len. Fail loud if a malformed
# partial-backed node (key longer than its host backup) ever reaches
# here -- it would otherwise feed split_pages>num_pages to the L2 split
# and drive child.host_len negative.
if split_len > child.host_len:
raise RuntimeError(
"[CP_HICACHE_FAILFAST][split_exceeds_host_backup] radix split at "
f"split_len={split_len} exceeds the node's backed host_len="
f"{child.host_len} (a partial-backed node must not exist); "
f"node_id={getattr(child, 'id', None)}"
)
# Leave committed for the pre-split metadata BEFORE reassigning
# child.cp_hicache, so a first-touch lazy rebuild of the running
# counts still sees the intact, in-tree child as committed.
# child is provably committed here (a pending split raises at the
# top of _split_node; _node_backuped requires host_len>0 & not
# pending).
self._cp_account_leave_committed(_old_cp_meta)
if self._is_cp_shared_l2_metadata(_old_cp_meta):
parent_object_key = self.cache_controller._cp_shared_l2_object_key(
new_node.id
)
child_object_key = getattr(_old_cp_meta, "object_key", None)
if not child_object_key:
child_object_key = self.cache_controller._cp_shared_l2_object_key(
child.id
)
split_pages = split_len // self.page_size
allocator = getattr(
self.cache_controller, "cp_shared_l2_page_allocator", None
)
split_object = getattr(allocator, "split_committed_object", None)
if split_object is None:
raise RuntimeError(
"CP shared physical L2 split requires allocator "
"split_committed_object"
)
split_object(
child_object_key,
split_pages_by_payload={
payload_kind: split_pages
for payload_kind in _old_cp_meta.object_ranges
},
parent_object_key=parent_object_key,
child_object_key=child_object_key,
)
new_node.cp_hicache, child.cp_hicache = _old_cp_meta.split(
split_len,
parent_object_key=parent_object_key,
child_object_key=child_object_key,
)
else:
new_node.cp_hicache, child.cp_hicache = _old_cp_meta.split(
split_len
)
new_node.host_len = split_len
child.host_len = child.host_len - split_len
# Re-enter each half's OWN metadata object (count-neutral overall
# -- .split() partitions page_owners) so the later
# evict-subtraction uses the correct per-half counts.
self._cp_account_enter_committed(new_node.cp_hicache)
self._cp_account_enter_committed(child.cp_hicache)
if self.enable_cp_l3:
# 3.1 proactive spill: the parent half (new_node) is a FRESH committed object
# (parent_object_key) that the commit frontier never sees -> enqueue it here, else its
# pages would never become L3-durable (the durable prefix would be holed after a split).
# The child half KEEPS the original object_key (split: child_key=_old_cp_meta.object_key)
# so its existing FIFO entry stays valid -- don't double-enqueue it. Rank-uniform: splits
# replay identically on every rank (replicated radix), so new_node.id/parent_object_key
# and the enqueue order match across ranks. split_committed_object marked both keys
# committed, so the maintainer's is_committed(parent_object_key) check passes.
self._cp_l3_enqueue_spill(parent_object_key, new_node)
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:].compacted()
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
prepared_cp_backup = params.cp_hicache_prepared_backup
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():
parent_node = node
node = node.children[child_key]
node.last_access_time = TreeNode.next_access_time()
node.priority = max(node.priority, priority)
raw_prefix_len = self.key_match_fn(node.key, key)
prefix_len = self._cp_floor_exact_valid_tail_extension_len(
node,
raw_prefix_len,
len(key),
floor_exact_key=prepared_cp_backup is not None,
)
stop_after_page_floor = prefix_len != raw_prefix_len
if prefix_len <= 0:
if stop_after_page_floor:
node = parent_node
break
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_ += self._node_device_resident_len(node)
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
prune_stale_tail_after_split = stop_after_page_floor
unfloored_prefix_len = prefix_len
prefix_len = self._cp_floor_backed_partial_split_len(
node, prefix_len
)
prune_stale_tail_after_split = (
prune_stale_tail_after_split
or prefix_len != unfloored_prefix_len
)
if prefix_len <= 0:
break
if self._cp_node_split_still_pending(node):
return self._cp_defer_insert_for_pending_backup_split(
node,
total_prefix_length,
reason="node_pending_backup",
)
if (
prune_stale_tail_after_split
and self._cp_stale_tail_prune_still_blocked(node)
):
return self._cp_defer_insert_for_pending_backup_split(
node,
total_prefix_length,
reason="stale_tail_unprunable",
)
new_node = self._split_node(node.key, node, prefix_len)
if prune_stale_tail_after_split:
self._cp_prune_stale_tail_after_page_floor(new_node, node)
# 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_ += self._node_device_resident_len(new_node)
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 stop_after_page_floor:
break
if len(key):
use_prepared_cp_backup = (
prepared_cp_backup is not None
and getattr(prepared_cp_backup, "logical_len", -1) == len(value)
)
new_node = TreeNode(
id=(
prepared_cp_backup.node_id
if use_prepared_cp_backup
else None
),
priority=priority,
)
new_node.parent = node
new_node.key = key.compacted()
new_node.value = value.clone()
node.children[child_key] = new_node
self.evictable_size_ += self._node_device_resident_len(new_node)
self._update_leaf_status(node)
self._update_leaf_status(new_node)
# Compute hash_value if storage or kv events are enabled, or for the CP L3 disk tier (the
# per-page content hash is L3's key). Clean-boot only (default-off+additive): a fresh process
# with --enable-cp-l3 computes the unbroken parent-chained hash from the root; flipping it at
# runtime would leave pre-existing nodes None-seeded (R1 HIGH-3).
if self.enable_storage or self.enable_kv_cache_events or self.enable_cp_l3:
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":
if use_prepared_cp_backup:
self._attach_prepared_cp_backup(new_node, prepared_cp_backup)
else:
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)
Reference in New Issue View Git Blame Copy Permalink