“如果你想体验极致的数学规整美感和爽感，就去实现Paged Attention；如果你想体验抽象分层设计原则和一致性维持的究极长痛，去实现Prefix Caching。”

​ ——我自己

======================================================================

这么多天没更新，今天终于更新了。之所以没有和之前一样2~3天一更新，是因为这中间的间隔，几乎全部在对Prefix Caching进行架构设计、实现迭代和反复的DEBUG。个人的体感是，在所有的推理框架实现当中，Prefix Caching绝对是除了分布式中最复杂的部分（当然，这部分非常非常大，但对于许多feature来说，其实并不是和分布式本身强绑定的）以外的皇冠上的明珠，是实现最为复杂的、同时也是相当重要的架构级feature，其挑战在于其并非是局部的算子或者数据结构设计，而是同时牵动存储分配、请求生命周期、缓存一致性的全局性架构。尤其是，如果需求足够好的兼容性，也就是需求既有radix-tree的token级别的粒度和相应可扩展性，又需求底层的Block Manager能够天然、不需要额外处理的自然同时兼容radix-tree模式，会引入非常复杂的系统架构设计需求和要求非常小心的实现。因此，这里会完整的把我的开发工程、主要最后定稿的想法和设计模式、途中遇到的一些问题，尽可能在遗忘和丢失之前。如果能够帮助到后来者和读者的话，那就太好了。

架构设计：在 vLLM 的地基上建 SGLang 的房子

LLM 推理的一个核心优化是前缀复用，这已经被多篇论文证明在特定的工作场景下非常有效。它的思想很简单：多个请求共享相同的 system prompt 时，只需计算一次 KV Cache，后续请求直接复用，通过指向相同的KV Cache，实现在无复制、无额外空间占用的情况下完成相同功能的目的。从逻辑上来说，它是对Paged Attention和分页、页表思想的自然延伸。从工程上来说。SGLang和vLLM分别用两种截然不同的方式实现了这个优化。下面的描述中就可以看出，vLLM的架构设计模式更多是在naive Paged Attention上的一个优美的最小侵入改动以部分的模拟前缀匹配，而SGLang则就真的是毫无妥协的、原生的基于完整一体化设计的Prefix Caching。

在我自己进行实现的时候，一开始其实并没有意识到一个重要问题：两者的差距其实并不小，而且其实是基于各自的先决条件（Block Size的不同默认），也没有意识到为什么vLLM要使用哈希表这种不足够优美的讨巧的“模拟radix tree”的路径。我在刚开始这个feature实现的想法是：vLLM应该是出于性能原因，或者架构太难以改动，因此做出了设计上的妥协，我应该完全能够在SGLang风格下对我自己的框架进行修改。当中途我发现这个过程里面的坑多的令人头皮发麻的时候，代码已经改的太多，无路可退了。因此，我本人几乎是非常头铁的一路撞到了底，硬生生把东西全都做了出来。不过，如果让我自己再做一遍的话，我可能应该会直接采用vLLM模式：除了可扩展性略差之外，在性能上并没有实质性的差距，而且实现要简单的多。

vLLM模式

vLLM的设计采用的是block粒度的hash table。以block_size（通常16）个token为一组，计算hash后存入字典。匹配只发生在block边界上。如果两个prompt在block中间分叉（只要block的id是不一样的），整个block都被计为miss，不会进入前缀匹配的范围。

SGLang模式

SGLang用token粒度的radix tree。每个token都是树的一个路由单元，KV Cache的存储粒度（page_size）也是1个token。树和存储粒度一致，匹配可以精确到任意token位置，不存在对齐问题。

粒度不一致的桥接

Pico-vLLM的底层存储是vLLM 风格的BlockManager（block_size=16），配合手写的Triton kernel实现PagedAttention。但前面也提到，我想要的SGLang风格的匹配模式。也就是说，token级别的radix tree，任意位置分叉都能在radix tree级别正确处理。

问题是：token粒度的索引和block粒度的存储，天然不兼容。这个问题的处理占据了我大量的开发时间，也引入了相当多的bug。不过，我最终得以找到一个相对优美的解决方案，如下所示。

三层架构

最终的设计是一个三层结构，中间有一个显式的桥接层：

