使用Transformer模型：从训练到推理
内容提要
本文介绍Transformer模型从训练到推理的差异。训练处理批量固定长度序列并更新权重，推理则固定权重逐token生成。推理分为预填充和解码两阶段，预填充一次性处理提示词，解码逐token生成。KV缓存存储先前token的键值，避免重复计算，将解码注意力成本从二次方降为线性。文章还讨论了KV缓存的内存占用及管理优化。
延伸解读
推理与训练的本质差异
训练阶段模型处理固定长度的批量序列，通过反向传播更新权重；推理阶段权重固定，模型逐token生成输出。这种差异导致性能瓶颈完全不同：训练受限于大规模矩阵乘法和反向传播，而推理则受限于重复的前向传播、内存移动以及为下一个token保留历史注意力键值。理解这一区别是优化推理性能的基础。
预填充与解码阶段
推理被分为预填充和解码两个阶段。预填充一次性处理整个提示词，计算所有token的键值并生成第一个输出token；解码阶段则每次只处理最新生成的token，利用缓存的键值计算注意力。预填充影响首token延迟，解码影响流式输出速度。短提示长回答主要受解码限制，长提示短回答则受预填充限制，而长对话历史两者都受影响。
KV缓存的计算与内存权衡
KV缓存通过存储历史token的键值，将解码阶段的注意力计算复杂度从二次方降为线性，但代价是内存占用。例如，一个32层、32个KV头、头维度128、BF16精度、序列长度4096的模型，单请求的KV缓存约需2GiB内存。实际部署中，多用户并发时KV缓存内存成为瓶颈，因此需要高效的内存管理策略，如预分配或分页缓存。
Q&A
Transformer模型在训练和推理阶段的主要区别是什么？
训练时，模型处理批量固定长度的token序列并更新权重；推理时，权重固定，模型逐token生成新token。训练主要由大型矩阵乘法和反向传播主导，而推理由重复的前向传播、内存移动以及保留先前注意力键值对主导。
什么是自回归生成？Transformer模型如何逐token生成文本？
自回归生成是指模型根据之前生成的token预测下一个token，每个新token依赖于之前生成的token。生成循环包括：对提示词进行分词，运行模型获得logits，选择token，将其追加到输入，重复直到满足停止条件。
Transformer推理中的prefill和decode阶段分别是什么？
Prefill阶段一次性处理整个提示词，计算所有提示词token的隐藏状态和注意力键值，并生成第一个新token的logits。Decode阶段则逐token生成，每次只处理最新生成的token，并更新KV缓存。
为什么需要KV缓存？它如何提高推理效率？
KV缓存存储先前token的键和值，避免在生成每个新token时重新计算整个序列的注意力。这使得解码阶段的每个token注意力成本从O(N^2)降低到O(N)，显著提高推理速度。
如何实现一个简单的KV缓存？
在自注意力模块中，当提供past_kv时，将新的键和值拼接到缓存的键值上。在生成循环中，先进行prefill处理整个提示词，然后每次解码只输入最新token，并传递缓存。
KV缓存的内存占用如何计算？
内存占用公式为：2 * 层数 * 批次大小 * 序列长度 * KV头数 * 头维度 * 每个元素字节数。例如，32层、32个KV头、头维度128、BF16精度、批次1、序列长度4096时，缓存占用约2 GiB。
KV缓存管理在推理系统中面临哪些挑战？
KV缓存占用大量内存，多用户并发时成为瓶颈。简单实现使用torch.cat反复分配和复制，效率低。生产系统需要预分配或分页管理，如PagedAttention，以减少内存碎片并支持更大批次。
