跳到主要内容

数据截至 (上游 commit f775db03aaa8)

03 · 连续批调度器

本章讲什么: Scheduler 每步怎么决定「这一步 GPU 算哪些请求」——waiting 队列排序、prefill 准入预算、chunked prefill、以及显存不足时的 retract(让路)。

3.1 它要解决的小问题​

静态批处理(攒够 N 条一起算、全算完再下一批)在真实流量下两头浪费:先到的请求干等,短请求陪长请求坐牢。

连续批处理(continuous batching)的答案是把调度粒度从「一批」降到「一步」:每生成一步就重新组批,谁的输出结束谁出列,新请求见缝插针。但代价是调度逻辑本身成了热路径,而且它必须在显存这座固定大小的仓库里玩俄罗斯方块:

  • 新请求要进来:它的 prefill 长度 + 预期生成长度放得下吗?
  • 正在跑的在变长:每步每人多一个 token 的 KV,装不下怎么办?

3.2 思路:prefill 优先、预算准入、不够就让路​

SGLang 的组批哲学三句话:

  1. 能 prefill 就先 prefill(get_next_batch_to_run 里 new_batch 优先于 decode)——首 token 延迟是用户体验,decode 批晚一步只是吞吐小降。
  2. 准入按预算说话,不按条数。 一个请求能不能进批,看三笔账:本批 prefill token 上限、显存余量、chunk 上限(见 3.4)。
  3. 显存满了不报错,让路。 把「代价最小」的请求逐出 running、已算的前缀写回 radix 树、请求塞回 waiting——它下一轮还能带着缓存前缀重新进来。

3.3 主循环与组批流程​

 waiting_queue ──► ① SchedulePolicy.calc_priority 排序(FCFS/LPM/…)
                        │
                        ▼
              ② PrefillAdder 逐个试装(三笔预算)
                        │ 装得下                      │ 装不下且是长输入
                        ▼                           ▼
              ③ 新 prefill batch             chunked prefill(切块下步继续)
                        │
                        ▼
              ④ 与 running_batch 合并 → forward → 出首 token 的请求转入 running
                        │
                        ▼
              ⑤ decode 步前 check_decode_mem:不够 → retract_decode 让路

get_next_batch_to_run(python/sglang/srt/managers/scheduler.py:3193)是这张图的代码本体。它每步做三件事:

  1. 合并上一批。 上一轮 prefill 完的请求(去掉刚结束的、去掉还在切块的)merge 进 running_batch(scheduler.py:3164-3190)。
  2. 尝试组新 prefill。 _get_new_batch_prefill_raw(scheduler.py:3299)从 waiting 队列挑人,组成 new_batch。
  3. 否则 decode。 没有新批就 update_running_batch(scheduler.py:3624)跑 decode——先 filter_batch 删掉已结束的,再查显存、必要时 retract,最后 prepare_for_decode 摆好张量。

调度循环本身的两种形态(normal / overlap)见第 1 章;本章聚焦「组批」这步。

3.4 准入:PrefillAdder 的三笔预算​

PrefillAdder(python/sglang/srt/managers/schedule_policy.py:511)每轮 prefill 组批时新建一个实例,核心方法是 add_one_req(schedule_policy.py:1208)。一个请求要进批,得同时过这几关:

预算代码里的变量超限时的返回
显存总账rem_total_tokens = KV 池余量(经 new_token_ratio 放大估算)AddReqResult.NO_TOKEN
本批 prefill 上限rem_input_tokens(默认 max_prefill_tokens=16384,server_args.py:867)OTHER(本批已满)
chunk 上限rem_chunk_tokens(chunked_prefill_size,server_args.py:852)切块,下步继续

new_token_ratio 是这里最不显然的机制。 准入不能只看「现在要多少显存」,还得给「这个请求未来要生成的 token」留位置——但未来的长度未知。SGLang 的解法是维护一个比率:NewTokenRatioTracker(python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py:13)从接近 1.0 的保守值起步,随 decode 步数衰减;显存估算按 max_new_tokens × ratio 记账。retract 发生后用已生成/上限的实际比例重新估计(estimate_new_token_ratio_after_retract),等于「被现实教育一次后调低乐观度」。

