leavelet
38ef664b4c
Heavily-forked CP HiCache produces KV corruption only after an L1->L2->L1 round-trip (eviction + reload), scaling with cached volume. Static reading and upstream-diff are exhausted, so trace the data flow and check four invariants directly, gated by SGLANG_CP_HICACHE_KV_TRACE (0=off, 1=structural lifecycle, 2=+compose/free/ack), off by default and cheap when off. Trace points (keyed by node_id; rid_map ties client rid->node_id): - backup_reserve / backup_d2h: host<->phys slots written (H2/H4) - split: prefix+suffix partition vs parent (H1) - evict: write_pending at device free (H4) - reload_node: host slots read on reload (H2) - reload_assign: reloaded device locs -> forward - dev_free / write_ack: device free vs backup-complete ordering (H4) - compose / remap: per-forward row-id cache reuse + unmapped count (H3) New mem_cache/cp_hicache_trace.py (cptrace/rng helper; values rendered space-free for offline parsing). The companion analyzer joins these by node_id; the decisive, request-independent check is H2: any node whose reload host-slot fingerprint differs from its backup fingerprint is the corruption source. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
4505 lines
180 KiB
Python
4505 lines
180 KiB
Python
from __future__ import annotations
|
|
|
|
import atexit
|
|
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.mem_cache.cp_hicache_trace import cptrace, rng, trace_enabled
|
|
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.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_hicache_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, NSATokenToKVPool):
|
|
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 (
|
|
kv_cache.kv_cache_dim * dtype_size * kv_cache.layer_num
|
|
+ indexer_size_per_token
|
|
* index_layer_count
|
|
* NSATokenToKVPool.index_k_with_scale_buffer_dtype.itemsize
|
|
)
|
|
if isinstance(kv_cache, MLATokenToKVPool):
|
|
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 _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 _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)
|
|
)
|
|
|
|
|
|
@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,
|
|
):
|
|
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
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
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()
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
|
|
self.tp_group = params.tp_cache_group
|
|
self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
|
|
self.pp_rank = params.pp_rank
|
|
self.pp_size = params.pp_size
|
|
self.enable_storage = server_args.hicache_storage_backend is not None
|
|
self.enable_storage_metrics = self.enable_storage and params.enable_metrics
|
|
self.extra_metric_labels = server_args.extra_metric_labels
|
|
|
|
(
|
|
extra_config,
|
|
prefetch_threshold,
|
|
prefetch_timeout_base,
|
|
prefetch_timeout_per_ki_token,
|
|
hicache_storage_pass_prefix_keys,
|
|
) = self._parse_storage_backend_extra_config(
|
|
server_args.hicache_storage_backend_extra_config
|
|
)
|
|
# TODO: support more timeout check functions
|
|
self.is_prefetch_timeout = self._prefetch_timeout_check_linear_func
|
|
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
|
|
|
|
self.load_cache_event = threading.Event()
|
|
cp_shared_kv_layout = None
|
|
if self._uses_cp_hicache:
|
|
cp_shared_kv_layout = CpSharedKVLayout(
|
|
page_size=self.page_size,
|
|
cp_size=params.token_to_kv_pool_allocator.cp_size,
|
|
cp_rank=params.token_to_kv_pool_allocator.cp_rank,
|
|
)
|
|
self.cache_controller = HiCacheController(
|
|
params.token_to_kv_pool_allocator,
|
|
self.token_to_kv_pool_host,
|
|
self.page_size,
|
|
self.tp_group,
|
|
load_cache_event=self.load_cache_event,
|
|
write_policy=server_args.hicache_write_policy,
|
|
io_backend=server_args.hicache_io_backend,
|
|
storage_backend=server_args.hicache_storage_backend,
|
|
prefetch_threshold=prefetch_threshold,
|
|
model_name=server_args.served_model_name,
|
|
storage_backend_extra_config=extra_config,
|
|
pp_rank=self.pp_rank,
|
|
pp_size=self.pp_size,
|
|
enable_storage_metrics=self.enable_storage_metrics,
|
|
cp_shared_kv_layout=cp_shared_kv_layout,
|
|
)
|
|
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] = {}
|
|
# 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)
|
|
|
|
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
|
|
|
|
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
|
|
)
|
|
committed_target = self._cp_zero_counts(cp_size)
|
|
committed_draft = self._cp_zero_counts(cp_size)
|
|
pending_target = self._cp_zero_counts(cp_size)
|
|
pending_draft = 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 getattr(node, "id", None) in pending_node_ids:
|
|
continue
|
|
counts = self._cp_assert_metadata_counts(
|
|
metadata, context=f"snapshot_committed node_id={getattr(node, 'id', '?')}"
|
|
)
|
|
committed_target = self._cp_add_counts(committed_target, counts)
|
|
if getattr(metadata, "draft_host_indices", None) is not None:
|
|
committed_draft = self._cp_add_counts(committed_draft, counts)
|
|
|
|
for node_id, pending_backup in pending.items():
|
|
metadata = pending_backup.metadata
|
|
counts = self._cp_assert_metadata_counts(
|
|
metadata, context=f"snapshot_pending node_id={node_id}"
|
|
)
|
|
pending_target = self._cp_add_counts(pending_target, counts)
|
|
if getattr(metadata, "draft_host_indices", None) is not None:
|
|
pending_draft = self._cp_add_counts(pending_draft, counts)
|
|
|
|
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_evict_key(self, node: TreeNode):
|
|
return (
|
|
int(getattr(node, "priority", 0)),
|
|
int(getattr(node, "id", 0)),
|
|
)
|
|
|
|
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 _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] = []
|
|
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}"
|
|
)
|
|
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}"
|
|
)
|
|
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 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
|
|
) -> Tuple[int, ...]:
|
|
metadata = getattr(node, "cp_hicache", None)
|
|
if metadata is not None:
|
|
return metadata.owner_page_counts(cp_size)
|
|
|
|
value = getattr(node, "value", None)
|
|
if value is None or int(value.numel()) == 0:
|
|
return tuple(0 for _ in range(cp_size))
|
|
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)
|
|
return tuple(int((owners == owner).sum().item()) for owner in range(cp_size))
|
|
|
|
def _cp_load_back_ancestor_unlock_contribution(
|
|
self,
|
|
node: TreeNode,
|
|
deficits: List[int],
|
|
planned_evicted_nodes: set,
|
|
cp_size: int,
|
|
) -> 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)
|
|
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
|
|
|
|
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)
|
|
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
|
|
)
|
|
)
|
|
if contribution <= 0 and unlock_contribution <= 0:
|
|
continue
|
|
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.")
|
|
|
|
def _apply_storage_runtime_config(
|
|
self,
|
|
*,
|
|
storage_backend: Optional[str],
|
|
prefetch_threshold: int,
|
|
prefetch_timeout_base: float,
|
|
prefetch_timeout_per_ki_token: float,
|
|
hicache_storage_pass_prefix_keys: bool,
|
|
enable_storage: bool,
|
|
enable_storage_metrics: bool,
|
|
extra_metric_labels: Optional[Dict[str, str]],
|
|
) -> None:
|
|
prefetch_timeout_per_page = (
|
|
self.page_size / 1024 * prefetch_timeout_per_ki_token
|
|
)
|
|
|
|
self.enable_storage = enable_storage
|
|
self.prefetch_threshold = prefetch_threshold
|
|
self.prefetch_timeout_base = prefetch_timeout_base
|
|
self.prefetch_timeout_per_page = prefetch_timeout_per_page
|
|
self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys
|
|
self.enable_storage_metrics = enable_storage_metrics
|
|
|
|
if self.enable_storage_metrics:
|
|
labels = {
|
|
"storage_backend": storage_backend,
|
|
"tp_rank": self.cache_controller.tp_rank,
|
|
"dp_rank": self.cache_controller.dp_rank,
|
|
"pp_rank": self.cache_controller.pp_rank,
|
|
"pp_size": self.cache_controller.pp_size,
|
|
}
|
|
if extra_metric_labels:
|
|
labels.update(extra_metric_labels)
|
|
existing_collector = getattr(self, "storage_metrics_collector", None)
|
|
if existing_collector is None:
|
|
self.storage_metrics_collector = StorageMetricsCollector(labels=labels)
|
|
elif set(existing_collector.labels.keys()) == set(labels.keys()):
|
|
existing_collector.labels = labels
|
|
else:
|
|
logger.warning(
|
|
"Storage metrics labels changed (%s -> %s). Keep existing labels to "
|
|
"avoid duplicate metric registration.",
|
|
sorted(existing_collector.labels.keys()),
|
|
sorted(labels.keys()),
|
|
)
|
|
|
|
def attach_storage_backend(
|
|
self,
|
|
storage_backend: str,
|
|
storage_backend_extra_config_json: Optional[str] = None,
|
|
served_model_name: Optional[str] = None,
|
|
hicache_storage_prefetch_policy: Optional[str] = None,
|
|
hicache_write_policy: Optional[str] = None,
|
|
) -> tuple[bool, str]:
|
|
"""Attach (enable) storage backend at runtime.
|
|
|
|
This will start storage threads inside `HiCacheController` and enable
|
|
prefetch/backup paths. Caller must ensure there are no running/queued
|
|
requests to avoid races.
|
|
"""
|
|
if self._uses_cp_hicache:
|
|
return (
|
|
False,
|
|
"CP shared KV HiCache does not support attaching a storage backend at runtime.",
|
|
)
|
|
|
|
# Validate inputs first (no side effects).
|
|
if hicache_storage_prefetch_policy is not None:
|
|
allowed = ["best_effort", "wait_complete", "timeout"]
|
|
if hicache_storage_prefetch_policy not in allowed:
|
|
return (
|
|
False,
|
|
f"Invalid hicache_storage_prefetch_policy: {hicache_storage_prefetch_policy!r}. "
|
|
f"Expected one of {allowed}.",
|
|
)
|
|
|
|
if hicache_write_policy is not None:
|
|
allowed = ["write_back", "write_through", "write_through_selective"]
|
|
if hicache_write_policy not in allowed:
|
|
return (
|
|
False,
|
|
f"Invalid hicache_write_policy: {hicache_write_policy!r}. "
|
|
f"Expected one of {allowed}.",
|
|
)
|
|
|
|
# If already enabled:
|
|
# - backend unchanged: treat as success, update policies only.
|
|
# - backend changed: treat as failure, do NOT update policies.
|
|
if self.enable_storage:
|
|
current_backend = self.cache_controller.storage_backend_type
|
|
|
|
if current_backend == storage_backend:
|
|
if hicache_storage_prefetch_policy is not None:
|
|
self.prefetch_stop_policy = hicache_storage_prefetch_policy
|
|
logger.info(
|
|
f"Set hicache_storage_prefetch_policy to {hicache_storage_prefetch_policy}"
|
|
)
|
|
if hicache_write_policy is not None:
|
|
self.cache_controller.write_policy = hicache_write_policy
|
|
self.write_through_threshold = (
|
|
1 if hicache_write_policy == "write_through" else 2
|
|
)
|
|
logger.info(f"Set hicache_write_policy to {hicache_write_policy}")
|
|
return (
|
|
True,
|
|
"HiCache storage backend already enabled with same backend; policies updated.",
|
|
)
|
|
|
|
return (
|
|
False,
|
|
f"HiCache storage backend is already enabled with backend '{current_backend}'. "
|
|
f"Cannot attach different backend '{storage_backend}'. Detach first.",
|
|
)
|
|
|
|
# Not enabled: update policies before controller attach so storage threads observe new values.
|
|
if hicache_storage_prefetch_policy is not None:
|
|
self.prefetch_stop_policy = hicache_storage_prefetch_policy
|
|
logger.info(
|
|
f"Set hicache_storage_prefetch_policy to {hicache_storage_prefetch_policy}"
|
|
)
|
|
|
|
if hicache_write_policy is not None:
|
|
self.cache_controller.write_policy = hicache_write_policy
|
|
self.write_through_threshold = (
|
|
1 if hicache_write_policy == "write_through" else 2
|
|
)
|
|
logger.info(f"Set hicache_write_policy to {hicache_write_policy}")
|
|
|
|
logger.info(f"Attaching HiCache storage backend: {storage_backend}")
|
|
try:
|
|
(
|
|
extra_config,
|
|
prefetch_threshold,
|
|
prefetch_timeout_base,
|
|
prefetch_timeout_per_ki_token,
|
|
hicache_storage_pass_prefix_keys,
|
|
) = self._parse_storage_backend_extra_config(
|
|
storage_backend_extra_config_json
|
|
)
|
|
except Exception as e:
|
|
logger.exception(f"Failed to parse storage_backend_extra_config_json: {e}")
|
|
return (
|
|
False,
|
|
f"Failed to parse storage_backend_extra_config_json '{storage_backend_extra_config_json}': {e}",
|
|
)
|
|
|
|
try:
|
|
self.cache_controller.attach_storage_backend(
|
|
storage_backend=storage_backend,
|
|
prefetch_threshold=prefetch_threshold,
|
|
model_name=served_model_name,
|
|
storage_backend_extra_config=extra_config,
|
|
)
|
|
except Exception as e:
|
|
logger.exception(
|
|
f"Failed to attach storage backend '{storage_backend}': {e}"
|
|
)
|
|
return False, f"Failed to attach storage backend '{storage_backend}': {e}"
|
|
|
|
self._apply_storage_runtime_config(
|
|
storage_backend=storage_backend,
|
|
prefetch_threshold=prefetch_threshold,
|
|
prefetch_timeout_base=prefetch_timeout_base,
|
|
prefetch_timeout_per_ki_token=prefetch_timeout_per_ki_token,
|
|
hicache_storage_pass_prefix_keys=hicache_storage_pass_prefix_keys,
|
|
enable_storage=True,
|
|
enable_storage_metrics=self._enable_metrics_flag,
|
|
extra_metric_labels=self.extra_metric_labels,
|
|
)
|
|
return True, "Attached HiCache storage backend successfully."
|
|
|
|
def detach_storage_backend(self) -> tuple[bool, str]:
|
|
"""Detach (disable) storage backend at runtime.
|
|
|
|
Caller must ensure there are no running/queued requests to avoid races.
|
|
"""
|
|
try:
|
|
# Drain any pending control queues before tearing down storage threads/backend.
|
|
# IMPORTANT: this must happen before we clear `ongoing_*`, otherwise acks/releases
|
|
# cannot be matched to nodes and may leak host pages / locks.
|
|
self._drain_storage_control_queues_local()
|
|
# Idempotent detach: always ask controller to best-effort cleanup, even if
|
|
# `self.enable_storage` is already False (may be leftover state from a
|
|
# previous partial detach).
|
|
self.cache_controller.detach_storage_backend()
|
|
except Exception as e:
|
|
logger.exception("Failed to detach storage backend.")
|
|
# Do NOT crash the server for admin operations. Return failure with detail.
|
|
return False, f"Failed to detach HiCache storage backend: {e}"
|
|
|
|
# Best-effort cleanup of any leftover bookkeeping.
|
|
self._drain_storage_control_queues_local()
|
|
# After controller threads are fully stopped, it's safe to force-release any
|
|
# leftover pending ops (e.g., async prefetch/backup that didn't get a revoke/ack).
|
|
self._force_release_pending_storage_ops()
|
|
|
|
self.enable_storage = False
|
|
self.enable_storage_metrics = False
|
|
return True, "Detached HiCache storage backend successfully."
|
|
|
|
def _force_release_pending_storage_ops(self):
|
|
"""Force release any leftover pending prefetch/backup bookkeeping.
|
|
|
|
This is a safety net for detach/shutdown paths. It assumes storage threads
|
|
have been stopped already (via controller.detach), so no concurrent access
|
|
to these structures should happen.
|
|
"""
|
|
cc = self.cache_controller
|
|
|
|
# Force release leftover prefetch ops: free pre-allocated host pages and
|
|
# drop the host protection on the matched prefix node.
|
|
try:
|
|
for req_id, info in list(self.ongoing_prefetch.items()):
|
|
try:
|
|
last_host_node, token_ids, host_indices, _operation = info
|
|
except Exception:
|
|
# Unexpected shape; just drop it.
|
|
self.ongoing_prefetch.pop(req_id, None)
|
|
continue
|
|
|
|
try:
|
|
if host_indices is not None:
|
|
cc.mem_pool_host.free(host_indices)
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to free host indices for prefetch %s", req_id
|
|
)
|
|
|
|
try:
|
|
last_host_node.release_host()
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to release host protection for prefetch %s", req_id
|
|
)
|
|
|
|
try:
|
|
cc.prefetch_tokens_occupied -= len(token_ids)
|
|
if cc.prefetch_tokens_occupied < 0:
|
|
cc.prefetch_tokens_occupied = 0
|
|
except Exception:
|
|
pass
|
|
|
|
self.ongoing_prefetch.pop(req_id, None)
|
|
except Exception:
|
|
logger.exception("Force release pending prefetch ops failed.")
|
|
|
|
# Force release leftover backup ops: drop host protection on nodes.
|
|
try:
|
|
for ack_id, node in list(self.ongoing_backup.items()):
|
|
try:
|
|
node.release_host()
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to release host protection for backup op %s", ack_id
|
|
)
|
|
self.ongoing_backup.pop(ack_id, None)
|
|
except Exception:
|
|
logger.exception("Force release pending backup ops failed.")
|
|
|
|
def _drain_storage_control_queues_local(self):
|
|
"""Drain storage control queues without TP synchronization.
|
|
|
|
This is intended for shutdown/detach paths where we want to make best-effort
|
|
cleanup even if queue sizes temporarily differ across ranks.
|
|
"""
|
|
self._drain_storage_control_queues_impl(
|
|
n_revoke=None,
|
|
n_backup=None,
|
|
n_release=None,
|
|
log_metrics=False,
|
|
)
|
|
|
|
def _drain_storage_control_queues_impl(
|
|
self,
|
|
n_revoke: Optional[int],
|
|
n_backup: Optional[int],
|
|
n_release: Optional[int],
|
|
log_metrics: bool,
|
|
):
|
|
cc = self.cache_controller
|
|
|
|
def _drain_queue(q, limit: Optional[int]):
|
|
drained = 0
|
|
while limit is None or drained < limit:
|
|
try:
|
|
item = q.get_nowait()
|
|
except Empty:
|
|
break
|
|
drained += 1
|
|
yield item
|
|
|
|
def _drain_revoke():
|
|
for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke):
|
|
info = self.ongoing_prefetch.pop(req_id, None)
|
|
if info is not None:
|
|
last_host_node, token_ids, _, _ = info
|
|
last_host_node.release_host()
|
|
cc.prefetch_tokens_occupied -= len(token_ids)
|
|
if cc.prefetch_tokens_occupied < 0:
|
|
cc.prefetch_tokens_occupied = 0
|
|
|
|
def _drain_backup():
|
|
for operation in _drain_queue(cc.ack_backup_queue, n_backup):
|
|
ack_id = operation.id
|
|
entry = self.ongoing_backup.pop(ack_id, None)
|
|
if entry is not None:
|
|
entry.release_host()
|
|
if log_metrics and self.enable_storage_metrics:
|
|
self.storage_metrics_collector.log_backuped_tokens(
|
|
operation.completed_tokens
|
|
)
|
|
|
|
def _drain_release():
|
|
host_indices_list = []
|
|
for host_indices in _drain_queue(cc.host_mem_release_queue, n_release):
|
|
host_indices_list.append(host_indices)
|
|
if host_indices_list:
|
|
host_indices = torch.cat(host_indices_list, dim=0)
|
|
cc.mem_pool_host.free(host_indices)
|
|
|
|
_drain_revoke()
|
|
_drain_backup()
|
|
_drain_release()
|
|
|
|
def _parse_storage_backend_extra_config(
|
|
self, storage_backend_extra_config: Optional[str]
|
|
):
|
|
"""
|
|
Parse storage backend extra config JSON and extract specific parameters.
|
|
|
|
Args:
|
|
storage_backend_extra_config: JSON string containing extra configuration
|
|
|
|
Returns:
|
|
tuple: (extra_config_dict, prefetch_threshold, prefetch_timeout_base, prefetch_timeout_per_ki_token, hicache_storage_pass_prefix_keys)
|
|
"""
|
|
# Parse extra config if provided. Extra config can be a JSON string or a json/toml/yaml file path prefixed with "@".
|
|
extra_config = {}
|
|
if storage_backend_extra_config:
|
|
try:
|
|
if storage_backend_extra_config.startswith("@"):
|
|
# Read config from a json/toml/yaml file
|
|
path = storage_backend_extra_config[1:]
|
|
ext = os.path.splitext(path)[1].lower()
|
|
with open(path, "rb" if ext == ".toml" else "r") as f:
|
|
if ext == ".json":
|
|
extra_config = json.load(f)
|
|
elif ext == ".toml":
|
|
import tomllib
|
|
|
|
extra_config = tomllib.load(f)
|
|
elif ext in (".yaml", ".yml"):
|
|
import yaml
|
|
|
|
extra_config = yaml.safe_load(f)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported config file {path} (config format: {ext})"
|
|
)
|
|
else:
|
|
# read config from JSON string
|
|
extra_config = json.loads(storage_backend_extra_config)
|
|
except Exception as e:
|
|
logger.error(f"Invalid backend extra config JSON: {e}")
|
|
raise e
|
|
|
|
prefetch_threshold = extra_config.pop("prefetch_threshold", 256) # tokens
|
|
prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", 1) # seconds
|
|
prefetch_timeout_per_ki_token = extra_config.pop(
|
|
"prefetch_timeout_per_ki_token", 0.25
|
|
) # seconds per 1024 tokens
|
|
hicache_storage_pass_prefix_keys = extra_config.pop(
|
|
"hicache_storage_pass_prefix_keys", False
|
|
)
|
|
|
|
if not isinstance(prefetch_threshold, int):
|
|
raise ValueError(
|
|
f"prefetch_threshold must be int, got {type(prefetch_threshold).__name__}"
|
|
)
|
|
if not isinstance(prefetch_timeout_base, (int, float)):
|
|
raise ValueError(
|
|
f"prefetch_timeout_base must be number, got {type(prefetch_timeout_base).__name__}"
|
|
)
|
|
if not isinstance(prefetch_timeout_per_ki_token, (int, float)):
|
|
raise ValueError(
|
|
f"prefetch_timeout_per_ki_token must be number, got {type(prefetch_timeout_per_ki_token).__name__}"
|
|
)
|
|
if not isinstance(hicache_storage_pass_prefix_keys, bool):
|
|
raise ValueError(
|
|
"hicache_storage_pass_prefix_keys must be bool, got "
|
|
f"{type(hicache_storage_pass_prefix_keys).__name__}"
|
|
)
|
|
|
|
return (
|
|
extra_config,
|
|
prefetch_threshold,
|
|
float(prefetch_timeout_base),
|
|
float(prefetch_timeout_per_ki_token),
|
|
hicache_storage_pass_prefix_keys,
|
|
)
|
|
|
|
def reset(self):
|
|
TreeNode.counter = 0
|
|
self.cache_controller.reset()
|
|
self.token_to_kv_pool_host.clear()
|
|
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.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 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.pop(node_id)
|
|
node = pending.node
|
|
node.host_len = pending.logical_len
|
|
node.cp_hicache = pending.metadata
|
|
node.host_value = None
|
|
if pending.locked:
|
|
self.dec_node_lock_ref(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.pop(node_id)
|
|
node = pending.node
|
|
# 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
|
|
if any(v > 0 for v in eviction_plan.remaining_deficit):
|
|
logger.warning(
|
|
"[CP_HICACHE_FALLBACK][cp_host_reservation_plan_insufficient] "
|
|
"reason=insufficient_evictable_host_slots node_id=%d phase=%s "
|
|
"required=%s target_avail=%s draft_avail=%s deficit=%s "
|
|
"planned_freed=%s remaining_deficit=%s victims=%s",
|
|
node_id,
|
|
phase,
|
|
admission.required_by_owner,
|
|
admission.target_available_by_owner,
|
|
admission.draft_available_by_owner,
|
|
admission.deficit_by_owner,
|
|
eviction_plan.planned_freed,
|
|
eviction_plan.remaining_deficit,
|
|
[getattr(node, "id", None) for node in 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)
|
|
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_indices_no_collective(
|
|
self,
|
|
device_indices: torch.Tensor,
|
|
node_id: int,
|
|
*,
|
|
admission_checked: bool = False,
|
|
):
|
|
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,
|
|
)
|
|
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,
|
|
prepared.metadata.owned_positions.numel(),
|
|
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, **kwargs):
|
|
details = " ".join(f"{key}={value}" for key, value in kwargs.items())
|
|
if details:
|
|
details = " " + details
|
|
logger.warning(
|
|
"[CP_HICACHE_FALLBACK][%s] reason=%s%s",
|
|
fallback_name,
|
|
reason,
|
|
details,
|
|
)
|
|
|
|
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:
|
|
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, "is_chunked", 0) > 0 and self.page_size > 1:
|
|
# Chunked prefill should still populate CP HiCache, but only for
|
|
# pages whose KV has become fully available. Leaving the current
|
|
# sub-page tail out of the prepared write avoids creating a
|
|
# scheduler-visible stale tail that the next chunk has to prune
|
|
# while the same request still owns/locks it.
|
|
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",
|
|
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),
|
|
result.metadata.owned_positions.numel(),
|
|
)
|
|
if trace_enabled(1):
|
|
md = result.metadata
|
|
cptrace(
|
|
1,
|
|
"rid_map",
|
|
rid=getattr(req, "rid", "<unknown>"),
|
|
node_id=node_id,
|
|
logical_len=len(kv_indices),
|
|
owned=md.owned_positions.numel(),
|
|
host_idx=rng(getattr(md, "host_indices", None)),
|
|
page_owners=getattr(md, "page_owners", None).numel()
|
|
if getattr(md, "page_owners", None) is not None
|
|
else 0,
|
|
)
|
|
|
|
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
|
|
if len(candidates) > 1:
|
|
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",
|
|
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:
|
|
return node.cp_hicache.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",
|
|
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,
|
|
)
|
|
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),
|
|
result.metadata.owned_positions.numel(),
|
|
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())
|
|
if trace_enabled(2) and finish_count > 0:
|
|
finished_ids = [
|
|
aid
|
|
for entry in self.cache_controller.ack_write_queue[:finish_count]
|
|
for aid in entry[2]
|
|
]
|
|
cptrace(2, "write_ack", finished=finish_count, node_ids=finished_ids)
|
|
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.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
|
|
_trace_wp = self._node_host_write_pending(node) if trace_enabled(1) else None
|
|
_trace_dev = rng(node.value) if trace_enabled(1) else None
|
|
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),
|
|
)
|
|
if trace_enabled(1):
|
|
cptrace(
|
|
1,
|
|
"evict",
|
|
node_id=node.id,
|
|
freed_len=freed_len,
|
|
device_resident=device_resident_len,
|
|
backed=self._node_backuped(node),
|
|
write_pending=_trace_wp,
|
|
dev=_trace_dev,
|
|
)
|
|
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, synchronize_across_ranks: bool = False
|
|
) -> int:
|
|
synchronize_across_ranks = (
|
|
synchronize_across_ranks
|
|
and getattr(self, "tp_group", None) is not None
|
|
and getattr(self, "tp_world_size", 1) > 1
|
|
)
|
|
if required_host_slots <= 0 and not synchronize_across_ranks:
|
|
return 0
|
|
|
|
leaves = list(self.evictable_host_leaves)
|
|
logger.debug(
|
|
"[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d sync=%s leaves=%d",
|
|
required_host_slots,
|
|
synchronize_across_ranks,
|
|
len(leaves),
|
|
)
|
|
eviction_heap = [
|
|
(self.eviction_strategy.get_priority(node), node) for node in leaves
|
|
]
|
|
heapq.heapify(eviction_heap)
|
|
|
|
num_evicted = 0
|
|
|
|
def all_ranks_done() -> bool:
|
|
local_done = int(num_evicted >= required_host_slots)
|
|
if not synchronize_across_ranks:
|
|
return bool(local_done)
|
|
done = torch.tensor(local_done, dtype=torch.int, device="cpu")
|
|
self._cp_hicache_all_reduce(
|
|
done,
|
|
op=torch.distributed.ReduceOp.MIN,
|
|
group=self.tp_group,
|
|
tag="host_evict_done_min",
|
|
)
|
|
return bool(done.item())
|
|
|
|
while len(eviction_heap) and not all_ranks_done():
|
|
_priority, x = heapq.heappop(eviction_heap)
|
|
if x == self.root_node:
|
|
break
|
|
if not x.evicted:
|
|
continue
|
|
parent_key = self.get_child_key_fn(x.key)
|
|
if x.parent is None or x.parent.children.get(parent_key) is not x:
|
|
continue
|
|
|
|
if not self._node_backuped(x):
|
|
if len(x.children) == 0:
|
|
parent = self._remove_host_leaf(x)
|
|
if (
|
|
len(parent.children) == 0
|
|
and parent.evicted
|
|
and self._node_backuped(parent)
|
|
):
|
|
new_priority = self.eviction_strategy.get_priority(parent)
|
|
heapq.heappush(eviction_heap, (new_priority, parent))
|
|
continue
|
|
|
|
if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
|
|
self._clear_pin(x)
|
|
|
|
if x.host_ref_counter > 0:
|
|
continue
|
|
|
|
host_indices = self._node_host_evict_indices(x)
|
|
physical_count = len(host_indices) if host_indices is not None else 0
|
|
|
|
self._record_remove_event(x)
|
|
if physical_count > 0:
|
|
if self._uses_cp_hicache and hasattr(
|
|
self.cache_controller, "evict_cp_host"
|
|
):
|
|
num_evicted += self.cache_controller.evict_cp_host(x.cp_hicache)
|
|
else:
|
|
num_evicted += self.cache_controller.evict_host(host_indices)
|
|
x.host_len = 0
|
|
x.cp_hicache = None
|
|
x.host_value = None
|
|
|
|
parent = self._remove_host_leaf(x)
|
|
if len(parent.children) == 0 and parent.evicted:
|
|
new_priority = self.eviction_strategy.get_priority(parent)
|
|
heapq.heappush(eviction_heap, (new_priority, parent))
|
|
|
|
return num_evicted
|
|
|
|
def evict_host(self, num_tokens: int):
|
|
if self._uses_cp_hicache:
|
|
return self._evict_host_for_physical_slots(num_tokens)
|
|
|
|
leaves = list(self.evictable_host_leaves)
|
|
eviction_heap = [
|
|
(self.eviction_strategy.get_priority(node), node) for node in leaves
|
|
]
|
|
heapq.heapify(eviction_heap)
|
|
|
|
num_evicted = 0
|
|
while num_evicted < num_tokens and len(eviction_heap):
|
|
_priority, x = heapq.heappop(eviction_heap)
|
|
if x == self.root_node:
|
|
break
|
|
# only evict the host value of evicted nodes
|
|
if not x.evicted:
|
|
continue
|
|
|
|
# Expire stale pins before checking host_ref_counter
|
|
if x.pin_expiry > 0 and time.monotonic() > x.pin_expiry:
|
|
self._clear_pin(x)
|
|
|
|
# node is protected from eviction as it has ongoing prefetch, backup, or pin
|
|
if x.host_ref_counter > 0:
|
|
continue
|
|
|
|
# Block deleted entirely (GPU already evicted, now CPU freed) --
|
|
# emit BlockRemoved so the router removes this block from its index.
|
|
self._record_remove_event(x)
|
|
num_evicted += self.cache_controller.evict_host(
|
|
self._node_host_evict_indices(x)
|
|
)
|
|
|
|
parent = self._remove_host_leaf(x)
|
|
|
|
if len(parent.children) == 0 and parent.evicted:
|
|
new_priority = self.eviction_strategy.get_priority(parent)
|
|
heapq.heappush(eviction_heap, (new_priority, parent))
|
|
|
|
def load_back(
|
|
self, node: TreeNode, mem_quota: Optional[int] = None
|
|
) -> Optional[torch.Tensor]:
|
|
|
|
if self._uses_cp_hicache:
|
|
start_time = time.perf_counter()
|
|
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()
|
|
if trace_enabled(1):
|
|
cptrace(
|
|
1,
|
|
"reload_assign",
|
|
node_id=loaded_node.id,
|
|
host_len=host_len,
|
|
offset=offset,
|
|
value=rng(loaded_node.value),
|
|
)
|
|
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 check_hicache_events(self):
|
|
self.writing_check()
|
|
self.loading_check()
|
|
if self.enable_storage:
|
|
self.drain_storage_control_queues()
|
|
if self.enable_storage_metrics:
|
|
self.storage_metrics_collector.log_storage_metrics(
|
|
self.cache_controller.storage_backend.get_stats()
|
|
)
|
|
|
|
def drain_storage_control_queues(self):
|
|
"""
|
|
Combine prefetch revoke, backup ack, and host mem release checks
|
|
to minimize TP synchronization and Python overhead.
|
|
"""
|
|
cc = self.cache_controller
|
|
|
|
qsizes = torch.tensor(
|
|
[
|
|
cc.prefetch_revoke_queue.qsize(),
|
|
cc.ack_backup_queue.qsize(),
|
|
cc.host_mem_release_queue.qsize(),
|
|
],
|
|
dtype=torch.int,
|
|
)
|
|
if self.tp_world_size > 1:
|
|
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
|
|
):
|
|
node.last_access_time = time.monotonic()
|
|
if len(key) == 0:
|
|
return 0
|
|
|
|
child_key = self.get_child_key_fn(key)
|
|
|
|
matched_length = 0
|
|
while len(key) > 0 and child_key in node.children.keys():
|
|
node = node.children[child_key]
|
|
node.last_access_time = time.monotonic()
|
|
# Refresh pin TTL on host insert hit
|
|
if self._is_pinned(node):
|
|
node.pin_expiry = time.monotonic() + node.pin_ttl
|
|
prefix_len = self.key_match_fn(node.key, key)
|
|
key = key[prefix_len:]
|
|
host_value = host_value[prefix_len:]
|
|
hash_value = hash_value[prefix_len // self.page_size :]
|
|
matched_length += prefix_len
|
|
|
|
if prefix_len < len(node.key):
|
|
new_node = self._split_node(node.key, node, prefix_len)
|
|
node = new_node
|
|
|
|
if len(key):
|
|
child_key = self.get_child_key_fn(key)
|
|
|
|
if len(key):
|
|
new_node = TreeNode(priority=node.priority)
|
|
new_node.parent = node
|
|
new_node.key = key.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 = time.monotonic()
|
|
child_key = self.get_child_key_fn(key)
|
|
value = []
|
|
|
|
while len(key) > 0 and child_key in node.children.keys():
|
|
child = node.children[child_key]
|
|
child.last_access_time = time.monotonic()
|
|
# Refresh pin TTL on cache hit
|
|
if self._is_pinned(child):
|
|
child.pin_expiry = time.monotonic() + child.pin_ttl
|
|
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):
|
|
if trace_enabled(1):
|
|
_orig = child.cp_hicache
|
|
_orig_owned = _orig.owned_positions.numel()
|
|
_orig_host = rng(getattr(_orig, "host_indices", None))
|
|
_orig_po = (
|
|
_orig.page_owners.numel()
|
|
if getattr(_orig, "page_owners", None) is not None
|
|
else 0
|
|
)
|
|
new_node.cp_hicache, child.cp_hicache = child.cp_hicache.split(
|
|
split_len
|
|
)
|
|
new_node.host_len = split_len
|
|
child.host_len = child.host_len - split_len
|
|
if trace_enabled(1):
|
|
pf, sf = new_node.cp_hicache, child.cp_hicache
|
|
cptrace(
|
|
1,
|
|
"split",
|
|
orig_id=child.id,
|
|
prefix_id=new_node.id,
|
|
suffix_id=child.id,
|
|
split_len=split_len,
|
|
orig_owned=_orig_owned,
|
|
orig_host=_orig_host,
|
|
orig_po=_orig_po,
|
|
pre_owned=pf.owned_positions.numel(),
|
|
pre_host=rng(getattr(pf, "host_indices", None)),
|
|
pre_po=pf.page_owners.numel()
|
|
if getattr(pf, "page_owners", None) is not None
|
|
else 0,
|
|
suf_owned=sf.owned_positions.numel(),
|
|
suf_host=rng(getattr(sf, "host_indices", None)),
|
|
suf_po=sf.page_owners.numel()
|
|
if getattr(sf, "page_owners", None) is not None
|
|
else 0,
|
|
)
|
|
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 = time.monotonic()
|
|
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
|
|
if self.enable_storage or self.enable_kv_cache_events:
|
|
new_node.hash_value = compute_node_hash_values(new_node, self.page_size)
|
|
|
|
# Emit BlockStored so the router indexes this block.
|
|
self._record_store_event(new_node)
|
|
|
|
if self.cache_controller.write_policy != "write_back":
|
|
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)
|