专栏名称: GiantPandaCV
专注于机器学习、深度学习、计算机视觉、图像处理等多个方向技术分享。团队由一群热爱技术且热衷于分享的小伙伴组成。我们坚持原创,每天一到两篇原创技术分享。希望在传播知识、分享知识的同时能够启发你,大家一起共同进步(・ω<)☆
目录
相关文章推荐
GiantPandaCV  ·  【ml-engineering ... ·  4 天前  
GiantPandaCV  ·  图解OpenRLHF中基于Ray的分布式训练流程 ·  6 天前  
51好读  ›  专栏  ›  GiantPandaCV

SGLang 后端原文解析

GiantPandaCV  · 公众号  · 3D  · 2024-11-21 20:32

主要观点总结

本文介绍了SGLang实现的后端/Runtime框架,包括LM Program的使用及面临的问题,SGLang在前后端做的改进,以及RadixAttention在KV Cache管理中的应用。文章还讨论了删除节点的方法,Request Scheduling后端请求调度,以及前端请求发送顺序的重要性。

关键观点总结

关键观点1: SGLang实现的后端/Runtime框架概述

文章介绍了SGLang在实现后端/Runtime框架时面临的挑战,包括语言模型的输出不稳定和KV Cache复用方法的改进。

关键观点2: LM Program的使用及问题

文章阐述了LM Program在推理端的主要面对的问题,包括语言模型的输出高度不稳定和现有的KV Cache复用方法需要改进。

关键观点3: SGLang在前后端的改进

针对上述问题,SGLang在前后端做了相应的改进。采用Compressed Finite State Machine解决语言模型的输出不稳定问题,基于RadixAttention优化后端Runtime来改进KV Cache的复用方法。

关键观点4: RadixAttention的应用

文章详细解释了RadixAttention在KV Cache管理中的应用,包括其基于树状结构的管理方式、LRU删除节点的方法、内存管理、request调度等。

关键观点5: 请求调度与前端部分

文章讨论了后端的请求调度,通过利用Radix Tree进行KV Cache reuse,提高cache hit rate,并介绍了前端部分的调度策略以及存在的问题。


正文



作者丨Chayenne Zhao
来源丨https://zhuanlan.zhihu.com/p/716543182
编辑丨GiantPandaCV


众所周知,麻辣系统博大精深。今天浅浅捋一捋 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 而言,在推理端主要面对的问题:

  1. 语言模型的输出高度不稳定:依据我的经验,类似 Vicuna-7b 这样的模型,在结尾多出一个 \n 都会极大的影响模型的输出。因此,我们需要有更加严格的 rule-strict generation,目前可以把这种 rule 简化为 format,譬如强制生成 Json 格式。

  2. 现有的 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/articles/Radix_tree

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 时,这些请求终究还是会被发送到后端的。

后端部分的解析到此结束,前端部分等我玩明白了再来。


- The End -


GiantPandaCV

长按二维码关注我们

本公众号专注:

1. 技术分享;

2. 学术交流

3. 资料共享

欢迎关注我们,一起成长!