跳转至

vLLM-Omni Prefix Cache 设计与优化笔记

在 vLLM-Omni 里，prefix cache 不只是一个“复用 KV cache”的性能优化点。它更像是在 vLLM 原生 KV prefix cache 旁边，为 Omni 多模态流水线补上的一层语义张量缓存：KV cache 解决 attention 计算复用，Omni prefix cache 则解决下游 stage 或后处理仍然需要完整 hidden states、多模态输出时的数据恢复问题。

1. 为什么原生 KV prefix cache 不够

vLLM 原生 prefix cache 的核心目标很清晰：当多个请求共享相同 prompt prefix 时，调度器只需要让模型计算未命中的 token，已经命中的完整 block 可以直接复用 KV cache。这个机制对纯文本自回归生成非常有效，因为后续生成主要依赖 attention KV 状态。

但 vLLM-Omni 面对的是多模态、多 stage 的 pipeline。某些模型不只需要 KV cache，还需要前序 token 对应的 hidden states 或 per-token multimodal outputs。例如 Qwen3-Omni 这类模型中，thinker stage 的输出可能会被 talker stage 消费，后处理也可能需要完整序列级别的中间表示。

问题就出在这里：如果 prefix 命中后，当前 forward 只计算了 suffix token，那么模型本轮自然只产出 suffix 部分的 hidden states。可是下游 stage 需要的是“完整 prompt 对应的输出”。因此只靠 KV cache，计算可以省掉，但语义输出不完整。

所以 Omni 需要额外回答一个问题：

prefix cache 命中后，如何把“已缓存 prefix 的语义张量”和“本轮新计算 suffix 的语义张量”重新拼成完整输出？

这就是 OmniTensorPrefixCache 的核心职责。

2. 核心设计：复用 vLLM 的 block 体系，旁路缓存 Omni 张量

Omni prefix cache 没有重新设计一套 token 到 block 的映射系统，而是复用 vLLM 已经维护好的 block / slot 管理。

它的关键思想是：

  1. vLLM scheduler 仍然负责判断 prefix cache hit，以及当前请求哪些 token 需要计算。
  2. vLLM KV cache 仍然负责保存 attention 所需的 KV 状态。
  3. Omni 额外维护一份 tensor cache，用同样的 slot mapping 记录 hidden states 和 per-token multimodal outputs。
  4. 当后处理或下游 stage 需要完整序列输出时，再根据 block table 把 cached prefix 与新计算 suffix 拼回来。

这样设计的好处是非侵入式。Omni 没有破坏 vLLM 原有的 prefix caching 逻辑，也没有让模型执行层理解一套新的缓存寻址规则，而是把自己需要的语义张量挂在 vLLM 的 block 管理体系上。

可以把它理解成两层缓存并行工作：

Text Only
1
2
3
4
5
vLLM KV prefix cache
  -> 复用 attention KV，减少模型计算

Omni tensor prefix cache
  -> 复用 hidden states / multimodal outputs，恢复完整语义输出

二者服务的目标不同，但它们共享同一套 block 命中和 slot 映射语义。

3. 写入路径：forward 后按 slot 写入缓存

Omni prefix cache 的写入发生在模型 forward 之后。模型执行完成后，runner 会从输出中抽取 hidden states 和 multimodal outputs，然后交给 prefix cache 更新。

核心链路可以简化为：

Text Only
1
2
3
execute_model
  -> extract_multimodal_outputs
  -> update_omni_tensor_prefix_cache

这里最重要的不是函数调用本身，而是写入方式：Omni 不按 request 维度随意保存张量，而是按 vLLM 的 slot mapping 写入 block cache。

简化后的逻辑类似：

Python
1
flat_cache[slot_mapping] = hidden_states

这行思路非常关键：slot mapping 是 vLLM 调度器对“当前 token 应该落在哪个 cache slot”的统一描述。Omni 用它写 hidden states，就能保证 tensor cache 与 KV cache 在 block 布局上对齐。

多模态输出也采用类似思路，但只缓存满足 per-token 语义的 tensor。例如某个 multimodal output 的第 0 维与序列长度一致，它才适合进入 prefix tensor cache。否则它更可能是请求级 metadata、配置或结构化信息，不应该被拆进 token slot，而应该通过 payload 原样传递。

这体现了一个很重要的边界：

prefix tensor cache 只缓存可以按 token/block 对齐恢复的内容，不把所有多模态信息都塞进缓存。

4. 命中恢复：cached prefix + new suffix

当请求命中 prefix cache 时，vLLM 会跳过已经命中的完整 block，只调度剩余 token 参与 forward。因此本轮模型输出只覆盖 suffix 部分。

