Hidden states#

The HiddenStateStore caches per-token hidden-state tensors (intermediate activations from an inference forward pass) alongside the KV cache, keyed by the same chunk keys. It exists to support any system where a downstream stage needs an upstream stage’s intermediate activations for the cached prefix and cannot reconstruct them from KV alone — for example a vLLM-Omni thinker -> talker pipeline, where the talker consumes the thinker’s hidden states for tokens whose KV the engine restored without re-running prefill.

Without this store, restoring KV alone causes downstream stages to see truncated or wrong activations for the cached prefix. With it, the talker (or any equivalent consumer) gets a contiguous [num_cached_prefix_tokens, hidden_dim] tensor matching the KV prefix LMCache restored.

Enabling it#

Set three keys in the LMCache YAML (LMCACHE_CONFIG_FILE):

chunk_size: 256

enable_hidden_state_cache: true
max_hidden_state_cpu_size: 4          # GiB, pinned CPU pool size
# hidden_state_layers: [0, 1]         # optional allowlist (see below)

  • enable_hidden_state_cache — master switch. When false (the default), engine.hidden_state_store is None and integrators must skip all HS APIs on that worker.

  • max_hidden_state_cpu_size — pinned-CPU pool size in GiB, dedicated to hidden states. Must be > 0 when the store is enabled. The pool is independent of the KV pool; HS allocator pressure never evicts KV.

  • hidden_state_layers — optional allowlist of layer_idx values accepted on store_hidden_states. Leave unset to accept all layer indices. Use this only when you know exactly which hook indices the worker writes (e.g. matching a fork’s hybrid multimodal hook list).

Using it from a worker#

The store is exposed on the engine (not as engine-level methods):

hs = engine.hidden_state_store
if hs is None:
    return  # HS caching disabled; nothing to do

# Store: hidden_states corresponds to token_ids[token_offset:]
hs.store_hidden_states(
    token_ids,                # full prefix (same as used for KV)
    hidden_states,            # [len(token_ids) - token_offset, hidden_dim]
    layer_idx=0,
    token_offset=num_computed,  # 0 for non-incremental callers
)

# Retrieve: contiguous prefix tensor or None on full miss
restored = hs.retrieve_hidden_states(token_ids, layer_idx=0)

Notes:

  • token_ids is always the full prefix so chunk keys align with KV exactly. token_offset lets incremental callers (such as vLLM-Omni) pass only the newly computed rows in hidden_states without zero-padding the already-cached prefix.

  • layer_idx is the storage layer index. Callers that need several intermediate tensors per request (e.g. multimodal hook outputs plus the main text hidden state) call store_hidden_states once per distinct layer_idx and retrieve per layer on restore.

  • retrieve_hidden_states is prefix-strict: it returns the longest contiguous CPU float32 prefix where every chunk has both KV and HS for the requested layer_idx, and stops at the first chunk where either is missing.

Eviction model#

The store implements a coupled-but-asymmetric eviction rule:

  • KV evicted ⇒ HS evicted. On every retrieve_hidden_states the store asks the KV StorageManager whether KV is still present for each chunk key. If KV is gone, the orphan HS entry is dropped and the returned prefix ends there.

  • HS evicted ⇏ KV evicted. Pressure on the HS pinned pool only evicts HS entries (its own LRU), never KV.

  • Restore stops at first missing HS or KV chunk for the requested layer.

Design notes (user-visible)#

  • Separate pinned pool from KV. KV tensors and hidden-state tensors have heterogeneous shapes and dtypes, so they get independent allocators. max_hidden_state_cpu_size sizes the HS pool only.

  • Same chunk keys as KV. The store reuses the engine’s TokenDatabase so HS chunks share the exact CacheEngineKey as the matching KV chunks. This is what enables the lazy coupled eviction check.

  • Lazy coupled-eviction check. No callbacks or shared indices are required between KV and HS — the store asks storage_manager.contains(key) per chunk on retrieve. Works uniformly for any backend that implements contains().

The full internal design (class layering, code paths, follow-up work) lives at docs/design/v1/hidden_state_store.md in the source tree.