众所周知，麻辣系统博大精深。今天浅浅捋一捋 SGLang 实现的后端/ Runtime 框架，前端部分留作后文。
LM Program
Using a program to schedule and control the generation processes of LLMs.
这个概念听上去很抽象，但是可以界定为 Program Controled LLM Generation，而 program 的类型非常多样。最典型的例子譬如 Tree of Thought Planning，需要用树状结构和状态机来控制模型的步进、回滚和剪枝。
对于 LM Program 而言，在推理端主要面对的问题：
- 语言模型的输出高度不稳定：依据我的经验，类似 Vicuna-7b 这样的模型，在结尾多出一个
\n都会极大的影响模型的输出。因此，我们需要有更加严格的 rule-strict generation，目前可以把这种 rule 简化为 format，譬如强制生成 Json 格式。 - 现有的 KV Cache 复用方法还有待改进：仍旧用 Tree of Thought 来举例，vllm 利用了近似传统操作系统中的 page cache 的方法来缓存。每次查找 cache 所在的 chunk 需要先对 prefix 进行 hash 映射，然后查找对应的 chunck。对于存在前缀关系的 prefix，很可能被完全随机的映射到毫无关系的不同 chunk，降低了 cache 的处理速度。而基于前缀树来管理 KV Cache 天然就很适合处理有前缀性质的 KV Cache reuse。
基于此，SGLang 在前后端做了相应的改进。对于第一个问题，采用的方法以 Compressed Finite State Machine（压缩有限状态自动机）为主，且待下回分解。而对于第二个问题，主要通过基于 RadixAttention（前缀树的一种）的方法来优化后端 Runtime。
RadixAttention
https://www.wikiwand.com/en/a...
Radix Tree（也称为 基数树 或 压缩前缀树）是一种基于树的数据结构，常用于高效地存储和查找字符串或键值对。它是一种空间优化的 Trie 树，通过合并具有公共前缀的节点来减少存储空间。每个节点代表一个字符串的前缀，路径上的节点共同构成完整的字符串或键。
而 RadixAttention 可以粗糙概括为基于 Radix Tree 管理的 KV Cache。实际上将树状结构引入到 KV Cache 管理中是个很自然的想法，毕竟已经有了利用哈希表来管理 KV Cache 的工作。我的朋友实现过树状结构管理 IO 的工作，略作推荐：
https://arxiv.org/abs/2404.00242
树状结构的增删查（没有改）是非常自然的，并且计算开销几乎可以忽略不记（在原文的 ablation study 中有论述）。这里先放上主图，一目了然：
从原文摘录的 RadixAttention 的增删查改示意图。
值得提出的是，删除节点采用的方法是 LRU（leaset recently used），在进行 batch inference 时，需要对正在使用的节点添加引用计数，避免因为其他查询请求没有用上而删去。类似这样的的实现细节自然是繁琐且工作量巨大的，开发实在不易。
此外，RadixAttention 的 running request 和 KV Cache 使用了同一块动态管理的内存。当 cache hit rate 较高时，相应 running request 的处理速度更快，batch size 也可以更大。在 system prompt 相当长时，cache hit rate 高达 99.94%，request 的处理速度非常之快。
Request Scheduling
后端对于给定的请求序列会利用 Radix Tree 进行 KV Cache reuse，可想而知在前端可以调度 request 发送到后端的顺序，提高临近 requests 之间的 prefix 重叠率，进一步提高 cache hit rate。这种调度算法给出的 request 顺序称为 longest-shared-prefix-first order。
作者在原文中证明了在离线和在线模式下，longest-shared-prefix-first order 是趋于最优的。而在分布式条件下，由于同一棵树的增删查改是确定的，因此只需要在多张 GPU 上初始化同一颗树，此后的增删查改完全一致，不需要为了 KV Cache 进行多卡间的通讯。
值得强调的是，这种 scheduling request 的方法可能会导致在请求数量非常大的情况下，存在某些和其他请求几乎没有 shared-prefix 的 request 一直没法被处理，也即可能导致这些请求饥饿（starvation）。目前这一块问题还未解决，然而实际上几乎不会遇到这种情况，因为 queue request 为 0 时，这些请求终究还是会被发送到后端的。
后端部分的解析到此结束，前端部分等我玩明白了再来。
END
作者：Chayenne Zhao
来源：GiantPandaCV
推荐阅读
- Ultrylytics 官宣: YOLO11 全新发布！
- AlphaFold3 重磅开源，诺奖级 AI 颠覆世界！GitHub 斩获 1.8k星，本地即可部署
- 小白视角：利用 vllm serve 新的 Embedding Model
- 小白视角：vllm 迁移到 SGLang 的体验与收获
欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