chunked prefill 是第三笔预算的直接产物:一条 32K 的长输入装不进一个 prefill 批,就切成 chunked_prefill_size 大小的块,分多步算。每算完一块,已算的 KV 立刻 cache_unfinished_req 写回 radix 树(第 2 章)——切块不丢共享。这条链路把「长 prompt」从「阻塞整个 decode」降级成「占用几步 prefill 额度」。

3.5 排序:缓存感知的 SchedulePolicy​

SchedulePolicy.calc_priority(schedule_policy.py:237)决定 waiting 队列里谁先被试装。策略分两类:

类别策略干什么
缓存感知LPM(默认形态)按「与树共享的前缀长度」排序,共享多的先进批,整批命中率最高
缓存感知DFS_WEIGHT按树上 DFS 权重排序,同前缀的请求聚在一起
缓存无关FCFS / LOF / RANDOM / ROUTING_KEY先到先服务 / 最长输出优先 / 随机 / 按路由键

LPM 的实现有个讨巧之处:它用一棵模拟树 waiting_queue_radix_tree(schedule_policy.py:234,RadixCache.create_simulated())做 in-batch 前缀匹配——不光和已缓存的比,还和「同一批里其他等待请求」比,让互相共享前缀的请求挨着进批。

坑: LPM 的排序开销随队列长度增长,_determine_active_policy(schedule_policy.py:296)在队列超过 128 条时直接降级 FCFS——高负载下缓存感知调度自动放弃,前缀局部性收益随之下降。这是代码里明写的取舍,不是 bug。

3.6 让路:retract​

decode 途中显存不够,由 update_running_batch(scheduler.py:3624)触发:先 check_decode_mem(schedule_batch.py:2873)估这一步的 KV 需求,不够就 retract_decode(schedule_batch.py:2880):

sorted_indices = self._get_decode_retraction_order(self.reqs)  # 按让路代价排序
while not self.check_decode_mem(selected_indices=sorted_indices):
    if len(sorted_indices) == 1:
        break                      # 至少保住一个请求,绝不团灭
    idx = sorted_indices.pop()
    req = self.reqs[idx]
    if self.release_req(idx, len(sorted_indices), server_args):
        retracted_reqs.append(req)  # 状态写回 radix 树,回 waiting 队列

被 retract 的请求回到 waiting 队列(_add_request_to_queue(req, is_retracted=True),scheduler.py:3695-3696),下一轮靠 radix 树里的前缀「复活」,重新 prefill 的只是被逐出后丢失的部分。

几个诚实的边界:

  • beam search 组不能 retract,只能 abort。 因为组内成员共享 KV、无法单独让路(retract_decode 里的 FINISH_ABORT 分支)。
  • retract 也可能失败。 如果开了 host 备份池而池子也满了,请求直接被 abort 并报「Retraction host KV pool exhausted」。
  • TEST_RETRACT 是留好的混沌开关。 测试用它强制周期性 retract(scheduler.py:3633),把这条冷路径变成常被踩的热路径。

3.7 本章代码地图​

主题文件路径符号名
组批总入口python/sglang/srt/managers/scheduler.pyget_next_batch_to_run
prefill 组批python/sglang/srt/managers/scheduler.py_get_new_batch_prefill_raw
decode 批维护python/sglang/srt/managers/scheduler.pyupdate_running_batch
请求排序python/sglang/srt/managers/schedule_policy.pySchedulePolicy.calc_priority、_determine_active_policy
准入预算python/sglang/srt/managers/schedule_policy.pyPrefillAdder.add_one_req
显存预估比率python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.pyNewTokenRatioTracker
让路python/sglang/srt/managers/schedule_batch.pyScheduleBatch.retract_decode、check_decode_mem
批数据结构python/sglang/srt/managers/schedule_batch.pyScheduleBatch(2057)、Req(828)