Scheduler、批次与 KV Cache#
这章解决什么问题#
这一章解决的不是“请求从哪进来”,而是“请求已经进来之后,系统怎样决定谁先跑、谁继续跑,以及这些请求占用的 KV cache 怎样被映射、复用和回收”。如果没有这层理解,你读 scheduler.py 时会只看到大量条件分支;读 memory_pool.py 时会只看到张量分配;两边都看了,仍然不知道它们为什么必须一起工作。
这里最值得先抓住的事实,是 python/sglang/srt/managers/schedule_batch.py 文件头写出的那条数据流:ScheduleBatch -> ModelWorkerBatch -> ForwardBatch。这条注释说明调度阶段并不是直接操作 GPU forward 所需的最低层张量,而是先构造一个更高层的 batch 表示,再逐步降到执行层。调度与内存的配合,也正是围绕这条转换链组织的。
Scheduler 看的不是单个请求,而是“当前批次状态”#
python/sglang/srt/managers/scheduler.py 里的 Scheduler 初始化时会建立等待队列、running_batch、last_batch、tree_cache、req_to_token_pool 等状态。它不是简单地从队列里拿一个请求就调用模型,而是持续维护“现在已经在跑什么、下一轮还能塞什么、哪些请求应该进入 prefill、哪些请求应该继续 decode”。
get_next_batch_to_run() 很适合当作阅读入口。这个方法先处理 timeout、过滤完成请求、把上一轮 prefill batch 合并进 running_batch,然后决定本轮是取新的 prefill batch,还是推进已有 decode batch。也就是说,调度器真正管理的是 batch 生命周期,而不是单个 request 生命周期。
继续往下看 get_new_batch_prefill() 和 _get_new_batch_prefill_raw(),你会看到 PrefillAdder 被用来在 token 预算、batch 大小、LoRA 约束、priority scheduling 与 chunked prefill 之间做折中。这里的重点不是把每个条件背下来,而是理解:SGLang 把“能不能接新请求”建模成一次 batch 构造问题,而不是某个全局开关。
这部分如果只靠文字,很容易把 Scheduler、ScheduleBatch 和 cache 看成三块并列知识。下面这张图专门用来解释它们之间的状态流转关系:请求怎样从 waiting queue 进入 batch,batch 怎样触发前向,cache 怎样在这个过程中决定复用和回收。
flowchart TD
A["Waiting requests"] --> B["Scheduler.get_next_batch_to_run()"]
B --> C["PrefillAdder / prefill admission"]
B --> D["running_batch / decode continuation"]
C --> E["ScheduleBatch"]
D --> E
E --> F["ModelWorkerBatch"]
F --> G["ForwardBatch"]
G --> H["ModelRunner.forward(...)"]
B --> I["tree_cache<br/>RadixCache / ChunkCache"]
B --> J["ReqToTokenPool"]
J --> K["TokenToKVPoolAllocator / KVCache"]
I --> K
K --> G
H --> L["BatchTokenIDOutput"]
L --> M["Detokenizer / next round"]这张图比纯文字多解释了一层“时序上的先后”。tree_cache、ReqToTokenPool 和 TokenToKVPoolAllocator 不是独立并列概念,而是在 batch 被组织和推进时一起参与决定“这一轮到底能不能跑、还能不能继续收新请求”。
ScheduleBatch 是调度层与执行层之间的桥#
schedule_batch.py 文件头已经把职责分得很清楚:ScheduleBatch 由 Scheduler 管理,ModelWorkerBatch 是给 tp_worker.py::TpModelWorker 的子集,ForwardBatch 才是 model_runner.py::ModelRunner 真正消费的低层张量表示。这相当于告诉你,调度器并不直接操作 attention kernel 所需的张量布局,而是先维护一个更适合做策略判断的中间层对象。
这层桥接非常重要。因为 batch 里不只有 token,还有 sampling 参数、多模态 pad 信息、prefix 匹配结果、finish reason、时间统计等运行时元数据。如果把这些东西直接混进 ModelRunner,执行层就会被调度细节污染;而如果没有 ScheduleBatch 这层,调度器又很难在 CPU 侧高效维护请求状态。
为什么 KV cache 被拆成两级内存池#
python/sglang/srt/mem_cache/memory_pool.py 在文件头直接说明了设计:SGLang 有 two levels of memory pool。ReqToTokenPool 负责“request 到 token location”的映射,TokenToKVPoolAllocator 负责管理 KV cache 索引,而真正的 KVCache 对象持有物理缓存。
这个拆分解决的是两个不同问题。ReqToTokenPool 解决“一个逻辑请求当前占了哪些 token 槽位”;TokenToKVPoolAllocator 和 KVCache 解决“这些槽位在设备上的物理存储在哪里,以及还能不能继续分配”。如果把这两件事混成一个结构,调度器既难以做请求级回收,也难以在分页、稀疏、SWA 或 Mamba 等不同缓存实现之间复用逻辑。
从 ReqToTokenPool.alloc(...) / free(...) 的接口也能看出这点:它关心的是请求槽位,而不是具体的 attention head 布局。反过来,TokenToKVPoolAllocator、PagedTokenToKVPoolAllocator 和不同 KVCache 实现关心的是物理缓存页、数据类型和设备布局。这正是一个典型的“逻辑索引层”和“物理存储层”分离。
如果把这部分再压缩成一句话,就是:调度器需要的是“我这轮还能排什么请求”,而缓存层需要的是“这些请求占的状态还在不在”。两级内存池把这两个问题拆开,才让 scheduler policy 和 cache policy 都能独立演进。
tree_cache 决定的是复用策略,不是调度本身#
scheduler.py 初始化 cache 时,会根据配置在 ChunkCache、RadixCache、SWARadixCache、MambaRadixCache 以及分层缓存变体之间选择。这里有两个特别值得抓住的源码注释:chunk_cache.py 文件头写着 “Cache for chunked prefill, used when RadixCache is disabled.”;radix_cache.py 文件头写着 “The radix tree data structure for managing the KV cache.”
这两句注释揭示了一个容易被忽略的事实:调度器并不直接决定“是否做前缀复用”,它依赖 tree_cache 提供匹配、插入、回收与保护语义。也正因为这样,Scheduler 才能把注意力放在 batch 组织上,而把 prefix 命中、节点锁定、eviction policy 等缓存策略下沉到 mem_cache。
对于第一次阅读源码的人,一个很实用的判断标准是:只要逻辑在回答“这批请求现在能不能凑成 batch”,大概率还在 scheduler 视角;只要逻辑在回答“这些 token 对应的缓存块能不能复用或回收”,大概率已经进入 mem_cache 视角。
边界条件、退化路径与 tradeoff#
调度与内存这一层最容易被写成“设计总是最优”,但现实不是这样。只要 batch 组成、LoRA 约束、chunked prefill、prefix reuse、设备内存紧张这些因素发生变化,scheduler 和 cache 的平衡点就会移动。也就是说,这一层不是一套静态最优策略,而是一套在约束之间不断折中的 runtime 机制。
这正是为什么本章一直强调 PrefillAdder、ScheduleBatch、ReqToTokenPool、TokenToKVPoolAllocator 和 tree_cache 要合起来看。因为只要其中任何一层的约束变化了,最终表现出来的就可能是“吞吐下降”“不能继续收新请求”“前缀复用失效”或“缓存回收变得激进”,而不是某个单一局部函数的问题。
调试时先看什么#
如果系统表现为“请求进得来,但队列推进异常慢”,优先看 scheduler 视角:get_next_batch_to_run()、get_new_batch_prefill() 以及 batch 相关状态。因为这类问题首先表现为“这一轮到底有没有形成可执行 batch”。
如果系统表现为“batch 逻辑看上去正常,但上下文状态复用失效、缓存回收异常或 memory pressure 明显上升”,则更适合直接转向 mem_cache 视角:memory_pool.py、radix_cache.py、chunk_cache.py 和相关 allocator。调度与缓存在这里是一体两面,排障时最好先判断自己到底是在追“调度问题”还是“状态驻留问题”。
一张更实用的现象对照表#
把这一章真正用起来时,最有价值的不是记住更多类名,而是先把现象和阅读入口对上号:
现象:新请求迟迟进不了运行批次
优先看:Scheduler.get_next_batch_to_run() / PrefillAdder / batch 预算
现象:前缀复用效果变差或缓存命中异常
优先看:tree_cache / RadixCache / ChunkCache / prefix 匹配路径
现象:显存占用异常或缓存回收看起来过早
优先看:ReqToTokenPool / TokenToKVPoolAllocator / memory_pool.py这类对照表看起来不“高级”,但它非常接近优秀技术书的价值:让读者在真实问题面前,不必重新从零建立排障路线。
为什么这一章必须同时讲 scheduler 和 cache#
如果只讲 scheduler,不讲 cache,读者会误以为 batch 的推进只是一个排队问题;如果只讲 KV cache,不讲 scheduler,读者又会把缓存理解成单纯的内存优化层。SGLang 把这两件事放在一起看的必要性,在于它们共同决定了“系统还能不能继续往前走”。
最典型的例子是 prefix reuse。它看起来像纯缓存优化,但实际会反过来影响 scheduler 能否更积极地接纳新请求。反过来,如果 batch 构造策略更保守,缓存即使还能复用,吞吐表现也未必明显提升。把这两个机制分开写,读者会学到两个局部事实;把它们合起来写,读者才会理解 runtime 为什么会表现成现在这样。
本章对应哪些代码路径#
这一章最重要的文件级锚点包括 python/sglang/srt/managers/scheduler.py、python/sglang/srt/managers/schedule_batch.py、python/sglang/srt/mem_cache/memory_pool.py、python/sglang/srt/mem_cache/radix_cache.py、python/sglang/srt/mem_cache/chunk_cache.py 与 python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py。
要继续深挖,推荐的顺序是:先看 Scheduler.get_next_batch_to_run() 建立整体节奏,再看 ScheduleBatch 文件头理解数据结构边界,然后回到 memory_pool.py 把两级内存池读清楚,最后再看 RadixCache 与 ChunkCache 这类策略实现。这样读,比一上来就陷进某个 allocator 或 eviction policy 的细节更稳。
读完这一章之后,一个理想的收获是:你不再把“调度”和“缓存”看成两份平行机制,而会把它们理解成同一件事的两个侧面。前者决定当前批次怎样推进,后者决定这些批次背后的上下文状态怎样继续留在系统里。这种合起来看的能力,正是继续读执行模型和性能问题之前最需要补上的那一块。
叶王 © 2013-2026 版权所有。如果本文档对你有所帮助,可以请作者喝饮料。