SGLang 推理框架深度解析:请求的调度与生命周期
SGLang 推理框架深度解析:请求的调度与生命周期
在大语言模型(LLM)服务的核心,一个高效的调度系统是实现高吞-吐量与低延迟的关键。SGLang 作为先进的推理框架,其内部的调度器(Scheduler)通过一套精密设计的机制,系统化地管理着每一个请求从进入到完成的完整生命周期。这种管理不仅关乎顺序,更是一场在计算、内存与I/O之间的精妙平衡艺术。
核心理念:Prefill 与 Decode 的两阶段分离
要理解 SGLang 的调度逻辑,首先必须明白 LLM 推理的两个核心阶段:
-
Prefill(预填充)阶段:此阶段处理用户输入的完整提示(Prompt)。这是一个计算密集型过程,因为模型需要一次性为输入序列中的所有 Token 计算注意力并生成初始的键值缓存(KV Cache)。 首个 Token 的生成延迟(Time to First Token, TTFT)主要取决于此阶段的效率。
-
Decode(解码)阶段:在 Prefill 完成后,模型进入自回归生成阶段,即逐个生成新的 Token。此阶段属于访存密集型,因为每生成一个新 Token,都需要从显存中读取此前所有 Token 的 KV Cache。 尽管单步计算量较小,但其总时长决定了最终的生成速度(Tokens per second)。
SGLang 的调度器正是围绕这两个截然不同的阶段来组织请求,以最大化 GPU 的利用率。
请求状态的容器:四大核心批处理队列
SGLang 的调度器通过四个关键的列表或队列来追踪和管理不同状态下的请求。这些队列共同构成了请求流转的骨架。
-
waiting_queue
:这是一个等候区,所有新抵达的请求都会首先被放入此队列。 SGLang 允许根据不同策略对该队列进行排序,例如默认的 “lpm”(Longest Prefix Match,最长前缀匹配)策略会优先处理那些与其他请求有更长公共前缀的请求,以最大化 KV 缓存的复用率。 -
new_batch
:这是为 Prefill 阶段准备的批次。调度器会从waiting_queue
中挑选出符合当前内存和计算条件的请求,组成new_batch
。这些请求即将进入计算密集型的预填充处理。 -
running_batch
:此批次包含了所有已经完成 Prefill 阶段、正处于 Decode 阶段的请求。它们在等待逐个生成新的 Token,是连续批处理(Continuous Batching)的核心部分。 -
cur_batch
:代表当前调度周期内,即将被送入模型进行单次前向传播的最终批次。它的形成逻辑体现了 SGLang 的调度优先级:如果new_batch
不为空,那么cur_batch
就是new_batch
,优先处理预填充任务;否则,cur_batch
将由running_batch
中的请求构成,继续生成任务。
一个请求的完整生命周期
一个请求在 SGLang 系统中的旅程,是一个在上述队列间动态流转的循环过程。
1. 到达与入队
当一个新请求到达时,经过基本校验后,它被封装成一个内部 Req
对象,并被加入到 waiting_queue
中,静待调度器的“翻牌”。
2. 调度与批处理形成
调度器的核心循环会周期性地执行。在一个周期的开始,调度器首先会处理上一轮遗留的任务,然后尝试从 waiting_queue
中构建新的批次。
- 内存感知调度:调度器会检查当前可用的 GPU 显存,然后从
waiting_queue
的队首(根据优先级排序)开始,逐个将请求移入new_batch
,直到无法容纳下一个请求为止。 - 分块预填充(Chunking):如果一个请求的 Prompt 过长,导致其所需的 KV Cache 超出当前可用内存,SGLang 不会直接拒绝或长时间等待。相反,它会将该请求进行“分块”,只将能够处理的一部分放入
new_batch
,剩余部分则留在waiting_queue
中等待后续的 Prefill 周期。这确保了即使面对超长文本,系统也能保持运转,不会被单个大请求阻塞。
3. 执行:Prefill 与 Decode 的协作
批次形成后,cur_batch
被送入模型执行 run_batch
函数。
- Prefill 优先:如果
cur_batch
来自new_batch
,模型将执行forward_extend
操作。在这一步,系统会利用 RadixCache 寻找并复用公共前缀,然后为剩余的输入 Token 计算并存储 KV Cache。 - 持续 Decode:如果当前没有需要 Prefill 的新请求,
cur_batch
则由running_batch
构成,模型执行forward_decode
操作,为批次中的每个请求生成一个新 Token,并将其 KV Cache 追加到已有的缓存中。
4. 动态调整与容错:撤回机制(Retraction)
在资源极度紧张的情况下,SGLang 还设计了撤回机制。如果在为 Decode 请求分配下一个 Token 的缓存时发现内存不足,调度器可以根据特定策略,从 running_batch
中“撤回”(Retract)一个或多个请求,将它们重新放回到 waiting_queue
的队首。这会释放它们占用的 KV Cache,为其他请求腾出空间,保证了系统的稳定性。
5. 结果处理与迭代
run_batch
执行完毕后,调度器调用 process_batch_result
来更新每个请求的状态。
- 完成与释放:如果一个请求生成了终止符或达到了最大长度,它被标记为已完成。其完整的生成文本被缓存,占用的相关内存资源(如
req_to_token_pool
中的条目)被释放。同时,其最终的 KV Cache 序列会被插入到 RadixCache 中,供后续请求复用。 - 继续与迭代:未完成的请求则继续留在
running_batch
中,等待下一轮的 Decode。 - 循环往复:调度循环随即开始下一轮迭代,再次检查
waiting_queue
,形成新的cur_batch
,周而复始,直到所有请求处理完毕。
通过这种将请求划分为不同状态队列,并采用 Prefill 优先、分块处理和动态撤回等一系列精细化策略,SGLang 的调度器实现了对 GPU 资源的高效、灵活和稳健的管理。它像一位经验丰富的指挥家,确保了计算密集型和访存密集型任务能够和谐共存、交错执行,从而在复杂的并发场景下,依然能够提供卓越的推理性能。