推理工程的目标,是让生成式 AI 模型跑得更快、更便宜、更可靠,同时尽量不牺牲使其有价值的输出质量。要同时提升性能并守住质量,你需要对模型底层机制建立清晰直觉。
生成式 AI 模型由大型而复杂的神经网络构成。神经网络的历史可追溯到 20 世纪 50 年代,当时第一个用于简单二分类的感知机已经在硬件上实现。此后几十年里,感知机一度被放弃,后来又被重新发展为多层感知机,并引入了反向传播等关键概念。反向传播让隐藏状态和反复调整权重的学习过程成为可能。早期神经网络通常只有几层。到 21 世纪初,研究开始转向拥有数十层的深度神经网络。2012 年,AlexNet 证明了 GPU 对深度学习的实际价值,也催生了词嵌入模型、生成对抗网络(GAN)等新架构。
真正的转折点出现在 2017 年,Vaswani 等人发表《Attention Is All You Need》,提出了 Transformer。Transformer 是一种带有注意力机制的神经网络,能够学习序列中不同部分之间的关系。Transformer 是生成式 AI 的基石。它不只用于大语言模型,也为嵌入、语音、图像和视频生成等多种模态提供基础。
在各种模态中,有两种重要的基于 Transformer 的模型风格:
-
自回归令牌生成:从标记化的序列开始,预测最可能的下一个令牌。
-
迭代去噪:从随机噪声开始,通过扩散过程向最可能的输出进行细化。
本章探讨了大语言模型(自回归标记生成)和图像生成模型(迭代去噪)的架构细节。
2.1 神经网络
几代人对神经网络的研究构成了生成式人工智能的理论基础。要成为一名高效的推理工程师,你需要对神经网络中的基本概念有直观的理解。本节提供了一个高层级的介绍;第 8 章提供了进一步阅读的建议。
神经网络的基本单元是节点(又称神经元)。节点是一个小程序,它接收一个输入,将其乘以一些权重,加上一些偏置,然后返回结果。一组节点构成一个层。层内的节点彼此独立——它们各自进行计算。节点之间的连接,或者说神经网络中的“网络”,是在层与层之间进行的,其中一层中的节点接收上一层的输出。
大语言模型(LLM)背后的神经网络包含几十到几百层。共有三种类型的层:
-
输入层:第一层,负责接收并处理输入到神经网络的数据。
-
隐藏层:第一层和最后一层之间的每一层,它们迭代地转换输入以得出输出。
-
输出层:最后一层,负责返回网络的预测结果。
每一层产生的输出都会被下一层作为输入读取。对于隐藏层,这些输出被称为隐藏状态。
隐藏状态是神经网络内部数据表示的一种。内部表示的一个关键方面是其维度,即所用向量的实际大小。
文本输入的内部表示增加了维度,将文本块编码为数百或数千个数字的向量,以捕捉语义含义。但图像模型的内部表示将维度从数百万像素降低到可管理的尺寸。有用于创建这些内部表示的神经网络,也有用于使用它们的神经网络:
-
编码器:接收文本或图像等输入,并创建包含额外信息和语义含义的输入内部表示。
-
解码器:使用内部表示来生成文本或图像等输出。
神经网络是可组合的。您可以将多个神经网络组合成一个单一模型,或者按顺序使用它们来构建流水线。
现代大语言模型(LLM)大多采用仅解码器(decoder-only)架构,而仅编码器(encoder-only)模型在今天已较为罕见,BERT 系列的传统文本嵌入模型就是一个典型的例子。许多其他模态的模型则使用编码器-解码器(encoder-decoder)架构。Whisper 是一款流行的开源音频转录模型,它使用编码器处理音频输入,并使用解码器生成文本标记。
2.1.1 线性层与矩阵乘法
神经网络中最本质的操作是矩阵乘法(matmul)。矩阵乘法接收一个输入向量(数字列表)和一个矩阵(数字网格),并将向量与矩阵相乘,从而产生一个输出向量。在神经网络中,线性层是最简单的矩阵乘法形式。给定一个输入向量,线性层会应用一个权重矩阵并加上一个偏置向量:
任何给定线性层的权重都只是生成式 AI 模型总权重的一小部分,权重矩阵中的各个值是在训练过程中设定的。
2.1.2 激活函数
矩阵乘法是可组合的,这意味着将一个向量乘以两个矩阵,等同于将该向量乘以这两个矩阵的乘积。
这对多层神经网络来说是一个问题,因为一系列线性层(每层都是一个矩阵乘法)会合并成一个单一的层,所有矩阵相乘在一起。
深度多层神经网络之所以有价值,是因为更多层能够更有效地利用更多参数,并在隐藏状态中编码更多含义。神经网络通过引入激活函数来打破线性,从而把各层区分开来。激活函数必须是非线性的,否则可组合的矩阵乘法最终仍会把多层网络压缩成一层;同时它又必须可微,或至少在大多数区域上近似可微,以支持反向传播。
推理里最常见的激活函数之一是 ReLU,也就是修正线性单元。它非常简单:如果 x 大于 0,就返回 x;否则返回 0。激活函数的变体很多,例如因形状像耐克标志而得名的 Swish,但大多遵循同一模式:压低或截断负值,同时保留正值。
ReLU、SiLU、Swish 和 SwiGLU 等激活函数运行速度快,易于训练,并且能够有效打破线性,从而支撑多层神经网络。
2.2 LLM 推理机制
LLM 是自回归的 token 生成模型:它会基于前文,一次生成一个新的 token。这些 token 是语言模型处理文本的基本单位,也就是用来表示文本片段的数字。现代 LLM 普遍采用子词分词,因此一个 token 可能对应一个完整单词,也可能只是单词的一部分。
把文本转成 token,以及再把 token 还原成文本,并不需要任何神经网络。分词器本质上只是字符串与其数值 token 表示之间的映射。
语言模型的词表,则给出了 token 与字符串之间的完整对应关系。不同模型的词表和分词器并不相同;较新的模型通常采用更高效的分词方案,因为生成同样内容所需的 token 越少,端到端推理就越快。大多数模型的词表都包含超过 100,000 个 token。在其他模态中,例如语音合成,词表还会被扩展,以便用 token 表示音频波形等其他信息。
推理涉及两到三段 token 序列:
- 输入序列:传入 LLM 的提示词、聊天内容、上下文、函数以及其他输入。
- 推理序列:可选,对于推理模型而言,这是用于思考的中间输出序列。
- 输出序列:由 LLM 生成的响应。
这些序列合在一起,受模型上下文窗口限制,也就是模型在一次请求中能够处理和生成的 token 总数。你还可以用 max_tokens 参数,进一步限制输出序列的长度。
虽然最终输入会被拼成单个字符串,但 LLM 在训练时其实支持多种输入形式,例如带角色的多轮聊天、用于工具调用的函数签名,以及某些场景下的多模态输入。把这些输入组织成一个统一序列的工作,由聊天模板完成;不同模型在细节上略有差异,因此推理引擎必须正确实现。
将套用聊天模板后的输入序列分词,是推理的第零步。然后,推理主要有两个阶段:
-
预填充:处理输入序列,为每个输入 token 计算注意力,并将这些值存储在 KV 缓存中。
-
解码:通过模型执行前向传播,以自回归方式生成 token。
解码阶段的每次前向传播都必须生成一个 token。这需要额外几个步骤,因为神经网络输出的是向量,而不是 token。
LLM 解码中使用的神经网络输出层会生成一个 logits 向量。该向量的长度等于模型的词汇表大小。归一化后,这些 logits 表示词汇表中每个潜在 token 作为正确输出的概率。
输出 token 是依据归一化后的概率分布采样得到的。你可以通过推理参数对这一过程做细调:
-
Temperature:在归一化之前调整 logits 本身。
-
Top-k:归一化后选择最有可能的 k 个令牌,然后在它们之间重新归一化。
-
Top-p:归一化后选择概率之和达到 p 的最小令牌集合。
较低的 temperature、top-k 或 top-p,会让输出更可预测,因为模型只能在更高概率的 token 中进行选择。把 temperature 设为 0,或把 top-k 设为 1,都可以让 token 选择变成确定性过程。
在生成结构化输出时,例如要求输出符合 JSON 模式,还可以使用 logit bias 等额外工具进一步约束模型。这些工具会在每次前向传播之后生效,也支撑着工具调用等重要能力;实现得是否正确,对推理质量影响很大。
整个“生成 logits,再选择 token”的循环会一直持续到模型输出停止 token 为止,除非更早触达上下文窗口或 max_tokens 的限制。
因此,推理真正最耗时的两部分,是预填充阶段构建 KV 缓存,以及解码阶段持续生成输出 token。更准确地说,后者是在生成最终会被选为 token 的 logits 向量。由于这两步都依赖大型神经网络,它们占据了推理的大部分时间和资源。
2.2.1 大语言模型架构
Hugging Face(最大的开源模型仓库)上的每个大语言模型都包含一个 config.json 文件:该文件由数十行内容组成,详细描述了模型的架构。
模型架构是在训练过程中,针对模型各个组成部分的性质和形态所做的一系列决策的集合。在同一架构下,可能存在以下情况:
-
多种规模:不同参数量的模型,例如 Llama 8B 和 70B。
-
多种变体:某一模型的”基础”版本和”指令微调”版本共享相同的架构。
-
无限次微调:LoRA(低秩适配)等方法改变的是模型行为,而非架构。
这些架构之所以重要,是因为它们决定了运行时和引擎的支持。如果你针对某一特定架构进行了高度优化的部署,那么你可以部署同一架构的另一个变体,并享受相同的性能提升。
模型架构是大多数配置文件的第一行内容之一。以解析架构名称 Qwen3MoeForCausalLM 为例:
- Qwen:模型系列,或模型的品牌名称。
- 3:该系列中架构的主要版本号。
- MoE:表示混合专家模型(第 2.2.4 节)。
- CausalLM:表示因果语言模型。
因果语言模型根据前面的词元来预测序列中的下一个词元,与之相对的是掩码语言模型——后者根据左右两侧的上下文词元来填空。当今所有的生成式大语言模型都是因果语言模型。
除架构名称之外,config.json 文件还包含有关构成模型底层神经网络的各层的性质与维度,以及推理过程中流经这些层的向量的相关信息。
2.2.2 Transformer 块
LLM 的主体是由数十到数百个 Transformer 块组成的序列。这些块构成了一个大型神经网络的核心,包含三种层:
-
嵌入层:神经网络的输入层,接收词元并返回嵌入向量。
-
Transformer 块:网络中的隐藏层是用于生成预测的 Transformer 块。
-
输出层:也称为语言模型头(LMHead),将 Transformer 块输出的隐藏状态转换为 logits 向量,模型词汇表中的每个词元对应一个 logit 值。
在 Transformer 模块中,包含用于注意力机制、前馈神经网络和归一化的子层。
前馈神经网络是一种多层感知机。这些线性子层构成了大型语言模型(LLM)中可训练权重的主要部分,而注意力子层则是第二大组成部分。其他组件(如归一化和激活函数)在模型规模中所占的比例微乎其微,几乎可以忽略不计。
尽管线性子层占据了权重的最大份额,但在推理过程中,更复杂的操作是注意力机制。
2.2.3 注意力机制
注意力机制是 Transformer 用来把当前词元与序列中其他词元联系起来的方法。人类能够自然理解词与词之间的关系,而注意力机制让大语言模型具备了类似能力。
以句子 “I decided to write a book because I thought it would be easy, but it was actually hard.” 为例,注意力机制能够判断其中的 “it” 指的是“写一本书”这件事。
注意力的标准形式是缩放点积注意力,如以下公式所示。
注意力机制接受三种输入:
-
Q(查询):正在生成或更新的标记的嵌入表示。
-
K(键):所有先前token的表示。
-
V(值):所有先前token的注意力计算值。
模型中的注意力子层是多头的,其中每个头对应一次注意力运算。如果你将神经网络的架构可视化,这些头在同一子层上是相互并行的。每个头可以负责关注标记之间不同类型的关系,例如主谓一致和共指消解。
注意力机制主要有两种类型:
-
自注意力:Q、K 和 V 均来自同一序列。
-
交叉注意力:Q 来自与 K 和 V 不同的序列,将 Q 条件化于外部信息之上。
大型语言模型使用自注意力(配合因果掩码以防止前瞻序列),而图像生成和多模态模型还会使用交叉注意力(例如,在生成的图像与关联的文本提示之间)。
从定义上看,注意力需要比较当前词元与序列中所有先前词元之间的关系,因此它对序列长度呈二次复杂度。上下文越长,计算就越慢。
但在实际推理里,借助 KV 缓存,注意力实现通常表现为线性开销。KV 缓存会存储先前词元的键和值,从而避免每次都重新计算。
KV 缓存在大型语言模型的预填充阶段构建,在解码阶段被读取和更新,默认存放在 GPU 内存中。它的存储、访问和复用,是推理工程中的重要主题,第 5.3 节会详细展开。
2.2.4 混合专家模型
神经网络的密度由层间连接数决定。更密集的网络能保留更多信息,而更稀疏的网络在运行时所需的计算量和内存更少。
混合专家(MoE)是一种在线性层中引入稀疏性的架构优化方法。MoE 模型不再使用单一的大矩阵,而是拥有数百个较小的矩阵,也就是“专家”,并把每个输入路由到少量专家上处理。这一过程称为激活专家。
以 Qwen3-235B-A22B 为例,该模型共有 2350 亿参数,但每次请求只会激活其中的 220 亿参数。由于活跃参数数量更少,MoE 模型在单请求、本地推理场景里非常高效。
但在生产服务器上做批量推理时,情况会复杂得多。不同请求会激活不同专家;除非在大规模专家并行(第 5.4.2 节)中真正实现稀疏性,否则在任意时刻,几乎所有模型参数都可能处于活跃状态。
专家路由本身也是细粒度的。模型每次前向传播都会逐层计算并生成一个 token。路由器是 LLM 内部的一个小模型,负责在每一层决定该激活哪些专家。以 Qwen 为例,模型共有 128 个专家,路由器会在 94 层中的每一层,为每个生成的 token 选择 8 个专家。
MoE 架构在参数量超过 1000 亿的大模型中尤为常见,但也存在参数量仅为 200 亿到 300 亿的小型 MoE 模型。混合专家架构还带来了一种新的推理并行方式,称为专家并行,它能够在多 GPU 上实现大模型的高吞吐量推理。
参数量低于 320 亿、尤其是低于 80 亿的模型,往往更适合传统的密集架构。针对代码补全等任务的领域专用模型,也通常很难从 MoE 中获得太多收益,因为对它们来说,整个模型几乎就相当于一个专家。
2.3 图像生成推理机制
图像生成模型接收文本提示,并据此生成图像。这类模型的公开走红其实略早于 LLM;Midjourney 的闭源模型和开源 Stable Diffusion 都在 2022 年夏天首次发布。
图像生成模型并非像 LLM 那样的单体模型,而是由多个模型协同工作的流水线,共同完成图像生成任务。在图像生成的基础模型中,有三个核心组件:
-
文本编码器:将文本提示转换为图像生成模型可以理解的指令。
-
去噪模型:模型的核心,根据提示从噪声迭代生成图像。
-
变分自编码器(VAE):将模型输出从潜在空间转换为像素空间。
这种流水线理念贯穿于整个图像模型生态系统。除基础模型外,图像推理通常还包括:
-
LoRA:轻量级微调,用于改变风格和提升质量。
-
ControlNet:通过轮廓和边缘引导输出图像,使其匹配大致的形状和颜色。
围绕图像生成模型的开源生态系统庞大而丰富,包括 ComfyUI 等工具,可用于构建复杂的基础模型与适配器流水线,通过替换组件来生成独特的输出。
整个图像生成流水线在潜在空间中运行。一张普通的高清图像可能为 1024×1024 像素,总像素数超过一百万。由于去噪模型需要并行计算整张图像的注意力,在像素空间中进行操作是不可行的。
潜在空间是图像的低维表示。图像的潜在空间矩阵可能为 128×128,约为其所代表的像素空间总数值的百分之一。
潜在空间被初始化为随机值,也就是噪声。去噪模型再根据文本提示,通过一系列步骤把噪声逐步精炼成图像。每一步都会更新整个潜在空间,这与大语言模型逐个处理 token 的方式不同。大多数图像生成模型需要 30 到 50 个步骤,才能生成高质量图像。
在每个步骤里,模型都会执行两次前向传播:一次带条件,也就是带文本提示;一次不带条件。随后再根据引导比例把两者合并。也就是说,一个 50 步的生成过程,实际需要 100 次前向传播。
这一过程以及图像生成的其他关键部分,通过推理参数在逐请求的基础上进行控制。最重要的参数包括:
- 提示词:描述图像应呈现的内容。
- 负面提示词:单独描述图像中不应出现的风格或对象。
- 步骤数:通过去噪步骤的数量来权衡速度与质量,大多数模型为30到50步。
- 引导比例:控制创意与提示词遵循度之间的平衡,通常为整数值,约为4。
- 图像尺寸:从固定的分辨率和宽高比菜单中选择输出图像的规格。
尽管这些核心机制在各类图像生成模型中普遍存在,但其架构在过去几年中已发生了相当大的演变。
2.3.1 图像生成模型架构
图像生成模型建立在 Transformer 之上,更具体地说,是扩散 Transformer。它与 LLM 使用的 Transformer 很相似,只是处理的不是离散 token 的嵌入,而是图像数据。
扩散 Transformer 以图像块为单位处理图像。在训练文本到图像模型时,训练图像会被切分为 2×2 或 4×4 的重叠图像块,并嵌入到潜在空间中。推理时则反过来:图像生成完成后,再把潜在空间还原为像素。
图像生成系统通常是多模型流水线,包括文本编码器、去噪模型和变分自编码器。Stable Diffusion XL(SDXL)是这类流水线的经典示例。尽管它已经不是新模型,但其架构至今仍很有参考价值。
SDXL的流水线包含两个用于去噪的扩散模型:基础模型和精炼模型。这两个模型针对不同任务进行训练:基础模型负责从纯噪声生成连贯图像,而精炼模型则负责添加细节并确保与提示词的一致性。
现代模型在提示词遵循、脸和手的细节准确性、文字渲染,以及图像到图像推理支持等方面,都明显超过了 SDXL。像 Qwen Image 这样的新模型,整体流程与 SDXL 相近,但几乎每个组件都更大、更强。
| 组件 | SDXL(2023 年) | Qwen Image(2025 年) |
|---|---|---|
| 文本编码器 | 基于 CLIP 的模型 | Qwen 2.5 VL(7B) |
| 去噪器 | 参数量不足 4B | 参数量 20B |
| VAE | 单重编码 | 双重编码 |
现代图像生成模型(如 Qwen Image)能力的大幅跃升,主要来自更大的组件模型和更复杂的流水线。可读文字渲染和更逼真的人脸,部分就得益于用完整 LLM 替换了微型 NLP 文本编码器,并把去噪器规模提升到原来的数倍。
这些更大的模型也更吃资源。虽然 GPU 性能在不断增长,但推理工程师不能只靠硬件进步来高效运行它们。
图像生成模型最新的研究方向,是将扩散 Transformer 架构与 LLM 架构相融合。
任何能被 token 化的内容,都可以用 LLM 建模。LLM 天然具备更强的文本理解能力,而且本来就是图像生成流水线中的关键组件。
LLM 也解决了扩散模型的一些固有问题。扩散模型只能生成固定尺寸的输出;LLM 则是自回归的,可以生成可变长度结果。扩散模型往往要经历多达 100 次前向传播才能生成一张图像,而 LLM 每次前向传播只生成一个 token。HunyuanImage-3.0 等模型,就在沿着这种 LLM 风格的方向演进。
尽管这仍是图像生成的前沿方向,但与此同时,业界也在继续探索其他现有架构,用于加速出图和把图像模型扩展到视频生成。
2.3.2 少步骤图像生成模型
图像生成中最耗时的部分,是 30 到 50 个去噪步骤。与其让每个步骤变得更快,是否有一种方法可以通过简单地减少步骤数来优化图像模型?
少步骤图像生成模型正是为此而训练的:以 8 个或更少的去噪步骤生成高分辨率图像。这些模型开箱即用就比传统图像生成模型快 80% 到 90%,但其输出质量也明显较低。
创建这些模型有两种主要方法:
-
潜在一致性:训练模型直接预测目标潜在图像向量,并重复预测两到四次以提升质量。
-
蒸馏:使用对抗蒸馏和/或渐进蒸馏,训练一个小模型在更少的推理步骤中模拟更大的模型。
如今,蒸馏比潜在一致性更为常见。当 FLUX 和 Qwen Image 等新图像模型发布时,开放图像模型社区的成员除了创建面向质量和风格的 LoRA 之外,还会制作蒸馏版本。
如果你的使用场景对延迟敏感、而对质量要求较低,例如实时生成式滤镜,就应优先考虑少步骤图像生成模型。
2.3.3 视频生成
视频生成模型在架构上与图像生成模型类似,只是规模更大。它们通常拥有 3 到 5 倍的参数量,并在潜在空间里编码 10 到 100 倍的信息。
视频生成的朴素方法是逐帧处理。早期的视频生成采用这种逐帧方式:先生成起始帧,再利用该帧生成下一帧,以此类推。
逐帧方法的问题在于误差累积。早期出现的小问题会随着每一帧不断叠加,导致视频很快就偏离正轨。
相比之下,现代视频模型将整个视频保存在潜在空间中,并在每个去噪步骤中对其进行修改。每一帧都与其他所有帧相互关注,并在每次前向传播中得到更新。
如果图像模型的潜在空间表示 X 和 Y 两个物理维度,那么视频模型的潜在空间则表示三个维度:X、Y 和 T(时间)。
这种方法的主要局限在于视频帧数是固定的,就像图像生成模型具有固定宽高比一样。现代视频生成模型生成的序列长度仅有几秒钟。
视频长度仍然受限于算力。即便使用最新 GPU,在大规模潜在空间上计算注意力也极其昂贵;生成一秒视频,往往就需要数秒推理时间。算力需求高到什么程度?视频模型通常以批量大小 1 运行,也就是说,一个拥有 8 块 GPU 的完整节点往往只能同时处理一个请求。
尽管每个去噪步骤上的注意力计算代价高昂,但视频生成模型的总步骤数与图像模型相同,通常约为 50 步。
与大语言模型和图像生成相比,视频生成是一种较为新兴的模态。其局限性与两年前大语言模型的局限性如出一辙。
| 大语言模型(2023 年末) | 视频生成模型(2025 年末) |
|---|---|
| 首 token 延迟高,生成速度低 | 生成速度慢 |
| 幻觉频发 | 物理效果不真实 |
| Ampere GPU 满负荷运转 | Blackwell GPU 满负荷运转 |
| 上下文窗口有限 | 视频输出时长短 |
如今,这些限制在大语言模型里大多已经被缓解。如何把它们从电影级视频生成中继续消除,以及如何在世界模型、3D 物体生成(XYZ 维,而非 XYT 维)和其他生成式渲染模型中推进,同样是当前非常活跃的研究方向。
其中一个重要方向,是重新引入某种自回归生成思路,类似于把大语言模型架构融入图像生成。像 Self Forcing 这样的新方法,并不是回到误差会迅速累积、几乎无法实用的纯逐帧生成,而是把全局质量控制与迭代式生成结合起来。
在视频模型中加入自回归成分,可以在一定程度上缓解注意力瓶颈。不过,注意力仍然是视频推理里最昂贵、也最关键的部分。
2.4 计算推理瓶颈
在一个完全优化的系统中,每种资源在任何时候都被充分利用。在 GPU 中,主要有两种资源:
-
计算能力:GPU 每秒能够执行的浮点运算次数。
-
内存带宽:GPU 每秒能够传输的字节数。
理想情况下,计算单元不会因为等待内存中的数据而空闲,内存带宽也不会因为等待计算完成而处于闲置状态。
然而在现实世界中,系统存在瓶颈:即一种资源空闲而另一种资源饱和的失衡状态。发现这些瓶颈是提升性能的第一步。如果某个操作的瓶颈在于内存带宽,那么再多的计算优化也无法加快系统速度,反之亦然。
在大多数情况下,推理系统存在以下瓶颈:
-
LLM 预填充(KV 缓存构建)受计算能力限制。
-
LLM 解码(Token 生成)受内存带宽限制。
-
图像和视频生成受计算能力限制。
在优化每个阶段的性能时,目标是降低瓶颈对整体系统性能的制约。例如,将多个请求合并批处理可以使 LLM 解码对内存带宽的依赖降低,因为处理一批请求在相同内存流量的情况下会使用更多的计算资源。
2.4.1 算子:字节比率与算术强度
每块 GPU 都有特定的计算速度(以每秒操作次数衡量)和内存带宽(以每秒千兆字节或太字节衡量)。通过比较这两个指标,可以确定特定 GPU 的算子:字节比率。
例如,H100 GPU 在 FP16 精度下可以执行 989 teraFLOPS 的密集计算,对应的内存带宽为 3.35 TB/s。由此得出算子:字节比率约为 295。
为了在 H100 GPU 上实现 FP16 推理的完美平衡(万物皆应如此),推理系统每访问一个字节的内存,就需要执行 295 次浮点运算。
要计算出这一比率,需要计算算法的算术强度。算术强度,也称为操作强度,是当前计算任务中工作量与内存访问量之间的比率。
算子:字节比率是在每秒的时间尺度上衡量的,而算术强度则是在单个函数或算法的执行过程中进行衡量的。
算术强度通常通过屋顶线模型进行可视化,该模型将性能与带宽上限(斜线)和性能上限(水平线)进行对比绘制。
通过与屋顶线模型对比,可以判断算法属于以下哪种情况:
-
计算瓶颈:当算术强度高于硬件的算子:字节比率并触及水平性能上限时,即为计算瓶颈。
-
内存瓶颈:当算术强度低于硬件的算子:字节比率并触及斜线带宽上限时,即为内存瓶颈。
要找到瓶颈所在,需要分析系统中开销最大的计算的算术强度。对于推理而言,注意力机制就是这样一种计算。
2.4.2 LLM 推理瓶颈
LLM 推理分为两个阶段:
-
预填充(Prefill):决定首个令牌的生成时间(TTFT),受计算能力限制。
-
解码(Decode):决定每秒生成的令牌数(TPS),受内存带宽限制。
对于每个阶段,可以通过比较最重要操作的算术强度与可用硬件的运算量/字节比(ops:byte ratio)来证明瓶颈的存在。
在预填充和解码两个阶段中,最耗资源的操作都是注意力机制(attention)。注意力机制的精确算术强度取决于模型架构(维度、注意力头数等)、输入序列长度以及注意力算法的实现方式。
两者的本质区别在于:预填充阶段并行处理整个输入序列,而解码阶段则逐个生成令牌。
在预填充阶段,模型权重仅加载一次,随后在输入矩阵与注意力矩阵之间进行一系列大规模矩阵乘法运算。相比单次内存读取,这涉及大量计算,因此算术强度较高。
在解码阶段,每生成一个令牌都需要加载一次模型权重,而生成过程使用的是计算量更小的向量-矩阵乘法。在这种情况下,相对于加载整个模型权重所需的数据量,浮点运算次数相对较少,因此算术强度较低。
以下是一个计算精确算术强度的示例:考虑一个注意力头维度为128(d=128)的模型,在长度为4096个令牌(N=4096)的序列上执行解码步骤。本分析使用标准注意力算法,不进行任何优化。
标准注意力实现可以概括为三步:
- 从 HBM 中按块加载 Q 和 K,计算
S = QK^T,再将 S 写回 HBM。 - 从 HBM 中读取 S,计算
P = softmax(S),再将 P 写回 HBM。 - 从 HBM 中按块加载 P 和 V,计算
O = PV,再将 O 写回 HBM。
根据本练习的参数,确定各矩阵的大小:
-
N:序列长度,设定为4096。
-
d:注意力头的维度,设定为128。
-
Q、K、V:大小为N×d,即4096×128。
-
S、P:计算得大小为N×N,即4096×4096。
-
O:计算得大小为N×d,即4096×128。
假设使用FP16推理,矩阵中每个值占两个字节。作为参考,一个4096×4096的矩阵大约为32 MiB,大致相当于一张高分辨率RAW格式单反照片的数据量。
注意力算法的三行代码遵循相同的模式:从内存加载数据、执行计算、将结果存储到内存。
将这三步分别拆成“读取、计算、写回”三个维度后,可以得到:
- 第一步:读取
Q = 2Nd和K = 2Nd;计算S = (2d) * (N^2);写回S = 2N^2。 - 第二步:读取
S = 2N^2;计算P = 3(N^2);写回P = 2N^2。 - 第三步:读取
P = 2N^2和V = 2Nd;计算O = (2N) * (Nd);写回O = 2Nd。
计算总内存移动量时,将第一列和第三列求和,分别追踪从GPU内存的读取和写入:
将三步中的读写量相加后,总内存搬运量为 8N^2 + 8Nd 字节。
计算总计算量时,将第二列求和:
将三步中的计算量相加后,总计算量为 4N^2d + 3N^2 次浮点运算。
计算算术强度时,将工作量(总计算量)与内存流量进行比较:
因此,算术强度等于总计算量除以总内存搬运量;在本例中,其结果约为 62 ops/byte。
在本示例中,算术强度为 62,远低于 H100 GPU 的运算量/字节比 295。具体数值因模型、序列长度和硬件的不同而有所差异,但此示例说明了一个普遍规律:解码阶段受内存带宽限制。
像这样计算算术强度是一种学术练习,并非推理工程师的日常任务。但亲自推导一遍有助于建立直觉认知。
2.4.3 图像生成推理瓶颈
图像和视频模型相对较小——其参数量仅为前沿语言模型的十分之一——但其注意力机制的计算开销同样巨大。
图像和视频生成模型采用迭代去噪方式,而非自回归式的令牌生成。
正如大语言模型预填充阶段的注意力机制需要一次性处理整个输入序列,图像和视频生成时的注意力机制也必须考虑以潜在空间表示的完整图像或视频对象。
与大语言模型预填充阶段类似,图像和视频生成模型的推理同样受计算资源瓶颈制约。针对这些模态优化推理的具体技术将在第 6.5 节和第 6.6 节中介绍。
2.5 优化注意力机制
对于大型语言模型(LLM)而言,注意力机制的计算复杂度随输入序列长度呈二次方增长。每次注意力计算都依赖于每个前驱词元的 K 值和 V 值。在实践中,由于 KV 缓存存储了前驱词元的键值计算结果,注意力在解码阶段是一个线性时间操作。
即便是线性扩展的算法也会变得非常昂贵。注意力机制是各类模型和架构中推理开销最大的部分之一,因此优化注意力机制自然成为一个重要且高度活跃的研究领域。
注意力是一个敏感过程,因为每个词元都依赖于之前的每一个词元。注意力中的细微误差会迅速累积,使得注意力优化成为一项精细的工作。
图 2.14 展示了注意力算法本身是简单直观的。然而,这种基础实现效率低下——中间矩阵 S 和 P 在某一步骤结束时被存储,随即在下一步骤中被加载。
优化注意力机制有两种策略:
-
实现改进:编写性能更高的内核,更高效地利用内存和算力。
-
新算法:创建在优于二次方时间内运行的注意力算法,同时将质量损失降至最低。
实现改进仍受限于注意力的二次方时间复杂度,但这类方法是无损的(不影响质量),并使当今硬件上处理长序列的推理成为可能。其他算法方法以质量换取时间和空间复杂度的降低,但训练技术可以将影响降至最低。
最著名的注意力实现是 FlashAttention 系列论文和内核。基础算法可以用寥寥数行代码实现,而 FlashAttention 则使用数万行代码,在专为特定 GPU 构建的手工融合内核中实现注意力——用于 H100 的 FlashAttention 与用于 B200 的 FlashAttention 代码不同。
FlashAttention 通过消除多余的内存读写操作,并将注意力算法的布局精确匹配 GPU 的能力来发挥作用。FlashAttention 对于计算密集型操作(如 LLM 预填充和视频生成)尤为有用。
另一个重要实现是 PagedAttention。KV 缓存增长迅速,会占满 GPU 内存并消耗大量读取时间。PagedAttention 将 KV 缓存划分为可通过查找表访问的块(页),这意味着 KV 缓存可以以碎片化内存的方式分散存储在 GPU 上,而无需占用单块连续内存。
尽管 FlashAttention 和 PagedAttention 是有价值的优化,但它们并没有改变注意力是二次方算法这一事实。注意力的新变体改善了底层的时间和空间复杂度:
-
滑动窗口注意力:对前 w 个词元的滑动窗口计算注意力,将注意力复杂度从 O(N²) 降至 O(Nw),其中 w 通常在 8K 到 32K 的范围内。
-
门控注意力:训练中引入的各类层允许以相对于块长度线性的时间近似计算特定上下文块的注意力。
-
线性注意力:用线性时间算法替换二次方的 softmax 方程,以近似计算注意力。
-
压缩注意力:定期压缩序列早期的上下文,注意力同时考虑压缩后的上下文和未压缩的近期词元。
-
多潜变量注意力:在低维潜在空间中近似计算注意力。
直觉上,序列中相邻词元之间的相互影响大于距离较远的词元,这是合理的。我现在写的这句话与上一句密切相关,但与本章开头的句子关联则少得多。
这种直觉可以通过训练加以延伸。滑动窗口注意力等算法在训练中应用时,所创建的模型在推理中使用相同技术时能够保持较高质量。
另一个研究方向是通过使用不同于 Transformer 的架构来完全避免注意力机制。Mamba 是一种选择性状态空间模型,用循环状态更新替代自注意力,实现了对序列长度的线性扩展。混合模型有时将 Mamba 风格的状态空间模型块与 Transformer 块混合使用。状态空间模型的应用仍然有限,但随着 NVIDIA Nemotron 3 Nano 等开放模型采用混合架构,混合模型正变得越来越流行。