Marconi: Prefix Caching for the Era of Hybrid LLMs

本文最后更新于：2025年6月30日

本文是 Marconi: Prefix Caching for the Era of Hybrid LLMs 的阅读笔记，原文链接。

Attention 的二次复杂度给计算和显存都带来了巨大的压力，所以有一些工作致力于寻找线性复杂度的 Attention(linear attention)。

Attention 记录的中间状态会随着序列长度增加而增加，Linear attention 只会维护一个固定大小的状态。具体参考苏神的博客

可以想像的是，linear attention 的“智商”不如传统 attention。所以现在普遍的做法是结合两者，即在若干个 linear 层之后加上标准的 attention，比如说Jamba，Mamba和Minimax，这让它们集两家之长：既有处理长上下文的效率，又不失注意力机制的记忆能力。

但这种新架构也引入了新的问题：前缀缓存（prefix caching）在这些新混合模型上效果不佳。

为何在混合式大语言模型上做缓存如此困难

前缀缓存的理念很简单：如果许多用户的请求都以相同的文本开头（例如系统提示或共享文档），你可以一次性计算该前缀的内部状态，将其缓存，并在所有后续请求中重复使用，从而节省大量冗余计算。

对于传统的Transformer模型来说，这很简单。它们的内部状态（KV缓存）可以被轻松地切分。如果你有了一个1000个 token 序列的缓存，你可以轻易地将其前500个 token 的部分用于另一个请求。

然而，混合模型中的循环SSM层则不同。它们按顺序处理 toekn 并就地更新（in-place update）其内部状态。你可以把它想象成一个滚动的摘要：第1000个 toekn 的状态包含了从第1到999个 toekn 的所有压缩信息。关键在于，你无法回滚这个状态，这意味着代表1000个 toekn 的状态不能被用来表示仅前500个 toekn 的状态。

这给缓存系统带来了一个“要么全有，要么全无”的困境：

最大化复用：为了确保能处理任何可能的前缀重叠，你必须非常频繁地保存SSM状态的检查点（例如，每32个 toekn 一次）。这被称为细粒度检查点（fine-grained checkpointing）。
带来的后果：这样做会产生大量巨大但效用低的缓存条目。SSM状态本身很大（通常是单个 toekn KV缓存大小的10到100倍）。如此频繁地保存它们，会使缓存中充满了极少被复用的条目，导致“缓存颠簸”（cache thrashing）现象——有用的数据不断被踢出缓存。

如论文所示，在使用细粒度检查点时，KV缓存块的复用率比SSM状态块的复用率高出65倍以上，这说明大量缓存空间被浪费在从未被命中的SSM状态上。

graph TD
    subgraph "细粒度缓存（问题所在）"
        direction LR
        Req1["请求: 'The quick brown fox jumps over the lazy dog'"] -- "生成大量检查点" --> Cache;
    end

    subgraph "缓存内存"
        direction TB
        State1["SSM状态<br>'The quick' (体积大)"];
        State2["SSM状态<br>'The quick brown' (体积大)"];
        State3["SSM状态<br>'The quick brown fox' (体积大)"];
        State4["...更多..."];
    end
    
    subgraph "新请求"
        Req2["请求: 'The quick brown cat...'"];
    end

    Req2 -- "需要" --> State2;
    State1 -- "从未被复用" --> Wasted1((浪费));
    State3 -- "从未被复用" --> Wasted2((浪费));

    style Wasted1 fill:#f99,stroke:#333,stroke-width:2px
    style Wasted2 fill:#f99,stroke:#333,stroke-width:2px
    style State2 fill:#9f9,stroke:#333,stroke-width:2px

Marconi的解决方案

Marconi通过一个双管齐下的方法解决了这个困境：限制准入缓存的内容，并对踢出缓存的内容更加智能。

智能准入策略

Marconi不会缓存所有东西，而是只缓存那些有很高复用可能性的SSM状态。它识别了两种关键的复用场景：

纯输入前缀：例如系统提示或指令，这些内容会在许多不同用户会话中共享。
输入和输出前缀：这通常是对话历史，新请求会附加在上一轮对话之后。

为了管理这一点，Marconi使用了一个基数树（radix tree）来追踪所有请求的历史。它的准入策略简单而有效：

对于对话历史，它总是在最后一个解码 toekn 之后设置检查点，因为这是最可能的续写点。
为了找到共享的“纯输入”前缀，它在处理新请求之前执行一次推测性插入（speculative insertion）。它会检查将新请求的输入添加到基数树中是否会创建一个新的分支点。如果会，那么这个分支就代表一个被至少两个不同序列共享的公共前缀，Marconi就知道应该在这里设置检查点。

通过这种方式，Marconi为每个序列最多只准入两个SSM状态，极大地减少了缓存的混乱，并确保了被准入的状态有很高的复用概率。

graph TD
    subgraph "Marconi的推测性准入"
        direction LR
        subgraph "已有的基数树"
            Root --> N1("系统提示<br>...");
            N1 --> N2("用户: 你好吗?");
        end

        subgraph "新请求"
            Req("系统提示<br>...<br>用户: 天气怎么样?");
        end
    end
    
    subgraph "Marconi的逻辑"
       A("1. 推测性插入") --> B{"新请求是否会在<br>'系统提示...'处创建分支?"};
       B -- "是的!" --> C("2. 决策:<br><b>为'系统提示...'创建<br>SSM状态检查点</b>");
    end
    
    A -- "检查" --> Req
    A -- "利用" --> N1

感知FLOP的驱逐策略

大多数缓存系统使用简单的“最近最少使用”（LRU）策略进行驱逐。Marconi认为，对于混合模型来说，这还不够。原因是不同的缓存条目相对于它们占用的内存，能提供的计算节省量是不同的。

Marconi引入了一个新指标，称为FLOP效率（FLOP efficiency）：即每字节缓存内存所节省的计算量（浮点运算次数）。

其关键洞见在于，对于混合模型，FLOP效率随序列长度的增加而增加。一个长前缀可以节省大量的计算，但其SSM状态占用的内存与一个短前缀完全相同。

Marconi的驱逐策略为每个缓存条目计算一个效用分数，该分数平衡了新近度（recency）和这个新指标：效用分 = 新近度分数 + α * FLOP效率分数。

权重 α 通过观察工作负载自动调整。这个策略天生就会用较短序列的命中率来换取较长序列的命中率——这是一个 tradeoff，因为混合模型处理短序列本已非常高效，而最大的性能提升空间正是在长序列上。

测试

该论文将Marconi与vLLM和SGLang等基线系统（经过扩展以支持混合模型）进行了评估。

命中率大幅提升：与采用细粒度检查点的系统（vLLM）相比，Marconi将 toekn 命中率平均提高了4.5倍至34.4倍。
仅凭其感知FLOP的驱逐策略，就比一个拥有智能准入但使用标准LRU驱逐策略的系统（SGLang+）的 toekn 命中率高出19%至219%。
更快的首 toekn 生成速度：这些命中率的提升直接转化为更低的时延，将95百分位的首 toekn 生成时间（TTFT）最多减少了71.1%（617毫秒）。
对于SSM层比例更高、SSM状态维度更大的模型，这些优势更加明显。

荷兰游记下一篇