这篇作品和十多年前的 Spark 的设计非常类似,都是 frontend (DSL) – runtime codesign 来实现全局的加速;

Frontend with DSL

Runtime (Radix Attention)

解决的问题

还是针对 KV cache 的计算角度,sglang 希望能解决不同 request 之间的 prefix 共享问题,尤其是在 system prompt 存在的情况下,很多时候的 system prompt 是大多数请求共享的,并不需要每次都计算,这样可以节省很多计算量;

不止如此,还有更多的 shared-prefix 的情况,在这些情况下应该能共享 KV cache, 这就是最核心的 motivation;

方法

因为要面向最大 prefix 建模,这里就需要一个 radix tree 的数据结构,要求在一个 prompt 的任意 token 位置都可以裂开分成多个 child 对应不同的 后续 prompt 内容;

这里由于要支持任意 token 位置断开,所以在 pagedattention 的框架下,必须要把每个 page 大小设置成一个 token 才能避免 internal fragmentation 的显存浪费;

为了支持 prompt 之间复用,理论上最好能支持连续提问的同一组问题复用,就像一个用户连续提问了多组,那么最好能设计一个生命周期 > decode 流程的算法,也就是 request 完全返回但是 kv cache 并不直接销毁,这里作者用的是 LRU 算法来回收过期冷门 token;

不过存在 lru 的 eviction 机制就需要 live kv cache (指存在 prompt 正在处理中的对应 kv) 不能被回收,也就是用一个类似 ping-page 的思路,但是可以通过 ref_cnt (本身 pagedattention 的 cow 机制也需要这个) 来实现追踪;

在有了这个 live kv cache 的 ping 保证,那么就可以灵活管理 memory 空间了,标准 CPU 的物理 LRU 是存在一个固定上限的,但是这里则是可以灵活控制一个逻辑容量上限,比如请求多了就自动回收一些远古 LRU 来回血,以支持新来的请求;

不过在 page_size=1 的设计也存在一定的局限性,联想传统 OS 中的 page 是和 page table 尺寸相关的, page 小说明 id 多,page table 的尺寸就会大很多,浪费比较多的空间用于 track;

Cache-aware scheduling