LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)

KV Cache

Premium

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

Get code access

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

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

生成是一步一步挤出来的

Log in to continue reading

This is premium content. Please log in to access the full article.

先回忆自回归生成的循环。给定已有的 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​,…,x

注意条件是从 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 才是循环的主体，也是所有推理优化真正要对付的部分。

Table of Contents

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