精心调♥教哈Gemi得到的文档,感觉可读性已经非常好了~
SGLang 整体执行流程
本文档旨在介绍 SGLang 的整体执行流程,从入口开始,涵盖关键模块的交互,并为后续深入理解推测解码等高级功能提供基础。
SGLang 宏观架构: Router-Worker 模式
SGLang 采用的是一个典型的“ API 网关 + 多 Worker”的分布式架构,以实现高吞吐量和可扩展性。
- API Server (Router) :
- 实现 : 这是一个用 Rust 编写的高性能
sgl-router进程,其入口点是sgl-router/src/main.rs。 - 职责 :
- Inference Worker (Python Worker) :
- 实现 : 这是一个 Python 进程,通过
python -m sglang.launch_server命令启动。其核心实现在python/sglang/srt/目录下,特别是entrypoints/http_server.py和entrypoints/engine.py。 - 职责 :
请求处理流程
一个典型的请求流程如下:
- 用户客户端向 SGLang 的 API Server (Router) 发送一个标准的 OpenAI 格式的 API 请求。
- Router 接收到请求,并根据其内部的负载均衡策略,从健康的 Worker 池中选择一个最合适的 Worker。
- Router 将请求转发给选定的 Worker。
- Worker 接收到请求,调用其内部的
engine,执行模型的前向传播。 - 在执行过程中,Worker 会与 GPU 上的 KV Cache 交互,读取历史 token 的 KV 值,并写入新生成 token 的 KV 值。
- Worker 将推理结果返回给 Router。
- Router 最终将结果返回给用户客户端。
Engine 初始化
SGLang 的运行时生命周期由 Engine 类(位于 python/sglang/srt/entrypoints/engine.py )的初始化过程驱动。这个过程精心编排了一系列子进程和服务的启动。
初始化流程概述
- 配置解析 (
ServerArgs) :Engine的构造函数接收所有配置参数,并将它们统一封装在ServerArgs对象中。这包括模型路径、并行化策略 (tp_size,pp_size,dp_size)、内存配置等。 - 环境设置 (
_set_envs_and_config) : 设置 NCCL、CUDA 等底层库所需的环境变量,并配置日志、资源限制(ulimit)等。 - 启动 Scheduler 进程 (
run_scheduler_process) : 这是最核心的步骤。
- 单机/单 DP 分片 : 根据
tp_size和pp_size启动多个Scheduler进程。每个进程负责一个 GPU,并通过tp_rank和pp_rank参数被告知其在张量/流水线并行中的角色。 - 多 DP 分片 (
dp_size > 1) : 启动一个DataParallelController进程,该进程内部会为每个 DP rank 启动一组Scheduler进程。
- 启动 Detokenizer 进程 (
run_detokenizer_process) : 启动一个独立的进程,专门负责将模型输出的 token IDs 反序列化为文本。 - 初始化 Tokenizer (
_init_tokenizer_manager) : 在主进程中创建TokenizerManager,负责接收外部请求、应用聊天模板、并将 tokenized 的结果通过 ZMQ 发送给Scheduler。 - 等待与同步 : 主进程会等待所有
Scheduler进程加载完模型并发送 “ready” 信号后,才完成初始化。
Python 端多进程架构
Engine 初始化过程的核心是创建一组协同工作的 Python 进程。这种多进程架构旨在将不同的任务隔离开,以提高整体的吞吐量和鲁棒性。以下是主要的进程及其关系:
+----------------------------------------------------------------------------------------------------------+
| Python Worker Processes (TP > 1) |
| |
| +-----------------------+ (ZMQ, Broadcast to ALL) +---------------------------+ |
| | Main Process (Engine) |----------------------------------->| Scheduler Process (GPU 0) | |
| | - TokenizerManager | | Scheduler Process (GPU 1) | |
| +-----------------------+ | ... | |
| ^ | Scheduler Process (GPU N-1)| |
| | (forks) +---------------------------+ |
| | | |
| | (All Schedulers make IDENTICAL scheduling decisions) |
| | | |
| +---------v-------------+ v |
| | Detokenizer Process |<--(ZMQ from rank 0)------------------ (One ModelRunner thread per Scheduler) |
| | - Converts IDs to text| (Sync via NCCL All-Reduce during forward) |
| +-----------------------+ |
| |
+----------------------------------------------------------------------------------------------------------+
-
Main Process (Engine) : 这是
python -m sglang.launch_server命令启动的父进程。它不直接参与模型计算,主要职责是: -
Scheduler Process : 这是推理工作的核心。
敏锐的读者可能已经从上图中注意到一个关键细节:当
tp_size > 1时,系统中存在多个并行的
Scheduler进程。那么,SGLang 是如何确保在没有中央协调者的情况下,所有这些并行的
Scheduler能够做出完全一致的决策呢?
关键在于:
每个
Scheduler都在独立地进行决策,但它们总能得到相同的结果。因为Scheduler使用了确定性调度算法。这保证没有scheduler瓶颈阻塞整个系统。更重要的是,相比于Scheduler -> N个worker,它省掉了一次进程间通信。
- Detokenizer Process : 这是一个辅助进程。
这种分离的设计使得 CPU 密集型的任务(如 tokenization)和 GPU 密集型的任务(模型推理)可以并行进行,最大化了硬件利用率。
模型与权重加载
Scheduler->TpModelWorker->ModelRunner: 这是模型权重加载的核心链条。ModelRunner:
核心组件与执行流程
SGLang 的执行流程主要围绕以下几个核心组件展开:
Engine: 作为系统的总协调器,负责接收请求、管理调度器和 worker,并返回最终结果。Scheduler: 位于python/sglang/srt/managers/scheduler.py,负责对传入的请求进行批处理(batching)和调度。TpModelWorker: 位于python/sglang/srt/managers/tp_worker.py,代表一个持有模型分片(TP/PP shard)的工作进程。它在ModelRunner的帮助下,在 GPU 上执行实际的模型推理。EAGLEWorker: 位于python/sglang/srt/speculative/eagle_worker.py,是TpModelWorker的一个特殊版本,专门用于执行 EAGLE 推测解码。
单机执行流程
- 请求接收 :
Engine接收来自用户的GenerationRequest。 - 调度 :
Engine将请求发送给Scheduler。Scheduler将多个请求组合成一个批次(ScheduleBatch)。 - 模型推理 :
Scheduler将批处理后的请求分发给TpModelWorker(或EAGLEWorker)。 - 前向传播 :
TpModelWorker在其持有的模型分片上执行前向传播。 - 返回结果 :
Engine从Scheduler获取最终的生成结果,并以流式方式返回给用户。
Scheduler 核心调度循环: handle_requests
Scheduler 的核心是一个位于 handle_requests 方法中的 while True 无限循环。这个循环是 SGLang 推理服务的心脏,它持续不断地执行以下任务,驱动所有请求从接收到完成的整个生命周期。
这个核心循环的伪代码如下:
# in sglang.srt.managers.scheduler.Scheduler.handle_requests
def handle_requests(self):
while True:
# 1. 接收新请求
# 尝试从 Router 非阻塞地接收一个或多个请求
new_req_inputs = self._recv_new_requests()
# 2. 处理新请求
# 将接收到的原始请求输入 (e.g., TokenizedGenerateReqInput)
# 转换为内部的 Req 对象,并加入到 self.waiting_queue
for req_input in new_req_inputs:
self._request_dispatcher.dispatch(req_input)
# 3. 调度并执行模型
# 这是最关键的步骤,决定下一个要执行的批次
self.schedule()
# 4. 处理模型输出
# schedule() 内部会调用模型并处理其输出,
# 包括更新请求状态、发送 token 回复等。
请求生命周期详解
让我们跟踪一个生成请求( TokenizedGenerateReqInput )的完整生命周期,以理解上述循环是如何工作的:
- 请求到达 :
- Router 将一个来自用户的 HTTP 请求转换为
TokenizedGenerateReqInput对象,并通过 ZMQ Socket 发送给一个 Worker 的Scheduler进程。 - 在
handle_requests循环的第 1 步,_recv_new_requests方法通过self.router_port.recv_pyobj()接收到这个对象。
- 请求入队 :
- 循环的第 2 步调用
self._request_dispatcher,它会匹配到handle_generate_request方法。 handle_generate_request方法进行一系列检查(如长度校验),然后创建一个核心的Req对象。这个对象封装了请求的所有信息,包括 prompt token、采样参数、状态等。- 最后,这个
Req对象被添加到self.waiting_queue列表中,等待被调度。
- 调度决策 (
schedule方法) :
- 循环的第 3 步调用
self.schedule(),这是调度的核心。 schedule()方法首先会尝试将self.waiting_queue中的新请求(prefill 请求)和self.running_batch中正在运行的请求(decode 请求)合并。- 它使用
self.policy.schedule()策略函数来决定哪些请求可以被调度。策略会考虑当前系统的负载、KV Cache 的剩余空间、请求优先级等因素。 - 调度的结果是一个
ScheduleBatch对象,它包含了所有被选中要进行 prefill 或 decode 的请求。
- 模型执行 :
schedule()方法根据ScheduleBatch的内容,决定是执行 prefill 还是 decode。- Prefill : 如果批次中包含新的请求,
schedule()会调用self.model_worker.forward_prefill(batch)。 - Decode : 如果批次中只包含正在运行的请求,则调用
self.model_worker.forward_extend(batch)。 - 这些
forward_*方法会触发底层的ModelRunner在 GPU 上执行实际的模型前向传播。
- 输出处理与状态更新 :
- 模型执行完成后,
schedule()方法会调用self._process_outputs(output_dict)来处理返回的结果。 _process_outputs会遍历批次中的每一个请求,更新它们的状态:
- 资源释放 :
- 对于已完成的请求,
schedule()方法会调用self._free_request(req)。 _free_request会释放该请求占用的所有资源,最重要的是调用self.tree_cache.free(req)来归还其占用的 KV Cache 物理块。这些物理块随后可以被新的请求重新使用。
通过这个 handle_requests -> schedule -> model_worker.forward -> _process_outputs 的闭环, Scheduler 得以高效地处理大量并发请求,实现了持续的批处理(Continuous Batching)。
并行化与 Tensor 形状
SGLang 主要使用张量并行( Tensor Parallelism , TP)来加速单次 forward 的计算。为了避免 padding 带来的计算浪费,SGLang 的底层 Kernel 采用“摊平”的物理视图,通过 qo_indptr([B+1]) 索引来区分批次中的不同序列。
- 定义 :
关键 Tensor 的形状变化 (TP > 1)
- 核心思想 :