自回归推理、Prefill 与 Decode

一条请求的准确状态机

给定 token ids $u_{0:T_p-1}$ ，decoder-only 模型先执行 prefill：整段 prompt 以 $[B,T_p]$ 输入，各层同时处理 $T_p$ 个位置并建立 KV Cache，最后一个 prompt 位置产生 logits $z_{T_p}$ 。采样器从它选出第一个输出 token $u_{T_p}$ 。

此后第一个 decode 调用把刚选出的 $u_{T_p}$ 作为 shape $[B,1]$ 的新输入，追加它的 K/V，并产生下一个位置的 logits $z_{T_p+1}$ 。循环为：

z_{T_p}=f(u_{0:T_p-1}),\qquad u_{T_p}\sim\pi(z_{T_p})

z_{t+1}=f(u_t;K_{0:t-1},V_{0:t-1}),\qquad u_{t+1}\sim\pi(z_{t+1})

$\pi$ 是 greedy、temperature、top-k/top-p 等采样策略。模型每次前向产生下一个 token 的分布，不是把当前输入 token 原样输出。

两个阶段的 shape 与硬件特征

项目	Prefill	单步 Decode
token ids	$[B,T_p]$	$[B,1]$
hidden states	$[B,T_p,H]$	$[B,1,H]$
attention Q	$[B,H_q,T_p,D]$	$[B,H_q,1,D]$
attention K/V source	当前 $T_p$ 个位置	历史 cache 加当前新位置，长度约 $S$
score 逻辑 shape	$[B,H_q,T_p,T_p]$	$[B,H_q,1,S]$
GEMM 形态	较大的 token 矩阵，更易形成高并行度	token 维为 1，常接近 GEMV/小 GEMM
cache	批量写入 prompt K/V	每层追加一个 K/V 位置并读取历史
常见服务指标	主导 TTFT 的模型执行部分	主导 TPOT / inter-token latency

Prefill 往往更偏 compute-bound，但长 prompt 的 attention IO、内存容量和 chunking 也可能成为瓶颈；decode 在小 batch 下常受权重与 KV 读取带宽、kernel launch 和 CPU 调度影响。它们不是硬性分类，batch size、sequence length、量化、fusion、硬件和并行策略都会移动瓶颈。

TTFT、TPOT 与吞吐不是同一个量

从服务端接收请求到客户端看到首 token 的 Time to First Token 可拆成：

TTFT\approx T_{queue}+T_{tokenize/template}+T_{prefill}+T_{sample}+T_{stream}

这些项可能流水或重叠，上式是核对边界，不保证可机械相加。Time per Output Token 常统计首 token 之后相邻输出 token 的间隔；若请求总延迟为 $T_{e2e}$ 、TTFT 为 $T_{first}$ 、输出 $G>1$ 个 token，一种请求级定义是：

TPOT=\frac{T_{e2e}-T_{first}}{G-1}

生产指标必须说明是平均、p50/p99、仅 completed requests，还是逐 step 直方图。吞吐（tokens/s 或 requests/s）还受到 continuous batching、并发与调度影响；提高批量可能增加总吞吐，却恶化某个请求的 TTFT 或 TPOT。

可运行的小模型阶段计时

实验问题：不用网络或模型下载，能否完整走过 token ids → prefill → 首 token → 增量 decode，并分别测量模型侧 TTFT 与 steady TPOT？

import math, platform, statistics, time
import torch

torch.manual_seed(31)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DTYPE = torch.float32

class TinyDecoder:
    def __init__(self, vocab=128, hidden=256, q_heads=4, kv_heads=2, head_dim=64):
        self.vocab, self.hidden = vocab, hidden
        self.q_heads, self.kv_heads, self.head_dim = q_heads, kv_heads, head_dim
        scale = hidden**-0.5
        shapes = [(vocab, hidden), (hidden, q_heads*head_dim),
                  (hidden, kv_heads*head_dim), (hidden, kv_heads*head_dim),
                  (q_heads*head_dim, hidden), (hidden, vocab)]
        self.w = [torch.randn(s, device=DEVICE, dtype=DTYPE)*scale for s in shapes]

    def project(self, x):
        b, t, _ = x.shape
        q = (x@self.w[1]).view(b,t,self.q_heads,self.head_dim).transpose(1,2)
        k = (x@self.w[2]).view(b,t,self.kv_heads,self.head_dim).transpose(1,2)
        v = (x@self.w[3]).view(b,t,self.kv_heads,self.head_dim).transpose(1,2)
        return q, k, v

    def attend(self, q, k, v, causal):
        repeat = self.q_heads // self.kv_heads
        k = k.repeat_interleave(repeat, 1); v = v.repeat_interleave(repeat, 1)
        scores = q @ k.transpose(-2,-1) / math.sqrt(self.head_dim)
        if causal:
            t = q.shape[-2]
            scores += torch.triu(torch.full((t,t), -torch.inf, device=DEVICE), 1)
        context = torch.softmax(scores, -1) @ v
        return context.transpose(1,2).reshape(q.shape[0],q.shape[2],-1) @ self.w[4]

    def prefill(self, token_ids):
        x = self.w[0][token_ids]; q,k,v = self.project(x)
        hidden = x + self.attend(q,k,v,causal=True)
        return hidden[:,-1] @ self.w[5], (k,v)

    def decode(self, token_ids, cache):
        x = self.w[0][token_ids]; q,k_new,v_new = self.project(x)
        k = torch.cat((cache[0],k_new),2); v = torch.cat((cache[1],v_new),2)
        hidden = x + self.attend(q,k,v,causal=False)
        return hidden[:,-1] @ self.w[5], (k,v)

