手写 mini-sglang 阶段一 · 算得对：把骨架换成正式结构

阶段一的里程碑：同样的输出，正式的结构¶

阶段一的六章把第 0 步那 126 行里的每一个朴素部件换成正式的模块：Req / Batch / Context、不依赖 nn.Module 的算子层、流式加载的模型、分页的 KV 池、可插拔的注意力后端、带采样器的 Engine。这一页先把阶段结束时的样子跑给你看——同一个提示词、同样的 16 个 token，外加骨架做不到的一件事：4 个请求放进一个 batch，一步算完。六章读完，回来对照这里的数字。

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

一个 Engine：给它一个填好字段的 Batch，返回每个请求的下一个 token；
还没有调度器——batch 的字段靠自己填（正是阶段二要自动化的那几行）；
单个请求和第 0 步一样快，多个请求一起算时吞吐是第 0 步的三倍；
每一步的输出仍然和 Hugging Face 逐 token 一致。

先跑起来¶

cd minisgl && python examples/stage1_milestone.py

examples/stage1_milestone.py

"""阶段一的里程碑：同样的提示词、同样的 16 个 token，但骨架已经换成正式结构。

和第 0 步比两件事：单个请求的 tokens/s（正式结构在这里并不占便宜），
以及骨架根本做不到的——把 4 个请求放进一个 batch，一步算完。
"""

import time

import torch
from minisgl.core import Batch, Req, SamplingParams
from minisgl.distributed import DistributedInfo
from minisgl.engine import Engine, EngineConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

path = "models/Qwen3-0.6B"
engine = Engine(EngineConfig(model_path=path, tp_info=DistributedInfo(0, 1), dtype=torch.float32,
                             max_running_req=4, num_page_override=1024, max_seq_len_override=256))
tok = AutoTokenizer.from_pretrained(path)
print(f"正式结构：{engine.device}，注意力后端 {engine.config.attention_backend}，"
      f"KV 池 {engine.num_pages} 页，{torch.get_num_threads()} 个线程（Engine 自己设的）")


def make_reqs(prompts, n):
    reqs = []
    for row, p in enumerate(prompts):
        engine.page_table[row, :256] = torch.arange(row * 256, row * 256 + 256)   # 手工"分配"KV 位置
        reqs.append(Req(input_ids=torch.tensor(tok(p).input_ids, dtype=torch.int32), table_idx=row,
                        cached_len=0, output_len=n, uid=row, sampling_params=SamplingParams(max_tokens=n),
                        cache_handle=None))
    return reqs


def step(reqs, phase):
    """没有调度器：自己把 batch 的字段填好，前向 + 采样，再把新 token 写回每个请求。"""
    batch = Batch(reqs=reqs, phase=phase)
    engine.graph_runner.pad_batch(batch)
    batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len) for r in reqs]).int()
    batch.input_ids = torch.cat([r.input_ids[r.cached_len:r.device_len] for r in reqs])
    batch.out_loc = torch.cat([engine.page_table[r.table_idx, r.cached_len:r.device_len] for r in reqs])
    engine.attn_backend.prepare_metadata(batch)
    out = engine.forward_batch(batch, engine.sampler.prepare(batch))
    for i, r in enumerate(reqs):
        r.append_host(out.next_tokens_cpu[i:i + 1])
    return out.next_tokens_cpu.tolist()


def generate(prompts, n=16):
    reqs = make_reqs(prompts, n)
    outs = [[] for _ in reqs]
    t0 = time.perf_counter()
    for i, t in enumerate(step(reqs, "prefill")):
        outs[i].append(t)
    t1 = time.perf_counter()
    for _ in range(n - 1):
        for i, t in enumerate(step(reqs, "decode")):
            outs[i].append(t)
    t2 = time.perf_counter()
    return outs, t1 - t0, (t2 - t1) / (n - 1)


# ------------------------------------------------------------------ 1. 和第 0 步同一个请求
prompt = "The capital of France is"
[out], t_prefill, t_step = generate([prompt])
print(f"单个请求：{tok.decode(out)!r}")
print(f"  prefill {t_prefill * 1000:.0f} ms，decode 每步 {t_step * 1000:.0f} ms（{1 / t_step:.1f} tokens/s）")

# ------------------------------------------------------------------ 2. 骨架做不到的：4 个请求一个 batch
prompts = [prompt, "List three prime numbers:", "def fibonacci(n):", "Once upon a time,"]
outs, t_prefill, t_step4 = generate(prompts)
print(f"4 个请求一个 batch：prefill {t_prefill * 1000:.0f} ms，decode 每步 {t_step4 * 1000:.0f} ms"
      f"（每步出 4 个 token，合计 {4 / t_step4:.1f} tokens/s；一步只比单请求慢 {(t_step4 / t_step - 1) * 100:.0f}%）")
written = sum(len(tok(p).input_ids) + 15 for p in prompts)               # 每个请求：提示词 + 15 个送回去的新 token
print(f"KV 池写了 {written} / {engine.num_pages} 页：每个 token 只写自己那一页，没有任何拷贝")

