手写 mini-sglang 阶段二 · 排得好：同时服务很多请求

阶段二的里程碑：调度器接管¶

阶段一结束时，batch 的字段还要自己填：哪些请求进这一步、它们的 KV 放在哪几页、位置从几开始。阶段二的五章把这几行变成一个调度器——连续批处理、准入控制、页的分配与回收、Radix Cache、分块 prefill、重叠调度——对外只剩一个离线接口 LLM.generate。这一页先把阶段结束时的样子跑给你看：三组数字，每组对应骨架做不到的一件事。

这一阶段结束时你手里有什么

LLM.generate(prompts)：一次吞下任意多个请求，长短不一也行，先完成的先退出 batch，后来的随时补进来；
同一个系统提示词只算一遍 KV，后面的请求直接复用；
长提示词被切成固定大小的块，别的请求可以插在块之间；
CPU 上准备下一步的时间藏在计算后面（重叠调度，本页先关掉它，数字只反映调度器本身）。

先跑起来¶

cd minisgl && python examples/stage2_milestone.py

examples/stage2_milestone.py

"""阶段二的里程碑：调度器接管了阶段一里手工填 batch 的那几行。

三件事：离线接口一次吞下 8 个请求（连续批处理、准入、页分配全自动）；
同一个系统提示词的 4 个问题，Radix Cache 让 prefill 少算四分之三；
一个长提示词被分块 prefill 切成几步。
"""

import time

import torch
from minisgl.core import SamplingParams
from minisgl.env import ENV
from minisgl.llm import LLM

ENV.DISABLE_OVERLAP_SCHEDULING.value = True                   # 先关掉重叠调度，数字只反映调度器本身


def make(**kw):
    return LLM("models/Qwen3-0.6B", dtype=torch.float32, max_running_req=8, num_page_override=2048,
               max_seq_len_override=512, **kw)


def count_prefill(llm):
    """把引擎的前向包一层，数一数 prefill 实际算了多少个 token、分了几步"""
    computed = []
    orig = llm._forward
    llm._forward = lambda fi: (computed.append(sum(r.extend_len for r in fi.batch.reqs))
                               if fi.batch.is_prefill else None, orig(fi))[1]
    return computed


# ------------------------------------------------------------------ 1. 8 个请求，一次 generate
prompts = ["The capital of France is", "List three prime numbers:", "def fibonacci(n):", "Once upon a time,",
           "1 + 1 =", "The quick brown fox", "import torch\n", "Roses are red,"]
llm = make()
t = time.perf_counter()
outs = llm.generate(prompts, SamplingParams(max_tokens=16, ignore_eos=True))
dt = time.perf_counter() - t
print(f"8 个请求 x 16 个 token：{dt:.2f} s，合计 {8 * 16 / dt:.1f} tokens/s（阶段一手工组 batch 是 70，第 0 步是 25.6）")
print(f"  第一个：{outs[0]['text']!r}")
llm.shutdown()

# ------------------------------------------------------------------ 2. 同一个系统提示词：Radix Cache
system = "You are a helpful assistant. Answer concisely and accurately. " * 4
questions = ["What is 2 + 2?", "Name a primary color.", "What is the capital of Japan?", "Spell cat."]
for cache_type in ("naive", "radix"):
    llm = make(cache_type=cache_type)
    computed = count_prefill(llm)
    llm.generate([system + questions[0]], SamplingParams(max_tokens=4))          # 第一个请求先进来
    t = time.perf_counter()
    llm.generate([system + q for q in questions[1:]], SamplingParams(max_tokens=4))
    dt = time.perf_counter() - t
    total = sum(len(llm.tokenizer(system + q).input_ids) for q in questions)
    print(f"{cache_type:<5}：4 个提示词共 {total} 个 token，prefill 实际计算 {sum(computed)} 个，后 3 个请求用时 {dt * 1000:.0f} ms")
    llm.shutdown()

# ------------------------------------------------------------------ 3. 长提示词：分块 prefill
long_prompt = "Explain, step by step, why the sky is blue. " * 25
llm = make(max_extend_tokens=96)                              # 每一步 prefill 最多 96 个 token
computed = count_prefill(llm)
llm.generate([long_prompt], SamplingParams(max_tokens=2))
print(f"分块 prefill：{len(llm.tokenizer(long_prompt).input_ids)} 个 token 的提示词分了 {len(computed)} 步，"
      f"每步 {computed}；别的请求可以插在这些步之间")
