diff --git a/csrc/engine/pybind_sequence.cpp b/csrc/engine/pybind_sequence.cpp index 7d628dc..cd183ec 100644 --- a/csrc/engine/pybind_sequence.cpp +++ b/csrc/engine/pybind_sequence.cpp @@ -116,6 +116,8 @@ void pybind_engine_sequence(pybind11::module &engine) { bind_sequence_readonly(sequence, "sampling_params", [](const Sequence &seq) -> const auto & { return seq.generation.params; }); bind_sequence_property(sequence, "multi_modal_data", [](Sequence &seq) -> auto & { return seq.input.multi_modal_data; }); bind_sequence_property(sequence, "cached_hit_tokens", [](Sequence &seq) -> auto & { return seq.cache.hit_tokens; }); + // 只读：host prefix cache 落盘时需要 seq 当前占用的 block id + bind_sequence_readonly(sequence, "block_table", [](const Sequence &seq) -> const auto & { return seq.cache.block_table; }); bind_sequence_property(sequence, "last_output_token", [](Sequence &seq) -> auto & { return seq.mtp.last_output_token; }); bind_sequence_property(sequence, "last_accepted_token", [](Sequence &seq) -> auto & { return seq.mtp.accepted_tokens; }); bind_sequence_property(sequence, "last_next_n_token", [](Sequence &seq) -> auto & { return seq.mtp.draft_tokens; }); diff --git a/deepinfer/_C/engine.pyi b/deepinfer/_C/engine.pyi index 91aa482..f500ad7 100644 --- a/deepinfer/_C/engine.pyi +++ b/deepinfer/_C/engine.pyi @@ -455,6 +455,9 @@ class Sequence: def cached_hit_tokens(self, arg1: typing.SupportsInt | typing.SupportsIndex) -> None: ... @property + def block_table(self) -> list[int]: + ... + @property def chunk_checkpoints(self) -> list[int]: ... @chunk_checkpoints.setter diff --git a/deepinfer/hicache/__init__.py b/deepinfer/hicache/__init__.py new file mode 100644 index 0000000..9461a6f --- /dev/null +++ b/deepinfer/hicache/__init__.py @@ -0,0 +1,3 @@ +from deepinfer.hicache.host_cache import HostPrefixCache + +__all__ = ["HostPrefixCache"] diff --git a/deepinfer/hicache/host_cache.py b/deepinfer/hicache/host_cache.py new file mode 100644 index 0000000..dd80146 --- /dev/null +++ b/deepinfer/hicache/host_cache.py @@ -0,0 +1,642 @@ +"""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 and self.l3.readable(): + 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 + self.l3.breaker.trip("L3 fetch raised") + 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 diff --git a/deepinfer/hicache/kernels.py b/deepinfer/hicache/kernels.py new file mode 100644 index 0000000..9cc372c --- /dev/null +++ b/deepinfer/hicache/kernels.py @@ -0,0 +1,26 @@ +"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions).""" + +from __future__ import annotations + +from pathlib import Path +from types import ModuleType + +_module = None + + +def copy_pages_module() -> ModuleType: + global _module + if _module is None: + from torch.utils.cpp_extension import load + + # Do not touch TORCH_CUDA_ARCH_LIST here: the engine's own JIT kernels (deep_jit / TileLang) run + # in this process and must keep their arch-specific (sm_1xxa) targets. torch's loader derives the + # arch for this trivial kernel from the visible device. + here = Path(__file__).resolve().parent + _module = load( + name="deepinfer_hicache_page_copy", + sources=[str(here / "page_copy.cpp"), str(here / "page_copy.cu")], + extra_cuda_cflags=["-O3"], + verbose=False, + ) + return _module diff --git a/deepinfer/hicache/mooncake_store.py b/deepinfer/hicache/mooncake_store.py new file mode 100644 index 0000000..6271ad7 --- /dev/null +++ b/deepinfer/hicache/mooncake_store.py @@ -0,0 +1,224 @@ +"""Mooncake Store as an L3 tier under HostPrefixCache (GPU -> pinned host pool -> Mooncake Store). + +Enabled by DEEPINFER_HICACHE_MOONCAKE_STORE: a JSON object (or "@/path/to/file.json") with Mooncake's +`MooncakeDistributedStore.setup(config)` keys, e.g. + + {"local_hostname": "127.0.0.1", "metadata_server": "http://127.0.0.1:18080/metadata", + "master_server_addr": "127.0.0.1:50551", "protocol": "tcp", "global_segment_size": 0} + +Objects + page groups dsv4kv//g/ G consecutive KV pages (G * page_bytes) + states state// the SWA stored region (state_bytes) + versions the layout (checkpoint config, page/state sizes, env-dependent state layout), so a layout change + is a clean miss. The state namespace follows the contract discussed in kvcache-ai/Mooncake#4803. Pages hold all + layers of a block, so they are independent of the DP rank: any rank or instance can reuse them. + Pages are grouped (DEEPINFER_HICACHE_L3_GROUP_PAGES, default 4 = 512 tokens) because per-object overhead dominates + for single 114 KB pages (0.17 GB/s vs 1.7 GB/s at 4 pages over TCP on one host). + +Data moves between the Store and the (registered) pinned host pools with batch_put_from_multi_buffers / +batch_get_into_multi_buffers, so no extra copies. Puts run on a background thread (write-behind); the engine thread +learns about completions through drain_done(). Gets are synchronous (called from try_resume). + +Failure handling. A Store transfer that cannot complete blocks for Mooncake's fixed 60 s transfer timeout, so one bad +get stalls the engine thread that long. A circuit breaker therefore opens on the first transfer failure (any put or +get error other than "not found" / "already exists") that persists after one retry, or on a get slower than +DEEPINFER_HICACHE_L3_SLOW_S (default 5 s). The retry absorbs sporadic transfer errors (seen at ~0.015% of EFA puts). +While it is open, reads are skipped (host-only lookup) and queued puts are dropped. It closes after +DEEPINFER_HICACHE_L3_COOLDOWN_S (default 30 s, doubled on each consecutive trip up to 10 min). + +TCP transport limits. Mooncake <= 0.3.13 queues at most 1024 transfers per peer, and one long prompt stores more +buffers than that in a single batch put; the rejected transfers hang until the 60 s timeout. Unless they are already +set, MC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER and MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER default to 65535, which is the +default since kvcache-ai/Mooncake#3710. +""" + +from __future__ import annotations + +import collections +import hashlib +import json +import os +import queue +import threading +import time +from typing import Any + +from deepinfer.logger import init_logger + +logger = init_logger(__name__) + +ENV = "DEEPINFER_HICACHE_MOONCAKE_STORE" +FORMAT_VERSION = 1 +OBJECT_NOT_FOUND, OBJECT_ALREADY_EXISTS = -704, -705 +TCP_ENV_DEFAULTS = {"MC_TCP_MAX_QUEUED_TRANSFERS_PER_PEER": "65535", "MC_TCP_MAX_PENDING_ADMISSIONS_PER_PEER": "65535"} + + +def config_from_env() -> dict | None: + raw = os.environ.get(ENV, "").strip() + if not raw: + return None + if raw.startswith("@"): + with open(raw[1:]) as f: + raw = f.read() + return json.loads(raw) + + +def layout_fingerprint(parts: dict[str, Any]) -> str: + """Version byte + short hash of everything that determines the bytes of a page / state.""" + blob = json.dumps(parts, sort_keys=True).encode() + return f"v{FORMAT_VERSION}-{hashlib.sha256(blob).hexdigest()[:16]}" + + +class _Breaker: + """Opens on a transfer failure; closes after a cooldown that doubles on consecutive trips.""" + + def __init__(self, cooldown_s: float, max_cooldown_s: float = 600.0) -> None: + self.base, self.max = cooldown_s, max_cooldown_s + self.cooldown = cooldown_s + self.open_until = 0.0 + self.trips = 0 + self._lock = threading.Lock() + + def allow(self) -> bool: + return time.monotonic() >= self.open_until + + def trip(self, why: str) -> None: + with self._lock: + now = time.monotonic() + if now < self.open_until: # already open (e.g. the put thread and a get failed together) + return + if now - self.open_until > self.cooldown: # healthy for a while since the last trip: reset the backoff + self.cooldown = self.base + self.open_until = now + self.cooldown + self.trips += 1 + logger.warning("HostPrefixCache L3: %s; Store disabled for %.0f s", why, self.cooldown) + self.cooldown = min(self.cooldown * 2, self.max) + + +class MooncakeL3: + def __init__(self, config: dict, fingerprint: str, group_pages: int) -> None: + for k, v in TCP_ENV_DEFAULTS.items(): # read by Mooncake's TCP transport when setup() creates it + os.environ.setdefault(k, v) + from mooncake.store import MooncakeDistributedStore # optional dependency (pip install mooncake-transfer-engine) + + cfg = dict(config) + cfg.setdefault("global_segment_size", 0) # memory is contributed by a standalone mooncake_client by default + self.store = MooncakeDistributedStore() + rc = self.store.setup(cfg) + if rc != 0: + raise RuntimeError(f"MooncakeDistributedStore.setup failed: {rc}") + self.fingerprint = fingerprint + self.group_pages = group_pages + self._page_prefix = f"dsv4kv/{fingerprint}/g{group_pages}/" + self._state_prefix = f"state/{fingerprint}/" + self._puts: queue.Queue = queue.Queue() + self._done: collections.deque = collections.deque() # tokens of finished puts (appended by the worker thread) + self.put_objects = self.put_failed = self.put_bytes = self.put_dropped = self.put_retried = self.get_retried = 0 + self.put_seconds = 0.0 + self.slow_get_s = float(os.environ.get("DEEPINFER_HICACHE_L3_SLOW_S", "5")) + self.breaker = _Breaker(float(os.environ.get("DEEPINFER_HICACHE_L3_COOLDOWN_S", "30"))) + threading.Thread(target=self._put_worker, name="hicache-l3-put", daemon=True).start() + logger.info("HostPrefixCache L3: Mooncake Store ready (fingerprint %s, %d pages per object)", fingerprint, group_pages) + + def register(self, ptr: int, nbytes: int) -> None: + rc = self.store.register_buffer(ptr, nbytes) + if rc != 0: + raise RuntimeError(f"Mooncake register_buffer({nbytes} bytes) failed: {rc}") + + def group_key(self, last_page_hash: bytes) -> str: + return self._page_prefix + last_page_hash.hex() + + def state_key(self, boundary_hash: bytes) -> str: + return self._state_prefix + boundary_hash.hex() + + # ---------------------------------------------------------------- write-behind (any thread may submit) + + def submit_put(self, keys: list[str], ptrs: list[list[int]], sizes: list[list[int]], token: object) -> None: + """Queue a batch put; `token` comes back from drain_done() when the put has returned (success or not).""" + self._puts.put((keys, ptrs, sizes, token)) + + def drain_done(self) -> list[object]: + out = [] + while self._done: + out.append(self._done.popleft()) + return out + + def _put_worker(self) -> None: + while True: + keys, ptrs, sizes, token = self._puts.get() + if not self.breaker.allow(): # Store unhealthy: drop the put, release the host entries right away + self.put_dropped += len(keys) + self._done.append(token) + continue + t0 = time.perf_counter() + try: + rcs = list(self.store.batch_put_from_multi_buffers(keys, ptrs, sizes)) + # retry the failed objects once: sporadic transfer errors are not a sick Store + redo = [i for i, rc in enumerate(rcs) if rc not in (0, OBJECT_ALREADY_EXISTS)] + if redo: + self.put_retried += len(redo) + again = self.store.batch_put_from_multi_buffers( + [keys[i] for i in redo], [ptrs[i] for i in redo], [sizes[i] for i in redo] + ) + for i, rc in zip(redo, again, strict=True): + rcs[i] = rc + # an object that already exists (another rank or instance stored it first) is not a failure + bad = [rc for rc in rcs if rc not in (0, OBJECT_ALREADY_EXISTS)] + self.put_objects += len(keys) - len(bad) + self.put_failed += len(bad) + self.put_bytes += sum(sum(s) for s, rc in zip(sizes, rcs, strict=True) if rc == 0) + if bad: + self.breaker.trip(f"batch put failed for {len(bad)}/{len(keys)} objects (rc {bad[0]})") + except Exception: # noqa: BLE001 - a Store problem must not take the engine down + logger.exception("HostPrefixCache L3: batch put failed") + self.put_failed += len(keys) + self.breaker.trip("batch put raised") + self.put_seconds += time.perf_counter() - t0 + self._done.append(token) + + # ---------------------------------------------------------------- read (engine thread) + + def readable(self) -> bool: + """False while the breaker is open: skip the Store and use the host cache only.""" + return self.breaker.allow() + + def exists(self, keys: list[str]) -> list[bool]: + if not keys: + return [] + return [rc == 1 for rc in self.store.batch_is_exist(keys)] + + def get(self, keys: list[str], ptrs: list[list[int]], sizes: list[list[int]]) -> list[bool]: + """Read each object into its list of buffers; True where the whole object arrived.""" + if not keys: + return [] + t0 = time.perf_counter() + rcs = list(self.store.batch_get_into_multi_buffers(keys, ptrs, sizes)) + redo = [i for i, (rc, s) in enumerate(zip(rcs, sizes, strict=True)) if rc != sum(s) and rc != OBJECT_NOT_FOUND] + if redo: # one retry, as for puts + self.get_retried += len(redo) + again = self.store.batch_get_into_multi_buffers([keys[i] for i in redo], [ptrs[i] for i in redo], [sizes[i] for i in redo]) + for i, rc in zip(redo, again, strict=True): + rcs[i] = rc + dt = time.perf_counter() - t0 + ok = [rc == sum(s) for rc, s in zip(rcs, sizes, strict=True)] + # "not found" just means evicted since exists(); anything else is a transfer problem + bad = [rc for rc, good in zip(rcs, ok, strict=True) if not good and rc != OBJECT_NOT_FOUND] + if bad: + self.breaker.trip(f"batch get failed for {len(bad)}/{len(keys)} objects (rc {bad[0]}, {dt:.1f} s)") + elif dt > self.slow_get_s: + self.breaker.trip(f"batch get of {len(keys)} objects took {dt:.1f} s") + return ok + + def stats_dict(self) -> dict[str, int | float]: + return dict( + l3_put_objects=self.put_objects, + l3_put_failed=self.put_failed, + l3_put_gib=round(self.put_bytes / 2**30, 3), + l3_put_seconds=round(self.put_seconds, 3), + l3_put_queue=self._puts.qsize(), + l3_put_dropped=self.put_dropped, + l3_put_retried=self.put_retried, + l3_get_retried=self.get_retried, + l3_breaker_trips=self.breaker.trips, + l3_breaker_open=int(not self.breaker.allow()), + ) diff --git a/deepinfer/hicache/page_copy.cpp b/deepinfer/hicache/page_copy.cpp new file mode 100644 index 0000000..4a20e78 --- /dev/null +++ b/deepinfer/hicache/page_copy.cpp @@ -0,0 +1,9 @@ +#include + +void copy_pages(int64_t src_base, int64_t dst_base, torch::Tensor src_ids, torch::Tensor dst_ids, int64_t page_bytes); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("copy_pages", ©_pages, py::arg("src_base"), py::arg("dst_base"), py::arg("src_ids"), py::arg("dst_ids"), + py::arg("page_bytes"), + "Copy pages src_base[src_ids[i]] -> dst_base[dst_ids[i]] on the current CUDA stream. Either base may be pinned host memory."); +} diff --git a/deepinfer/hicache/page_copy.cu b/deepinfer/hicache/page_copy.cu new file mode 100644 index 0000000..a445421 --- /dev/null +++ b/deepinfer/hicache/page_copy.cu @@ -0,0 +1,45 @@ +// Page copy between two pools addressed by UVA pointers: either pool may live in GPU memory or in +// pinned host memory (the "kernel" transfer method of SGLang's HiCache). One thread block (row of the +// grid) per page; pages are copied as 16-byte vectors on the current CUDA stream. +#include +#include +#include +#include + +#include +#include + +namespace { + +__global__ void copy_pages_kernel(const uint8_t *__restrict__ src_base, uint8_t *__restrict__ dst_base, + const int32_t *__restrict__ src_ids, const int32_t *__restrict__ dst_ids, int64_t page_bytes) { + const int64_t n16 = page_bytes >> 4; + const int4 *src = reinterpret_cast(src_base + static_cast(src_ids[blockIdx.x]) * page_bytes); + int4 *dst = reinterpret_cast(dst_base + static_cast(dst_ids[blockIdx.x]) * page_bytes); + const int64_t stride = static_cast(gridDim.y) * blockDim.x; + for (int64_t i = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; i < n16; i += stride) { + dst[i] = src[i]; + } +} + +} // namespace + +void copy_pages(int64_t src_base, int64_t dst_base, torch::Tensor src_ids, torch::Tensor dst_ids, int64_t page_bytes) { + TORCH_CHECK(page_bytes > 0 && page_bytes % 16 == 0, "page_bytes must be a positive multiple of 16, got ", page_bytes); + TORCH_CHECK(src_ids.is_cuda() && dst_ids.is_cuda(), "page ids must be CUDA tensors"); + TORCH_CHECK(src_ids.scalar_type() == torch::kInt32 && dst_ids.scalar_type() == torch::kInt32, "page ids must be int32"); + TORCH_CHECK(src_ids.is_contiguous() && dst_ids.is_contiguous(), "page ids must be contiguous"); + TORCH_CHECK(src_ids.numel() == dst_ids.numel(), "src/dst page id count mismatch"); + const int64_t n = src_ids.numel(); + if (n == 0) { + return; + } + constexpr int threads = 256; + const int64_t n16 = page_bytes / 16; + const unsigned blocks_y = static_cast(std::min((n16 + threads - 1) / threads, 16)); + const dim3 grid(static_cast(n), blocks_y); + auto stream = c10::cuda::getCurrentCUDAStream(); + copy_pages_kernel<<>>(reinterpret_cast(src_base), + reinterpret_cast(dst_base), src_ids.data_ptr(), dst_ids.data_ptr(), page_bytes); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/tests/hicache/test_host_cache.py b/tests/hicache/test_host_cache.py new file mode 100644 index 0000000..77ffc4e --- /dev/null +++ b/tests/hicache/test_host_cache.py @@ -0,0 +1,190 @@ +"""HostPrefixCache logic against a fake engine: hashing/matching, the keep-tokens cap, snapping to a stored +state, store dedupe, LRU eviction and deferral. The real state slice kernels are replaced by plain copies; +the page transfer kernel is the real one. Needs one GPU.""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest +import torch + +os.environ["DEEPINFER_YOCO_PREFILL_TRUNCATION_MODE"] = "compact" +os.environ["DEEPINFER_YOCO_PREFILL_KEEP_TOKENS"] = "140" +os.environ["DEEPINFER_PREFILL_STORE_EXTRA_STATE"] = "1" + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +PAGE = 114048 +STATE = 4096 +BLOCK = 128 + + +class FakeSeq: + def __init__(self, request_id: str, tokens: np.ndarray) -> None: + self.request_id = request_id + self.tokens = tokens + self.block_table: list[int] = [] + self.state_id = -1 + self.num_cached_tokens = 0 + self.cached_hit_tokens = 0 + self.chunked_rest_len = 0 + self.data = self # SequenceData stand-in + + def get_len(self) -> int: + return len(self.tokens) + + def get_token_ids(self) -> list[int]: + return self.tokens.tolist() + + +class FakeScheduler: + def __init__(self, num_blocks: int, num_states: int) -> None: + self.free_blocks = list(range(num_blocks)) + self.free_states = list(range(num_states)) + self.refuse = False + + def try_allocate_seq(self, seq: FakeSeq) -> bool: + n = (len(seq.tokens) + BLOCK - 1) // BLOCK + if self.refuse or len(self.free_blocks) < n or not self.free_states: + return False + seq.block_table = [self.free_blocks.pop(0) for _ in range(n)] + seq.state_id = self.free_states.pop(0) + return True + + def get_num_unfinished_seqs(self) -> int: + return 0 + + +class _Obj: + pass + + +@pytest.fixture +def engine(monkeypatch): + from deepinfer.worker.cache_engine import GPUCacheInterface + + monkeypatch.setattr(GPUCacheInterface, "cache_info_ready", True, raising=False) + monkeypatch.setattr(GPUCacheInterface, "state_size_to_disk", STATE, raising=False) + monkeypatch.setattr(GPUCacheInterface, "swa_block_size_kvstore", 134 + 512 + 140, raising=False) + monkeypatch.setattr(GPUCacheInterface, "swa_block_size_inference", 134, raising=False) + monkeypatch.setattr( + GPUCacheInterface, + "select_state_slice_to_store", + classmethod(lambda cls, sid, pos, cache, state_tensor=None, TO_CE=False: state_tensor.copy_(cache[sid, :STATE])), + ) + monkeypatch.setattr( + GPUCacheInterface, + "load_state_from_store_to_prefill_cache", + classmethod(lambda cls, sid, pos, tensor, cache: cache[sid, :STATE].copy_(tensor)), + ) + e = _Obj() + e.scheduler = FakeScheduler(num_blocks=512, num_states=8) + e.cache_config = _Obj() + e.cache_config.block_size = BLOCK + e.driver_worker = _Obj() + e.driver_worker.cache_engine = _Obj() + e.driver_worker.cache_engine.gpu_blocks = torch.zeros((512, PAGE), dtype=torch.uint8, device="cuda") + e.driver_worker.cache_engine.gpu_states = torch.zeros((8, STATE), dtype=torch.uint8, device="cuda") + return e + + +def _prefill(engine, cache, seq: FakeSeq) -> None: + """Simulate the engine: allocate (if the cache did not), write random KV/state, run the store hook.""" + if not seq.block_table: + assert engine.scheduler.try_allocate_seq(seq) + blocks = engine.driver_worker.cache_engine.gpu_blocks + states = engine.driver_worker.cache_engine.gpu_states + start = seq.num_cached_tokens // BLOCK + for b in seq.block_table[start:]: + blocks[b] = torch.randint(0, 256, (PAGE,), dtype=torch.uint8, device="cuda") + states[seq.state_id] = torch.randint(0, 256, (STATE,), dtype=torch.uint8, device="cuda") + cache.on_prefill_step([seq], []) + torch.cuda.synchronize() + cache.poll() + + +def _tokens(n: int, seed: int) -> np.ndarray: + return np.random.default_rng(seed).integers(0, 100000, n, dtype=np.int32) + + +def test_hit_snaps_to_stored_state_and_restores_bytes(engine): + from deepinfer.hicache import HostPrefixCache + + cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + blocks = engine.driver_worker.cache_engine.gpu_blocks + states = engine.driver_worker.cache_engine.gpu_states + + a = FakeSeq("a", _tokens(1000, 1)) + assert cache.try_resume(a, "a", a.tokens) == 0 # cold + _prefill(engine, cache, a) + # 896 tokens = 7 pages; states at the last 4 boundaries (896, 768, 640, 512) + assert cache.stats.stores == 1 and cache.stats.pages_stored == 7 and cache.stats.states_stored == 4 + a_pages = [blocks[b].clone() for b in a.block_table[:7]] + a_state = states[a.state_id, :STATE].clone() + + # Same prefix plus a new turn: hit at 896 (= 7 full pages, and >= keep_tokens remain to recompute). + b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)])) + assert cache.try_resume(b, "b", b.tokens) == 896 + torch.cuda.synchronize() + cache.poll() + assert b.num_cached_tokens == 896 and b.cached_hit_tokens == 896 + assert all(torch.equal(blocks[b.block_table[i]], a_pages[i]) for i in range(7)) + assert torch.equal(states[b.state_id, :STATE], a_state) + assert cache.stats.hits == 1 and cache.stats.hit_tokens == 896 + + # Same 950-token prefix but a different tail at 1000 tokens: pages match up to 896, the keep cap allows + # only 768, and a state exists there -> hit at 768 (the lookback states make short extensions hit). + c = FakeSeq("c", np.concatenate([a.tokens[:950], _tokens(50, 3)])) + assert cache.try_resume(c, "c", c.tokens) == 768 + torch.cuda.synchronize() + cache.poll() + _prefill(engine, cache, c) # pages 0..6 and all four states already exist: nothing new is stored + assert cache.stats.stores == 1 and cache.stats.pages_stored == 7 and cache.stats.states_stored == 4 + + # b's prefill stores its remaining pages (1500 -> 11 pages, 4 new) and states at 1408..1024. + _prefill(engine, cache, b) + assert cache.stats.pages_stored == 11 and cache.stats.states_stored == 8 + d = FakeSeq("d", np.concatenate([b.tokens, _tokens(300, 4)])) + assert cache.try_resume(d, "d", d.tokens) == 1408 + + +def test_eviction_and_capacity(engine): + from deepinfer.hicache import HostPrefixCache + + cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=2) + a = FakeSeq("a", _tokens(7 * BLOCK + 10, 1)) + _prefill(engine, cache, a) + assert len(cache.pages) == 7 + + big = FakeSeq("big", _tokens(60 * BLOCK, 2)) # does not fit at all + _prefill(engine, cache, big) + assert cache.stats.store_skipped == 1 and len(cache.pages) == 7 + + e = FakeSeq("e", _tokens(40 * BLOCK, 3)) # fits without eviction (43 free) + _prefill(engine, cache, e) + assert cache.stats.evicted_pages == 0 and len(cache.pages) == 47 + + f = FakeSeq("f", _tokens(20 * BLOCK, 4)) # needs 20, 3 free -> evict 17 LRU pages (a's 7 first) + _prefill(engine, cache, f) + assert cache.stats.evicted_pages == 17 and len(cache.pages) == 50 + assert ( + cache.try_resume(FakeSeq("a2", np.concatenate([a.tokens, _tokens(300, 5)])), "a2", np.concatenate([a.tokens, _tokens(300, 5)])) == 0 + ) + # Only two state slots: each store keeps the two newest boundaries that fit, evicting LRU (a's, then e's). + assert len(cache.states) == 2 and cache.stats.evicted_states == 4 + + +def test_deferred_when_gpu_cannot_allocate(engine): + from deepinfer.hicache import HostPrefixCache + + cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=2) + a = FakeSeq("a", _tokens(1000, 1)) + _prefill(engine, cache, a) + engine.scheduler.refuse = True + b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)])) + assert cache.try_resume(b, "b", b.tokens) == -1 + assert b.block_table == [] and b.num_cached_tokens == 0 and cache.stats.deferred == 1 + engine.scheduler.refuse = False + assert cache.try_resume(b, "b", b.tokens) == 896 diff --git a/tests/hicache/test_host_cache_l3.py b/tests/hicache/test_host_cache_l3.py new file mode 100644 index 0000000..6e1ce07 --- /dev/null +++ b/tests/hicache/test_host_cache_l3.py @@ -0,0 +1,193 @@ +"""HostPrefixCache L3 (Mooncake Store) logic against a dict-backed fake store: write-behind of page groups and +states, a second cache (another rank / instance) extending its empty host pool from L3, no re-puts, a fetch that +loses an object between exists() and get(), and the circuit breaker after a transfer failure. Same fake engine as +test_host_cache.py. Needs one GPU.""" + +from __future__ import annotations + +import ctypes +import time + +import numpy as np +import pytest +import torch + +from deepinfer.hicache.mooncake_store import _Breaker +from tests.hicache.test_host_cache import BLOCK, PAGE, STATE, FakeSeq, _prefill, _tokens, engine # noqa: F401 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +GROUP = 4 + + +class FakeL3: + """In-memory stand-in for MooncakeL3 with the same interface; puts complete synchronously.""" + + def __init__(self) -> None: + self.group_pages = GROUP + self.objects: dict[str, bytes] = {} + self.puts = 0 + self.drop_on_get: set[str] = set() # keys that vanish between exists() and get() + self.fail_gets = False # every get is a transfer failure (trips the breaker, like MooncakeL3.get) + self.gets = 0 + self.breaker = _Breaker(30.0) + self._done: list[object] = [] + + def readable(self) -> bool: + return self.breaker.allow() + + def register(self, ptr: int, nbytes: int) -> None: + pass + + def group_key(self, h: bytes) -> str: + return "dsv4kv/test/g4/" + h.hex() + + def state_key(self, h: bytes) -> str: + return "state/test/" + h.hex() + + def submit_put(self, keys, ptrs, sizes, token) -> None: # noqa: ANN001 + for k, ps, ss in zip(keys, ptrs, sizes, strict=True): + if k not in self.objects: + self.objects[k] = b"".join(ctypes.string_at(p, s) for p, s in zip(ps, ss, strict=True)) + self.puts += 1 + self._done.append(token) + + def drain_done(self) -> list[object]: + out, self._done = self._done, [] + return out + + def exists(self, keys: list[str]) -> list[bool]: + return [k in self.objects for k in keys] + + def get(self, keys, ptrs, sizes) -> list[bool]: # noqa: ANN001 + self.gets += 1 + if self.fail_gets: + self.breaker.trip("test transfer failure") + return [False] * len(keys) + ok = [] + for k, ps, ss in zip(keys, ptrs, sizes, strict=True): + data = None if k in self.drop_on_get else self.objects.get(k) + if data is None or len(data) != sum(ss): + ok.append(False) + continue + off = 0 + for p, s in zip(ps, ss, strict=True): + ctypes.memmove(p, data[off : off + s], s) + off += s + ok.append(True) + return ok + + def stats_dict(self) -> dict: + return {"l3_put_objects": self.puts} + + +@pytest.fixture +def l3(monkeypatch): + from deepinfer.hicache import HostPrefixCache + + store = FakeL3() + monkeypatch.setattr(HostPrefixCache, "_make_l3", lambda self: store) + return store + + +def _settle(cache) -> None: # noqa: ANN001 + torch.cuda.synchronize() + cache.poll() # retires D2H copies and submits the planned L3 puts + cache.poll() # drains the put completions (refs back to 0) + + +def test_second_cache_hits_from_l3_with_exact_bytes(engine, l3): # noqa: F811 + from deepinfer.hicache import HostPrefixCache + + blocks = engine.driver_worker.cache_engine.gpu_blocks + states = engine.driver_worker.cache_engine.gpu_states + a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + a = FakeSeq("a", _tokens(1000, 1)) + _prefill(engine, a_cache, a) + _settle(a_cache) + # 7 pages -> one full group (pages 0..3) + states at 896/768/640/512 + assert len([k for k in l3.objects if k.startswith("dsv4kv/")]) == 1 + assert len([k for k in l3.objects if k.startswith("state/")]) == 4 + assert all(e.refs == 0 for e in list(a_cache.pages.values()) + list(a_cache.states.values())) + a_pages = [blocks[b].clone() for b in a.block_table[:4]] + a_state = states[a.state_id, :STATE].clone() + + # Another rank / instance: empty host pool, same prefix + a new turn. Host alone would miss; L3 gives the group + # and the state at 512 (the only boundary inside the stored group). + b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)])) + assert b_cache.try_resume(b, "b", b.tokens) == 4 * BLOCK + torch.cuda.synchronize() + b_cache.poll() + assert all(torch.equal(blocks[b.block_table[i]], a_pages[i]) for i in range(4)) + assert torch.equal(states[b.state_id, :STATE], a_state) + s = b_cache.stats + assert s.l3_lookups == 1 and s.l3_hits == 1 and s.l3_pages_loaded == 4 and s.l3_states_loaded == 1 + assert all(e.refs == 0 for e in list(b_cache.pages.values()) + list(b_cache.states.values())) + assert len(b_cache.free_pages) + len(b_cache.pages) == b_cache.num_pages + + +def test_no_reput_of_known_objects(engine, l3): # noqa: F811 + from deepinfer.hicache import HostPrefixCache + + cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + a = FakeSeq("a", _tokens(1000, 1)) + _prefill(engine, cache, a) + _settle(cache) + puts = l3.puts + again = FakeSeq("a-again", a.tokens.copy()) + _prefill(engine, cache, again) # identical prompt: nothing new on the host, nothing new for L3 + _settle(cache) + assert l3.puts == puts + + +def test_object_lost_between_exists_and_get_falls_back(engine, l3): # noqa: F811 + from deepinfer.hicache import HostPrefixCache + + a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + a = FakeSeq("a", _tokens(1000, 1)) + _prefill(engine, a_cache, a) + _settle(a_cache) + l3.drop_on_get = {k for k in l3.objects if k.startswith("dsv4kv/")} + + b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)])) + assert b_cache.try_resume(b, "b", b.tokens) == 0 # miss, sequence untouched + assert b.block_table == [] and b_cache.stats.l3_get_failed == 1 + assert len(b_cache.pages) == 0 and len(b_cache.free_pages) == b_cache.num_pages + + +def test_transfer_failure_opens_the_breaker(engine, l3): # noqa: F811 + from deepinfer.hicache import HostPrefixCache + + a_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + a = FakeSeq("a", _tokens(1000, 1)) + _prefill(engine, a_cache, a) + _settle(a_cache) + + l3.fail_gets = True + b_cache = HostPrefixCache(engine, size_bytes=50 * PAGE, num_state_slots=8) + b = FakeSeq("b", np.concatenate([a.tokens, _tokens(500, 2)])) + assert b_cache.try_resume(b, "b", b.tokens) == 0 and l3.gets == 1 + assert not l3.readable() and l3.breaker.trips == 1 + # while open, later requests do not touch the Store at all (no 60 s stalls), even once it is healthy again + l3.fail_gets = False + c = FakeSeq("c", np.concatenate([a.tokens, _tokens(500, 3)])) + assert b_cache.try_resume(c, "c", c.tokens) == 0 and l3.gets == 1 + # after the cooldown the Store is used again + l3.breaker.open_until = 0.0 + d = FakeSeq("d", np.concatenate([a.tokens, _tokens(500, 4)])) + assert b_cache.try_resume(d, "d", d.tokens) == 4 * BLOCK and l3.gets == 2 + + +def test_breaker_backoff() -> None: + br = _Breaker(10.0, max_cooldown_s=25.0) + assert br.allow() + br.trip("x") + assert not br.allow() and br.trips == 1 + br.trip("y") # already open: no second trip + assert br.trips == 1 + for want in (20.0, 25.0): # consecutive trips right after reopening double the cooldown, capped + br.open_until = time.monotonic() - 1e-3 + br.trip("z") + assert abs(br.open_until - time.monotonic() - want) < 1.0