hf = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32).eval()
for p, o in zip(prompts, outs):
    ids = tok(p, return_tensors="pt").input_ids
    with torch.no_grad():
        ref = hf.generate(ids, max_new_tokens=16, do_sample=False)[0, ids.shape[1]:].tolist()
    print(f"  {p!r:30} -> {tok.decode(o)[:32]!r:36} 与 HF 一致: {o == ref}")
print(f"单独跑和放进 4 个请求的 batch 里跑，第一个请求的 16 个 token 完全相同：{outs[0] == out}")
engine.shutdown()

运行结果

正式结构：cpu，注意力后端 torch，KV 池 1024 页，16 个线程（Engine 自己设的）
单个请求：' Paris. The capital of Italy is Rome. The capital of Spain is Madrid.'
  prefill 62 ms，decode 每步 42 ms（23.7 tokens/s）
4 个请求一个 batch：prefill 96 ms，decode 每步 57 ms（每步出 4 个 token，合计 70.0 tokens/s；一步只比单请求慢 36%）
KV 池写了 79 / 1024 页：每个 token 只写自己那一页，没有任何拷贝
  'The capital of France is'     -> ' Paris. The capital of Italy is '   与 HF 一致: True
  'List three prime numbers:'    -> ' 2, 3, 5, 7, 11,'                   与 HF 一致: True
  'def fibonacci(n):'            -> '\n    if n == 0:\n        return 0' 与 HF 一致: True
  'Once upon a time,'            -> ' there were 3000 people in a tow'   与 HF 一致: True
单独跑和放进 4 个请求的 batch 里跑，第一个请求的 16 个 token 完全相同：True

读这几个数字¶

单个请求：23.7 tokens/s，第 0 步是 25.6。 正式结构在单请求上不占便宜，甚至慢了一点。这不意外：分页 KV 的写入是按位置散写的，注意力后端要为每个请求切出自己的那段 KV，Req / Batch 的记账也都是 Python。这些开销换来的东西在单请求的数字里看不到——它们是为下面那一行准备的。

4 个请求一个 batch：每步 57 ms，只比单请求的 42 ms 慢 36%，合计 70 tokens/s。 这是第 0 步的骨架做不到的事。decode 是访存密集的：一步要把 2.4 GB 的 fp32 权重从内存读一遍，矩阵乘的"批"是 1 行还是 4 行，读权重的代价几乎一样。所以把请求攒成 batch，每一步多出来的只是那点计算，权重只读一次。这就是连续批处理的全部道理，推理系统的吞吐都从这里来。

KV 池写了 79 页，没有任何拷贝。 第 0 步每一步都 torch.cat 把整段 KV 复制一遍；现在每个 token 只写自己那一页，page table 记着它在哪。单独跑和放进 batch 里跑，第一个请求的 16 个 token 完全相同——批处理不改变任何一个请求的结果，这一点后面每一章的测试都会反复确认。

骨架的每一块换成了什么¶

第 0 步	阶段一	换掉的理由
一个 `ids` 张量、一个 `pos` 张量	`Req`、`Batch`、`Context`	多个请求各有自己的长度和状态；模型从全局 `Context` 读当前 batch，`forward()` 不用传参
四个函数 + 一个权重 `dict`	算子层：`BaseOP` 与各种 `Layer`	每个 `Layer` 自己管权重的形状、加载和切分，张量并行时才切得动
`safe_open` 一次读完	模型与权重加载	按需流式加载；同一个 decoder 文件装下 Qwen3、Qwen2.5、Llama 3
`list` + `torch.cat`	KV 池、page table 与 token pool	预分配、按页写、零拷贝；多个请求共用一个池
`einsum` + 显式掩码	注意力后端与参考实现	变长 batch、分页 KV、前缀复用的元数据；GPU 上换成 FlashInfer / FlashAttention
`argmax`	Engine 与采样器	温度、top-k、top-p 按请求设置，一次批量采样；`Engine` 把上面所有东西按正确的顺序初始化

怎么读这六章¶

按顺序读，每章先读自测题、自己写一版、再对照本书和官方的实现、最后跑本章的测试。读完Engine 与采样器，回来重跑这一页的脚本：step() 里那几行"自己填 batch 字段"的代码，你会发现每一行都对应阶段二调度器的一个职责——positions 和 out_loc 从哪来（页分配）、cached_len 为什么会变（前缀缓存）、谁决定这一步哪些请求进 batch（准入控制）。阶段二做的就是把这几行自动化。

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

不看书，能画出 Req 的三个长度字段在 prefill 前后各是多少；
能解释 page table 的一行是什么、out_loc 是怎么算出来的；
换成 page_size=4 重跑本页脚本，输出不变（第 4 章的练习）；
说得出为什么 4 个请求一步只比 1 个请求慢三成。

小结¶

阶段一把骨架的六个朴素部件换成六个正式模块，接口与官方一致。
单请求不比骨架快；多请求一个 batch 才是正式结构的意义——decode 读一次权重，算几行都差不多。
KV 按页写、零拷贝；批处理不改变任何请求的结果。
step() 里手工填的那几行就是阶段二调度器要做的事。

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