手写 mini-sglang
阶段一 · 算得对：把骨架换成正式结构
Engine 与采样器¶
Engine 是一个 TP rank 上的计算引擎：它按固定的顺序初始化模型、KV 池、page table、注意力后端、采样器和 CUDA Graph，然后对外只提供一个方法 forward_batch——给一个准备好的 batch，返回每个请求的下一个 token。这一章写完它，并在没有调度器的情况下手工驱动它完成 prefill 和 decode，结果与 Hugging Face 逐 token 一致。"算得对"这一部分到此完成。
自测：能答上来就可以跳过本章
- 引擎怎样决定 KV 池有多少页？为什么要在加载模型之前先量一次可用显存？
forward_batch返回的 token 为什么有 GPU 和 CPU 两份？- 一个 batch 里有的请求贪心、有的请求
temperature=0.7, top_p=0.9，采样器怎么在一次调用里处理？ - top-k 和 top-p 同时设置时，先做哪个？
自测参考答案（先自己答，再展开对照）
- 页数 =（
memory_ratio× 加载模型前的空闲显存 − 模型占用）÷ 每页的字节数，TP 下各 rank 取最小值。先量一次是为了算出模型本身占了多少：加载前后空闲显存的差就是模型占用。 - GPU 上的一份直接写回 token pool，作为下一轮的输入，不需要经过 CPU；CPU 上的一份（异步拷贝）给调度器判断是否结束、交给反分词。
- 把每个请求的参数排成张量，一次批量处理：全是贪心时直接 argmax；混合时，贪心的请求用一个极小的温度（效果等同于 argmax），和其他请求一起走采样的路径。
- 先 top-k 再 top-p。
本章要写的文件：engine/config.py、engine/engine.py、engine/sample.py、utils/device.py（设备抽象）。
到这一章为止的文件：46 个（新建 5，修改 0，提前引入 3）
这一步的 main：minisgl/
├── __init__.py
├── core.py
├── attention/ （4 个文件，不变）
├── distributed/ （3 个文件，不变）
├── engine/
│ ├── __init__.py ← 本章新建
│ ├── config.py ← 本章新建
│ ├── engine.py ← 本章新建
│ ├── sample.py ← 本章新建
│ ├── graph.py ← 提前引入（第 18 章讲）
├── kernel/ （2 个文件，不变）
├── kvcache/ （4 个文件，不变）
├── layers/ （8 个文件，不变）
├── models/ （7 个文件，不变）
├── moe/
│ ├── __init__.py ← 提前引入（第 20 章讲）
│ ├── base.py ← 提前引入（第 20 章讲）
├── scheduler/ （2 个文件，不变）
├── utils/
│ ├── __init__.py
│ ├── misc.py
│ ├── hf.py
│ ├── torch_utils.py
│ ├── registry.py
│ ├── logger.py
│ ├── device.py ← 本章新建
examples/ch06_engine.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。
配置¶
minisgl/engine/config.py · EngineConfig
大部分字段对应命令行参数。@dataclass(frozen=True)
class EngineConfig:
model_path: str
tp_info: DistributedInfo
dtype: torch.dtype
device: str = "auto" # auto：有 GPU 用 cuda:<rank>，否则用 cpu（官方只支持 CUDA）
max_running_req: int = 256
attention_backend: str = "auto"
moe_backend: str = "auto"
cuda_graph_bs: List[int] | None = None
cuda_graph_max_bs: int | None = None
page_size: int = 1
memory_ratio: float = 0.9
cpu_kv_cache_bytes: int = 2 << 30 # CPU 上 KV 池默认最多 2 GiB
distributed_timeout: float = 60.0
distributed_port: int = 2333
use_dummy_weight: bool = False
max_seq_len_override: int | None = None
num_page_override: int | None = None
@cached_property
def hf_config(self):
return cached_load_hf_config(self.model_path)
@cached_property
def model_config(self) -> ModelConfig:
from minisgl.models import ModelConfig
return ModelConfig.from_hf(self.hf_config)
@property
def max_seq_len(self) -> int:
if self.max_seq_len_override is not None:
return self.max_seq_len_override
return self.model_config.rotary_config.max_position
@property
def max_forward_len(self) -> int:
return self.max_seq_len
@property
def distributed_addr(self) -> str:
return f"tcp://127.0.0.1:{self.distributed_port}"
device 和 cpu_kv_cache_bytes 是我们加的；model_config 是 cached_property，第一次访问时从 config.json 解析（第 3 章）。
初始化的顺序¶
minisgl/engine/engine.py · Engine.__init__
顺序是有讲究的：
def __init__(self, config: EngineConfig):
set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size)
self.config = config = _adjust_config(config)
self.device = _pick_device(config)
if is_cuda(self.device):
torch.cuda.set_device(self.device)
else:
_set_cpu_threads(config.tp_info.size)
torch.manual_seed(42)
self.stream = create_stream(self.device)
set_stream(self.stream)
self.dtype = config.dtype
self.ctx = Context(config.page_size)
set_global_ctx(self.ctx)
self.tp_cpu_group = self._init_communication(config)
init_free_memory = self._sync_get_memory()[1]
logger.info_rank0(f"Free memory before loading model: {_gib(init_free_memory)}")
# ======================= 模型：先在 meta 设备上建空壳，再把权重"装"进去
set_rope_device(self.device)
with torch.device("meta"), torch_dtype(config.dtype):
self.model = create_model(config.model_config)
self.model.load_state_dict(self._load_weight_state_dict(config))
# ======================= KV 池：剩下的显存都给它
self.num_pages = self._determine_num_pages(init_free_memory, config)
num_tokens = self.num_pages * config.page_size
self.ctx.kv_cache = self.kv_cache = create_kvcache_pool(
model_config=config.model_config,
num_pages=self.num_pages + 1, # 多出的 1 页给 dummy 请求，它的写入落在这里
page_size=config.page_size,
device=self.device,
dtype=self.dtype,
)
# ======================= page table：每行对应一个请求，存每个 token 的 KV 位置
self.max_seq_len = min(config.max_seq_len, num_tokens)
aligned_max_seq_len = _align_up_32(self.max_seq_len)
self.ctx.page_table = self.page_table = torch.zeros(
(config.max_running_req + 1, aligned_max_seq_len), # 最后一行给 dummy 请求
dtype=torch.int32, device=self.device,
)
# ======================= 注意力后端、MoE 后端、采样器
self.ctx.attn_backend = self.attn_backend = create_attention_backend(
config.attention_backend, config.model_config)
if config.model_config.is_moe:
self.ctx.moe_backend = self.moe_backend = create_moe_backend(config.moe_backend)
self.sampler = Sampler(self.device, config.model_config.vocab_size)
# ======================= dummy 请求与 CUDA Graph
self.dummy_req = Req(
input_ids=torch.tensor([0], dtype=torch.int32),
table_idx=config.max_running_req,
cached_len=0,
output_len=1,
uid=-1,
sampling_params=None, # type: ignore[arg-type]
cache_handle=None, # type: ignore[arg-type]
)
self.page_table[self.dummy_req.table_idx].fill_(num_tokens) # 指向 dummy 页
self.graph_runner = GraphRunner(
stream=self.stream,
device=self.device,
model=self.model,
attn_backend=self.attn_backend,
cuda_graph_bs=config.cuda_graph_bs,
cuda_graph_max_bs=config.cuda_graph_max_bs,
free_memory=init_free_memory,
max_seq_len=aligned_max_seq_len,
vocab_size=config.model_config.vocab_size,
dummy_req=self.dummy_req,
)
- TP 信息和设备。之后构建的每一层都要读 TP 信息决定自己的切分。
- stream 和全局 Context。在 CUDA 上，引擎的所有计算都在自己的 stream 上（第 11 章的重叠调度需要调度器和引擎各用一条 stream）。
- 通信。TP>1 时建立进程组（第 16 章）。
- 量一次可用显存，然后加载模型。加载前后的差值就是模型占用的显存。
- KV 池。用剩下的显存算出页数（见下一节），多分配一页给 dummy 请求。
- page table，最后一行给 dummy 请求，并让它指向 dummy 页（
fill_(num_tokens)：dummy 页的第一个位置）。 - 注意力后端和 MoE 后端。注意力后端在构造时要读 KV 池（
get_global_ctx().kv_cache），所以必须在 KV 池之后。 - 采样器和 CUDA Graph。CUDA Graph 在捕获时会真的跑模型，必须放在最后。
显存规划¶
minisgl/engine/engine.py · Engine._determine_num_pages
每页的字节数是 def _determine_num_pages(self, old_free_memory: int, config: EngineConfig) -> int:
cache_per_page = (
2 # K 和 V
* config.model_config.head_dim
* div_even(config.model_config.num_kv_heads, config.tp_info.size, allow_replicate=True)
* config.page_size
* self.dtype.itemsize
* config.model_config.num_layers
)
num_pages = config.num_page_override
if num_pages is None:
new_free_memory = self._sync_get_memory()[1]
model_memory = old_free_memory - new_free_memory
available = int(config.memory_ratio * old_free_memory) - model_memory
if not is_cuda(self.device):
available = min(available, config.cpu_kv_cache_bytes)
num_pages = available // cache_per_page
assert num_pages > 1, "Not enough memory for KV cache"
logger.info_rank0(f"KV cache: {num_pages * config.page_size} tokens, "
f"{_gib(num_pages * cache_per_page)}")
return num_pages
2（K、V）× head_dim × 本 rank 的 KV 头数 × page_size × dtype 字节数 × 层数。可用显存的算法：
memory_ratio 默认 0.9，留出 10% 给激活、CUDA Graph 和各种工作区。张量并行时，各 rank 的可用显存取最小值（_sync_get_memory 用一次 all-reduce），保证所有 rank 的 KV 池一样大——调度器在每个 rank 上做完全相同的决策，前提是资源完全相同。
CPU 上"空闲内存"是整台机器的，按比例算会得到几十 GB 的 KV 池，所以我们额外加了一个上限 cpu_kv_cache_bytes（默认 2 GiB）；测试里通常直接用 num_page_override 指定页数。
forward_batch¶
minisgl/engine/engine.py · Engine.forward_batch
- 在 def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
with self.ctx.forward_batch(batch):
if self.graph_runner.can_use_cuda_graph(batch):
logits = self.graph_runner.replay(batch)
else:
logits = self.model.forward()
for req in batch.reqs:
req.complete_one()
next_tokens_gpu = self.sampler.sample(logits[: batch.size], args).to(torch.int32)
next_tokens_cpu = next_tokens_gpu.to("cpu", non_blocking=True)
copy_done_event = create_event(self.device)
copy_done_event.record(self.stream)
return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event)
forward_batch 上下文里执行模型（或 replay CUDA Graph），得到 [batch_size, vocab] 的 logits；
- 推进所有请求的状态：complete_one()。注意这发生在采样之前、且在 CPU 上立即完成——调度器此后就可以把这些请求当作"已经多了一个 token"来调度下一轮，而不必等 GPU 算完（第 11 章）；
- 采样得到 next_tokens_gpu，再 non_blocking 地拷一份到 CPU，并记录一个 event。调用方需要 CPU 上的值时，先 copy_done_event.synchronize()。
两份 token 各有用途：GPU 上的那份由调度器写回 token pool，作为下一轮的输入，全程不离开 GPU；CPU 上的那份用来判断是否结束、发给 detokenizer。
设备抽象¶
官方在这里直接用 torch.cuda.Stream()、torch.cuda.Event() 和 pin_memory=True。我们把它们换成几个小函数：
minisgl/utils/device.py · create_stream
def create_stream(device: torch.device) -> Any:
return torch.cuda.Stream(device=device) if is_cuda(device) else NullStream()
minisgl/utils/device.py · NullEvent
CPU 上所有计算都是同步完成的：class NullEvent:
"""CPU 上的"事件"：记录和等待都是空操作（record 时计算早已完成）。"""
def record(self, stream: Any = None) -> None:
pass
def synchronize(self) -> None:
pass
def query(self) -> bool:
return True
record 的时候结果早已算好，synchronize 当然不需要等。锁页内存只在 CUDA 上有意义（它让主机到设备的拷贝可以异步），CPU 上 pin(device) 返回 False。
采样器¶
minisgl/engine/sample.py · Sampler
采样分两步：
@dataclass
class Sampler:
device: torch.device
vocab_size: int
def prepare(self, batch: Batch) -> BatchSamplingArgs:
"""在调度器一侧（CPU）为整个 batch 准备采样参数张量。"""
params = [r.sampling_params for r in batch.reqs]
if all(p.is_greedy for p in params):
return BatchSamplingArgs(temperatures=None)
MIN_T = MIN_P = 1e-6
# 混合 batch 里的贪心请求：温度取极小值，softmax 退化成 one-hot，效果等于 argmax
ts = [max(0.0 if p.is_greedy else p.temperature, MIN_T) for p in params]
top_ks = [p.top_k if p.top_k >= 1 else self.vocab_size for p in params]
top_ps = [min(max(p.top_p, MIN_P), 1.0) for p in params]
top_k = top_p = None
if any(k != self.vocab_size for k in top_ks):
top_k = make_device_tensor(top_ks, torch.int32, self.device)
if any(p < 1.0 for p in top_ps):
top_p = make_device_tensor(top_ps, torch.float32, self.device)
return BatchSamplingArgs(make_device_tensor(ts, torch.float32, self.device), top_k, top_p)
def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor:
if args.temperatures is None:
return torch.argmax(logits, dim=-1)
impl = sample_torch
if logits.is_cuda:
try:
import flashinfer.sampling # noqa: F401
impl = sample_flashinfer
except ImportError:
pass
return impl(logits.float(), args.temperatures, args.top_k, args.top_p)
prepare(batch)在调度器一侧、前向之前调用，把每个请求的采样参数整理成张量并异步拷到设备上。整个 batch 都是贪心时返回temperatures=None，sample直接argmax，省掉 softmax。sample(logits, args)在前向之后调用。
一个 batch 里的请求可以有各自的参数。混合 batch 里的贪心请求被赋予一个极小的温度（1e-6），softmax 之后退化成 one-hot，采样结果就是 argmax——这样整个 batch 可以走同一条代码路径。top-k 值为 -1（不限制）时换成词表大小；只有 batch 里至少一个请求真的设置了 top-k（或 top-p），才生成对应的张量。
PyTorch 参考实现：
minisgl/engine/sample.py · sample_torch
与 FlashInfer 的默认行为一样，先做 top-k、再在剩下的概率上做 top-p。top-p 保留"排在它前面的概率之和小于 p"的 token，所以概率最大的那个永远保留。GPU 上改用 FlashInfer 的采样 kernel（def sample_torch(logits: torch.Tensor, temperatures: torch.Tensor, top_k: torch.Tensor | None,
top_p: torch.Tensor | None) -> torch.Tensor:
"""每个请求可以有各自的温度、top-k、top-p。与 FlashInfer 一样先做 top-k、再做 top-p。"""
probs = torch.softmax(logits / temperatures.unsqueeze(1), dim=-1)
if top_k is not None or top_p is not None:
sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True)
if top_k is not None:
rank = torch.arange(probs.shape[1], device=probs.device).unsqueeze(0)
sorted_probs = sorted_probs.masked_fill(rank >= top_k.unsqueeze(1), 0.0)
if top_p is not None:
sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
before = sorted_probs.cumsum(dim=-1) - sorted_probs # 排在它前面的概率之和
sorted_probs = sorted_probs.masked_fill(before >= top_p.unsqueeze(1), 0.0)
probs = torch.zeros_like(probs).scatter_(1, sorted_idx, sorted_probs)
return torch.multinomial(probs, num_samples=1).squeeze(1)
sample_flashinfer），它用拒绝采样避免排序，结果的分布相同。
手工驱动引擎¶
没有调度器时，我们自己扮演调度器：给每个请求"分配"一段 KV 位置（直接写 page table），每一轮填好 batch 的四个字段，调用 forward_batch，再把新 token 追加到请求上。
examples/ch06_engine.py
"""第 6 章：没有调度器，手工驱动 Engine 做 prefill + decode，与 HF generate 对比。"""
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))
print(f"设备 {engine.device}，注意力后端 {engine.config.attention_backend}，KV 池 {engine.num_pages} 页，"
f"page table {tuple(engine.page_table.shape)}")
tok = AutoTokenizer.from_pretrained(path)
prompts = ["The capital of France is", "def fibonacci(n):"]
reqs = []
for row, p in enumerate(prompts):
ids = tok(p).input_ids
engine.page_table[row, :256] = torch.arange(row * 256, row * 256 + 256) # 手工"分配"KV 位置
reqs.append(Req(input_ids=torch.tensor(ids, dtype=torch.int32), table_idx=row, cached_len=0,
output_len=8, uid=row, sampling_params=SamplingParams(max_tokens=8),
cache_handle=None))
outs = [[] for _ in reqs]
phase = "prefill"
for step in range(8):
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)) # 前向 + 采样，并推进 Req 状态
for i, r in enumerate(reqs):
r.append_host(out.next_tokens_cpu[i:i + 1])
outs[i].append(int(out.next_tokens_cpu[i]))
if step == 0:
print("prefill 输入 token 数:", len(batch.input_ids), " 之后 decode 每轮:", len(reqs))
phase = "decode"
hf = AutoModelForCausalLM.from_pretrained(path, dtype=torch.float32)
for p, o in zip(prompts, outs):
ids = tok(p, return_tensors="pt").input_ids
ref = hf.generate(ids, max_new_tokens=8, do_sample=False)[0, ids.shape[1]:].tolist()
print(f"{p!r}: {tok.decode(o)!r} 与 HF 一致: {o == ref}")
engine.shutdown()
运行结果
第一轮是 prefill：两个请求共 9 个 token 拼成一维送进模型；之后每轮 decode 两个请求各 1 个 token。8 个 token 与 Hugging Face 的 设备 cpu，注意力后端 torch，KV 池 1024 页，page table (5, 256)
prefill 输入 token 数: 9 之后 decode 每轮: 2
'The capital of France is': ' Paris. The capital of Italy is Rome' 与 HF 一致: True
'def fibonacci(n):': '\n if n == 0:\n' 与 HF 一致: True
generate 完全相同。下一部分的调度器要做的，就是把这段手工代码里的"分配 KV 位置""填 batch 字段""追加 token"自动化，并且让请求可以随时加入和离开。
官方实现
与官方的差异
- 设备：官方只支持 CUDA；我们通过
utils/device.py同时支持 CPU。 - 通信：官方即使 TP=1 也会初始化一个 gloo 进程组；我们只在 TP>1 时初始化，测试里可以在同一个进程中反复创建引擎（
shutdown会清掉全局状态）。 - CPU 线程数：我们在 CPU 上把 PyTorch 的线程数设为物理核数除以 TP 数，否则 decode 的小矩阵乘会被线程同步拖慢几十倍（见导读）。
测试¶
tests/test_ch06_engine.py 有三个测试：prefill logits 与 HF 的最大误差小于 1e-4；手工循环的贪心输出与 HF 逐 token 相同；采样器的 top-k、top-p、混合贪心的行为和概率（top-k=2 时第一名被抽中的频率接近 e⁵ / (e⁵ + e⁴) ≈ 0.731）。
tests/test_ch06_engine.py · test_sampler_greedy_and_filters
def test_sampler_greedy_and_filters():
sampler = Sampler(torch.device("cpu"), vocab_size=6)
assert sampler.prepare(_batch([SamplingParams(), SamplingParams()])).temperatures is None
logits = torch.tensor([[5.0, 4.0, 3.0, 0.0, 0.0, 0.0]] * 3)
args = sampler.prepare(_batch([SamplingParams(temperature=1.0, top_k=2),
SamplingParams(temperature=1.0, top_p=0.5),
SamplingParams()])) # 第三个请求是贪心
torch.manual_seed(0)
draws = torch.stack([sampler.sample(logits, args) for _ in range(2000)])
assert set(draws[:, 0].tolist()) == {0, 1} # top-k=2：只可能是前两个
assert set(draws[:, 1].tolist()) == {0} # 第一名概率约 0.67 > 0.5，top-p 只留它
assert set(draws[:, 2].tolist()) == {0} # 混合 batch 里的贪心请求
p0 = (draws[:, 0] == 0).float().mean().item() # top-k 后重新归一化：e^5/(e^5+e^4) = 0.731
assert abs(p0 - 0.731) < 0.03
怎么讲清楚
讲引擎：初始化顺序是 TP 信息 → stream 与上下文 → 通信 → 加载模型 → KV 池 → page table → 注意力后端 → 采样器 → CUDA Graph。KV 池的大小 = 显存比例 × 加载模型前的空闲显存 − 模型占用，换算成页数，TP 下各 rank 取最小值（要在加载前先量一次，才知道模型占了多少）。forward_batch 前向、推进请求状态、采样，返回 GPU 上的 token（直接写回 token pool）和异步拷贝到 CPU 的一份（给调度器判断是否结束）。采样器一次处理参数各不相同的请求：全贪心时直接 argmax，混合时贪心请求用极小的温度，先 top-k 再 top-p。
练习¶
- 在一张 24 GB 的卡上跑 Qwen3-0.6B（bf16），模型占约 1.2 GB，
memory_ratio=0.9，KV 池大约有多少个 token？ - 给采样器加上
min_p：保留概率不低于"最大概率 × min_p"的 token。它应该放在 top-k、top-p 之前还是之后？ - 如果
forward_batch在采样之后才调用complete_one()，对正确性有影响吗？对第 11 章的重叠调度呢？
参考答案
- 可用约 0.9 × 24 − 1.2 ≈ 20.4 GB，每 token 112 KiB，约 18 万个 token（实际还要扣掉 CUDA Graph 等占用）。
- min_p 的阈值取决于最大概率，通常作用在温度缩放后的完整分布上：可以放在 top-k、top-p 之前计算掩码，再与它们的掩码取交集。
- 对正确性没有影响：
complete_one只改 CPU 上的计数。对重叠调度也没有本质影响，因为调度器在forward_batch返回之后才组下一轮；关键是它不依赖采样结果的值，所以 CPU 不需要等 GPU。
小结¶
-
Engine的初始化顺序：TP 信息 → stream 与上下文 → 通信 → 加载模型 → KV 池 → page table → 后端 → 采样器 → CUDA Graph。 - KV 池的大小 =
memory_ratio × 加载前空闲显存 − 模型占用，按每页字节数换算成页数；TP 下各 rank 取最小值。 -
forward_batch前向、推进请求状态、采样，返回 GPU 上和（异步拷贝中的）CPU 上两份 token。 - 采样器支持每个请求不同的参数：全贪心时 argmax；混合时贪心请求用极小温度；先 top-k 后 top-p。
发现错误或有更好的写法？在 GitHub 上编辑此页提交问题