RadixTree（token 粒度）
  ↑ match_prefix: 在 token 级别遍历，返回 block-aligned 的结果
  ↑ insert: 在 token 级别分裂/合并，block_ids 按边界分配
  |
PrefixCache（桥接层）
  ↑ 向下取整对齐：matched_len = (raw_len // block_size) * block_size
  ↑ 双层引用计数协调
  |
BlockManager（block 粒度）
  ↑ allocate / free / inc_ref / dec_ref
  ↑ swap_in / swap_out（offload 预留）

RadixTree 的节点存储的 cached_blocks 是逻辑 block id 列表。一条边可能有 90 个 token 但只覆盖 5 个完整 block（80 token），剩余 10 个 token 的 KV 在第 6 个 block 的前 10 个 slot 里，但这个 block 不完整，不能被 prefix cache 复用。

match_prefix 在遍历完所有匹配的边之后，在返回前做一次统一对齐：

1
2
3

aligned_len = (matched_len // self.block_size) * self.block_size
aligned_blocks = matched_blocks[:aligned_len // self.block_size]
return aligned_blocks, aligned_len, last_node

这个设计把对齐逻辑集中在一个点上，无论树的内部结构怎么分裂合并，返回给外部的永远是 block-aligned 的结果。Engine 不需要知道 radix tree 内部的 token 粒度。如果后续实现写时复制，这个部分就可以进行改变，此刻以预留的形式存在。

insert 的 block 归属

insert 触发 split 时，需要把原节点的 cached_blocks 正确分配给 split_node 和 child。一条 96-token 的边在第 90 个 token 处分裂：

split 前：node [t0..t95]  cached_blocks = [B0, B1, B2, B3, B4, B5]  (6 blocks, 96 tokens)

split 后：
  split_node [t0..t89]  cached_blocks = [B0, B1, B2, B3, B4]  (5 blocks, 80 aligned tokens)
  child      [t90..t95] cached_blocks = [B5]                   (1 block, 但只填了 6 token)

block 归属的计算：num_split_blocks = (i + match_len) // block_size - i // block_size。这个公式处理了分裂点不在 block 边界上的情况。前半段拿走所有完整 block，后半段继承剩余的。

语义设计：处理边界情况和分层解耦

（以下内容有ai辅助排版）

两层引用计数的必要性

SGLang 只有一层引用计数（lock_ref），因为它的 page 分配器和 radix tree 是一体的。vLLM 的 hash table 也只需要一层 block 引用。但 Pico-vLLM 的 radix tree 和 BlockManager 是独立的模块，而它们各自需要回答不同的问题，因其抽象层级和最细粒度同时不同。具体来说，其被设计的按如下职责工作：

BlockManager.logical_ref_count的含义是"这个block此刻有几个持有者"

持有者只有两类：正在运行的请求（通过 kv_cache 持有）和 RadixTree（作为一个整体实体持有）。ref归0意味着block立刻回收到 free pool。这是决定"block 是否可回收"的唯一依据。

RadixTreeNode.lock_ref的含义是"这个prefix段还有几个活跃请求在用"

lock_ref > 0意味着有请求正在使用这段前缀的 KV Cache。驱逐这个节点不会导致恶性的数据丢失（因为请求的持有阻止block被释放），但会引发额外的处理复杂性。在目前策略下，lock_ref归0意味着节点变得可以最小代价驱逐，进入LRU队列等待清理。

核心不变量：

logical_ref_count[block] =
    #{ 活跃请求 r : block ∈ r.kv_cache.logical_block_ids }
  + 𝟙[ 存在 radix 节点 n : block ∈ n.cached_blocks ]

这和文件系统的 inode + hard link 模型完全同构——BlockManager 是 inode table，RadixTree 是目录树。"block 什么时候可以物理回收"和"目录项什么时候可以删除"是两回事。RadixTree 无论内部有多少个节点引用了同一个 block（比如 split 前后的不同节点），在 BlockManager 看来都只算一个持有者。

lock_ref 的 node 指针方案

这是整个设计中最关键的决策，也是踩了最多坑之后才得出的结论。

最初的实现用 token 路径做 inc/dec：match 时沿 tokens 遍历 inc，close 时沿同样的 tokens 遍历 dec。看起来对称——但 insert 会改变树结构。

Bug 1：release 用 match 重新找 block

三个请求使用相同的 prompt。req4 首先 prefill，insert 把 blocks [A, B, C] 加入 radix tree。req5 在 submit 时 tree 为空（所有 submit 发生在 step 之前），matched_blocks = []，自己 allocate 了 [D, E, F]。insert 时发现路径已在，newly_held = []。

req5 close 时，release 调用 radix_tree.match(tokens) 重新找 block——找到了 req4 的 [A, B]，然后对不属于自己的 block 执行 dec_ref。A、B 的 ref 从 1 降到 0，被释放。req6 close 时又 match 到已经释放的 [A, B]，继续 dec，ref 变成 -1。

req5 close:
  match 返回 [A, B]（属于 req4 的 block）
  block_mgr.dec_ref([A, B]) → A=0 free, B=0 free  ← 错误释放！

req7 submit:
  match 返回 [A, B]（tree 指向已释放的 block）
  adopt_blocks 读到 block_mapping[A] = (NONE, -1) → 崩溃

修复：release 不再 match，而是接收 held_blocks = kv_cache.logical_block_ids（请求自己实际持有的 block 列表）。

Bug 2：split 后边长不对齐 block_size

请求 A 的 prompt 有 144 token，insert 建了一条 144-token 的边。请求 B 的 prompt 前 125 个 token 相同，第 126 个开始不同。insert 触发 split：

split_node [t0..t124]  len=125  ← 不是 16 的倍数！

请求 C 来了，match 走到 split_node，完整匹配 125 个 token。返回 matched_len=125。Engine 调用 adopt_blocks(blocks, 125)，assert 125 % 16 == 0 失败。

修复：match 返回前统一做 aligned_len = (matched_len // block_size) * block_size。

Bug 3：token 路径遍历的根本缺陷

即使修了前两个 bug，第三个更深的问题出现了。请求 R 在 submit 时 match 到 112 个 token（对齐后），inc_ref(tokens[:112]) 沿路径遍历。此时树有一条 144-token 的边，边长 > 112，inc_ref 走不到这条边，没有计数。

然后 R 的 insert 触发 split，把 144-token 的边劈成 split_node（125 token）+ child（19 token）。split_node.ref 复制了原节点的 ref。

R close 时，radix_held_len = 112 + ext * 16 = 144。dec_ref(tokens[:144]) 这次能走到 split_node（125 ≤ 144），对 split_node 执行 dec。但 inc 时没碰过 split_node——inc 了 0 次，dec 了 1 次，underflow。

inc_ref(tokens[:112]): 边长 144 > 112 → break，不 inc
  （insert 触发 split）
dec_ref(tokens[:144]): split_node 边长 125 ≤ 144 → dec → 0-1 = -1 💥

根本原因在于用token路径重新遍历做 inc/dec，但树的结构在match和close之间被insert改变了。 同样的token序列在两次遍历中走到了不同的节点。

修复方案：node指针方案

SGLang的lock_ref用的是node指针+parent chain，而不是token路径重遍历。最终采用了这个方案：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

def match_prefix(self, tokens):
    ...
    last_node = self.root    # 初始指向 root
    while ...:
        if match_len == len(edge):
            last_node = child    # 只在完整匹配边时更新
            ...
    return matched_blocks, aligned_len, last_node

def inc_lock_ref(self, node):
    curr = node
    while curr is not None and curr is not self.root:
        curr.lock_ref += 1
        curr = curr.parent

def dec_lock_ref(self, node):
    curr = node
    while curr is not None and curr is not self.root:
        curr.lock_ref -= 1
        curr = curr.parent

请求对象上存的是 request.last_node（指针）而不是token路径。close时dec_lock_ref(request.last_node)通过parent chain向上遍历以正确处理引用数减少问题。

**这个设计方案的主要目的是为了确保引用技术操作在结构变化下永远对称。**split 只是在 node 和 root 之间插入了新的中间节点（split_node），Python 对象的 parent 指针会被正确更新。从 child 走到 root 的 parent chain 自动包含了新插入的 split_node，不需要重新遍历 token。

split 前：root → child [t0..t143]  lock_ref=K
  inc_lock_ref(child): child.lock += 1

split 后：root → split_node [t0..t124] → child [t125..t143]
  split_node.lock = child.lock（纯复制）
  child.parent = split_node（指针更新）

dec_lock_ref(child): child.lock -= 1 → split_node.lock -= 1  ← 自动走过新节点，对称！

insert部分不改变lock_ref的计数。新节点创建时 lock_ref=0，分裂时 split_node.lock 纯复制自 child.lock。lock_ref 的增减完全由 match 和 close 控制，insert 是透明的。

权责清单

最终的修改入口和权责划分：

关键设计原则在于每一行的两列操作完全独立，互不依赖。BlockManager不知道RadixTree的存在，RadixTree不直接调BlockManager的free，PrefixCache是唯一协调两者的层。其中，加粗部分是可以根据不同的驱逐策略进行联合设计和改变的部分，只要确保可驱逐和操作成对合法，完全可以对这些部分进行改变，只要正确处理node废弃之后的状态、指针引用即可。

实际上，如果驱逐策略变得很复杂，两层引用计数的重要性将会开始凸显。这是因为实际上Radix Tree的引用计数本身不是释放是否合法的硬性指标（它是虚拟层级的最上层），而是“判断是否应该驱逐、应该驱逐哪些”的参照性标准。从技术本质上来说，在radix tree层面上驱逐一个计数不等于0的块是完全合法的，因为Block Manager此时并不会真的立刻释放它，而是给了Block Manager一个可以尽快在合法条件下按需将其释放的信号（或者，反过来说，radix tree是阻止block Manager将相应块释放的控制机制），这就赋予了极大的自由度（和处理被驱逐的废弃节点的相应活跃请求的复杂性）。对于需要考虑换入换出成本的radix tree的offload相关的驱逐策略，这还会变得更重要。因此，许多策略本身可以通过影响和控制相关节点的引用计数量，同时依然参照引用计数的大小来制定驱逐策略，在未来的扩展性设计中，这将节省大量开发成本。

设计模式总结，SGLang vs vLLM vs Pico-vLLM

SGLang 的优势：page_size=1 意味着树和存储粒度完全一致，不存在对齐问题。split 后每条边的 token 都有独立的 page，不会出现"半满 block"的归属问题。代价是 page table 更大（每个 token 一个 entry），管理开销更高。

vLLM 的优势：hash table 没有分裂合并，O(1) 查找，实现简单。block 粒度的管理开销最低。代价是匹配精度受限于 block 边界——两个 prompt 在 block 中间分叉时，整个 block 都 miss，浪费已计算的 KV Cache。

Pico-vLLM 的定位：在 vLLM 的 block-level 存储架构上获得 SGLang 的 token-level 匹配精度。代价是需要一个显式的桥接层处理粒度转换和双层引用计数。可以预计，这应该不是性能上最高效的方案，因为增加抽象层级和转换必然引入开销。不过，个人认为，这样引入的内存布局的兼容性可以在其他地方的可扩展性层面发挥重大作用。以我个人的经验而言，至少有一个地方非常需求这个层面的自由度调优兼容性：PD分离和多卡TP通信。在这两者的通信当中，NCCL需要拷贝后发送，而NIXL的零拷贝则需要按块为粒度进行发送，后者在block size过小的情况下会引入严重的launch overhead，从而极大拖累通信性能。因此，block size的自由调整和radix tree风格的前缀匹配的适配是非常重要且有意义的。

Benchmark和性能测试

好不容易把这么复杂的系统设计完之后，不做benchmark，看看Prefix Caching技术到底有何优劣，就太可惜了。设计两个不同长度的Prompt评测情景，对其进行性能分析。

短 Prompt（144 token）

3 种 system prompt（assistant / coder / analyst），每种 3 轮请求。第 1 轮 cold start，第 2-3 轮 warm。

hit_rate = 62.8%

短prompt下GPU prefill只需几毫秒，CPU overhead（Python侧tensor创建、scheduler调度、slot_mapping计算）占主导。prefix cache节省的GPU时间被CPU开销掩盖，加速效果有限。

实际上，对于小模型、非常短的prompt来说，由于前缀匹配引入了额外的CPU overhead的开销，用时不仅不会减少，反而会明显增加。在实测过程中，对于这个量级的prompt，大约会拖慢3040%的执行速度。换言之，Prefix Caching本身所引入的额外CPU overhead大约是5ms量级左右。如果Prefix Caching能够节省的内存量超过这个量级，那么Prefix Caching的策略就是值得的。

如果反向思考的话，这本身可能暗示了另一种策略的产生：如果对话不够长、模型足够小，那么也许直接通过某种分离策略来bypass正常的Prefix流程，按照之前的naive paged cache方法进入控制流，反而是更好的做法。

长 Prompt（2083 token shared + ~35 token variable）

固定一个 2083-token 的超长 system prompt（技术参考文档），8 个不同的 user query（各 30-50 token）。共享前缀占比约 98.5%。

共享前缀: 2083 tokens
可变部分: 28-48 tokens

                OFF(ms)     ON(ms)   speedup
Cold (req0):    48.92       41.59    1.18x
Warm (avg 7):   41.17       16.06    2.56x
Warm (best):    41.71       12.08    3.45x

2000+ token 的 prefill 让 GPU 计算成为瓶颈（约 30-40ms），CPU overhead 占比降到 20% 以下。prefix cache 省掉了 98.5% 的 prefill 计算，warm 请求的 TTFT 降到 cold 的 30-40%。

理论上界分析：prefix cache 的最大加速比是 total_tokens / new_tokens ≈ 2083 / 35 ≈ 60x。实际上的加速比只达到3.45x。这实际上是因为CPU侧的overhead实际上占据了相当大的比例，几乎有~5ms之多，这部分是无法消除的，其开销完全不会随prefix cache命中率的提高而减少（tensor创建、block table构造等仍需要执行）。此外，CUDA Graph 有固定的launch overhead，Python GIL下的 scheduler调度会产生额外开销，adopt_blocks 需要更新 GPU 上的 block table同样会引入额外开销，这些都会在我们评测使用的小模型上产生不可逾越的加速比gap上限。阿姆达定律告诉我们，最终的加速比不可能超过这些部分占的时间比例的倒数。考虑到这一点，Prefix Caching已经在职责范围内很好的完成了它的任务了。

总结和心得

到这里，所有Pico-vLLM计划内的所有主要计划内feature就已经都完整实现了，而且对于单卡的性能部分、PD分离部分的延迟抖动和异构并行度的正确性验证、Prefix Caching部分都有了完整的benchmark和Profiling。在工作过程中，也产生了若干个明显可以后续优化的点，包括了1、TP模式中的通信异步化和层间架构的重设计以允许通算重叠，2、PD分离的时候传递的后端替换为NIXL以评测性能变化，3、Scheduler的Prefill的Chunk策略和更具体调度策略的设计，4、Prefix共享block的写时复制（COW）语义的实现，5、GPU to CPU的换入换出的offload驱逐策略的实现，6、驱逐策略和Radix Tree结构的分离解耦，以及其他CPU侧代码的整体性能优化，都是未来可以不断扩充、实现和改善的地方。这些改善将在后续不断的进行，但集中的开发工程本身即到此为止。总体来说，这一个月的集中开发所取得的成果是令人感到满意的。

然后是心得。其实框架做到最后已经有sys的意思了。分层抽象，协同设计，性能和兼容性的取舍，每一个“传统”的sys设计思想，在推理框架里都有明确的体现。例如，怎么在开始动笔之前设计一个抽象层级，和它的精确的语义？这个很重要，语义设计不清就会权责不明，权责不明就产生耦合，产生耦合后续维护性就变差。有时候，这个问题比性能本身还重要。宁可性能损失3%，换取可维护性，大多数时候可能也是值得的。在这次的例子中，对于每一层的设计来说，应该具有哪些函数、其语义是什么，应该管理的东西有哪些、引用计数的来源和去向，预测的断言和边界条件情况，如果能够先尽可能的设计好再动手开发，可能反而会比现在更省心思、更省时间，也更不容易产生按下一个bug，又连带造成另一个bug，最后在修修补补中逐渐缠绕的一团乱麻的情况。在完善的开发流程中，最好都应该在纸面上先仔细想好、确保边界条件思考的尽可能全面再动笔，不仅可以省下后续大量的debug的无用时间，更可以提升最终成品的质量。本质上，在软件工程里，这种思想已经重复了无数次，但说到底绝知此事要躬行。此处教训相当宝贵，是应当牢牢记住的。