4.1 Scheduler、批次与 KV Cache

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 视角。

本章对应哪些代码路径#

这一章最重要的文件级锚点包括 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 的细节更稳。

读完这一章之后，一个理想的收获是：你不再把“调度”和“缓存”看成两份平行机制，而会把它们理解成同一件事的两个侧面。前者决定当前批次怎样推进，后者决定这些批次背后的上下文状态怎样继续留在系统里。这种合起来看的能力，正是继续读执行模型和性能问题之前最需要补上的那一块。