Transformers 大模型推理优化实战：低精度量化、Flash Attention 与架构创新-平芜编程栈

Transformers 大模型推理优化实战：低精度量化、Flash Attention 与架构创新

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

本文基于 Transformers 仓库中的官方优化指南（docs/source/ar/llm_tutorial_optimization.md）编写，系统讲解如何从"低精度量化、Flash Attention、架构创新（相对位置编码 + KV Cache 优化）"三个维度降低大语言模型（LLM）的推理内存占用并提升推理效率。读完后，你将掌握 VRAM 需求估算方法、8-bit/4-bit 量化的完整实操代码、KV Cache 与多轮对话加速的实现原理，以及 MQA/GQA 等模型架构选择的决策依据，并能对照仓库源码（如src/transformers/utils/quantization_config.py、src/transformers/cache_utils.py）理解每个优化手段的落地位置。

1. 为什么 LLM 部署难：内存与长上下文的挑战

GPT3/4、Falcon、Llama 这类大语言模型在文本理解与生成上接近人类水平，但要在真实任务中部署它们，存在两大障碍：

参数量巨大：要表现出接近人类的文本理解与生成能力，LLM 通常需要数十亿乃至数千亿参数（见 Kaplan et al. 的 scaling 研究），这直接推高了推理时的内存需求；
长上下文：许多真实任务需要给模型提供大量背景信息，要求模型在推理时能处理超长输入序列。

这两点共同指向一个核心问题：如何增强 LLM 的计算与内存能力，尤其是在处理海量输入序列时。本指南围绕三类已被验证有效的技术展开：

低精度（Lower Precision）：研究表明，以 8-bit / 4-bit 等更低的数值精度运行，可以在不明显损害模型效果的前提下获得计算与内存收益（详见仓库文档docs/source/en/main_classes/quantization.md）；
Flash Attention：对注意力算法的改造版本，不仅内存更高效，还通过优化 GPU 显存（VRAM）利用方式带来速度提升；
架构创新（Architectural Innovations）：由于 LLM 在推理时总是以"长输入上下文 + 自回归生成"的方式运行，研究者提出了专门的模型结构来提升推理效率，代表性进展包括 ALiBi、旋转位置编码（RoPE）、多查询注意力（MQA）与分组查询注意力（GQA）。

整个指南从"张量视角"分析自回归生成过程，通过一系列可运行的实践示例，逐一验证每种优化带来的实际收益。

2. 低精度：从 VRAM 估算到 8-bit / 4-bit 量化

2.1 先算账：加载一个 LLM 需要多少显存？

把 LLM 看作一组权重矩阵和向量，把文本输入看作一串向量序列（下文中"权重"指模型中所有权重矩阵与向量）。每个参数都是一个如4.5689的十进制数，通常以 float32、bfloat16 或 float16 格式存储，由此可以估算把模型加载进内存所需的显存：

加载一个 X 十亿参数的模型权重，在 float32 精度下大约需要 4 * X GB 显存。

但如今模型很少以完整 float32 训练，更多是 bfloat16，较少为 float16，因此经验法则变为：

加载一个 X 十亿参数的模型权重，在 bfloat16/float16 精度下大约需要 2 * X GB 显存。

对于较短的文本输入（少于 1024 个 token），推理的内存需求几乎完全由加载权重主导，因此可以近似认为推理内存需求 ≈ 把模型装入 GPU 显存的内存需求。按 bfloat16 估算一些知名模型的显存需求：

模型	参数量	bfloat16 显存估算
GPT3	175B	2 × 175 =350 GB
BLOOM	176B	2 × 176 =352 GB
Llama-2-70b	70B	2 × 70 =140 GB
Falcon-40b	40B	2 × 40 =80 GB
MPT-30b	30B	2 × 30 =60 GB
bigcode/starcoder	15.5B	2 × 15.5 =31 GB

写这篇文档时，市场上最大的 GPU 芯片是提供 80 GB 显存的 A100 与 H100。上表多数模型仅加载就需要超过 80 GB，因此必然需要张量并行（tensor parallelism）和/或流水线并行（pipeline parallelism）。指南原文指出，当时 Transformers 不原生支持张量并行，因为它要求以特定方式编写模型结构（可参考 text-generation-inference 库中的自定义建模写法）；而朴素流水线并行开箱即用。当前仓库的英文版同篇文档已更新说明：对于配置类中定义了base_tp_plan的受支持模型，Transformers 现已支持张量并行——从这一演进可以看出仓库在此方向上的持续投入。

