Principles

Tokenization Basics BPE Algorithm GPT Tokenizers BPE Training Engineering

Model Architecture

From token ids to logits Embedding and LM Head

Attention Mechanisms

From Self-Attention to GQA Attention Sink

Position Encoding

Position Encoding Basics RoPE Math Derivation RoPE Implementation Length Extrapolation

GPU Programming Basics

GPU Architecture Basics Tensor Layout Triton Basics: Vector Add

Flash Attention Principles From Naive Implementation to Auto-Tuning Block Pointers and Multi-Dim Support Causal Masking Optimization Grouped Query Attention Backward Pass Implementation

Distributed Training

Data Parallelism ZeRO Optimizer Fully Sharded Data Parallel Tensor Parallelism Pipeline Parallelism Multi-Dimensional Hybrid Parallelism

Hands-on Training

Pretraining Data Tokenizer Training Model Architecture Data Pipeline Training Loop Monitoring and Validation

PagedAttention

Premium

按最大长度预留连续显存只用上了两三成，用操作系统的分页思路把利用率做到 96%

Get code access

前两章我们做了两件事：用 KV cache 省掉重复计算，用 continuous batching 让 GPU 不再空转。但每一章的结尾都指向同一个没解决的问题：这些缓存到底该怎么在显存里摆放。

这一章要回答的就是它。我们会先算清楚显存到底被谁吃掉了，看到一个相当难堪的事实：现有系统里真正用于存 token 的显存只有两三成。然后引入一个四十多年前操作系统就用过的办法，把利用率拉到 96%，顺便白送一个意想不到的能力。

显存到底被谁吃掉了

先做一道算术题，感受一下 KV cache 有多占地方。

Log in to continue reading

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

拿 OPT-13B 举例，算一个 token 要存多少： $K$ 和 $V$ 各一份，每份的宽度是 hidden size 5120，每一层都要存一份、共 40 层，FP16 每个数 2 字节。乘起来：

2 \times 5120 \times 40 \times 2\,\text{B} = 800\ \text{KB}

一个 token 就要 800 KB。OPT 支持生成到 2048 个 token，所以一个请求的 KV cache 最多能吃掉 1.6 GB。

再看这块卡上还剩多少地方。一张 A100 40GB，装完 13B 模型的 FP16 权重（约 26 GB，占 65%）后，留给 KV cache 的空间只有十几 GB。按每请求 1.6 GB 算，同时只能服务十来个请求。

这就是第一层认识：KV cache 不是什么边角开销，它和模型权重是同一个量级的对手。而上一章刚讲过，batch 越大吞吐越高，所以能同时装下多少请求，直接决定了这台机器值多少钱。

Table of Contents

显存到底被谁吃掉了

装得下的比你以为的还少

这个问题，操作系统四十年前就解决过

PagedAttention：把缓存切成块

注意力怎么在散落的块上算

分页白送的能力：共享

块该切多大

代价：这不是免费的抽象

一个后续：prefix caching 的演进