推理运行时¶
推理引擎把一组长度不同、动态到达的请求映射为 GPU batch。核心状态不是一个输入 tensor,而是请求生命周期、KV block、调度预算、采样器和流式输出的组合。
推理运行时与服务给出从 Orca 到 PagedAttention、chunked prefill 与 P/D 分离的因果链;vLLM / PagedAttention 进一步把逻辑 token、物理 block、block table 与 copy-on-write 落成可执行模型。
请求状态机¶
一个请求至少经历:
被抢占的请求还可能进入 swapped 或 recompute 状态。每次转移都要同时更新队列、KV 所有权、token 数、采样状态和客户端流。
核心转移不变量¶
最小转移函数接收当前状态、目标状态和当前 KV block 数,返回新的状态与资源数;任何终态都释放请求所有权,非法边直接失败。
ALLOWED = {
"waiting": {"admitted", "cancelled", "failed"},
"admitted": {"prefilling", "cancelled", "failed"},
"prefilling": {"decoding", "cancelled", "failed"},
"decoding": {"finished", "cancelled", "failed"},
}
TERMINAL = {"finished", "cancelled", "failed"}
def apply_transition(state, target, kv_blocks):
if not isinstance(kv_blocks, int) or kv_blocks < 0:
raise ValueError("kv_blocks must be a non-negative integer")
if target not in ALLOWED.get(state, set()):
raise ValueError(f"illegal transition: {state} -> {target}")
return target, 0 if target in TERMINAL else kv_blocks
state, blocks = apply_transition("waiting", "admitted", 0)
state, blocks = apply_transition(state, "prefilling", 2)
state, blocks = apply_transition(state, "cancelled", blocks)
assert state == "cancelled" and blocks == 0
assert "decoding" not in ALLOWED.get(state, set())
try:
apply_transition("waiting", "admitted", -1)
except ValueError:
pass
else:
raise AssertionError("negative KV ownership must be rejected")
状态机的输入是当前状态、目标状态与资源变化,输出是一个新的可枚举状态;终态必须同步清空 KV 所有权。完整请求类比单个转移函数更长,因此默认折叠,但其中的状态集合、合法边和终态回收仍是正文语义的一部分。
请求状态转移与资源回滚
from dataclasses import dataclass
from enum import Enum, auto
class State(Enum):
WAITING = auto()
ADMITTED = auto()
PREFILLING = auto()
DECODING = auto()
FINISHED = auto()
CANCELLED = auto()
FAILED = auto()
TERMINAL = {State.FINISHED, State.CANCELLED, State.FAILED}
ALLOWED = {
State.WAITING: {State.ADMITTED, State.CANCELLED, State.FAILED},
State.ADMITTED: {State.PREFILLING, State.CANCELLED, State.FAILED},
State.PREFILLING: {State.DECODING, State.CANCELLED, State.FAILED},
State.DECODING: TERMINAL,
}
@dataclass
class Request:
request_id: str
state: State = State.WAITING
kv_blocks: int = 0
computed_tokens: int = 0
def transition(self, target):
if target not in ALLOWED.get(self.state, set()):
raise ValueError(f"illegal transition: {self.state} -> {target}")
self.state = target
if target in TERMINAL:
self.kv_blocks = 0
def reserve_blocks(self, count):
if self.state not in {State.ADMITTED, State.PREFILLING, State.DECODING}:
raise ValueError("request does not own runtime resources")
if not isinstance(count, int) or count <= 0:
raise ValueError("reservation count must be a positive integer")
self.kv_blocks += count
r = Request("r0")
r.transition(State.ADMITTED)
r.reserve_blocks(2)
try: r.reserve_blocks(0)
except ValueError: pass
else: raise AssertionError("non-positive reservation must be rejected")
r.transition(State.PREFILLING)
r.transition(State.DECODING)
r.computed_tokens = 8
r.transition(State.FINISHED)
assert r.kv_blocks == 0 and r.computed_tokens == 8
try:
r.transition(State.DECODING)
raise AssertionError("terminal request was revived")
except ValueError:
pass
reference 把非法跳转和终态复活变为显式错误,并保持 computed_tokens 作为已提交历史。生产引擎还需让队列移动、block 分配、采样状态和客户端事件在同一事务边界提交;尤其不能把“字段已改”误当作 GPU 使用已经结束。block table、连续批处理与回滚的组合实现见手撕:推理引擎。
最小请求对象包括:
request and tenant id
prompt / generated token ids
priority, deadline and arrival time
sampling and grammar state
model, adapter and tokenizer version
logical block table
computed-token count
finish reason and cancellation token
Continuous batching¶
静态 batch 等所有序列结束再换批,短请求会等待最长请求。continuous batching 在每个模型迭代边界移除已完成请求、加入新请求。Orca 给出了迭代级调度的代表性设计。
调度器每轮决定:
- 哪些 waiting 请求做 prefill;
- 哪些 running 请求做一个或多个 decode token;
- 本轮总 token budget;
- 是否抢占或换出低优先级请求;
- 是否有足够 KV block;
- 哪些请求因 SLO 或公平性优先。
“batch size”因此是动态的;性能报告应同时给出 scheduled sequences 与 scheduled tokens。
build_batch baseline¶
输入是一轮开始时的请求快照,以及 token、序列和 prefill chunk 三个预算;输出为 (request_id, phase, scheduled_tokens) 列表。baseline 先给每个 decode 请求一个 token,再用剩余预算放入 waiting / prefill 请求。
def build_batch(requests, token_budget, max_sequences, prefill_chunk):
if min(token_budget, max_sequences, prefill_chunk) <= 0:
raise ValueError("scheduler budgets must be positive")
request_ids = [request["id"] for request in requests]
if len(request_ids) != len(set(request_ids)):
raise ValueError("request ids must be unique within a scheduling round")
batch, used = [], 0
decode = [r for r in requests if r["phase"] == "decode"]
prefill = [r for r in requests if r["phase"] in {"waiting", "prefill"}]
for request in decode + prefill:
if len(batch) == max_sequences or used == token_budget:
break
available = token_budget - used
tokens = 1 if request["phase"] == "decode" else min(
request["prompt_left"], prefill_chunk, available
)
if tokens > 0 and tokens <= available:
batch.append((request["id"], request["phase"], tokens))
used += tokens
return batch
requests = [{"id": "d0", "phase": "decode", "prompt_left": 0},
{"id": "p0", "phase": "prefill", "prompt_left": 6},
{"id": "d1", "phase": "decode", "prompt_left": 0}]
batch = build_batch(requests, token_budget=5, max_sequences=3, prefill_chunk=3)
assert [item[0] for item in batch] == ["d0", "d1", "p0"]
assert sum(item[2] for item in batch) == 5
assert all(tokens == 1 for _, phase, tokens in batch if phase == "decode")
try:
build_batch(requests + [requests[0]], 5, 4, 3)
except ValueError:
pass
else:
raise AssertionError("one request cannot be scheduled twice in a round")
每个请求在一轮中最多出现一次,decode 恰好消费一个 token,且两个预算都不得越界。函数不修改请求状态,也不承诺公平性;生产调度还必须联合 KV / workspace reservation、age、deadline、cache affinity、抢占和原子提交。状态推进与更多预算断言见手撕:推理引擎 · Continuous batching。
分页 KV¶
连续为每个请求预留最大上下文会浪费显存。分页运行时把逻辑 token 区间映射到固定大小物理 block:
block table 让一个请求的 KV 不必物理连续。PagedAttention 进一步让 attention kernel 按映射读取 K/V。
Block 大小¶
- 大 block:映射少、访问规整,但最后一块内部碎片大;
- 小 block:碎片小,但元数据、分配和 kernel 查表更多。
选择必须使用真实 prompt/output 长度与并发分布,而不是只测固定长度。
Prefix cache 与 copy-on-write¶
共享 system prompt、文档前缀或 beam 分支可指向相同只读 block。当前缀末尾的 block 尚未填满,新分支继续写入前必须 copy-on-write,否则会修改其他请求的历史。
cache key 至少包含:
exact token ids
model and weights
adapter
position and rope configuration
attention/cache dtype
tenant or permission boundary
字符串相同但 tokenization、adapter 或权限不同都不能复用。跨租户共享即使内容碰巧相同,也要经过明确的数据隔离设计。
Chunked prefill¶
长 prompt 的一次 prefill 可能占据整个 GPU,阻塞交互式 decode。chunked prefill 将其分为多个 token 块,与 decode 请求拼成每轮预算:
Sarathi-Serve 展示了用 chunked prefill 控制 throughput–latency 权衡的路线。chunk 越大,prefill GEMM 越高效;越小,decode stall 越低但 launch 和重复调度更多。
抢占¶
显存不足时有两种常见选择:
- recompute:释放 KV,恢复时重新 prefill 已处理 token;
- swap:把 KV 移到 CPU 或更低层存储,恢复时传回。
recompute 消耗计算,swap 消耗链路带宽和主机内存。选择取决于上下文长度、带宽、优先级和预期等待时间。无论哪种方式,都必须保留 sampler、grammar 与 position 状态。
Prefill 与 decode 输入¶
运行时常把不同请求的 token 展平成一条 token array,并用 metadata 描述每个序列的 query range、context length 和 block table。kernel 正确性依赖这些 offset:
- prefill 请求可能一次有多个 query token;
- decode 请求通常只有一个新 query;
- chunked prefill 已有历史 cache;
- prefix hit 使“prompt 长度”不等于“本轮计算长度”;
- speculative verify 一次有多个候选 query。
只根据 tensor shape 推断语义容易在混合 batch 中出错。
Streaming 与 detokenization¶
token 生成完成不等于客户端已收到文本。运行时还要:
- 增量解码 byte/subword;
- 缓冲潜在 stop sequence;
- 处理 Unicode 未完成片段;
- 发送 backpressure;
- 客户端断开后及时取消;
- 保留 finish reason;
- 避免慢客户端阻塞 GPU 调度线程。
tokenizer 和网络发送通常放在独立线程或进程,但状态顺序必须与请求 token 序列一致。
引擎正确性¶
最小回归矩阵包括:
- 单请求 prefill + decode 对齐框架 reference;
- 不同长度 continuous batch;
- prefix hit/miss 与 copy-on-write;
- block 分配、回收和抢占;
- chunked prefill 与 decode 混排;
- beam/speculative 分支;
- adapter、量化和不同 cache dtype;
- 取消、超时、OOM 与进程重启;
- 固定 seed 下的可接受重放边界。
策略与过载控制见调度与服务,跨集群阶段拆分见 Prefill–Decode 分离。