第 3 课：SGLang SRT 源码地图：Scheduler、ScheduleBatch 与 ModelRunner

本课目标与通过标准

本课产出：能沿 SGLang SRT 的真实 event loop 追踪请求从 waiting queue 到 prefill/decode batch、TP worker、ModelRunner 和结果处理，并与 vLLM V1 的职责切分做源码级比较。

通过要求：唯一代码填空题通过全部断言；Q1～Q3 都能沿源码对象给出因果链；能指出一个正确性不变量、一个性能边界和一个需要 benchmark 才能确认的结论；总分至少 8/10。

源码版本与阅读方法

本课程使用固定 commit 的永久链接保证行号可复现，同时链接 current docs 供核对最新变化。阅读时先画调用图和状态所有权，再进入分支；不要从大文件第一行机械顺读。

建议阅读顺序：

顺序	源码入口	关注对象
1	python/sglang/srt/managers/scheduler.py::run_event_loop	normal/overlap 两条主循环如何选择
2	get_next_batch_to_run、get_new_batch_prefill、update_running_batch	waiting、chunked request、running batch 怎样迁移
3	managers/schedule_batch.py::Req、ScheduleBatch、NextBatchPlan	请求状态与 GPU batch tensor 的边界
4	managers/tp_worker.py::TpModelWorker.forward_batch_generation	scheduler 如何调用模型 worker
5	model_executor/model_runner.py::forward、_forward_raw、sample	attention backend、CUDA graph、模型与 sampler 在哪里接合

源码快照：SGLang 9deb6952afa4。

核心心智模型

1. 它是什么，解决什么问题

SGLang SRT（SGLang Runtime）是由加州大学伯克利分校团队开发的高性能大语言模型推理引擎。与纯通用批处理推理系统不同，SGLang 最初专为复杂的结构化生成、多轮对话与智能体（Agent）复杂逻辑树探索而生。 在复杂大模型程序（如 ReAct Agent、Tree-of-Thought、自我纠错长思维链）中，系统面临严重的复杂前缀重复计算与小批量自回归的 CPU 驱动瓶颈：

• 传统系统仅能缓存固定的 System Prompt，无法高效捕捉在会话中途分叉、并发探索的多路径前缀；
• 在单步 Decode 过程中，PyTorch Eager 模式下的 Python 循环与 Kernel 发射延迟高达数百微秒，严重压制 GPU 算力利用率。

SGLang SRT 提出了以 RadixAttention（基数树动态前缀感知）与 多尺寸 CUDA Graph 捕获分发引擎为核心的运行时体系（图 3-1 与图 3-2）。它解决的核心工程问题是：通过树状结构全自动追踪任何历史 Token 前缀，并利用高效的分层合批容器（ScheduleBatch）与 CUDA Graph，将多轮交互与复杂 Agent 工作流的 TTFT 与推理吞吐推向极致。

2. 它如何工作

SGLang SRT 的分层协作与批处理推进时序如图 3-1 与图 3-2 所示：

TokenizerManager (多进程)

并行 Tokenizer 编解码

SRT Scheduler

ScheduleBatch 动态合批容器

Radix Tree 索引空间

ModelRunner (TP / CUDA Graph)

CUDA Graph 捕获执行图

FlashInfer / PagedAttention 算子

图 3-1：SGLang SRT 运行时分层交互与 ScheduleBatch 容器拓扑

接收批次请求

Radix Tree 快速最长前缀检索

构造混合批次 (Prefill + Decode ScheduleBatch)

命中捕获尺寸的 CUDA Graph?

重放 CUDA Graph (超低调度开销)

回退 Eager 动态形状前向执行

多项式采样与树节点更新

是否

图 3-2：SGLang SRT 混合批处理与 CUDA Graph 捕获分发时序流程

1. 三大核心抽象分工解耦（图 3-1）：
   • Req（单请求持久业务实体）：拥有单请求完整的跨步语义状态，包括输入/输出 Token 列表、采样参数、停止条件，以及在 Radix Tree 中的当前挂载节点；
   • ScheduleBatch（调度器层批量控制视图）：由 Scheduler 在 CPU 侧负责编排维护的动态批次容器。负责请求优先级排序、Prefill 与 Decode 区分聚合，以及向物理内存池请求分配 KV Slot；
   • ForwardBatch（设备端执行描述符）：由 TpModelWorker / ModelRunner 消费的 GPU 物理张量集合。包含打平的 Input Tokens、序列长度偏移数组 seq_lens、PagedAttention 物理查找表指针，以及与 CUDA Graph 对齐的批次元数据。
