mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
2ff094fddc
Stored page groups (DEEPINFER_HICACHE_L3_GROUP_PAGES consecutive pages, default 4) and SWA states are written behind to a Mooncake Store straight from the registered pinned pools (batch_put_from_multi_buffers on a background thread; host slots stay pinned until the put returns). A host miss is extended from the Store: batch_is_exist finds the longest group run and the deepest boundary with a state (host or Store), batch_get_into_multi_buffers fills fresh host slots, and the existing H2D path runs unchanged. Keys are namespaced and carry a versioned layout fingerprint (pages under dsv4kv/<fp>/g<G>/, states under state/<fp>/, as discussed in kvcache-ai/Mooncake#4803). Pages are rank-independent, so prefixes are shared across DP ranks, instances and restarts. Any Store error falls back to the host-only path. tests/hicache/test_host_cache_l3.py: a second cache hits from L3 with exact bytes, no re-puts, an object lost between exists() and get() falls back cleanly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
642 lines
29 KiB
Python
642 lines
29 KiB
Python
"""Host-memory prefix cache for DeepInfer, in the style of SGLang's HiCache.
|
|
|
|
Layout
|
|
Two pinned host pools, both organized in pages:
|
|
- KV pages: one page = one paged-KV block (block_size tokens, all layers), same bytes as a GPU block.
|
|
- State slots: the SWA "stored region" of a sequence (the part DeepSeek persists), taken at a
|
|
128-token-aligned position with GPUCacheInterface.select_state_slice_to_store.
|
|
Pages are indexed by a hash chain over the prompt tokens (page i's key covers tokens [0, (i+1)*block_size)),
|
|
so a lookup is a walk from the first page until the first miss. Eviction is LRU, separately per pool.
|
|
|
|
Transfer
|
|
"kernel" method only: a small CUDA kernel copies pages between a GPU pool and the pinned host pool through
|
|
UVA pointers, on the engine's current stream, so it is ordered with the model's own kernels and never
|
|
blocks the CPU. Completion is tracked with CUDA events polled from the engine loop.
|
|
|
|
Policy
|
|
Write-through at the end of every prefill: all full blocks of the prompt plus the state at the last
|
|
128-aligned position are stored (pages already present are only touched). On a new request, the longest
|
|
page-chain match is capped so that at least `keep_tokens` (the decoder bounded-replay rectangle) remain to
|
|
be recomputed, then snapped down to the nearest position that has a stored state; the sequence is
|
|
pre-allocated, its blocks and state are filled from the host, and it enters the scheduler with
|
|
num_cached_tokens set. Requires DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE=compact (the decoder SWA state is
|
|
rebuilt from the recomputed suffix) and DEEPINFER_PREFILL_STORE_EXTRA_STATE=1 (ring slack for the store).
|
|
|
|
Not covered: tokens generated during decode are not stored (they are re-prefilled as part of the next turn's
|
|
prompt and stored then); the cache is per engine process (route a conversation's turns to the same rank).
|
|
|
|
L3 (optional): with DEEPINFER_HICACHE_MOONCAKE_STORE set, every stored page group and state is also written behind
|
|
to a Mooncake Store (deepinfer/hicache/mooncake_store.py), and a host miss is extended from the Store, so prefixes
|
|
are shared across DP ranks, server instances and restarts.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import time
|
|
from collections import OrderedDict, deque
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from deepinfer import envs
|
|
from deepinfer.logger import init_logger
|
|
from deepinfer.worker.cache_engine import KVSTORE_GRAIN, GPUCacheInterface
|
|
|
|
if TYPE_CHECKING:
|
|
from deepinfer._C.engine import Sequence
|
|
from deepinfer.engine.llm_engine import LLMEngine
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class _Entry:
|
|
__slots__ = ("key", "slot", "refs", "ready")
|
|
|
|
def __init__(self, key: bytes, slot: int) -> None:
|
|
self.key = key
|
|
self.slot = slot
|
|
self.refs = 0 # in-flight transfers (store or load) touching this slot
|
|
self.ready = False # host copy complete and hit-able
|
|
|
|
|
|
@dataclass
|
|
class _Chain:
|
|
tokens: np.ndarray
|
|
hashes: list[bytes] = field(default_factory=list)
|
|
counted: bool = False # lookup stats recorded (a deferred request is looked up repeatedly)
|
|
l3_tried: bool = False # L3 consulted once for this request
|
|
|
|
|
|
@dataclass
|
|
class Stats:
|
|
lookups: int = 0
|
|
hits: int = 0
|
|
hit_tokens: int = 0
|
|
lookup_tokens: int = 0
|
|
deferred: int = 0 # hit found but GPU allocation not possible yet
|
|
stores: int = 0
|
|
store_skipped: int = 0
|
|
pages_stored: int = 0
|
|
pages_loaded: int = 0
|
|
states_stored: int = 0
|
|
states_loaded: int = 0
|
|
evicted_pages: int = 0
|
|
evicted_states: int = 0
|
|
errors: int = 0
|
|
l3_lookups: int = 0
|
|
l3_hits: int = 0
|
|
l3_pages_loaded: int = 0
|
|
l3_states_loaded: int = 0
|
|
l3_get_failed: int = 0
|
|
l3_get_seconds: float = 0.0
|
|
l3_errors: int = 0
|
|
|
|
def as_dict(self) -> dict[str, int]:
|
|
return dict(self.__dict__)
|
|
|
|
|
|
class HostPrefixCache:
|
|
def __init__(self, engine: LLMEngine, *, size_bytes: int, num_state_slots: int, staging_slots: int = 16) -> None:
|
|
assert envs.DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE == "compact", "host prefix cache needs decoder bounded replay (compact)"
|
|
assert envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE, "host prefix cache needs DEEPINFER_PREFILL_STORE_EXTRA_STATE=1"
|
|
assert GPUCacheInterface.cache_info_ready, "cache engine not initialized"
|
|
|
|
self.engine = engine
|
|
self.scheduler = engine.scheduler
|
|
cache_engine = engine.driver_worker.cache_engine
|
|
self.gpu_blocks: torch.Tensor = cache_engine.gpu_blocks
|
|
self.gpu_states: torch.Tensor = cache_engine.gpu_states
|
|
assert self.gpu_blocks.dim() == 2 and self.gpu_blocks.is_contiguous() and self.gpu_blocks.dtype == torch.uint8
|
|
self.device = self.gpu_blocks.device
|
|
self.block_size: int = engine.cache_config.block_size
|
|
assert self.block_size == KVSTORE_GRAIN, (self.block_size, KVSTORE_GRAIN)
|
|
self.page_bytes: int = self.gpu_blocks.size(1)
|
|
self.state_bytes: int = GPUCacheInterface.state_size_to_disk
|
|
assert self.page_bytes % 16 == 0 and self.state_bytes % 16 == 0, (self.page_bytes, self.state_bytes)
|
|
self.keep_tokens: int = envs.DEEPINFER_YOCO_PREFILL_KEEP_TOKENS
|
|
# The state kernels gather [b - window_ring, b); keep b at least two blocks in so that range never goes negative.
|
|
self.min_hit_tokens = 2 * self.block_size
|
|
# A follow-up prompt can only resume at a position <= its length - keep_tokens, so a state stored only
|
|
# at the end of a prompt is unusable by any prompt that extends it by fewer than keep_tokens (+ up to a
|
|
# block). Store the state at the last few 128-boundaries instead; how far back is bounded by the SWA
|
|
# ring: the window ending at b must still be in the ring, b - ring_inference >= L - ring_kvstore.
|
|
self.state_lookback_blocks = int(os.environ.get("DEEPINFER_HICACHE_STATE_LOOKBACK", "4"))
|
|
self.ring_slack = GPUCacheInterface.swa_block_size_kvstore - GPUCacheInterface.swa_block_size_inference
|
|
|
|
self.num_pages = max(0, size_bytes // self.page_bytes)
|
|
self.num_state_slots = num_state_slots
|
|
self.staging_slots = staging_slots
|
|
assert self.num_pages > 0 and num_state_slots > 0
|
|
self.gpu_blocks_ptr = self.gpu_blocks.data_ptr()
|
|
# Host pools, staging ring and the copy kernel are allocated lazily at the first store (nothing can be
|
|
# cached before a prefill completes, and an unused cache then costs no pinned memory).
|
|
self._allocated = False
|
|
self.host_pages = self.host_states = None
|
|
self.host_pages_ptr = self.host_states_ptr = 0
|
|
self._staging: list[torch.Tensor] = []
|
|
self._staging_idx = 0
|
|
self.kernel = None
|
|
|
|
self.pages: OrderedDict[bytes, _Entry] = OrderedDict()
|
|
self.states: OrderedDict[bytes, _Entry] = OrderedDict()
|
|
self.free_pages: list[int] = list(range(self.num_pages - 1, -1, -1))
|
|
self.free_states: list[int] = list(range(num_state_slots - 1, -1, -1))
|
|
self._inflight: deque[tuple[torch.cuda.Event, list[_Entry]]] = deque()
|
|
self._chains: dict[str, _Chain] = {}
|
|
self.stats = Stats()
|
|
self.enabled = True
|
|
self.l3 = self._make_l3()
|
|
self.l3_group = self.l3.group_pages if self.l3 is not None else 0
|
|
self._l3_uploaded: OrderedDict[str, None] = OrderedDict() # recently put/seen L3 keys (skip re-puts)
|
|
# DEEPINFER_HICACHE_DEBUG_SYNC=1: synchronize after every launch and log it (pinpoints a faulting kernel).
|
|
self._debug_sync = os.environ.get("DEEPINFER_HICACHE_DEBUG_SYNC", "0") == "1"
|
|
logger.info(
|
|
"HostPrefixCache: %d pages x %d B (%.1f GiB) + %d state slots x %d B (%.2f GiB), allocated at first store; keep_tokens=%d",
|
|
self.num_pages,
|
|
self.page_bytes,
|
|
self.num_pages * self.page_bytes / 2**30,
|
|
num_state_slots,
|
|
self.state_bytes,
|
|
num_state_slots * self.state_bytes / 2**30,
|
|
self.keep_tokens,
|
|
)
|
|
|
|
def _make_l3(self): # noqa: ANN202
|
|
from deepinfer.hicache import mooncake_store
|
|
|
|
config = mooncake_store.config_from_env()
|
|
if config is None:
|
|
return None
|
|
model_cfg = os.path.join(str(self.engine.model_config.model), "config.json")
|
|
with open(model_cfg, "rb") as f:
|
|
cfg_hash = hashlib.sha256(f.read()).hexdigest()
|
|
fp = mooncake_store.layout_fingerprint(
|
|
dict(
|
|
config=cfg_hash,
|
|
page_bytes=self.page_bytes,
|
|
state_bytes=self.state_bytes,
|
|
block_size=self.block_size,
|
|
extra_state=int(envs.DEEPINFER_PREFILL_STORE_EXTRA_STATE),
|
|
extra_size=int(envs.DEEPINFER_PREFILL_STORE_EXTRA_SIZE),
|
|
backtracking=int(envs.DEEPINFER_MAX_BACKTRACKING_TOKENS),
|
|
swa_kvstore=int(GPUCacheInterface.swa_block_size_kvstore),
|
|
)
|
|
)
|
|
group = int(os.environ.get("DEEPINFER_HICACHE_L3_GROUP_PAGES", "4"))
|
|
return mooncake_store.MooncakeL3(config, fp, group)
|
|
|
|
def _allocate(self) -> None:
|
|
if self._allocated:
|
|
return
|
|
from deepinfer.hicache.kernels import copy_pages_module
|
|
|
|
t0 = time.perf_counter()
|
|
self.kernel = copy_pages_module()
|
|
self.host_pages = torch.empty((self.num_pages, self.page_bytes), dtype=torch.uint8, pin_memory=True)
|
|
self.host_states = torch.empty((self.num_state_slots, self.state_bytes), dtype=torch.uint8, pin_memory=True)
|
|
self.host_pages_ptr = self.host_pages.data_ptr()
|
|
self.host_states_ptr = self.host_states.data_ptr()
|
|
# Staging ring on the GPU for state slices; reuse is safe because every transfer runs on the same stream.
|
|
self._staging = [torch.empty(self.state_bytes, dtype=torch.uint8, device=self.device) for _ in range(self.staging_slots)]
|
|
if self.l3 is not None:
|
|
self.l3.register(self.host_pages_ptr, self.host_pages.numel())
|
|
self.l3.register(self.host_states_ptr, self.host_states.numel())
|
|
self._allocated = True
|
|
logger.info("HostPrefixCache: host pools pinned in %.1fs", time.perf_counter() - t0)
|
|
|
|
# ------------------------------------------------------------------ hashing
|
|
|
|
def _chain(self, request_id: str, tokens: np.ndarray | None) -> _Chain | None:
|
|
chain = self._chains.get(request_id)
|
|
if chain is None and tokens is not None:
|
|
chain = self._chains[request_id] = _Chain(np.ascontiguousarray(tokens, dtype=np.int32))
|
|
return chain
|
|
|
|
def _extend(self, chain: _Chain, num_pages: int) -> None:
|
|
hashes, toks, bs = chain.hashes, chain.tokens, self.block_size
|
|
prev = hashes[-1] if hashes else b""
|
|
for i in range(len(hashes), num_pages):
|
|
prev = hashlib.blake2b(prev + toks[i * bs : (i + 1) * bs].tobytes(), digest_size=16).digest()
|
|
hashes.append(prev)
|
|
|
|
def forget(self, request_id: str) -> None:
|
|
self._chains.pop(request_id, None)
|
|
|
|
# ------------------------------------------------------------------ transfers
|
|
|
|
def _ids(self, ids: list[int]) -> torch.Tensor:
|
|
return torch.from_numpy(np.asarray(ids, dtype=np.int32)).pin_memory().to(self.device, non_blocking=True)
|
|
|
|
def _next_staging(self) -> torch.Tensor:
|
|
t = self._staging[self._staging_idx]
|
|
self._staging_idx = (self._staging_idx + 1) % len(self._staging)
|
|
return t
|
|
|
|
def _track(self, entries: list[_Entry], uploads: list | None = None) -> None:
|
|
ev = torch.cuda.Event()
|
|
ev.record()
|
|
self._inflight.append((ev, entries, uploads or []))
|
|
|
|
def _checkpoint(self, what: str) -> None:
|
|
if self._debug_sync:
|
|
from deepinfer.utils import allow_cuda_sync
|
|
|
|
with allow_cuda_sync():
|
|
torch.cuda.synchronize()
|
|
logger.info("hicache debug: ok after %s", what)
|
|
|
|
def poll(self) -> None:
|
|
"""Retire completed transfers (events on one stream complete in order, so only the head is checked)."""
|
|
while self._inflight and self._inflight[0][0].query():
|
|
_, entries, uploads = self._inflight.popleft()
|
|
for e in entries:
|
|
e.refs -= 1
|
|
e.ready = True
|
|
for up in uploads:
|
|
self._l3_submit(*up)
|
|
if self.l3 is not None:
|
|
for held in self.l3.drain_done():
|
|
for e in held:
|
|
e.refs -= 1
|
|
|
|
def _evict(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> int:
|
|
"""Evict `need` LRU entries that are idle; evicts nothing (returns 0) if that many cannot be found."""
|
|
victims, busy = [], []
|
|
for key, e in pool.items():
|
|
if len(victims) >= need:
|
|
break
|
|
if e.refs or not e.ready:
|
|
busy.append(key)
|
|
continue
|
|
victims.append(key)
|
|
for key in busy:
|
|
pool.move_to_end(key)
|
|
if len(victims) < need:
|
|
return 0
|
|
for key in victims:
|
|
free.append(pool.pop(key).slot)
|
|
return len(victims)
|
|
|
|
def _reserve(self, pool: OrderedDict[bytes, _Entry], free: list[int], need: int) -> bool:
|
|
if need > len(free) + len(pool):
|
|
return False
|
|
if len(free) < need:
|
|
evicted = self._evict(pool, free, need - len(free))
|
|
if pool is self.pages:
|
|
self.stats.evicted_pages += evicted
|
|
else:
|
|
self.stats.evicted_states += evicted
|
|
return len(free) >= need
|
|
|
|
# ------------------------------------------------------------------ load (hit)
|
|
|
|
def try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int:
|
|
"""Fill `seq` (not yet in the scheduler) from the host cache.
|
|
|
|
Returns the number of cached tokens (0: miss, sequence untouched; -1: hit but GPU blocks/state not
|
|
allocatable right now, sequence untouched, retry later).
|
|
"""
|
|
if not self.enabled:
|
|
return 0
|
|
try:
|
|
return self._try_resume(seq, request_id, prompt_ids)
|
|
except Exception: # noqa: BLE001 - never take the engine down for a cache problem
|
|
logger.exception("host prefix cache: resume failed, request %s continues uncached", request_id)
|
|
self._fail()
|
|
return 0
|
|
|
|
def _try_resume(self, seq: Sequence, request_id: str, prompt_ids: np.ndarray) -> int:
|
|
bs = self.block_size
|
|
n = len(prompt_ids)
|
|
chain = self._chain(request_id, prompt_ids)
|
|
if not chain.counted:
|
|
chain.counted = True
|
|
self.stats.lookups += 1
|
|
self.stats.lookup_tokens += n
|
|
limit = ((n - self.keep_tokens) // bs) * bs # the replay rectangle [n - keep, n) must be recomputed
|
|
if limit < self.min_hit_tokens:
|
|
return 0
|
|
max_pages = limit // bs
|
|
self._extend(chain, max_pages)
|
|
if self.l3 is not None and not chain.l3_tried:
|
|
chain.l3_tried = True
|
|
try:
|
|
self._l3_fill(chain, max_pages, limit)
|
|
except Exception: # noqa: BLE001 - fall back to the host-only lookup
|
|
logger.exception("host prefix cache: L3 fetch failed for %s", request_id)
|
|
self.stats.l3_errors += 1
|
|
matched = 0
|
|
for h in chain.hashes[:max_pages]:
|
|
e = self.pages.get(h)
|
|
if e is None or not e.ready:
|
|
break
|
|
matched += 1
|
|
b = min(matched * bs, limit)
|
|
while b >= self.min_hit_tokens:
|
|
st = self.states.get(chain.hashes[b // bs - 1])
|
|
if st is not None and st.ready:
|
|
break
|
|
b -= bs
|
|
else:
|
|
return 0
|
|
|
|
try:
|
|
allocated = self.scheduler.try_allocate_seq(seq)
|
|
except RuntimeError: # can never fit (prompt longer than the whole KV pool): let the scheduler reject it
|
|
return 0
|
|
if not allocated:
|
|
self.stats.deferred += 1
|
|
return -1
|
|
num_pages = b // bs
|
|
page_entries = [self.pages[chain.hashes[i]] for i in range(num_pages)]
|
|
for e in page_entries:
|
|
e.refs += 1
|
|
self.pages.move_to_end(e.key)
|
|
st.refs += 1
|
|
self.states.move_to_end(st.key)
|
|
|
|
block_table = seq.block_table
|
|
self.kernel.copy_pages(
|
|
self.host_pages_ptr,
|
|
self.gpu_blocks_ptr,
|
|
self._ids([e.slot for e in page_entries]),
|
|
self._ids(list(block_table[:num_pages])),
|
|
self.page_bytes,
|
|
)
|
|
self._checkpoint(f"load: page copy H2D ({num_pages} pages, b={b})")
|
|
staging = self._next_staging()
|
|
self.kernel.copy_pages(self.host_states_ptr, staging.data_ptr(), self._ids([st.slot]), self._ids([0]), self.state_bytes)
|
|
self._checkpoint("load: state copy H2D")
|
|
GPUCacheInterface.load_state_from_store_to_prefill_cache(seq.state_id, b, staging, self.gpu_states)
|
|
self._checkpoint("load: load_state_to_cache")
|
|
self._track(page_entries + [st])
|
|
|
|
seq.num_cached_tokens = b
|
|
seq.cached_hit_tokens = b
|
|
self.stats.hits += 1
|
|
self.stats.hit_tokens += b
|
|
self.stats.pages_loaded += num_pages
|
|
self.stats.states_loaded += 1
|
|
return b
|
|
|
|
# ------------------------------------------------------------------ store (write-through at prefill end)
|
|
|
|
def on_prefill_step(self, scheduled: list[Sequence], finished: list[Sequence]) -> None:
|
|
"""Engine hook (pe_do_save_kvcache): called after a prefill step's GPU work has completed."""
|
|
if not self.enabled:
|
|
return
|
|
try:
|
|
for seq in scheduled:
|
|
if seq.chunked_rest_len == 0: # last chunk of this prompt
|
|
self._store(seq)
|
|
except Exception: # noqa: BLE001
|
|
logger.exception("host prefix cache: store failed")
|
|
self._fail()
|
|
|
|
def _store(self, seq: Sequence) -> None:
|
|
chain = self._chains.get(seq.request_id)
|
|
if chain is None: # request did not go through try_resume: hash the prompt from the sequence itself
|
|
chain = self._chain(seq.request_id, np.asarray(seq.data.get_token_ids(), dtype=np.int32))
|
|
self._allocate()
|
|
bs = self.block_size
|
|
b = (seq.get_len() // bs) * bs
|
|
if b < self.min_hit_tokens:
|
|
return
|
|
num_pages = b // bs
|
|
self._extend(chain, num_pages)
|
|
missing = [i for i in range(num_pages) if chain.hashes[i] not in self.pages]
|
|
for i in range(num_pages):
|
|
if i not in missing:
|
|
self.pages.move_to_end(chain.hashes[i])
|
|
# State boundaries: b, b - bs, ... while the window ending there is still in the ring.
|
|
seq_len = seq.get_len()
|
|
new_states = []
|
|
for k in range(self.state_lookback_blocks):
|
|
pos = b - k * bs
|
|
if pos < self.min_hit_tokens or pos < seq_len - self.ring_slack:
|
|
break
|
|
key = chain.hashes[pos // bs - 1]
|
|
if key in self.states:
|
|
self.states.move_to_end(key)
|
|
else:
|
|
new_states.append((pos, key))
|
|
uploads = self._l3_plan(chain, num_pages, [key for _, key in new_states])
|
|
if not missing and not new_states:
|
|
for up in uploads: # everything is already on the host: upload right away
|
|
self._l3_submit(*up)
|
|
return
|
|
if not self._reserve(self.pages, self.free_pages, len(missing)):
|
|
self.stats.store_skipped += 1
|
|
return
|
|
# States are best effort: keep the newest boundaries that fit.
|
|
while new_states and not self._reserve(self.states, self.free_states, len(new_states)):
|
|
new_states.pop()
|
|
if not missing and not new_states:
|
|
self.stats.store_skipped += 1
|
|
return
|
|
|
|
block_table = seq.block_table
|
|
entries: list[_Entry] = []
|
|
if missing:
|
|
new = []
|
|
for i in missing:
|
|
e = _Entry(chain.hashes[i], self.free_pages.pop())
|
|
e.refs = 1
|
|
self.pages[e.key] = e
|
|
new.append(e)
|
|
self._checkpoint(f"store: before page copy ({len(new)} pages, blocks {block_table[:3]}..., b={b}, state_id={seq.state_id})")
|
|
self.kernel.copy_pages(
|
|
self.gpu_blocks_ptr,
|
|
self.host_pages_ptr,
|
|
self._ids([block_table[i] for i in missing]),
|
|
self._ids([e.slot for e in new]),
|
|
self.page_bytes,
|
|
)
|
|
self._checkpoint("store: page copy D2H")
|
|
entries += new
|
|
self.stats.pages_stored += len(new)
|
|
for pos, key in new_states:
|
|
st = _Entry(key, self.free_states.pop())
|
|
st.refs = 1
|
|
self.states[key] = st
|
|
staging = self._next_staging()
|
|
GPUCacheInterface.select_state_slice_to_store(seq.state_id, pos, self.gpu_states, state_tensor=staging)
|
|
self._checkpoint(f"store: select_state_slice at {pos}")
|
|
self.kernel.copy_pages(staging.data_ptr(), self.host_states_ptr, self._ids([0]), self._ids([st.slot]), self.state_bytes)
|
|
self._checkpoint("store: state copy D2H")
|
|
entries.append(st)
|
|
self.stats.states_stored += 1
|
|
self._track(entries, uploads)
|
|
self.stats.stores += 1
|
|
|
|
# ------------------------------------------------------------------ L3 (Mooncake Store)
|
|
|
|
def _l3_seen(self, key: str) -> bool:
|
|
if key in self._l3_uploaded:
|
|
self._l3_uploaded.move_to_end(key)
|
|
return True
|
|
return False
|
|
|
|
def _l3_mark(self, key: str) -> None:
|
|
self._l3_uploaded[key] = None
|
|
while len(self._l3_uploaded) > 1_000_000:
|
|
self._l3_uploaded.popitem(last=False)
|
|
|
|
def _l3_plan(self, chain: _Chain, num_pages: int, new_state_keys: list[bytes]) -> list:
|
|
"""Page groups of this prompt and its new states the Store has not seen yet: [(page hashes, groups, state keys)]."""
|
|
if self.l3 is None:
|
|
return []
|
|
g = self.l3_group
|
|
groups = [i for i in range(num_pages // g) if not self._l3_seen(self.l3.group_key(chain.hashes[(i + 1) * g - 1]))]
|
|
state_keys = [k for k in new_state_keys if not self._l3_seen(self.l3.state_key(k))]
|
|
if not groups and not state_keys:
|
|
return []
|
|
return [(chain.hashes[:num_pages], groups, state_keys)]
|
|
|
|
def _l3_submit(self, hashes: list[bytes], groups: list[int], state_keys: list[bytes]) -> None:
|
|
"""Pin the host slots of the planned groups/states and queue a write-behind put. Entries that are gone (evicted)
|
|
or not ready are skipped."""
|
|
g = self.l3_group
|
|
keys, ptrs, sizes, held = [], [], [], []
|
|
for i in groups:
|
|
es = [self.pages.get(h) for h in hashes[i * g : (i + 1) * g]]
|
|
if any(e is None or not e.ready for e in es):
|
|
continue
|
|
keys.append(self.l3.group_key(hashes[(i + 1) * g - 1]))
|
|
ptrs.append([self.host_pages_ptr + e.slot * self.page_bytes for e in es])
|
|
sizes.append([self.page_bytes] * g)
|
|
held += es
|
|
for k in state_keys:
|
|
st = self.states.get(k)
|
|
if st is None or not st.ready:
|
|
continue
|
|
keys.append(self.l3.state_key(k))
|
|
ptrs.append([self.host_states_ptr + st.slot * self.state_bytes])
|
|
sizes.append([self.state_bytes])
|
|
held.append(st)
|
|
if not keys:
|
|
return
|
|
for e in held:
|
|
e.refs += 1
|
|
for k in keys:
|
|
self._l3_mark(k)
|
|
self.l3.submit_put(keys, ptrs, sizes, held)
|
|
|
|
def _l3_fill(self, chain: _Chain, max_pages: int, limit: int) -> None:
|
|
"""Extend the host cache with page groups and a state from the Store, so the host lookup that follows hits
|
|
deeper. No-op when the Store adds nothing beyond what the host already has."""
|
|
bs, g, l3 = self.block_size, self.l3_group, self.l3
|
|
self._allocate() # pools are lazy; a fresh instance can hit L3 before it ever stored anything
|
|
host_pages = 0
|
|
for h in chain.hashes[:max_pages]:
|
|
e = self.pages.get(h)
|
|
if e is None or not e.ready:
|
|
break
|
|
host_pages += 1
|
|
first_group = host_pages // g
|
|
n_groups = max_pages // g - first_group
|
|
if n_groups <= 0:
|
|
return
|
|
self.stats.l3_lookups += 1
|
|
group_keys = [l3.group_key(chain.hashes[(first_group + i + 1) * g - 1]) for i in range(n_groups)]
|
|
present = l3.exists(group_keys)
|
|
k = 0
|
|
while k < n_groups and present[k]:
|
|
k += 1
|
|
if k == 0:
|
|
return
|
|
l3_pages = (first_group + k) * g # pages available in host + L3
|
|
# deepest boundary <= min(l3_pages * bs, limit) beyond the host-only reach that has a state (host or L3)
|
|
top = min(l3_pages * bs, limit) // bs
|
|
candidates = [p for p in range(top, host_pages, -1) if p * bs >= self.min_hit_tokens]
|
|
if not candidates:
|
|
return
|
|
host_state = {p: (self.states.get(chain.hashes[p - 1]) is not None) for p in candidates}
|
|
need_l3 = [p for p in candidates if not host_state[p]]
|
|
in_l3 = dict(zip(need_l3, l3.exists([l3.state_key(chain.hashes[p - 1]) for p in need_l3]), strict=True))
|
|
pick = next((p for p in candidates if host_state[p] or in_l3.get(p)), None)
|
|
if pick is None:
|
|
return
|
|
state_from_l3 = not host_state[pick]
|
|
fetch_groups = list(range(first_group, (pick + g - 1) // g)) # groups covering pages [first_group*g, pick)
|
|
|
|
# Pin the host prefix so reserving slots below cannot evict it, then reserve.
|
|
prefix = [self.pages[h] for h in chain.hashes[:host_pages]]
|
|
for e in prefix:
|
|
e.refs += 1
|
|
try:
|
|
need_pages = len(fetch_groups) * g
|
|
if not self._reserve(self.pages, self.free_pages, need_pages):
|
|
return
|
|
if state_from_l3 and not self._reserve(self.states, self.free_states, 1):
|
|
return
|
|
slots = [self.free_pages.pop() for _ in range(need_pages)]
|
|
keys = [l3.group_key(chain.hashes[(i + 1) * g - 1]) for i in fetch_groups]
|
|
ptrs = [[self.host_pages_ptr + s * self.page_bytes for s in slots[j * g : (j + 1) * g]] for j in range(len(fetch_groups))]
|
|
sizes = [[self.page_bytes] * g for _ in fetch_groups]
|
|
st_slot = None
|
|
if state_from_l3:
|
|
st_slot = self.free_states.pop()
|
|
keys.append(l3.state_key(chain.hashes[pick - 1]))
|
|
ptrs.append([self.host_states_ptr + st_slot * self.state_bytes])
|
|
sizes.append([self.state_bytes])
|
|
t0 = time.perf_counter()
|
|
ok = l3.get(keys, ptrs, sizes)
|
|
self.stats.l3_get_seconds += time.perf_counter() - t0
|
|
if not all(ok): # evicted from the Store in between, or a transfer error: drop the whole fetch
|
|
self.stats.l3_get_failed += 1
|
|
self.free_pages.extend(slots)
|
|
if st_slot is not None:
|
|
self.free_states.append(st_slot)
|
|
return
|
|
loaded = 0
|
|
for j, gi in enumerate(fetch_groups):
|
|
for t in range(g):
|
|
page = gi * g + t
|
|
slot = slots[j * g + t]
|
|
h = chain.hashes[page]
|
|
if page >= pick or h in self.pages: # beyond the hit, or the host already holds this page
|
|
self.free_pages.append(slot)
|
|
continue
|
|
e = _Entry(h, slot)
|
|
e.ready = True
|
|
self.pages[h] = e
|
|
loaded += 1
|
|
if st_slot is not None:
|
|
st = _Entry(chain.hashes[pick - 1], st_slot)
|
|
st.ready = True
|
|
self.states[st.key] = st
|
|
self.stats.l3_states_loaded += 1
|
|
self.stats.l3_hits += 1
|
|
self.stats.l3_pages_loaded += loaded
|
|
finally:
|
|
for e in prefix:
|
|
e.refs -= 1
|
|
|
|
# ------------------------------------------------------------------ misc
|
|
|
|
def _fail(self) -> None:
|
|
self.stats.errors += 1
|
|
if self.stats.errors >= 3:
|
|
logger.error("host prefix cache: disabled after %d errors", self.stats.errors)
|
|
self.enabled = False
|
|
|
|
def stats_dict(self) -> dict[str, int | float]:
|
|
d = self.stats.as_dict()
|
|
d.update(
|
|
enabled=int(self.enabled),
|
|
pages_used=len(self.pages),
|
|
pages_total=self.num_pages,
|
|
states_used=len(self.states),
|
|
states_total=len(self.free_states) + len(self.states),
|
|
inflight=len(self._inflight),
|
|
hit_rate_tokens=(self.stats.hit_tokens / self.stats.lookup_tokens) if self.stats.lookup_tokens else 0.0,
|
|
)
|
|
if self.l3 is not None:
|
|
d.update(self.l3.stats_dict())
|
|
return d
|