mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +08:00
ee2fe708ef
The open-source engine has no KV reuse across requests, so every turn of a
multi-turn conversation prefills its whole prompt. This adds deepinfer/hicache:
- Pinned host pools organized in pages (one page = one 128-token KV block)
plus slots for the SWA state that the production KV store persists
(select_state_slice_to_store / load_state_from_store_to_prefill_cache).
Pages are indexed by a hash chain over the prompt, LRU evicted, written
through at the end of every prefill (engine hook pe_do_save_kvcache).
- Transfers use one small CUDA kernel copying pages between GPU and pinned
host memory through UVA pointers on the engine stream ("kernel" I/O
method, ~180 GiB/s on GB300); completion is tracked with CUDA events.
- A hit is capped so that at least yoco keep_tokens remain to recompute
(bounded replay rebuilds the decoder SWA state from that suffix), then
snapped to the nearest 128-aligned position with a stored state; states
are stored at the last few boundaries of a prompt so short follow-ups hit.
The sequence is pre-allocated, filled from host and scheduled with
num_cached_tokens set (HostPrefixCache.try_resume).
- Engine: read-only Sequence.block_table binding (to copy a finished
prefill's blocks out) and its stub.
Validated on 4x GB300 with DeepSeek-V4.1-Flash through an OpenAI-compatible
server (separate PR): greedy outputs byte-identical with and without the
cache over 8 multi-turn requests; InferenceX AgentX replay at concurrency 1
served 224 requests per hour vs 99, TTFT p50 0.59 s vs 8.9 s, with 89-96%
of prompt tokens served from the cache. tests/hicache/test_host_cache.py
covers the cache logic against a fake engine.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TNVW4sqoaPqbNtSrgpA2zK
433 lines
19 KiB
Python
433 lines
19 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).
|
|
"""
|
|
|
|
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
|