2. 两阶段批处理决策（Prefill 优先与 Decode 轮转，图 3-2）： 在调度器核心函数 get_next_batch_to_run() 中，状态机遵循严谨的转移次序：
   • 第一步：合并已就绪的 Prefill：上一个 step 中刚刚执行完前向 Prefill 的请求，在当前轮首先被合并进入活跃的 running_batch 集合；
   • 第二步：新 Prefill 准入判定：若 WAITING 队列中存在可接纳的新请求，利用 Radix Tree 检索最长公共前缀后，优先构造 prefill 批次；
   • 第三步：Decode 批次推进：若无新 Prefill，则将存活的 running_batch 请求打包为 decode 批次下发执行。
3. CUDA Graph 捕获与回放（图 3-2）： ModelRunner 在启动阶段预先针对一系列典型的批次大小（如 Batch Size = 1, 2, 4, 8, 16, 32, 64, 128）录制捕获了静态的 CUDA Graph。对 decode 批次，若请求数不等于某个捕获尺寸，会被填充（Padding）到下一个更大的捕获尺寸后重放（Replay）整张图，绕过 PyTorch 的 CPU 解释与底层驱动启动开销；只有批次大小超过最大捕获尺寸，或本轮是 extend/prefill 批次时，才回退至 Eager 动态前向模式。

3. 正确性条件与常见误区

• 禁止跨层合并三大对象（Req / ScheduleBatch / ForwardBatch）： 初学者常试图将三者合为一个大对象。这一做法是致命的：Req 是长周期 CPU 业务对象，而 ForwardBatch 包含了大量与具体 step 强绑定的物理 GPU Tensor。若将其混为一谈，会导致庞大的中间激活与张量随长请求一直驻留显存引发 OOM，且破坏了 Overlap 异步调度流水线的隔离性。

• 异步 Overlap 调度中的写后读（WAR）屏障： 在 SGLang 的 Overlap 模式中，CPU 主线程在准备当前轮 Batch 元数据时，GPU 正在并发读取上一个轮次的输入 Buffer。系统必须通过 Stream Barrier（CUDA Event / 内存栅障）保证 GPU 完全完成读取后，CPU 才能向共享内存复写新的输入张量，杜绝发生脏读与内存撕裂。

• Radix Tree 引用计数与修剪互斥不变量： 只要某个树节点对应的 KV 块正被某个活跃 Req 读取，其引用计数必须严格大于 0。LRU 驱逐算法只能修剪引用计数归零的闲置叶子节点，严禁误删活跃节点。

4. 性能与工程取舍

• CUDA Graph Padding 开销 vs CPU 发射延迟： 为了提高 CUDA Graph 命中率，系统通常需要把不成整数倍的请求向上对齐填充（Padding）到最近的静态尺寸桶。Padding 带来了一定程度的无用计算，但消除了数十微秒的 Python 解释开销，在小 Batch Decode 场景下利远大于弊。

• Prefill 优先策略对 TTFT 与 ITL 的博弈： get_next_batch_to_run 默认优先调度新到达的 Prefill 请求，这极大优化了首字到达时间（TTFT）；但代价是正在生成的 Decode 请求会被暂时挂起一轮，导致字间生成延迟（ITL）产生短暂尖峰。因此，SGLang 在长序列场景下引入了 Chunked Prefill 与动态优先级调节来平抑这种抖动。

具体演示：NextBatchPlan 批次选择决策演算

考虑调度循环中的一个经典决策时刻：

• 当前运行队列：running = ["d0"]（包含 1 个正在进行自回归的请求）；
• 上一轮刚刚完成 Prefill 计算的请求：completed_prefill = ["p_done"]；
• 等待队列中新到达、且通过 Radix 匹配准备执行 Prefill 的请求：new_prefill = ["p_new"]。

调度决策推导：

