SGLang Unified Memory：前缀保留量 1.9 倍，重复请求耗时降低 37.9%

前置要求：请阅读 Unified Radix Cache。

背景

模型在推理过程中，会保存key和value，我们称之为KV，供接下来生成答案时使用。那么KV Cache就是存储KV的容器。而在现代的混合模型中，有不同类型的attention层，它们保存KV的方式也不同：

• Full Attention层为每个 token 保存 key 和 value。 读入的 token 越多，KV 就越多。生成答案时，注意力会读取这些 KV，回看文档内容。
• Mamba Attention层不断更新一块大小固定的状态。 这类attention的KV会不断覆写入一个固定的状态存储里。每读入一个新 token，固定的状态里的内容会变化，但这块状态不会变大。因此，只有在精确定位到同一个时刻的快照（checkpoint）的情况下，才能复用KV，否则就会因为覆写，导致状态无法复用。因此每个活跃请求都需要自己的工作状态。如果要把某一时刻的状态另存一份，供以后复用，这份快照（checkpoint）还需要另占一个状态槽位。

下图只看一个请求。上下文变长时，KV 不断累积，循环状态则在原来的空间里更新。混合模型继续生成时，两者都需要。

这就带来了容量分配问题。大量短请求需要很多状态槽位，但每个请求的 KV 很少；少量长请求需要的工作状态较少，KV 却很多。如果 SGLang 在启动时就给两类数据各划一块固定大小的显存，一边可能已经放满，另一边却还有空位，无法同时应对两种情况。

Unified Memory

共用空闲空间

SGLang 的 Unified Memory 让两类数据使用同一块 GPU 显存。循环状态从一端向中间增长，KV 从另一端向中间增长。中间的空闲空间，哪一边需要就分给哪一边。整块显存的大小不变，两类数据能用到的容量随需求调整。

假设总容量是 24 个单位，每个请求需要 2 单位循环状态和 1 单位 KV。如果固定给状态 8 单位、给 KV 16 单位，只能放下四个请求：状态池已经满了，KV 池却还有 12 单位空闲。改成共享池后，图中的七个请求总共只用 21 单位，可以全部放下。

长上下文可能遇到相反的情况：KV 放不下了，状态池却还有空位。可以切到“长上下文”查看。共享池解决的是固定分界带来的容量浪费；如果两类数据加起来已经放满整个缓冲区，请求仍然需要等待，或者先驱逐一部分缓存。

这里的 Unified Memory，指的是同一张 GPU 上几类缓存共用容量，不涉及把数据搬到 CPU。HiCache负责另一件事：把缓存历史保存在 CPU 内存或外部存储里，需要时再加载回来。

把数据放成可以搬移的块

上一节的共享缓冲区是一整段连续的显存。分配器从两端往里分配。因此我们还需要做到：释放后的空间还回中间的空闲区；后面还会看到，它有时要挪动数据块来填补空洞。这些操作都有一个前提：分配器管理的每个块，在显存里都得连成一段。

KV 默认并不是这样存的。SGLang 按页分配 KV，一页存若干个 token；而每个 Full Attention 层都要为这些 token 各存一份 KV。默认每层有独立的缓冲区，于是同一页会按层拆成好几小块，分散在各层的缓冲区里。释放这一页，每层缓冲区里各留下一个小空洞。因此这些空洞都是分散在每一层的独立存储里。

所以 Unified Memory 改用按页存放的布局：同一页在所有层的 KV 紧挨着放，组成一段连续的空间，称为 page envelope。页里的数据完全不变，变的只是它们在显存里的位置。现在释放这一页，腾出的是一整段连续空间，可以还给共享空闲区；搬移这一页时，也只需复制这一整段，不用逐层去搬。

使用不同存放布局的代价则是：

• kernel 必须能按跨步方式读取每一层的数据，所以 Unified Memory 目前只支持一部分 attention 后端，也要求 K 和 V 的 head_dim 相同。
• 每次 forward 前，需要把虚拟编号换算成物理位置。

实测开销见文末“实验结果”。

下图用 3 层、2 页做对比，提供直观展示。

循环状态也按块存放，一个块包含各个循环层所需的状态。KV 页和状态块保留各自的大小，但都从同一个缓冲区取空间。因此，释放一个较大的状态块，就可能放下好几个较小的 KV 页。

实时搬运整合已释放空间

在Unified Memory中，释放的memory block可能夹在两个仍在使用的块之间。图中，状态块 A、B、C 占据左端。释放 B 后，虽然出现了一个空洞，C 却把它和中间的共享空闲区隔开了，导致无法完全利用上已经释放的空间。因此，SGLang 把 C 搬到 B 原来的位置，C 腾出的空间就能接入共享空闲区，用来分配新的 KV 页。

C 搬走后，请求还得找得到它。SGLang 给 C 一个不变的虚拟编号，再用一张表记录它当前的位置。请求仍然引用 C，表里的记录从“C 在位置 8”改成“C 在位置 4”。这样，搬移数据时不用逐个修改所有引用它的请求。

这个过程叫作压实（compaction）。SGLang 可以先留下空洞，等某次分配需要这段空间时再压实。搬移必须等到安全的时候才能进行，如果 GPU 计算或网络传输仍在使用这块数据，就不能在操作中途把它搬走。

共享池初版引入了这套分配和位置映射机制。后续扩展让三类数据都能共用缓冲区：全注意力 KV、滑动窗口 KV 和循环状态。滑动窗口层只需要最近一段 token，其容量需求也可能与另外两类不同。

Unified Memory的显存驱逐策略

请求完成后，SGLang 可以保留文档的 KV，以及请求末尾的 checkpoint。另一个请求如果以完全相同的 token 前缀开头，就能复用这些历史，只计算新加的后缀。而显存有限的时候，驱逐这段缓存能立即腾出空间，但后续请求可能就得重新计算它。

共用显存以后，驱逐决策也要把所有能回收的容量算进去。假设一次新分配需要 6 单位 KV 空间。删掉一段缓存前缀，会释放 1 单位 KV 和 2 单位状态 checkpoint。删掉两段，再进行必要的压实，就已经有了（1 + 2）* 2 = 6 单位可用空间。如果只数释放了多少 KV，就只会算到 2 单位，进而继续删掉本来可以留下的前缀。

SGLang 的驱逐优化会在每次驱逐后，检查共享分配器能否满足当前分配。能满足就停止。这个检查也会计入另一类数据腾出且可以回收的空间，因此能留下更多有用的缓存前缀。

实验结果

下面的实验先缓存 28 段不同的前缀，再发送一个无关的长请求制造显存压力，最后重新请求原来的前缀。前缀还在，就能省下重复处理输入的计算；前缀被驱逐了，就需要重算。

两组都已经启用 Unified Memory。 比较的是为新分配腾空间时，驱逐在什么时候停止。采用改进后的规则，更多前缀得以保留，28 个返回的请求也更早完成。

这个收益有前提：一次分配要足够大。实验中，7,000 token 的长请求一次申请全部 KV 空间；如果按 2,048 token 分四次申请，新旧规则的结果一样，都保留 17 段前缀。整体吞吐也没有稳定提升。

Unified Memory 本身也有开销：维护位置表、压实时复制数据块，以及上文说的按页布局代价。在 5 个模型上，token 间延迟中位数从快 0.81% 到慢 1.32% 不等。它的收益是更好地利用显存，而不是让计算更快。