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..e073fc5 --- /dev/null +++ b/deepinfer/hicache/host_cache.py @@ -0,0 +1,432 @@ +"""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). +""" + +from __future__ import annotations + +import hashlib +import os +import time +from collections import OrderedDict, deque +from dataclasses import dataclass, field + +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 + +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) + + +@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 + + def as_dict(self) -> dict[str, int]: + return dict(self.__dict__) + + +class HostPrefixCache: + def __init__(self, engine, *, 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 + # 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 _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)] + 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]) -> None: + ev = torch.cuda.Event() + ev.record() + self._inflight.append((ev, entries)) + + 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 = self._inflight.popleft() + for e in entries: + e.refs -= 1 + e.ready = True + + 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, 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, 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) + 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, finished) -> 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) -> 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)) + if not missing and not new_states: + 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) + self.stats.stores += 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, + ) + return d diff --git a/deepinfer/hicache/kernels.py b/deepinfer/hicache/kernels.py new file mode 100644 index 0000000..0384de7 --- /dev/null +++ b/deepinfer/hicache/kernels.py @@ -0,0 +1,25 @@ +"""JIT build of the page-copy kernel (torch cpp_extension; cached under ~/.cache/torch_extensions).""" + +from __future__ import annotations + +from pathlib import Path + +_module = None + + +def copy_pages_module(): + 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/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