1. 合并上一轮成果：running.extend(completed_prefill) \implies running = ["d0", "p_done"]。此时持久运行集合已经原子扩充为 2 个请求；

2. 判定运行批次：

   • 检测到 new_prefill 非空，根据“新 Prefill 优先以压低首字时延”策略，本次调度计划选出执行 batch_to_run = ("prefill", ["p_new"])；
   • 最终返回计划字典：{"batch_to_run": ("prefill", ["p_new"]), "running_batch": ["d0", "p_done"]}。

3. 后继演进： 在下一个 step 中，若没有更新的 Prefill 进入，系统将执行 batch_to_run = ("decode", ["d0", "p_done"])，正式对两路请求并发进行自回归解码。状态机转移清晰无歧义。

实践任务：唯一代码填空题

补齐 NextBatchPlan 的教学版选择逻辑：先合并上轮完成的 prefill，再优先运行新 prefill，否则运行 decode。

只能修改 TODO/______ 位置，不得删除断言或放宽通过条件。

def choose_sglang_batch(running, completed_prefill, new_prefill):
    running = list(running)
    # TODO：完成 prefill 的请求进入 running；新 prefill 优先于 decode。
    ______
    if new_prefill:
        batch_to_run = ("prefill", list(new_prefill))
    elif running:
        batch_to_run = ("decode", list(running))
    else:
        batch_to_run = None
    return {"batch_to_run": batch_to_run, "running_batch": running}

assert choose_sglang_batch(["d0"], ["p_done"], ["p_new"]) == {
    "batch_to_run": ("prefill", ["p_new"]),
    "running_batch": ["d0", "p_done"],
}
assert choose_sglang_batch(["d0"], [], []) == {
    "batch_to_run": ("decode", ["d0"]),
    "running_batch": ["d0"],
}

检查方法

运行断言；补一个三者皆空的用例。再说明真实源码在哪些条件下不会简单优先新 prefill。

Q1

Req、ScheduleBatch、ForwardBatch 三者为何不能合并成一个类？

你的答案：

Q2

为什么 get_next_batch_to_run 返回 plan 时还要返回更新后的 running_batch？

你的答案：

Q3

如何从源码判断性能问题在 scheduler CPU、batch preparation 还是 ModelRunner/kernel？

你的答案：

评分规则

• 代码 4 分：主路径 2 分、边界条件 1 分、能映射回源码对象 1 分；
• Q1～Q3 各 2 分：必须包含对象、状态变化、正确性或成本链；
• 一票否决：把论文峰值写成普遍生产结论；把源码支持写成所有模型/硬件可用；混淆算法正确性与性能；只背类名而说不清状态所有权。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先独立完成。核对后请更换一个 batch、cache 或硬件条件重新推演，不能只复制。

def choose_sglang_batch(running, completed_prefill, new_prefill):
    running = list(running)
    running.extend(completed_prefill)
    if new_prefill:
        batch_to_run = ("prefill", list(new_prefill))
    elif running:
        batch_to_run = ("decode", list(running))
    else:
        batch_to_run = None
    return {"batch_to_run": batch_to_run, "running_batch": running}

assert choose_sglang_batch(["d0"], ["p_done"], ["p_new"]) == {
    "batch_to_run": ("prefill", ["p_new"]),
    "running_batch": ["d0", "p_done"],
}
assert choose_sglang_batch(["d0"], [], []) == {
    "batch_to_run": ("decode", ["d0"]),
    "running_batch": ["d0"],
}

Q1 参考答案

Req、ScheduleBatch 和 ForwardBatch 三者各自代表了完全不同的抽象层级、生命周期、设备存储位置和状态写入者，严格解耦是现代大模型推理引擎高吞吐与正确性的根基：

1. 生命周期与存储介质的分离：Req 是全生命周期驻留在 CPU 上的长周期状态机，包含原始 Prompt tokens、已生成的采样 tokens、采样参数（temperature/top_p）、停止条件、结构化约束（Regex/JSON schema）以及 Radix Tree 节点的引用引用，跨数十甚至上千个 step 存活；ScheduleBatch 是调度器单步决策产生的控制对象，驻留在 CPU 端，管理本步调度的一组请求列表、KV block 增减分配意图、内存配额检查与抢占/回退（preempt/retract）元数据；ForwardBatch 则是经过 flattened 和 tensorized 后的设备端结构体，通常在 TP worker / ModelRunner 层面生成，持有 GPU 上的平铺张量（如 input_ids, req_pool_indices, seq_lens, out_cache_loc, attention metadata 等），仅用于喂给 PyTorch / FlashInfer / Triton kernels 执行当轮 forward，生命周期仅为单次 forward。

