Introduction
Language Model Programs用于调度和控制LLM的生成过程,LM programs有两项共同的特点,(1) 包含了多个穿插在控制流中的LLM调用(对于完成复杂任务和提高整体质量十分关键); (2) 接受结构化的输入并产生结构化的输出(对于LM programs的组合以及集成到现有软件系统中十分重要)。
然而目前LM programs面临着两项重要的挑战:
- 因为LLM的非确定性,LM programs实现起来繁琐且困难。程序的字符串操作、复杂输入输出的处理以及对prompts的tuning都大大降低了可读性。
- LM programs对计算和访存的冗余使用大大降低了其性能。一方面是现有的LM programs(e.g., vLLM, TGI, TensoRT-LLM)没有利用到LLM调用之间存在的共同前缀,另一方面就是针对结构化输出的解码速度问题。
SGLang包括两部分,一个是前端语言,一个是后端运行时。
在前端部分,SGLang整体作为一个嵌入在Python中的DSL(Domain-Specific Language),提供了各类原语(fork, gen,select等)。
在后端部分,SGLang设计了RadixAttention, compressed finite state machine, API speculative execution三种机制。
Radix Attention
SGLang采取了Radix Tree的数据结构并结合LRU驱逐策略来存储KV Cache。具体来说,Radix Tree是一种更高效的前缀树,通过路径压缩来让一个节点可以存储变长的字符串从而有效缩短长度。SGLang在Radix Tree中存储tokens到相应的KV Cache的映射,将KV Cache存储在非连续的分页中,每一页的大小对应一个token。同时,每个节点存储了当前正在使用该节点信息的请求数,只有当请求数为0时才能驱逐该节点。
SGLang对于排队的请求采用了一种cache-aware的调度方式,按照请求与cache中匹配的长度来排序,优先处理那些重叠长度更长的请求来放大cache的效益。
$\bigstar$ 问题:公平性,这种调度方式很容易造成饥饿,并且对于用户体验是大打折扣的,当并发请求较多时,那些具有公共前缀的新请求很容易挤掉等待中的请求。这一问题SGLang本身并没有解决。
SGLang的RadixAttention是支持分布式的,采用张量并行(Tensor Parallelism)时,每个GPU维护一个共享的KV Cache。对于数据并行的情况,SGLang采取了每个worker使用一个sub-tree(存储具体的请求内容和KV Cache),同时router维护一个meta-tree(存储索引,记录worker中的路径前缀)。router进行调度时,考虑请求与worker的亲和度以及请求之间的亲和度(共享前缀长度越长,亲和度越高),从而提高复用率。worker和router之间采用异步更新方式,当worker进行驱逐时,并非直接写入meta-tree,而是先记录下来。在系统低负载时,router再集中更新meta-tree。
$\bigstar$ 问题:locality与并行效率的权衡,采用上述调度方式,最大化了cache的复用,但容易伤害到数据并行的效率,即不同的worker之间的工作负载问题,如何在确保cache复用的同时保证不同worker之间的负载均衡也是一个问题。
Compressed Finite State Machine
设计压缩有限状态机(Compressed FSM)的目的在于让LLM的输出遵循特定格式(e.g., JSON, 代码),通常采用正则表达式来定义这种格式。
传统实现
将正则表达式转化为一个FSM,每个状态表示decode过程中的一个合法的位置,状态之间的连接边表示允许出现的下一个Token。
- 在生成每个Token时,系统维护一个当前FSM状态。
- 它检查FSM,找出从当前状态可以转移到的所有下一个状态,从而得到一个合法Token集合。
- 模型计算出的下一个Token的概率分布中,所有不在此合法集合内的Token概率会被置为零,然后从剩余Token中采样。
传统方法的瓶颈在于,当输出中包含长的、确定的字符序列时(例如JSON键 {"summary": "),传统方法显得极其低效。这个序列在Token化后可能对应多个Token,但在生成它的每一个步骤中,实际上只有一个Token是合法的(因为下一个字符是唯一确定的)。然而,系统仍然需要为每一个Token执行一次完整的前向传播,造成了大量的计算浪费,严重拖慢了生成速度。
SGLang设计的压缩有限状态机则是将那些连续的、确定性的路径压缩为一条单独的边,这条路径代表了一段token序列。SGLang进行解码时,在compressed FSM中遇到压缩边时,会直接一次性生成这一序列。考虑到后续的decode过程,这一连续序列会通过retokenization来生成若干相应的tokens。