llm.shutdown()

运行结果

8 个请求 x 16 个 token：1.23 s，合计 104.5 tokens/s（阶段一手工组 batch 是 70，第 0 步是 25.6）
  第一个：' Paris. The capital of Italy is Rome. The capital of Spain is Madrid.'
naive：4 个提示词共 231 个 token，prefill 实际计算 231 个，后 3 个请求用时 409 ms
radix：4 个提示词共 231 个 token，prefill 实际计算 73 个，后 3 个请求用时 244 ms
分块 prefill：302 个 token 的提示词分了 4 步，每步 [96, 96, 96, 14]；别的请求可以插在这些步之间

读这几个数字¶

8 个请求 1.23 秒，104.5 tokens/s。 第 0 步 25.6，阶段一手工组 4 个请求的 batch 是 70。两件事叠在一起：batch 更大了（8 个请求一步算完，decode 读一次权重出 8 个 token），而且不用自己填字段了——generate 里的每一步都是调度器在做阶段一 step() 里那几行的事：准入、分配页、算 positions 和 out_loc、收回结束请求的页。真实服务里请求是陆续到的、长短不一的，连续批处理的意义是batch 的成员随时在变而计算从不停下。

同一个系统提示词：prefill 从 231 个 token 降到 73 个，后 3 个请求 409 ms → 244 ms。 四个问题共享一段 158 个 token 的系统提示词。naive 缓存下每个请求都把它重算一遍；Radix Cache 在第一个请求算完后把前缀留在树里，后面三个只算各自的问题部分。这就是多轮对话和共享系统提示词场景下吞吐差几倍的来源——少算的那 158 × 3 个 token，一个也没有白省。

302 个 token 的提示词分了 4 步：96、96、96、14。 没有分块 prefill，一个长提示词进来会独占引擎一整步，正在 decode 的所有请求都得等它；分块之后每步最多算 max_extend_tokens 个，decode 请求插在中间，首 token 延迟稍长一点，换来的是其他请求的每 token 延迟不抖。

手工的那几行变成了什么¶

阶段一 `step()` 里你手工做的	阶段二	在哪一章
决定哪些请求进这一步的 batch	准入控制、连续批处理的主循环	调度器骨架与离线接口
`page_table[row, :256] = arange(...)`——手工"分配"KV 位置	`CacheManager`：按需分配、结束时回收、满了就等	CacheManager
`cached_len=0`——每个请求从头算	Radix Cache：匹配前缀、复用 KV、引用计数、LRU 淘汰	Radix Cache
整段提示词一步 prefill	分块 prefill：每步最多 `max_extend_tokens` 个 token	分块 prefill
算完这一步再准备下一步	重叠调度：准备下一步时 GPU 还在算这一步	重叠调度

怎么读这五章¶

调度器骨架先读——它定义主循环，后面四章都是往这个循环里加东西。读到 Radix Cache 时把本页脚本的 cache_type 来回切几次，看 prefill 实际计算 那个数怎么变；读到分块 prefill 时改 max_extend_tokens，看步数和每步的 token 数。最后一章重叠调度读完，把本页第一行的 ENV.DISABLE_OVERLAP_SCHEDULING.value = True 删掉重跑：CPU 上这项收益不大（没有 GPU 在另一头并行算），但在真卡上它是吞吐的最后一截。

验收：做到这些再进阶段三

画出一个请求从进入队列到释放 KV 页的状态变化，标出每一步调度器做的决定；
解释 Radix Cache 的树为什么要分裂节点，引用计数防的是什么；
说出分块 prefill 牺牲了什么、换来了什么；
把本页脚本的 max_running_req 改成 2，解释吞吐为什么掉、掉到哪。

小结¶

阶段二把阶段一手工填 batch 的那几行变成调度器，对外只剩 LLM.generate。
吞吐从 70 到 104.5 tokens/s：batch 更大、成员随时变、计算不停。
Radix Cache 让共享前缀只算一遍：231 个 prefill token 降到 73 个。
分块 prefill 让长提示词不再独占引擎：302 个 token 分 4 步，别的请求插在中间。

发现错误或有更好的写法？在 GitHub 上编辑此页提交问题