朴素流水线并行的用法非常简单：加载模型时传device_map="auto"，框架会自动把不同层放到可用的多块 GPU 上。注意朴素流水线并行虽然有效，但并不解决 GPU 空闲（idling）问题，需要更高级的流水线并行才能进一步利用硬件。如果你有一台 8 × 80GB A100 的节点，可以这样加载 BLOOM：

pip install transformers accelerate bitsandbytes optimum

from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("bigscience/bloom", device_map="auto", pad_token_id=0)

使用device_map="auto"后，注意力层会被均匀分布到所有可用 GPU 上。

2.2 实测基线：bigcode/octocoder 在 bfloat16 下需要约 29 GB

后续所有实验使用bigcode/octocoder（约 15.5B 参数），因为它能跑在单块 40 GB 的 A100 上。需要注意：本节所有内存与速度优化手段同样适用于需要模型/张量并行的更大模型。

模型以 bfloat16 加载，按上面的经验法则，预期推理显存需求约 31 GB。先加载模型与 tokenizer，再传给 Transformers 的pipeline对象：

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline import torch model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", dtype=torch.bfloat16, device_map="auto", pad_token_id=0) tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder") pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer:" result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):] result

输出：

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

生成的代码可以直接用来做字节到 GiB 的换算：

def bytes_to_giga_bytes(bytes): return bytes / 1024 / 1024 / 1024

调用torch.cuda.memory.max_memory_allocated()测量 GPU 峰值显存分配：

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出：

29.0260648727417

与粗略计算非常接近。数字不是完全精确，是因为字节到千字节换算应乘 1024 而非 1000，所以该经验公式也可以理解为"至多 X GB"的估算。如果改用完整 float32 运行，需要约 64 GB 显存。

如今几乎所有模型都以 bfloat16 训练。如果你的 GPU 支持 bfloat16，就没有理由用完整 float32 运行模型——float32 不会给出比训练精度更好的推理结果。

如果不確定模型权重在 Hub 上以什么格式存储，可以查看检查点的 config 文件中的"dtype"字段。建议用from_pretrained(..., dtype=...)加载时设置与 config 中一致的精度类型；除非原始类型是 float32，此时推理可以用float16或bfloat16。

为了准确测量后续实验的峰值显存，定义一个flush(...)函数释放全部已分配内存：

del pipe del model import gc import torch def flush(): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats()

flush()

Accelerate 库还提供了设备无关的辅助方法release_memory()，可一并释放模型占用的内存：

from accelerate.utils import release_memory # ... release_memory(model)

2.3 量化原理：为什么 8-bit / 4-bit 能跑，且不伤精度？

如果 GPU 显存不到 32 GB 怎么办？研究发现，模型权重可以被量化到 8-bit 甚至 4-bit 而几乎没有性能损失（Dettmers et al. 的工作）；GPTQ 论文进一步表明量化到 3-bit / 2-bit 也能在可接受的损失下运行。

不深入细节，量化方案的共同目标是：降低权重的存储精度，同时尽量保持模型推理结果准确（即尽可能接近 bfloat16 的推理结果）。量化对文本生成尤其有效，因为生成任务只关心选择"最可能的下一批 token"，并不真正关心下一 tokenlogit分布的精确数值——只要logit分布大体不变，argmax或topk的结果就不变。

尽管具体技术各有不同，所有量化方法的工作流程都可以概括为三步：

把所有权重量化到目标精度；
加载量化后的权重，输入序列仍以 bfloat16 格式传入；
在计算时把权重动态反量化回 bfloat16，再与 bfloat16 输入向量做运算。

也就是说，每个矩阵乘法

$$ Y = X * W $$

（其中 $X$ 是输入、$W$ 是权重矩阵、$Y$ 是输出）都变成了：

$$ Y = X * \text{dequantize}(W) $$

反量化与重量化在所有权重矩阵上按输入穿过网络图的顺序依次执行。因此，使用量化权重通常不会降低推理时间，反而会使其变长——这一点在实测中会被反复验证。

2.4 8-bit 量化实操：显存从 29 GB 降到 15 GB

用 Transformers 做量化需要安装bitsandbytes库：