2. 状态写入者与并发所有权（Ownership & Concurrency）：Req 仅由 CPU 调度器事件循环独占更新；ForwardBatch 是纯粹的只读设备视图，由底层 C++/CUDA 算子直接解包并高效消费。若将三者混为一体，会让持久的请求对象长期悬挂大尺寸临时显存张量导致显存泄露，或者迫使模型执行层反向侵入调度器业务逻辑，破坏关注点分离。

3. Overlap 流水线掩盖机制：在 SGLang overlap 调度模式下，当 GPU 正在执行 Step $t$ 的 ForwardBatch 时，CPU Scheduler 已经在准备 Step $t+1$ 的 ScheduleBatch 并投递 ForwardBatch。三者严格解耦是保证 CPU 调度逻辑与 GPU 计算完全流式掩盖（latency hiding）的前提。

Q2 参考答案

get_next_batch_to_run 在挑选本步批次时也会改动 running_batch：

1. 合并与清理：上一轮完成 prefill 的请求并入 running_batch；已结束（遇到停止符、达到最大长度或被 abort）的请求移出并释放 Radix 树引用；因显存不足被 retract 的请求退回等待队列。
2. 所有权推进：即使本步运行的是新的 prefill batch，暂停的 decode 请求集合也已经变化。返回更新后的 running_batch，主循环用一次赋值接管它，避免某个分支漏掉就绪请求或留下已结束的请求。

Q3 参考答案

定位 SGLang SRT 性能瓶颈需要结合代码分段高精度计时与 Nsight Systems / PyTorch Profiler 全链路 GPU timeline 分析：

1. Scheduler CPU 与 Batch Preparation 瓶颈：在 run_event_loop、get_next_batch_to_run 以及 prepare_for_decode / prepare_for_extend 注入微秒级打点。如果在 Nsys timeline 上观察到 CUDA Compute Stream 上不同 step 的 kernel 之间存在明显的空洞（GPU 气泡，通常超过 $200\ \mu\mathrm{s}$），且 CPU 单核利用率达到 100%，则说明瓶颈在于 Python 侧的调度循环、Radix Tree 节点遍历查找开销或 Python 对象的 tensorization 拼接开销。

2. CUDA Graph 捕获与回退瓶颈：若 GPU 单步执行时延显著高于预期，排查是否频繁触发 Eager 模式回退。在 Nsys 中比对单步 launch 的 kernel 个数：正常命中 CUDA Graph 时整张网络仅有数个 Graph Launch 节点；若 decode 步出现成百上千个零碎 kernel 串行发射，说明并发 batch size 超过了最大捕获尺寸（未命中的较小尺寸只会被填充到下一档，不会回退），或该步是 extend/prefill 批次，导致算子发射开销暴露在关键路径上。必要时调大 cuda_graph_max_bs。

3. ModelRunner 与算子/通信瓶颈：若 batch 已快速入队且 CUDA Graph 稳定生效，但单步 forward kernel 持续时间极长，需进一步下钻：在多卡 TP 部署下，检查 ncclKernel_AllReduce 通信占空比是否反常偏高（排查跨 NUMA 节点传输、NVLink 拓扑降级或卡间负载不均）；在单卡 kernel 内部，使用 NCU 排查 FlashInfer / PagedAttention 算子的 memory-bound 程度、L2 缓存命中率与 KV cache 读取带宽是否已逼近硬件物理瓶颈。

参考资料

• SGLang Scheduler source
• SGLang ScheduleBatch source
• SGLang TP worker source
• SGLang ModelRunner source
• SGLang paper

源码链接固定到本课审阅 commit；current docs、支持矩阵和默认参数会变化，面试或部署前必须按目标版本重新核对。

← 上一课
第 2 课：vLLM V1 统一调度器、KV… 下一课 →
第 4 课：SGLang RadixAtte…