LogoCookLLM文档
LogoCookLLM文档
首页CookLLM

原理精讲

词元化
Tokenization 基础BPE 算法详解GPT 系列 TokenizerBPE 训练工程化
模型架构
Transformer LM
从 token ids 到 logitsEmbedding 与 LM Head
Attention 机制
Self-Attention 到 GQAAttention Sink
位置编码
位置编码基础RoPE 数学推导RoPE 代码实现长度外推
GPU 编程基础
GPU 架构基础张量布局Triton 入门:向量加法
FlashAttention
Flash Attention 原理详解从朴素实现到 Auto-TuningBlock Pointer 与多维支持Causal Masking 优化Grouped Query Attention反向传播实现
分布式训练
数据并行ZeRO 优化器全分片数据并行张量并行流水线并行多维混合并行
推理优化
KV CacheContinuous BatchingPagedAttention

动手训练

概述
预训练
预训练数据Tokenizer 训练模型架构数据流水线训练循环监控与验证
X (Twitter)
系统工程推理优化

KV Cache

会员专享

decode 每步重算历史 token 的 K/V 是纯浪费,KV cache 把它省掉,代价是显存

在权益中心获取代码

你已经会让 Transformer 做一次完整的前向传播:给定一段 token,一次算出每个位置的输出(见 Transformer 前向传播)。但真正生成文本时,模型不是一次吐出一整段,而是一个 token 一个 token 地挤出来。这个"逐 token"的过程藏着一笔巨大的浪费:每生成一个新 token,朴素实现都要把整段历史重新过一遍模型,而历史 token 的 K 和 V 从第一次算出来那一刻起就再也不会变。

这一章就来消除这份重复劳动。我们要搞清楚四件事:这笔浪费到底有多大、为什么历史的 K/V 一定不变、怎么把它们缓存下来复用、以及这份缓存会吃掉多少显存。KV cache 是几乎所有推理系统的地基,理解它,才能看懂后面 batching、显存管理、量化这些优化在对付什么。

生成是一步一步挤出来的

先回忆自回归生成的循环。给定已有的 token 序列,每一步做四件事:

  1. 把当前整个序列喂进模型,做一次前向传播
  2. 取最后一个位置的 logits
  3. 对词表做 softmax,选出下一个 token(贪心、采样、top-k/top-p 都行)
  4. 把这个 token 拼回序列末尾,回到第 1 步

写成公式,模型每一步都在估计一个条件分布:

P(xt∣x1,x2,…,xt−1)P(x_t \mid x_1, x_2, \ldots, x_{t-1})P(xt​∣x1​,x2​,…,xt−1​)

注意条件是从 111 到 t−1t-1t−1 的每一个 token,不只是上一个。Transformer 预测下一个词时会注意到全部历史,这正是它能在几千个 token 上保持连贯的原因,也正是每一步都得把整段序列重新过一遍的根源。

这个循环其实分成性质完全不同的两段:

  • Prefill(预填充):把整个 prompt 一次性喂进去,所有位置并行算完,产出第一个新 token。假设 prompt 有 NNN 个 token,这一步一次处理 NNN 个位置。
  • Decode(解码):之后每一步只新增一个 token,逐个往外挤。

Generation Flow

BentoLM generation uses one prefill pass, then token-by-token decode with KV cache.

1
Prompt tokens
plain text continuation or ChatML prompt
input_ids
2
Prefill
run the full prompt once and build KV cache
past_key_values
3
Decode loop
feed only the newest token at each step
logits
4
Sampling
temperature, top-k, top-p and repetition penalty
next token
5
Stop
finish on max_new_tokens or EOS
output ids

生成 100 个 token,等于 1 次 prefill 加 99 次 decode。prefill 只发生一次,decode 才是循环的主体,也是所有推理优化真正要对付的部分。

登录以继续阅读

这是一篇付费内容,请登录您的账户以访问完整内容。

推理优化

从 KV cache 到批处理调度,理解 LLM 推理为什么慢、怎么加速

Continuous Batching

静态批处理让 GPU 大量空转,迭代级调度、selective batching 与 chunked prefill 怎么把它填满

目录

生成是一步一步挤出来的
decode 每步在重算什么
关键观察:K 和 V 算过就不变
KV Cache 机制
每一层每个头都有一份缓存
为什么第一个 token 慢,decode 却受显存带宽限制
KV cache 有多大:算一算内存
什么时候不划算
总结