隐藏状态#

HiddenStateStore 在 KV Cache 之外缓存每个 token 的隐藏状态张量（推理前向传播过程中的中间激活），并使用相同的块键作为索引。它的存在是为了支持这样一类系统：下游阶段需要上游阶段为已缓存前缀计算的中间激活，且无法仅凭 KV 重建这些激活——例如 vLLM-Omni 的 thinker -> talker 流水线，其中 talker 会消费 thinker 为某些 token 生成的隐藏状态，而引擎在恢复这些 token 的 KV 时并未重新运行 Prefill。

如果没有这个存储，仅恢复 KV 会导致下游阶段读到的已缓存前缀激活被截断或出错。有了它，talker（或任何等价的消费方）就能获得一个连续的 [num_cached_prefix_tokens, hidden_dim] 张量，与 LMCache 恢复的 KV 前缀相匹配。

启用方法#

在 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 —— 总开关。当为 false（默认值）时，engine.hidden_state_store 为 None，集成方必须在该 worker 上跳过所有 HS API。

  • max_hidden_state_cpu_size —— 专用于隐藏状态的固定 CPU 内存池大小（单位 GiB）。启用该存储时必须 > 0。该内存池独立于 KV 内存池；HS 分配器的压力绝不会逐出 KV。

  • hidden_state_layers —— store_hidden_states 接受的 layer_idx 值的可选白名单。不设置则接受所有层索引。仅当您确切知道 worker 写入哪些 hook 索引时才使用它（例如匹配某个 fork 的混合多模态 hook 列表）。

在 worker 中使用#

该存储暴露在引擎上（而非作为引擎级方法）：

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)

注意事项：

  • token_ids 始终是完整前缀，因此块键能与 KV 精确对齐。token_offset 允许增量调用方（例如 vLLM-Omni）只传入 hidden_states 中新计算的行，而无需对已缓存的前缀进行零填充。

  • layer_idx 是存储层索引。每个请求需要多个中间张量的调用方（例如多模态 hook 输出加上主文本隐藏状态）会为每个不同的 layer_idx 各调用一次 store_hidden_states，并在恢复时按层检索。

  • retrieve_hidden_states 是前缀严格的：它返回最长的连续 CPU float32 前缀，其中每个块对于所请求的 layer_idx 都同时具备 KV 和 HS，并在第一个缺少任一者的块处停止。

逐出模型#

该存储实现了一种耦合但非对称的逐出规则：

  • KV 被逐出 ⇒ HS 被逐出。 在每次 retrieve_hidden_states 时，该存储会向 KV 的 StorageManager 查询每个块键对应的 KV 是否仍然存在。如果 KV 已不存在，则丢弃对应的孤立 HS 条目，返回的前缀到此结束。

  • HS 被逐出 ⇏ KV 被逐出。 HS 固定内存池的压力只会逐出 HS 条目（其自身的 LRU），绝不会逐出 KV。

  • 对于所请求的层，恢复会在第一个缺失的 HS 或 KV 块处停止。

设计说明（用户可见）#

  • 与 KV 分离的固定内存池。 KV 张量和隐藏状态张量具有异构的形状和数据类型，因此它们使用各自独立的分配器。max_hidden_state_cpu_size 仅设置 HS 内存池的大小。

  • 与 KV 相同的块键。 该存储复用引擎的 TokenDatabase，因此 HS 块与匹配的 KV 块共享完全相同的 CacheEngineKey。正是这一点使得惰性耦合逐出检查成为可能。

  • 惰性耦合逐出检查。 KV 与 HS 之间不需要任何回调或共享索引——该存储在检索时为每个块调用 storage_manager.contains(key)。对于任何实现了 contains() 的后端，其行为都完全一致。

完整的内部设计（类的分层、代码路径、后续工作）位于源码树中的 docs/design/v1/hidden_state_store.md。