def sync():
    if DEVICE.type == "cuda": torch.cuda.synchronize()

def measure(fn):
    sync(); start=time.perf_counter(); result=fn(); sync()
    return result, (time.perf_counter()-start)*1e3

def run_request(model, prompt, max_new_tokens):
    def prefill_and_sample():
        logits, cache = model.prefill(prompt)
        return logits.argmax(-1,keepdim=True), cache
    (next_token,cache), ttft = measure(prefill_and_sample)
    generated, decode_ms = [next_token], []
    for _ in range(max_new_tokens-1):
        (logits,cache), elapsed = measure(lambda: model.decode(next_token,cache))
        decode_ms.append(elapsed); next_token=logits.argmax(-1,keepdim=True)
        generated.append(next_token)
    return torch.cat(generated,1), cache, ttft, decode_ms

def main():
    model=TinyDecoder(); b,pt,g=2,128,16
    prompt=torch.randint(0,model.vocab,(b,pt),device=DEVICE)
    for _ in range(3): run_request(model,prompt,g)  # same-shape warmup
    generated,cache,ttft,decode_ms=run_request(model,prompt,g)
    assert generated.shape==(b,g)
    assert cache[0].shape==cache[1].shape==(b,model.kv_heads,pt+g-1,model.head_dim)
    assert torch.all((0<=generated)&(generated<model.vocab))
    print(f"python={platform.python_version()} torch={torch.__version__}")
    print(f"device={DEVICE} dtype={DTYPE}")
    print(f"prompt={tuple(prompt.shape)} generated={tuple(generated.shape)} final_KV={tuple(cache[0].shape)}")
    print(f"model_ttft_prefill_plus_argmax_ms={ttft:.3f}")
    print(f"first_decode_step_ms={decode_ms[0]:.3f}")
    print(f"steady_tpot_median_ms={statistics.median(decode_ms[1:]):.3f}")
    print(f"decode_step_range_ms=[{min(decode_ms):.3f}, {max(decode_ms):.3f}]")

if __name__ == "__main__":
    with torch.inference_mode(): main()

完整文件位于 examples/prefill_decode_reference.py。CUDA 计时前执行三次相同 shape warmup，每个阶段前后同步。2026-08-13 在 RTX 3080、PyTorch 2.13.0+cu130 上一次实际结果为：prompt=[2,128]、生成 16 token、最终 KV [2,2,143,64]，模型侧 prefill+argmax 0.175 ms，首个 decode step 0.137 ms，稳态 TPOT 中位数 0.133 ms。

该模型只有一层简化 attention，没有 MLP、norm、RoPE、真实 tokenizer、框架调度或网络 streaming；绝对数字不能代表生产服务。它验证的是阶段边界、shape、同步计时模板和 token/cache 状态机。教学实现的 torch.cat 与 GQA 展开也会扭曲真实性能。

成本如何随长度变化

Prefill 对 prompt 做大矩阵投影，dense attention 的逻辑 score 规模为 $O(T_p^2)$ 。借助 KV Cache，第 $t$ 个 decode step 不再重算历史层，只处理一个新 token；但它仍读取 $O(tH_{kv}D)$ 的 K/V，并计算 $O(H_qtD)$ attention。权重读取、KV 读取、并行 batch 与 kernel launch 共同决定实际 TPOT。

因此：

增加 prompt 长度通常主要推高 prefill 与 TTFT，也增加后续每步 decode 的历史读取；
增加输出长度主要增加 decode step 数和请求占用调度器的时间；
continuous batching 可用不同请求填充每个 iteration，提高设备利用率，但排队、batch 组成和 prefill/decode 干扰会改变单请求指标；
chunked prefill 把长 prompt 切入多个调度轮次，可能改善公平性，却改变 TTFT 与调度开销。

不解决的问题与生产边界

Prefill/Decode 是执行阶段，不是具体 backend；每阶段内部仍可能用 dense SDPA、FlashAttention、paged decode kernel、量化 GEMM 或 CUDA Graph。
max_new_tokens、stop token、stop string、EOS 与 cancellation 决定何时结束；模型前向本身不知道 HTTP 请求是否应停止。
Streaming 以 token id / byte 序列为边界，还要处理 UTF-8、byte fallback、special tokens 和 stop string，不能简单把每个 token 文本直接拼接。
Request Lifecycle把 queue/admission、模型执行、增量 detokenization、cancel/timeout 与终态资源回收串成完整状态机。
speculative decoding 与 MTP 会一次提出或验证多个候选，但最终接受协议仍要维持目标分布；它们不是把普通 decode 的单步语义直接删掉。

上下游关系

上游 Tokenizer、Chat Template 与 Special Tokens 决定 prompt token ids 和 $T_p$ ；模板变化会改变 TTFT、停止条件和 prefix cache key。
MHA / MQA / GQA 与缩放点积注意力定义阶段内部的 head shape 与 attention 数学。
KV Cache 保存 prefill 结果并支持增量 decode；KV 页的对照实验说明它省掉哪些历史重算。
Logits 与基础采样承载 logits 到 token 的稳定 softmax、temperature、top-k/top-p；本页只用 argmax 固定阶段状态机。
后续 Continuous Batching 页面应在同一请求 trace 上报告 queue、TTFT、TPOT、吞吐与调度公平性，而不是只测一个孤立请求。

参考资料

vLLM — Metrics — TTFT、inter-token latency、queue/prefill/decode 等生产指标定义。
NVIDIA — Mastering LLM Techniques: Inference Optimization — prefill/decode 的计算与带宽特征概述。
PyTorch — scaled_dot_product_attention — attention API 与 backend 边界。