Omni 在下游需要完整输出时，会根据当前请求的 block table 找到已缓存的 prefix blocks，再拿本轮新产生的 suffix tensors 进行拼接。

简化流程是：

Text Only
1
2
3
cached prefix tensors  <- 从 block cache 按 block ids 取出
new suffix tensors     <- 本轮 forward 新产生
full tensors           <- concat(prefix, suffix)

也就是：

Python
1
full_hidden_states = cat(cached_hidden_states, new_hidden_states)

这个动作看起来简单，但它解决了多模态 prefix cache 中最容易被忽略的问题：计算层面 prefix 已经复用，语义输出层面也必须补齐。否则后续 stage 看到的就不是完整上下文，而只是新计算的一段 suffix。

这里还有一个细节：Omni 只恢复完整 block 对应的 cached prefix。因为 vLLM 原生 prefix cache 的命中粒度就是完整 block，未满一个 block 的尾部 token 不会作为 prefix cache hit 直接复用。Omni tensor cache 与这个规则保持一致，避免出现 KV cache 与 hidden states cache 命中边界不一致的问题。

5. 多模态输入的特殊处理

多模态输入比纯文本更麻烦，因为它通常包含两类信息：

  1. 可以按 token 对齐的张量，比如某些 token 位置对应的 multimodal hidden states。
  2. 不适合按 token 对齐的 metadata，比如原始多模态数据、处理参数、结构化描述等。

Omni 的处理方式是分层的。

对 per-token tensor，它进入 OmniTensorPrefixCache，按 slot/block 维护，prefix 命中时参与拼接恢复。

对非 per-token 信息，它不会被强行写进 tensor cache，而是通过 payload 继续传递。这样可以避免缓存层承担过多语义，也避免把请求级对象错误地解释成 token 级状态。

这也是 prefix cache 模块设计中比较克制的一点：它不是“多模态万能缓存”，而是只缓存那些能与 vLLM block cache 对齐、且下游确实需要恢复的中间张量。

6. Deferred multimodal outputs：减少频繁同步

Omni prefix cache 里还有一个值得注意的优化：deferred multimodal outputs。

多模态输出如果每个 step 都立刻从 GPU 搬到 CPU，可能带来频繁同步和拷贝开销。尤其在自回归生成中，step 粒度很细，小张量频繁跨设备搬运会打断流水线节奏。

因此 Omni 支持先暂存 deferred multimodal outputs，再在合适时机统一 commit 到 prefix cache。这个设计的目的不是改变缓存语义，而是减少不必要的 per-step GPU 到 CPU 同步，让缓存写入更贴近批处理。

从优化思想上看，它是在做一件很工程化的事：

缓存需要完整性，但写入不一定要立刻同步完成；只要在后续读取前提交即可。

这类优化很适合 vLLM-Omni 这种长链路推理系统，因为系统性能往往不是单个算子决定的，而是由调度、拷贝、同步、stage 间传递共同决定。

7. 这个模块的设计取舍

Omni prefix cache 的设计有几个明显取舍。

第一，它选择复用 vLLM block/slot 体系，而不是另建一套缓存索引。这降低了复杂度，也保证 KV cache 与 tensor cache 的命中边界一致。

第二，它只缓存完整 block 对应的内容。这牺牲了一点更细粒度复用的可能性，但换来了与 vLLM scheduler 行为一致的简单模型。

第三，它区分 per-token tensor 和请求级 payload。这样缓存层保持清晰，不会变成多模态数据的“大杂烩”。

第四，它把一些张量放在 CPU 侧缓存，避免长期占用 GPU 显存。代价是恢复时可能需要拷贝回设备，但这通常比把所有中间语义张量常驻 GPU 更可控。

第五，它通过 deferred commit 减少同步开销，说明这个模块不只是功能补齐，也在关注推理链路中的实际性能瓶颈。

8. 总结

vLLM-Omni 的 prefix cache 可以用一句话概括：

它在 vLLM 原生 KV prefix cache 之外，增加了一套与 block/slot 对齐的 Omni tensor cache，用来恢复多模态、多 stage pipeline 所需的完整 hidden states 和 multimodal outputs。

KV cache 让模型少算 prefix，Omni tensor cache 让下游仍然拿到完整语义输出。前者解决计算复用，后者解决输出完整性。

这也是 vLLM-Omni 很典型的扩展方式：不推翻 vLLM 已有的高性能调度和缓存体系，而是在其旁边补充 Omni 场景所需的语义层能力。这个设计既克制，又很实用。对于理解 vLLM-Omni 的整体架构来说，prefix cache 是一个很好的切入点：它把调度、KV cache、多模态中间表示、CPU/GPU 数据搬运和 stage 间输出需求都串在了一起。

📚 相关文章推荐

评论