pip install bitsandbytes

8-bit 量化只需给from_pretrained加上量化配置。指南原文使用旧式写法load_in_8bit=True；在当前仓库中，推荐（也是当前代码路径实际消费的）写法是传入BitsAndBytesConfig。该类的定义位于 src/transformers/utils/quantization_config.py，从源码可以看到：

load_in_8bit：启用 LLM.int8() 的 8-bit 量化，默认False；
load_in_4bit：用bitsandbytes的 FP4/NF4 层替换 Linear 层实现 4-bit 量化，默认False；两者互斥（同时为True会抛出ValueError）；
llm_int8_threshold（默认 6.0）：LLM.int8() 论文中异常值检测的阈值，超过该阈值的隐藏状态值会以 fp16 计算；
bnb_4bit_compute_dtype：4-bit 反量化后参与计算的精度（如可设torch.bfloat16提速）；
bnb_4bit_quant_type：4-bit 量化数据类型，可选fp4或nf4；
bnb_4bit_use_double_quant：双重量化（对第一次量化的常数再做一次量化），进一步省内存。

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_8bit=True), pad_token_id=0)

重新运行示例并测量内存：

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer) result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):] result

输出：

Here is a Python function that transforms bytes to Giga bytes:\n\n```python\ndef bytes_to_giga_bytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single

与之前完全相同，精度无损。再看内存：

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出：

15.219234466552734

显存从约 29 GB 降到 15 GB 出头，这个模型因此可以跑在 4090 这类消费级 GPU 上。内存效率收益很大、输出几乎无退化，但可以观察到推理有轻微变慢——这正是 2.3 节所说的反量化开销。

删除模型并清空内存，进入下一组实验：

del model del pipe

flush()

2.5 4-bit 量化实操：15.5B 模型只占约 9.5 GB

4-bit 量化使用同一套 API，把load_in_8bit=True换成load_in_4bit=True：

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", quantization_config=BitsAndBytesConfig(load_in_4bit=True), pad_token_id=0) pipe = pipeline("text-generation", model=model, tokenizer=tokenizer) result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):] result

输出：

Here is a Python function that transforms bytes to Giga bytes:\n\n```\ndef bytes_to_gigabytes(bytes):\n return bytes / 1024 / 1024 / 1024\n```\n\nThis function takes a single argument

输出与之前几乎一致——只是代码块前少了一个python标注。内存需求：

bytes_to_giga_bytes(torch.cuda.max_memory_allocated())

输出：

9.543574333190918

只有 9.5 GB！对 15B+ 参数模型来说非常少。

虽然本例中精度退化极小，但实践中 4-bit 量化相比 8-bit 或完整 bfloat16 推理更容易产生不同结果，需要你自行在目标任务上验证。同时注意，这里的推理又比 8-bit 略慢，原因是 4-bit 使用了更激进的量化方式，quantize与dequantize在推理时耗时更长。

del model del pipe

flush()

小结：OctoCoder 以 8-bit 运行时，所需显存从约 32 GB 降到约 15 GB；以 4-bit 运行时进一步降到 9 GB 出头。4-bit 量化让模型可以跑在 RTX3090、V100、T4 等对大多数人更易获得的 GPU 上。如果想把显存压得比 4-bit 更低，可以研究仓库量化文档（docs/source/en/main_classes/quantization.md）中介绍的 GPTQModel 等实现。

结论性提醒：模型量化是用（有时还有）推理时间换内存效率的交易。如果显存不是你的约束，通常不必考虑量化；但很多 GPU 不量化根本无法运行 LLM，此时 4-bit / 8-bit 量化方案是极其有用的工具。

3. Flash Attention：把 $N^2$ 的显存压力降到线性

当今高性能 LLM 大多共享同一套基础结构：前馈层、激活层、层归一化层，以及最关键的自注意力层。

自注意力层让模型理解输入 token 间的上下文关系，但其 GPU 峰值显存消耗在计算与内存两个维度上都随输入 token 数（序列长度，记为 $N$）二次方增长。短序列（1000 token 以内）不明显，长序列（约 16000 token）就成为严重问题。

对长度为 $N$ 的输入 $\mathbf{X}$，自注意力层输出 $\mathbf{O}$ 的计算公式为：

$$ \textbf{O} = \text{Attn}(\mathbf{X}) = \mathbf{V} \times \text{Softmax}(\mathbf{QK}^T) \text{ with } \mathbf{Q} = \mathbf{W}_q \mathbf{X}, \mathbf{V} = \mathbf{W}_v \mathbf{X}, \mathbf{K} = \mathbf{W}_k \mathbf{X} $$

$\mathbf{X} = (\mathbf{x}1, ... \mathbf{x}{N})$ 是输入序列。投影 $\mathbf{Q}$ 和 $\mathbf{K}$ 各含 $N$ 个向量，于是 $\mathbf{QK}^T$ 的规模是 $N^2$。LLM 通常有多个注意力头并行做自注意力计算。假设 LLM 有 40 个注意力头、以 bfloat16 运行，存储 $\mathbf{QK^T}$ 矩阵的显存需求是 $40 \times 2 \times N^2$ 字节：

$N = 1000$：只需约 50 MB 显存；
$N = 16000$：需要 19 GB 显存；
$N = 100,000$：仅存储 $\mathbf{QK}^T$ 矩阵就接近 1 TB。

一句话：默认自注意力算法在大输入上下文下迅速变得内存昂贵。随着 LLM 能力增强，它们从处理几句话的翻译/摘要，走向处理整个页面的任务，长输入成为刚需。

出路在于一种不需要显式构造 $\mathbf{QK}^T$ 矩阵的新注意力算法——Tri Dao 等人提出的Flash Attention。

其核心思想是：把 $\mathbf{V} \times \operatorname{Softmax}(\mathbf{QK}^T)$ 的计算拆开，通过对多组 Softmax 分块步骤迭代，逐块计算输出的小片段：

$$ \textbf{O}i \leftarrow s^a{ij} * \textbf{O}i + s^b{ij} * \mathbf{V}{j} \times \operatorname{Softmax}(\mathbf{QK}^T{i,j}) \quad \text{for multiple } i, j \text{ iterations} $$

其中 $s^a_{ij}$ 与 $s^b_{ij}$ 是需要在每一组 $i, j$ 上重新计算的 Softmax 归一化统计量。（Flash Attention 的完整实现更复杂，这里做了大幅简化；深入细节请阅读 Flash Attention 原论文。）

关键要点：

通过追踪 Softmax 归一化统计量并配合巧妙的数学处理，Flash Attention 在显存成本只随 $N$ 线性增长的前提下，给出与默认自注意力层数值完全相同的输出。

从公式直觉上会以为 Flash Attention 更慢——它确实需要更多 FLOPs，因为 Softmax 归一化统计量必须不断重算。但推理中它远快于默认注意力：

Flash Attention 大幅降低了对 GPU 上更慢的高带宽显存（VRAM）的访问需求，转而利用芯片上更快的 SRAM。

本质上，Flash Attention 保证所有中间读写都发生在片上高速 SRAM，而不是访问慢速 VRAM 来计算输出向量 $\mathbf{O}$。实践建议很明确：只要可用，就没有任何理由不用 Flash Attention——它数学上输出相同，且更快、更省内存。

4. 架构创新：为长文本任务优化的模型结构

前面我们从两个方向改进了计算与内存效率：把权重转成低精度格式；用更高效版本替换自注意力算法。本节讨论第三个维度——让 LLM 的结构本身适配长文本任务，例如检索增强问答、摘要、聊天。注意"聊天"不仅要求处理长文本输入，还要求高效处理用户与助手之间的多轮对话（如 ChatGPT 场景）。

模型一旦训练完成，基础结构就很难改变，因此在设计阶段就要预判任务并相应优化架构。结构中有两个组件对长输入序列最先成为内存和/或性能瓶颈：

位置编码（positional embeddings）
键值缓存（key-value cache）

4.1 位置编码：从绝对编码到 RoPE / ALiBi

自注意力把每个 token 与其他 token 建立关联。对输入序列"Hello", "I", "love", "you"，$\operatorname{Softmax}(\mathbf{QK}^T)$ 矩阵可以形如：每个词 token 获得一组概率质量，用以关注所有其他词 token，从而与全部其他 token 建立关系。例如单词"love"以 5% 关注"Hello"、以 30% 关注"I"、以 65% 关注自身。

一个基于自注意力但没有位置编码的 LLM，会很难理解各输入文本的相对位置，因为 $\mathbf{QK}^T$ 计算出的概率得分以 $O(1)$ 计算把每个词 token 关联到每个其他词 token，与相对位置距离无关。对没有位置编码的 LLM 来说，每个 token 到所有其他 token 的"距离"看起来都一样，区分"Hello I love you"与"You love I hello"将非常困难。

为了让 LLM 理解句子顺序，需要额外线索，通常以位置编码（positional encodings，也称positional embeddings）的形式注入：位置编码把每个 token 的位置编码为数值表示，让 LLM 据此更好理解句子顺序。

《Attention Is All You Need》的 authors 引入了正弦位置嵌入$\mathbf{P} = \mathbf{p}_1, \ldots, \mathbf{p}_N$，其中每个向量 $\mathbf{p}_i$ 是其位置 $i$ 的正弦函数。位置编码直接加到输入序列向量上：

$$ \mathbf{\hat{X}} = \mathbf{\hat{x}}_1, \ldots, \mathbf{\hat{x}}_N = \mathbf{x}_1 + \mathbf{p}_1, \ldots, \mathbf{x}_N + \mathbf{p}_N $$

从而提示模型更好地学习句子顺序。另一种路线（如 BERT 的 Devlin et al.）使用学习式位置嵌入，即在训练中学习 $\mathbf{P}$。

正弦与学习式位置嵌入曾是编码句子顺序的主流方法，但发现了两个问题：

二者都是绝对位置嵌入，即为每个位置 id（$0, \ldots, N$）编码唯一嵌入。研究表明绝对位置嵌入导致 LLM 在长文本输入上表现差——对长文本，让模型学习 token 之间的相对位置距离比绝对位置更有优势；
使用学习式位置嵌入时，LLM 必须在固定输入长度 $N$ 上训练，难以外推到比训练长度更长的输入。

近期，能解决上述问题的相对位置嵌入更受青睐，其中最突出的是：

旋转位置嵌入（RoPE, Rotary Position Embedding）
ALiBi

RoPE 与 ALiBi 的共同观点是：句子顺序信息最好直接注入自注意力算法——因为 token 正是在那里建立相互关系的；具体而言，应通过修改 $\mathbf{QK}^T$ 的计算来注入顺序信息。

RoPE指出：位置信息可以编码进 query-key 向量对（如 $\mathbf{q}_i$ 与 $\mathbf{k}_j$）中，把每个向量分别旋转角度 $\theta \cdot i$ 与 $\theta \cdot j$（$i, j$ 是各向量的句子位置）：

$$ \mathbf{\hat{q}}_i^T \mathbf{\hat{k}}_j = \mathbf{q}i^T \mathbf{R}{\theta, i - j} \mathbf{k}_j $$

其中 $\mathbf{R}_{\theta, i - j}$ 是旋转矩阵。$\theta$不在训练中学习，而是设为一个依赖于训练时最大输入序列长度的预定义值。

这样处理后，$\mathbf{q}_i$ 与 $\mathbf{k}_j$ 之间的概率得分仅在 $i \ne j$ 时受影响，且只取决于相对距离 $i - j$，与两个向量的具体位置 $i$、$j$ 无关。

RoPE 被当今许多重要 LLM 采用，如 Falcon、Llama、PaLM。仓库中对应实现集中在 src/transformers/modeling_rope_utils.py（RoPE 频率计算与配置解析），各模型（如 Qwen、Llama 系列）的建模代码通过它构造旋转缓存。

ALiBi则提出更简单的相对位置编码方案：把输入 token 之间的相对距离作为"负整数 × 预定义缩放值m"，在计算 softmax 前直接加到 $\mathbf{QK}^T$ 矩阵的每个 query-key 条目上。ALiBi 论文表明，这种简单的相对位置编码让模型在超长文本输入序列上仍保持高性能。ALiBi 同样被 MPT、BLOOM 等重要模型采用。

外推能力对比：RoPE 与 ALiBi 都能外推到训练中未见过的输入长度，但开箱即用的外推效果 ALiBi 明显优于 RoPE：

对 ALiBi，只需把下三角位置矩阵的数值加大以匹配输入序列长度即可；
对 RoPE，沿用训练时的 $\theta$ 会在远长于训练长度的输入上得到糟糕结果；不过社区发现了一些调整 $\theta$ 的有效技巧（如缩放 RoPE 基频），让 RoPE 位置编码也能在extrapolated 输入上正常工作。

RoPE 与 ALiBi 都是不在训练中学习的相对位置编码，基于如下直觉：
关于文本输入的位置线索应直接给到自注意力层的 $\mathbf{QK}^T$ 矩阵；
应促使 LLM 学习 token 之间恒定的相对距离位置编码；
token 彼此距离越远，query-key 概率应越低。RoPE 通过增大 query-key 向量夹角来降低向量点积；ALiBi 通过给点积加大负数来降低。

结论：面向需要处理大文本输入任务的 LLM，最好用 RoPE 或 ALiBi 这类相对位置编码训练。注意即使模型只在固定长度（如 $N_1 = 2048$）上训练，借助位置编码外推，实际可以用远大于 $N_1$ 的输入（如 $N_2 = 8192$）。

4.2 KV Cache：自回归生成的加速器及其代价

LLM 的自回归文本生成方式是：反复输入序列 → 采样下一个 token → 把该 token 追加到输入序列 → 直到模型生成"结束"标记 token。

先用一段简码展示自回归生成的朴素形态（用torch.argmax取最可能的下一 token）：

input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda") for _ in range(5): next_logits = model(input_ids)["logits"][:, -1:] next_token_id = torch.argmax(next_logits, dim=-1) input_ids = torch.cat([input_ids, next_token_id], dim=-1) print("shape of input_ids", input_ids.shape) generated_text = tokenizer.batch_decode(input_ids[:, -5:]) generated_text

输出：

shape of input_ids torch.Size([1, 21]) shape of input_ids torch.Size([1, 22]) shape of input_ids torch.Size([1, 23]) shape of input_ids torch.Size([1, 24]) shape of input_ids torch.Size([1, 25]) [' Here is a Python function']

可以看到，每一步都把刚采样的 token 追加进输入序列。

除极少数例外，LLM 都用因果语言建模目标训练，因此注意力结果矩阵的上三角被掩码——token从不依赖后续 token：$\mathbf{q}_i$ 永远不会与 $j > i$ 的 $\mathbf{k}j, \mathbf{v}j$ 建立关系，只关注此前的 $\mathbf{k}{m < i}, \mathbf{v}{m < i}$。为了减少不必要的计算，可以缓存每一层在所有前序时间步的 key-value 向量。

在 Transformers 中，向forward传入use_cache标志即可取回 KV cache，并在下一步连同当前 token 一起传回：

past_key_values = None # past_key_values 就是 KV cache generated_tokens = [] next_token_id = tokenizer(prompt, return_tensors="pt")["input_ids"].to("cuda") for _ in range(5): next_logits, past_key_values = model(next_token_id, past_key_values=past_key_values, use_cache=True).to_tuple() next_logits = next_logits[:, -1:] next_token_id = torch.argmax(next_logits, dim=-1) print("shape of input_ids", next_token_id.shape) print("length of key-value cache", past_key_values.get_seq_length()) # past_key_values 形状为 [num_layers, 0 for k, 1 for v, batch_size, length, hidden_dim] generated_tokens.append(next_token_id.item()) generated_text = tokenizer.batch_decode(generated_tokens) generated_text

输出：

shape of input_ids torch.Size([1, 1]) length of key-value cache 20 shape of input_ids torch.Size([1, 1]) length of key-value cache 21 shape of input_ids torch.Size([1, 1]) length of key-value cache 22 shape of input_ids torch.Size([1, 1]) length of key-value cache 23 shape of input_ids torch.Size([1, 1]) length of key-value cache 24 [' Here', ' is', ' a', ' Python', ' function']

可以看到：使用 KV cache 时，输入 token 序列不再变长，始终是一个输入向量；而 KV cache 的长度在每步解码时增加 1。

使用 KV cache 意味着 $\mathbf{QK}^T$ 本质上退化成了 $\mathbf{q}_c\mathbf{K}^T$，其中 $\mathbf{q}_c$ 是当前传入 token 的 query 投影，永远只是单个向量。

KV cache 带来两大优势：

计算效率显著提升：相比计算完整 $\mathbf{QK}^T$ 矩阵，计算量更少，推理速度更快；
峰值内存不再随生成 token 数二次增长，而是线性增长。

应始终使用 KV cache：它结果相同，且对长输入序列有显著加速。Transformers 在使用 text pipeline 或generate方法时默认启用 KV cache。仓库中 src/transformers/cache_utils.py 里的DynamicCache是生成模型的默认缓存实现，逐层存储 key/value 张量，形状为[batch_size, num_heads, seq_len, head_dim]；若传入 config，还会根据滑动窗口/混合层结构把缓存张量压缩到[batch_size, num_heads, min(seq_len, sliding_window), head_dim]，进一步降低内存需求。更完整的缓存体系说明见仓库文档 KV Cache 指南。

⚠️ 注意：尽管建议使用 KV cache，启用它之后 LLM 输出可能与未启用时略有不同。这是矩阵乘法内核（kernel）本身的数值特性所致，属于已知行为。

4.2.1 多轮对话：KV cache 复用历史上下文

KV cache 对聊天这类需要多次自回归解码的应用尤其有用。看一个例子：

User: How many people live in France? Assistant: Roughly 75 million people live in France User: And how many are in Germany? Assistant: Germany has ca. 81 million inhabitants

在这段对话中，LLM 运行了两次自回归解码：

第一轮：KV cache 为空，输入 prompt 是"User: How many people live in France?"，模型自回归生成"Roughly 75 million people live in France"，每个解码步都在增长 KV cache；
第二轮：输入 prompt 本应是"User: How many people live in France? \n Assistant: Roughly 75 million people live in France \n User: And how many in Germany?"。得益于 cache，前两句话的 key-value 向量已经计算过，所以实际只需处理"User: And how many in Germany?"。处理这段缩短的 prompt 时，其计算出的 key-value 向量直接拼接进第一轮解码留下的 cache，随后基于"完整三句话"的 cache 自回归生成第二个回答"Germany has ca. 81 million inhabitants"。

有两点值得注意：

保留完整上下文对聊天 LLM 至关重要，否则模型无法理解对话前文。例如上例中，模型需要理解用户问 "And how many are in Germany" 时指的是人口；
KV cache 对聊天极有价值，它允许持续"生长"已编码的聊天历史，而不必从头重新编码（若用 encoder-decoder 结构就必须从头重编）。

在transformers中，当传入return_dict_in_generate=True（叠加默认的use_cache=True）时，generate调用会返回past_key_values。注意该能力尚不通过pipeline接口暴露：

# Generation as usual prompt = system_prompt + "Question: Please write a function in Python that transforms bytes to Giga bytes.\n\nAnswer: Here" model_inputs = tokenizer(prompt, return_tensors='pt') generation_output = model.generate(**model_inputs, max_new_tokens=60, return_dict_in_generate=True) decoded_output = tokenizer.batch_decode(generation_output.sequences)[0] # Piping the returned `past_key_values` to speed up the next conversation round prompt = decoded_output + "\nQuestion: How can I modify the function above to return Mega bytes instead?\n\nAnswer: Here" model_inputs = tokenizer(prompt, return_tensors='pt') generation_output = model.generate( **model_inputs, past_key_values=generation_output.past_key_values, max_new_tokens=60, return_dict_in_generate=True ) tokenizer.batch_decode(generation_output.sequences)[0][len(prompt):]

输出：

is a modified version of the function that returns Mega bytes instead. def bytes_to_megabytes(bytes): return bytes / 1024 / 1024 Answer: The function takes a number of bytes as input and returns the number of

很好，不再为注意力层重复计算相同的 key 与 value。但要注意一个代价：$\mathbf{QK}^T$ 矩阵的峰值内存需求被大幅降低了，把 KV cache 留在内存里对长输入序列或多轮聊天可能非常吃内存——cache 需要为所有自注意力层、所有注意力头，保存所有前序输入向量 $\mathbf{x}_i$（$i \in {1, \ldots, c-1}$）的 key-value 向量。

以bigcode/octocoder为例，计算 KV cache 需要存储的浮点数个数：2（k 和 v）× 序列长度 × 注意力头数 × 每头维度 × 层数。在假想的 16000 输入序列长度下：

config = model.config 2 * 16_000 * config.n_layer * config.n_head * config.n_embd // config.n_head

输出：

7864320000

约80 亿个浮点数！以float16存储这 80 亿个浮点数需要约 15 GB 内存——接近模型权重本身的一半！研究者提出了两种显著降低 KV cache 内存成本的方法，即下文两小节。

4.2.2 多查询注意力（MQA）

Noam Shazeer 在Fast Transformer Decoding: One Write-Head is All You Need中提出MQA。如标题所示：与其使用n_head组 key-value 投影权重，只需用一组跨所有注意力头共享的 value 投影权重，模型性能并不会显著退化。

使用单组 value 投影权重后，key-value 向量 $\mathbf{k}_i, \mathbf{v}_i$ 在所有注意力头中必然相同，因此 cache 中只需存 1 组 key-value 投影，而不是n_head组。

由于大多数 LLM 使用 20 到 100 个注意力头，MQA 能大幅降低 KV cache 的内存消耗。对本指南使用的 LLM，可以把 16000 序列长度下的 cache 内存需求从 15 GB 降到不到 400 MB。

除了省内存，MQA 还提升计算效率。自回归解码中，每步都需要把巨大的 key-value 向量重新读回、与当前 key-value 向量拼接，再送入 $\mathbf{q}_c\mathbf{K}^T$ 计算——持续重载带来的内存带宽需求可能成为严重的时间瓶颈。减小 key-value 向量体积即减少内存访问，从而缓解带宽瓶颈。

要理解的关键点：把 key-value 注意力头数降到 1，只有在配合 KV cache 使用时才有意义。不使用 KV cache 的单个前向传播峰值内存不变——每个注意力头仍有唯一 query 向量，每个头仍有不同的 $\mathbf{QK}^T$ 矩阵。

MQA 已被社区广泛采用，众多热门 LLM 都在使用，包括 Falcon、PaLM、MPT、BLOOM；本指南使用的bigcode/octocoder检查点同样采用 MQA。

4.2.3 分组查询注意力（GQA）

谷歌的 Ainslie et al. 提出的GQA发现：相比使用常规的多个 key-value 头投影，MQA 有时会带来质量退化。该工作认为可以通过"更温和地"减少 query 头投影权重数量来保住更多模型性能：不用 1 组 key-value 投影权重，而是用n < n_head组。把n选为远小于n_head的值（如 2、4、8），既能保住 MQA 绝大部分内存与速度收益，又只牺牲更少的模型容量，因此性能损失更小。

此外，GQA 作者发现：已存在的模型检查点可以用不到原始预训练 5% 的计算量uptrain成 GQA 结构。虽然 5% 的预训练计算量本身可能仍是天文数字，但 GQA uptraining 让既有检查点可以服务于更长的输入序列。

GQA 提出较晚，撰写本文档时采用率相对更低；其最著名的应用是 Llama-v2。

总结性建议：如果 LLM 以自回归解码方式部署、且需要处理大输入序列（如聊天场景），强烈建议使用 GQA 或 MQA。

5. 结语：把"最好的算法与架构"用足

研究社区不断为越来越大的 LLM 提出加速推理的新方法。例如很有前景的**推测解码（speculative decoding）**方向：由更小更快的语言模型生成"简单 token"，只有"困难 token"才交给大 LLM 生成。

GPT3/4、Llama-2-70b、PaLM 这类巨型 LLM 之所以能在 Hugging Face Chat、ChatGPT 等聊天界面中高速运行，很大程度上正得益于本文讨论的**精度（量化）、算法（Flash Attention）、架构（相对位置编码 + KV Cache + MQA/GQA）**三个层面的优化。未来 GPU、TPU 等加速器只会更快、内存更大，但无论如何，都应确保用上了当前可用的最佳算法与架构，把每一份算力花在刀刃上。

快速决策参考：

场景	首选手段	预期收益
显存装不下模型	8-bit / 4-bit 量化（`BitsAndBytesConfig`）	本例 29 GB → 15.2 GB（8-bit）→ 9.5 GB（4-bit）
多卡加载超大模型	`device_map="auto"`流水线并行 / 张量并行	突破单卡 80 GB 上限
长上下文（≥16k token）	Flash Attention	注意力显存从 $O(N^2)$ 降到 $O(N)$
长文本任务（摘要/检索/聊天）	RoPE 或 ALiBi 相对位置编码	支持外推到训练长度之外的输入
自回归解码 + 多轮聊天	始终启用 KV cache，考虑 MQA/GQA 架构	解码步数计算量下降；16k 序列 cache 从约 15 GB 降到 <400 MB（MQA）