wehub-resource-sync
117 lines
4.6 KiB
ReStructuredText
117 lines
4.6 KiB
ReStructuredText
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``):
|
|
|
|
.. code-block:: yaml
|
|
|
|
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):
|
|
|
|
.. code-block:: python
|
|
|
|
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 :file:`docs/design/v1/hidden_state_store.md` in the
|
|
source tree.
|