大型语言模型推理详解

最新推荐文章于 2024-10-12 11:02:31 发布

小马不会过河

最新推荐文章于 2024-10-12 11:02:31 发布

阅读量1.1k

点赞数 21

文章标签：语言模型人工智能自然语言处理开源数据库

本文链接：https://blog.csdn.net/m0_59163425/article/details/141504439

版权

一、引言

本文将深入探索大型语言模型（LLM）推理的各个方面及其面临的挑战。主要聚焦于基于仅解码器Transformer模型的token生成过程，因为这类模型在推理阶段存在着独特的挑战和优化方案。同时，本文介绍的许多概念和方法也同样适用于Transformer编码器模型的推理实践。

本文章的主要内容：

LLM推理中常用的术语，例如键值（KV）缓存、内存带宽限制等；
推理优化的多种策略和配置，包括量化、内核融合、模型架构调整、批量大小选择、GPU选型等；
关键性能指标（延迟、吞吐量、成本）与推理优化的关系。

目标是帮助读者建立一个清晰的思维模型，以便能够在配置和优化LLM服务解决方案时做出明智的决策，并在将LLM部署到实际应用中时少走弯路。

内容概览

首先，解析Transformer解码器在token生成过程中的两大环节：提示处理与多个自回归（autoregressive）步骤，这两者的硬件利用效率存在显著差异。在后续章节中，这一区别将被反复提及。

随后，解读LLM推理的首个挑战：注意力机制在总序列长度上的计算需求平方级增长。为此，将探讨一种关键的缓存策略：KV缓存（KV cache），并分析它如何有效解决这一难题。KV缓存至关重要，因为它为自回归步骤提供核心输入。然而，KV缓存也存在自身问题，带来一系列问题。将深入分析这些问题及其应对策略。

在掌握KV缓存的要点后，将深入剖析Transformer推理过程中的硬件利用情况。此时，将引入两个核心概念：算术强度（arithmetic intensity）与屋顶线模型（roofline model），并将它们与硬件的关键特性（如峰值FLOPS、内存带宽）及性能指标（延迟、吞吐量、成本）紧密相连。随后，将运用这些知识于Transformer推理实践，揭示如何更高效地利用硬件资源并提升性能指标。此阶段所获知识将为我们理解各种主要性能优化策略背后的逻辑打下坚实基础。

模型量化（Quantization）作为近年来提升性能的重要策略之一，将在本系列中得到专门探讨。虽然量化本身足以构成一个完整的研究领域，但我们将以本文为基础，提供基础知识，并解读量化算法的优势与局限。

最后，将探讨现代大型语言模型（LLM）服务框架的运作机制。仅仅优化模型本身并不足以在LLM推理中达到最佳性能。模型服务器在高效管理请求与硬件资源方面发挥着至关重要的作用，对于确保端到端的最佳性能至关重要。

Reference

[1]: Attention Is All You Need (Vaswani et al., 2017)

二、LLM响应背后的双阶段流程

重新回顾Transformer架构以及基于其解码器部分的文本生成核心原理。在此，先界定本文将要使用的专业术语，并特别标注本文采用的表述方式。接下来，将深入剖析文本生成过程中的两个关键阶段：初始化阶段（initiation phase）和生成（或称解码）阶段（generation (or decoding) phase）。

首先，简要回顾一下Transformer。为了简化说明，假设每次处理一个序列（即批量大小为1）。下图展示了基于Transformer解码器的基本架构概览（图1），它用于从一系列输入token中生成输出token。

在这里插入图片描述
图1.Transformer解码器模型概览

请注意，解码器本身并不直接输出token，而是输出与词汇表大小相对应的logits（对数几率）。顺便一提，输出logits的最后一层常被称作“语言模型头”（LM头，language model head）。从logits中提取token的任务则交给了一种称为（token）“搜索策略”、“生成策略”或“解码策略”的启发式方法，这些方法包括：

贪心解码：简单地选择logits值最大的token，有时在选择前会对logits进行如重复惩罚等转换。
采样解码：利用logits作为多项分布进行采样，即根据采样结果从词汇表中选择一个token。这之前，可以通过温度缩放、top-k和top-p等方法调整采样分布。
更复杂的启发式方法，如束搜索、对比解码等。

为了简化说明，假设解码策略是模型操作的一部分（图2）。在LLM服务解决方案的语境下，这种接收token序列作为输入并返回相应输出token的实体，常被称作执行引擎或推理引擎，这一心智模型在实际应用中极为有用。

图2.简化的Transformer解码器模型示意图

那么，如何生成多个token呢？使用基于Transformer的解码器从输入文本序列（常称为提示）生成文本（常称为补全）主要包括以下步骤：

1. 将模型权重加载到GPU上。

2. 在CPU上对提示进行Tokenizing处理，并将token张量传输到GPU（图3）。

图3.Tokenization步骤示意图

3. 通过网络运行tokenized的提示，生成补全的第一个token。

这一单步骤阶段通常被称作初始化阶段，在后文中，会看到它有时也被称为“预填充阶段”（pre-fill phase）。

4. 将生成的token添加到输入token序列中，并以其为新的输入来生成补全的第二个token。然后，重复此过程，直至生成停止序列（如单个序列结束（EOS）token）或达到配置的最大序列长度（图4）。

这一多步骤阶段则被称为生成阶段、解码阶段、自回归阶段或增量阶段。

步骤3和4的流程图如下所示（图4）。

在这里插入图片描述
图4.Token生成过程的初始化与解码阶段

5. 将生成的补全token传回CPU并进行反tokenize处理，以获取最终的生成文本（图5）。

在这里插入图片描述
图5.反tokenize步骤示意图

注意：近期的一些前沿技术，如推测采样或前瞻解码，可能不完全遵循上述简单算法，旨在实现更低的延迟。

此时，你或许会感到不解或困惑。可能会问：初始化阶段和解码阶段之间的实际差异何在？从表面上看，这种区分似乎有些牵强。初始化阶段确实类似于while循环的初始化步骤，而在这两个阶段中，本质上都在做同一件事：在每次迭代中，对一个token序列执行前向传递，并逐次增加一个token。

实际上，你的观察是正确的。就硬件上的计算执行而言，这两个阶段在本质上是相同的，并无特别之处。然而，正如将在后文中看到的，这种设置涉及大量计算（计算量随总序列长度呈二次方增长），其中许多计算实际上是重复的。为缓解这一问题，采用了缓存策略来避免不必要的重复计算，这种优化方法被称为KV缓存，而这也是关键差异所在。

Reference

[1]: A Contrastive Framework for Neural Text Generation (Su et al., 2022)

[2]: Fast Inference from Transformers via Speculative Decoding (Leviathan et al., 2022)

[3]: Breaking the Sequential Dependency of LLM Inference Using Lookahead Decoding (Fu et al. 2023)

三、深入解析KV缓存

在前文中，我们对Transformer解码器在文本生成过程中的两个阶段进行了概要性介绍：一个是单步的启动阶段（用于处理初始提示），另一个是多步的生成阶段（逐个生成补全的token）。

这里将深入探讨LLM推理所面临的首个重大挑战：注意力层（特别是自注意力层）的计算需求随总序列长度（即提示token与生成的补全token之和）的平方而急剧增长。幸运的是，这些计算中有很多是重复的，可以通过缓存关键结果来显著降低计算负担。这种缓存技术实际上将原本二次方增长的注意力层计算转变为线性增长。

Transformer注意力层简要回顾

首先，回顾一下基础Transformer架构中的多头注意力（MHA，multi-head attention）层是如何工作的（见图1）。

图1.Transformer解码器层（上）及其中两个自注意力层的详细展示，以输入序列长度为3为例

为简化说明，我们假设处理的是单个长度为t的序列（即批量大小为1）：

在整个处理过程中，输入序列（即提示）中的每个token都被转换为一个高维向量（图1中的浅黄色部分）。
注意力层的输入是由前一个解码器块输出的一系列高维向量，每个输入token对应一个。
对于每个输入向量，注意力层会生成一个等维度的输出向量（图1中的浅蓝色部分）。

具体到单个注意力头：

我们使用三种不同的线性变换（投影）为每个输入向量生成三个低维向量：查询（query）、键（key）和值（value）（图1最左侧浅灰色向量）。这样，我们就得到了t个查询向量、t个键向量和t个值向量。
对于每个查询向量，我们计算其与所有键向量的点积，得到注意力分数，并据此对值向量进行加权平均，从而生成该查询的输出向量。这一过程为序列中的每个token创建了一个包含其他token信息的上下文表示。
但在自回归解码中，我们不能使用序列中尚未生成的token的信息。因此，通过掩码技术，我们将后续token的值向量对应的注意力分数置为零。

最后，将各注意力头的输出拼接起来，并通过一个线性层得到最终输出。

注意力计算的二次扩展

让我们看看计算注意力分数所需的浮点运算数（FLOPs）。对于一个给定的注意力头，在一个大小为b的批量中，每个序列的总长度为t（包括提示和生成的补全），注意力分数矩阵是通过将形状为(t, d_head)的查询张量与形状为(d_head, t)的转置键张量相乘而创建的。

单次矩阵乘法需要多少FLOPs？将形状为(n, p)的矩阵与形状为(n, m)的矩阵相乘大约涉及2.m.n.p次操作。在我们的例子中，单头单序列的注意力分数计算大约需要2.d_head.t^2 FLOPs。总体而言，注意力分数计算需要2.b.n_layers.n_head.d_head.t^{2=2.b.n_layers.d_model.t}2 FLOPs。现在显然可以看出，计算量随着t的二次方增长。

看一些实际数据，例如在Meta的Llama2-7B模型中，n_layers=32且d_model=4096。

注意：带掩码的注意力分数矩阵与值张量相乘所需的FLOPs数量与上述计算相同。

那涉及模型权重的矩阵乘法呢？通过类似的分析，我们可以证明它们的计算复杂度是O(b.n_layers.d_model^2.t)，即计算需求随总序列长度t线性增长。

为了理解二次扩展的严重性，让我们来看一个例子。生成第1001个token时，模型必须执行比生成第101个token多100倍的FLOPs。这种计算量的指数增长显然很快就变得不可承受。幸运的是，由于掩码机制，在各个步骤之间实际上可以省去许多计算。

掩码在生成阶段引发的冗余计算

现在我们深入问题的核心。由于掩码机制的作用，在生成过程中，对于给定的token，其输出表示仅依赖于其前面所有token的表示。这意味着，在多次迭代中，只要前面的token不变，该特定token的输出表示也将保持不变，从而导致了大量的冗余计算。

以我们之前讨论过的token序列为例，假设我们已从输入序列“What color is the sky? The sky”中生成了“is”。在上一次迭代中，“sky”是输入序列的最后一个token，因此其输出表示是基于序列中所有token的表示生成的，包括“What”、“color”、“is”、“the”、“sky”、“?”、“The”和“sky”的值向量。

下一次迭代的输入序列将是“What color is the sky? The sky is ”，但由于掩码机制，从“sky ”的角度来看，输入序列似乎仍然是“What color is the sky? The sky ”。因此，为“sky ”生成的输出表示将与上一次迭代中的相同。

现在用图示例说明（图2），使用图1中的图。启动步骤应处理长度为1的输入序列。冗余计算的元素在浅红色和浅紫色中突出显示。浅紫色元素对应于冗余计算的键和值。

图2.生成阶段注意力层中的冗余计算

回到我们的例子，在输入序列变为“What color is the sky? The sky is ”的新迭代中，唯一需要新计算的表示是序列中新增的最后一个token“is ”。

具体来说，我们需要哪些材料来完成这项任务？

一个针对“is”的查询向量。
键向量用于计算注意力分数，包括“What”、“color”、“is”、“the”、“sky”、“?”、“The ”、“sky ”和“is ”。
值向量用于计算输出，包括“What”、“color”、“is”、“the”、“sky”、“?”、“The ”、“sky ”和“is ”。

对于键和值向量，它们在之前的迭代中已经为除“is ”以外的所有token计算过。因此，我们可以保存（即缓存）并重用之前迭代中的键和值向量。这种优化称为“KV缓存”。然后，为“is ”计算输出表示就变得非常简单：

计算“is ”的查询、键和值。
从缓存中获取“What”、“color”、“is”、“the”、“sky”、“?”、“The ”、“sky ”和“is ”的键和值向量，并将它们与刚刚计算的“is ”的键和值连接起来。
使用“is ”的查询和所有的键来计算注意力分数。
使用注意力分数和所有的值来计算“is ”的输出向量。

看一下输入，只要可以使用它们的键和值向量，实际上不再需要之前的token。使用KV缓存时，模型的实际输入是最后生成的token（而不是整个序列）和KV缓存。图3提供了这种在生成阶段运行注意力层的新方法的示例。

在这里插入图片描述
图3.启用KV缓存的生成步骤

回到上一章节中的两个阶段：

启动阶段基本不受KV缓存策略的影响，因为它没有前置步骤。
但在解码阶段，情况发生了显著变化。现在，我们不再需要整个序列作为输入，而只需最后一个生成的token和KV缓存。

在注意力机制的运行过程中，注意力层会一次性地处理所有输入提示的token，而非在解码过程中逐个处理。在学术文献中[1]，这种一次性处理方式被称为批量注意力（有时可能误导性地被称为并行注意力），而逐个处理的方式则被称为增量注意力。

当利用KV缓存时，初始阶段会预先计算并填充所有输入token的键和值到KV缓存中，这一步骤常被称作预填充阶段（pre-fill phase）。在实际应用中，预填充和启动阶段这两个概念经常可以相互替代，但从现在开始，我们更倾向于使用“预填充阶段”这一术语。

这种启动和生成阶段之间的新区别不仅仅是概念上的。现在，在每个生成步骤中，除了权重矩阵外，还必须从内存中获取不断增长的缓存，并且每个序列只处理一个token。注意，使用针对每个阶段优化的GPU内核比同时使用同一个内核处理预填充和启用KV缓存的解码阶段能带来更好的性能（例如，参见[2]的研究）。

KV缓存实现线性注意力扩展

注意力如何扩展？转置后的键张量仍然是形状为(t, d_head)。然而，查询张量现在的形状是(d_head, 1)。单头单序列的注意力分数计算因此需要2.d_head.t FLOPs，整体而言，注意力计算需要2.b.n_layers.d_model.t FLOPs。注意力现在随总序列长度线性扩展！

我们是否解决了二次扩展的问题？例如，如果丢弃了缓存并需要重新计算它，答案是否定的。想象一下，你开发了一个聊天机器人应用[3]，并在每次对话轮次之间将缓存保留在内存中。现在，一个客户端已经闲置了很长时间。由于GPU内存有限，你实施了缓存驱逐策略，丢弃陈旧的对话。不幸的是，客户端恢复了，因此必须重新计算整个历史记录的缓存。此重新计算会产生与总对话长度二次方成比例的计算成本。

上述例子凸显了KV缓存作为一种折衷方案的本质：它不是免费的午餐。我们通过增加内存使用和数据传输的代价来换取计算量的减少。正如后续文章将详细探讨的，缓存的内存占用可能会相当可观。

回到聊天机器人示例，设计一个高效的缓存驱逐策略具有挑战性，因为它需要在两种昂贵的选择之间进行权衡：要么消耗更多稀缺的资源（GPU内存和带宽），要么需要大量的计算。

KV缓存在实际中的应用：以HuggingFace Transformers为例

在实际应用中，KV缓存是如何工作的呢？我们能否启用或禁用它？以HuggingFace Transformers库为例，所有专为文本生成设计的模型类（如XXXForCausalLM类）都实现了一个generate方法，用于生成文本。这个方法通过use_cache布尔参数来控制是否启用KV缓存（默认值为True）。

深入一层，查看模型的forward方法（例如，LlamaForCausalLM.forward文档），我们发现use_cache布尔参数如预期存在。在启用KV缓存的情况下，我们有两个输入：最后生成的token和通过参数input_ids和past_key_values分别传递的KV缓存。新的KV值（即包括当前迭代中计算的KV值）作为forward方法输出的一部分返回，用于下一次迭代。

这些返回的KV值是什么样子的？让我们来数一数张量。在启用KV缓存的情况下，forward方法返回一对张量列表（一个用于键，一个用于值）。模型中有多少个解码块（通常称为解码层，记作n_layers），就有多少对张量。对于批次中的每个序列的每个token，每个注意力头有一个维度为d_head的键/值向量，因此每个键/值张量的形状为(batch_size,seq_length,n_heads,d_head)。

以Meta的Llama2–7B模型为例，它有32个解码层和32个注意力头，每个头的维度为128。尽管我们将在后续文章中详细探讨KV缓存的具体大小，但现在我们已经可以对其可能达到的规模有一个初步的认识。

结论

让我们回顾一下关于LLM推理的知识。注意力分数的计算量会随着总序列长度的平方而增长。然而，由于注意力机制中的掩码，在每个生成步骤中，实际上只需要重新计算与最后生成的token相关的键和值，而其他过去token的键和值可以重复使用。我们可以将每次计算的新键和值缓存到GPU内存中，以便将来重用，从而节省计算量。

这种策略的主要优势在于，它使自注意力机制的计算量与总序列长度呈线性关系，而不是平方关系。

但是，正如上文所述，KV缓存也带来了一些新的问题，我们将在接下来的文章中探讨：

KV缓存会消耗GPU内存，并且缓存大小可能会变得非常大。而GPU内存通常是稀缺的，即使对于小型LLM也是如此。因此，KV缓存是提高模型处理能力（例如上下文窗口大小或吞吐量）和成本效率的主要障碍之一。

虽然KV缓存可以显著减少数据传输量，但与我们在单个生成步骤中执行的操作数量相比，其影响仍然有限。因为在执行矩阵向量运算时，我们需要加载大量的权重矩阵和不断增长的KV缓存。然而，在现代硬件上，数据加载的时间往往比实际计算时间更长，这导致了GPU计算资源的低效利用，并降低了成本效益。

Reference

[1]: See for example Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)

[2]: For example, since its release 2.2.0, the reference implementation of the widely adopted Flash-Attention algorithm features a dedicated kernel for the decoding phase when KV caching is enabled (flash_attn_with_kvcache) also referred to as Flash-Decoding.

[3]: Blog post — Scaling ChatGPT: Five Real-World Engineering Challenges (Orosz, 2024)

四、深度剖析KV缓存

在前述内容中，我们已介绍了KV缓存作为优化LLM推理过程的重要手段。其核心在于，通过存储（自）注意力机制中已计算完成的键（Key）和值（Value）张量至GPU内存中，实现了计算需求的线性增长，而非传统方法的二次方增长。

具体而言，KV缓存机制巧妙地利用了生成过程中的历史信息，避免了在每一步迭代中重复计算相同的Key和Value对，从而节省了大量计算资源。然而，这一策略也伴随着权衡：即以内存空间换取计算效率。本章将深入探讨KV缓存的潜力边界、所面临的挑战，以及行业内应对这些挑战的常见策略。

KV缓存的扩张极限

这非常直观：对于批处理中每个序列的每个token，需要为每个注意力层中的每个注意力头存储两个大小为d_head的向量张量（一个键张量和一个值张量）。每个张量参数所需的空间取决于精度：全精度（FP32）每参数4字节，半精度（BF16、FP16）每参数2字节，8位数据类型（INT8、FP8）每参数1字节等。

假设b为批处理大小，t为总序列长度（提示 + 完成），n_layers为解码器块/注意力层的数量，n_heads为每个注意力层的注意力头数量，d_head为注意力层的隐藏维度，p_a为精度。多头注意力（MHA）模型KV缓存的每token内存消耗（以字节计）为：

注意：我们提醒在MHA模型中，n_heads.d_head=d_model，但为简化上述公式我们不使用此处。

因此，KV缓存的总大小（以字节计）为：

KV缓存面临的第一个挑战之一是：它随着批处理大小和尤其是总序列长度呈线性增长。由于它与总序列长度呈现增长趋势，KV缓存的大小实际上没有上限，而我们的GPU内存显然是有限的。更糟糕的是，由于无法预先知道总序列长度，KV缓存的内存需求因此变得未知，使得内存管理尤为具有挑战性。

让我们看一些流行的MHA模型的数字（表1），例如Meta的Llama-2 [1]和OPT [2]，MosaicML的MPT [3]以及BigScience的BLOOM [4]：

在这里插入图片描述
表1.流行的多头注意力（MHA）模型规格

假设参数以半精度（FP16、BF16）存储，我们选择了一个较小的模型（Llama-2-7B）和一个较大的模型（BLOOM-176B）。对于Llama-2-7B（或BLOOM-176B），KV缓存内存消耗约为每token0.5MB（或每token4MB）。

对于Llama-2-7B，使用半精度加载模型权重消耗大约14GB内存，与缓存28k个token的键和值相同。例如，28k个token可以对应于一批56个长度为512的序列，这并不是特别极端的情况。

从上述数字可以看出，KV缓存的内存消耗可能会变得非常大，甚至超过加载大序列模型权重所需的内存量。

将这些数字与常用的NVIDIA数据中心GPU的内存容量进行比较（表2）：

在这里插入图片描述
表2.用于训练和/或为LLM提供服务的常用NVIDIA数据中心GPU规格

选择性价比较高的A10 GPU，以Llama-2-7B为例，计算最大的KV缓存容量。一旦加载了模型权重，剩余的10GB用于KV缓存，即总容量约为20k个token，包括提示在内，显然不能在处理或生成长序列时提供大量并发请求。

我们现已认识到，KV缓存机制成为了一个限制因素，它阻碍了我们处理或生成极长的序列（即构成处理长上下文窗口的障碍），同时也限制了处理大批量数据的能力，进而无法充分发挥硬件的效率。

在此视角下，最大化我们的处理能力意味着尽可能为KV缓存留出更多空间，可以通过以下方式实现：

减少模型权重的内存占用（权重量化）
减少KV缓存的内存占用（详见下文）
通过模型分片在多个GPU上汇集内存，但需要通过网络通信（模型并行化），或者使用其他存储介质如CPU内存或硬盘（卸载）

由于模型权重和不断增长的KV缓存必须在每次前向传递中加载，解码步骤涉及非常大的数据传输，正如接下来要解读的那样，实际上是受到内存带宽的限制，即我们实际上花费更多时间在数据传输而不是进行有用的计算。在这种情况下，延迟只能通过增加内存带宽（即更好的硬件）或减少数据传输来改善。较小的模型权重和KV缓存释放出内存以供更多序列使用，从而增加吞吐量（和/或最大序列长度）。

在这方面，减少内存占用的策略非常有用，因为它们不仅能够增加我们的硬件利用率和成本效益，还能减少延迟并增加吞吐量。

扩展知识：为什么输入token会被收费？（表3）

在这里插入图片描述
表3.OpenAI的样本费率（截至2024年12月1日）

到了这里，你应该能够理解为什么你会同时为输入和输出token付费了。一旦输入提示被处理，即在预填充阶段结束时，我们已经消耗了GPU内存（用于存储每个输入token的键和值张量）和计算资源（用于通过模型传递提示token）。

让我们看一些实际的数字。假设一个P参数模型的前向传递总FLOPs数约为2.P FLOPs/token[5]，使用Llama-2-7B处理一个提示消耗约0.5MB/token的GPU内存（见上文），以及约14 TFLOPs/token的GPU计算。对于一个1000个token的提示（略少于两页），这就是大约500MB的内存和14 TFLOPs的计算资源，而此时我们还没有生成任何内容。

现在让我们看看如何通过采用上述公式并逐个查看每个项来减少KV缓存的内存占用：

减少批处理大小？

在大多数情况下，我们不希望减少批处理大小，因为虽然它有助于减少KV缓存的内存占用和因此减少延迟，但会降低我们的硬件利用率，从而影响成本效益。在接下文中，我们确实会看到相反的情况，我们希望尽可能增加批处理大小。

减少对总序列长度的依赖？

不存储序列中所有token的键和值的一个原因是，我们可以选择在每次迭代中重新计算丢失的键和值，因为这样做值得花费FLOPS，而不是消耗GPU内存（例如，因为我们受到内存带宽的限制，在自回归阶段时会发生这种情况）。据我所知，在实践中并没有这样的情况，因此我们不会深入探讨这个方向。

另一个视角是，我们可以选择不存储模型几乎不关注或关注非常少的token的键和值。例如，对于只训练以关注序列部分（例如Mistral AI的Mistral-7B）或在内存消耗和模型精度之间达成折衷的模型，这可能是一种情况。让我解释一下。

像Mistral-7B [6]这样的模型被训练为不关注整个序列。Mistral-7B的注意力层确实通过仅关注最后的（4096）相邻token来构建token表示。这种注意力机制的变体称为滑动窗口注意力（SWA）或局部注意力。设计上，局部注意力确保在KV缓存中每个序列最多存储窗口大小（例如4096）的张量对。

另一种方法是利用注意力层在序列中分配注意力的模式。事实上已知，注意力模块在序列中不成比例地并且一致地向少数token分配更多的注意力（见图1）。相比之下，许多token对输出贡献非常少，因此根本不必费心存储它们的键和值。

图1.来自StreamingLLM论文的注意力（热力）图示例：大量注意力始终分配给第一个token和最后的相邻token（局部注意力）

通过丢弃这些token，我们事实上将相应的注意力分数设为零，并用一个更稀疏的矩阵来近似注意力矩阵。成功的近似将最小化近似误差，从而减少对模型精度（例如使用困惑度进行衡量）的影响。

让我们来看看过去几个月出现的一些方法，这些方法可以在不需要重新训练或微调的情况下直接应用：StreamingLLM框架、H2O（重要token预测器）、Scissorhands和FastGen。据我所知，然而，目前没有任何流行的LLM推理框架支持这些方法。

针对使用有限长度上下文窗口训练的模型，StreamingLLM框架[7]基于这样一个观察：初始token收集了大量注意力。因此，该框架通过仅保留最初的位置token（“汇聚token”）和最后的相邻token（局部注意力）来构建一个滑动窗口。StreamingLLM的KV缓存因此具有固定的长度，包括一个固定部分（通常为1到4个token）和一个滑动部分。

类似的H2O [8]和Scissorhands [9]方法明确旨在通过设置最大缓存token数（预算）并在达到缓存预算时丢弃token来压缩KV缓存。H2O算法每次只丢弃一个token，而Scissorhands则根据目标压缩比例（例如30%的KV缓存大小减少）丢弃尽可能多的token。

这两种方法都基于这样一个观察：在给定步骤中具有影响力的token（“关键token”或“重要token”）在未来步骤中仍然具有影响力（Scissorhands作者称之为重要性持久性假设）。换句话说，我们确信被丢弃的低影响token在未来步骤中仍然会相对被忽略，因此可以安全地丢弃。

这两种算法的关键方面显然是缓存逐出策略。Scissorhands简单地保留最近的token和在历史窗口内具有最高注意力分数的token。H2O丢弃累计注意力分数最低的token，因此仅保留在迭代中始终获得高注意力分数的token。两个团队的作者已经表明，他们的算法可以实现高达80%的KV缓存大小减少，并且几乎没有模型精度损失。

FastGen方法[10]（与无关的DeepSpeed-FastGen不要混淆）仍然基于注意力模式，但采取了另一种方法，即不设置缓存预算，而是设置注意力矩阵的最大近似误差，从而专注于保持模型精度。

FastGen是一个两步方法：首先，在预填充阶段结束时，对模型的注意力层进行剖析，以确定满足误差目标的压缩策略集。与其他方法类似，它假设在未来生成步骤中，识别出的注意力模式将保持不变。压缩策略包括：保留特殊token、保留标点符号token、保留最后的相邻token（局部注意力）等（见图2）。如果误差目标过于严格无法达到，FastGen则回退到常规的KV缓存。然后，在每个生成步骤中应用所选的压缩策略到KV缓存中。

图2.FastGen论文中压缩策略示例：特殊token（绿色）+ 标点符号token（橙色）+ 局部注意力（蓝色）。被丢弃的token为灰色。

需要注意的是，与其他方法相反，FastGen针对每个提示构建了定制的压缩策略。FastGen的作者们展示了，在给定的KV缓存压缩比例下，它们比H2O和Scissorhands更能保持模型的准确性。

无论如何，摆脱对不可预测的总序列长度的依赖是一种解脱，因为它允许为每个序列分配一个内存预算，从而极大地简化了内存管理。由于数据传输是延迟的主要贡献者，因此没有随序列长度线性增长的KV缓存，特别是对于较长的序列长度，可以带来显著的加速。

减少层数？

在这里并没有太多可得。较小的模型通常层数较少（表4），因此如果较小的模型在你的用例中表现良好，只需选择它即可。

表4.Llama-2模型规格

减少注意力头的数量？

对于给定的模型架构，模型大小主要由层数和注意力头的数量控制，减少注意力头的数量可能意味着选择一个较小的模型（参见表4）。

然而，如果我们仔细观察，我们会注意到我们只需要减少键和值头的数量，查询头的数量并不影响KV缓存的大小。这正是多查询注意力（MQA）[11]和分组查询注意力（GQA）[12]架构背后的核心思想。这些变体的多头注意力（MHA）的唯一动机是减少KV缓存的大小。

MQA首次在2019年引入。在MQA中，所有查询头共享相同的单个键和值头。换句话说，所有查询头使用相同的键计算它们的注意力分数，并且所有头的输出都使用相同的值（但注意力分数不同）（图3）。

图3.多头注意力（上）vs. 多查询注意力（下）（两个注意力头）

然而，对于较大的模型来说，去除所有头相对更为激进。例如，从64个头减少到1个头在模型表示能力上的削减比从32个头减少到1个头更为显著。GQA通过提供一种中间解决方案来解决这个问题：与其使所有查询头共享相同的唯一KV头，我们将它们分成g组查询头，每组的查询头共享相同的唯一KV头。换句话说，与其从n_heads缩减到1个KV头，KV头的数量从n_heads减少到1<g<n_heads。

从这个角度来看，MHA和MQA都是GQA的特例（g=1和g=n_heads，分别）。QGA允许在模型准确性和KV缓存大小（与延迟和吞吐量相关）之间更平滑地进行折中。

考虑到这个新参数g，KV缓存大小的公式变为：

在实践中，MQA/GQA架构已被谷歌研究的PaLM [13]、TII的Falcon [14]模型、Meta的Llama-2 [1]（仅70B）和Mistral AI的Mistral-7B [7]（表5）显著实施。

表5.使用MQA或GQA的模型系列

注意力头的隐藏维度？

再次强调，如果你不准备选择另一个模型的话，这里并没有太多可以得到的。根据模型系列的不同，注意力头的隐藏维度可能在各个模型大小中保持不变（例如Llama-2、Falcon），因此选择同一系列的较小变体并不会有所帮助。

使用更少字节的参数？

量化KV缓存确实是大幅减少其大小的好方法。然而，仅量化权重的算法如AWQ [15]或GPTQ [16]在定义上并不会有所帮助。只有那些同时量化权重和“激活”（即非权重部分）的算法，例如LLM.int8()[17]或SmoothQuant [18]才会产生量化的KV缓存。

需要注意的是，同时对权重和激活进行量化的算法的一个目的是在较低精度下执行计算密集的矩阵乘法。这在像训练期间的计算绑定时会提高性能，但正如我们将在接下来的文章中看到的那样，推理的自回归阶段实际上是内存带宽受限的，因此能够更快地计算并没有带来太多价值。由于推理受到内存带宽的限制，我们实际上只关心内存占用的减少，因为这意味着数据传输更少。

从这个角度来看，像LLM.int8()或SmoothQuant这样的量化算法可能有些过火：在将缓存的张量量化后移入GPU内存，并在从GPU内存中获取这些张量时解量化它们（需付出额外开销），应该已经足够了。

一些LLM推理系统已经包含了这种KV缓存量化功能。例如，FlexGen [19]使用4位数据格式量化并存储KV缓存和模型权重。NVIDIA TensorRT-LLM能够将KV缓存量化为8位数据格式（INT8或FP8）。流行的vLLM框架自版本0.3.0起就支持KV缓存（FP8）的量化。由于量化是在每次迭代中动态执行的，因此不需要校准步骤。

关于高效内存管理的重要性

到目前为止，我们隐含地假设内存中没有浪费：所有预留的内存都用于存储token，并且所有可用的内存都可以被预留。在实践中，简单粗暴的内存管理策略会导致内存的大部分被浪费（PagedAttention论文[20]指出，实际有效的内存利用率可能低至20%，即浪费了80%！）

由于请求的总序列长度事先不知道，我们可能会预留连续的内存块，足以容纳最大的序列长度。这种分配的显著部分可能永远不会被使用，并且由于不可用于其他请求，因此会被浪费（内存内碎片化）。
即使序列长度事先已知，由于内存逐步消耗而内存块被请求的生命周期预留，较短的请求也不能使用仍未使用的内存块。
如果使用像束搜索这样的解码策略，多个候选序列实际上可以部分共享它们的KV缓存。如果我们不考虑这种情况，我们将不可避免地浪费内存，存储了本可以共享的重复KV条目。

这些缺点正是现在流行的PagedAttention算法旨在解决的问题。PagedAttention分配固定大小且相对较小的内存块，称为块。每个块可以包含固定数量的token，并且在需要时可以跨不同请求共享。按需分配和小块大小减轻了内部内存碎片化，而相同大小的块消除了外部内存碎片化。

总体而言，PagedAttention实现了接近零的KV缓存内存浪费（低于4%[21]）。先前被浪费的内存现在可以用来容纳更多请求，从而增加吞吐量。PagedAttention推出时的吞吐量提升效果与当时的内存浪费水平一样显著。

PagedAttention最初由vLLM推理系统实现，现在受到所有主要推理框架的支持（例如HuggingFace TGI、NVIDIA TensorRT-LLM、LMDeploy TurboMind等）。

另一种可能的优化策略，PagedAttention没有涵盖的是跨请求重用键值缓存。这种情况适用于提示共享公共前缀的场景，这在多轮对话和代理人使用提示模板时常见（图4）。

图4.SGLang论文中的KV缓存共享示例（多轮对话），总共四个生成请求。蓝色框表示可共享的提示部分

能够跨请求重用KV缓存将显著改善延迟（特别是第一个token的延迟）和吞吐量（通过大幅减少具有共享前缀的并发请求的内存占用）。

这种KV缓存重用的例子是通过LMSYS SGLang论文中引入的RadixAttention技术实现的。

RadixAttention算法不是在完成生成请求后丢弃KV缓存，而是将其保留在GPU内存中，并将新条目添加到专用数据结构（基数树）中，将token序列映射到其KV缓存张量。当新请求到达时，调度程序使用基数树进行前缀匹配。如果有缓存命中，调度程序就会重用缓存的KV张量来满足请求。

由于GPU内存有限，缓存的KV张量不能永久保留。因此，RadixAttention算法包括一个淘汰策略（例如最近最少使用的淘汰策略）。最佳的缓存重用可能与先到先服务等调度安排不兼容。因此，RadixAttention配备了修改后的调度器，优先考虑与缓存前缀匹配的请求（缓存感知调度）。

注：PagedAttention和RadixAttention的命名有点误导，因为与人们可能认为的不同，它们并不是模型注意力层的优化（如FlashAttention），而是在模型服务器级别操作（它们帮助服务应用程序更好地管理主机上的KV缓存）。

如果GPU内存不足，为什么不“只是”使用多个GPU？或者将负载转移到CPU内存甚至硬盘上？

这是两种不同但有效的方法。

首先是将负载转移到容量更大但速度较慢的存储介质（CPU内存和硬盘）。并非所有推理框架都支持此功能，如HuggingFace Accelerate、DeepSpeed-Inference和更先进的FlexGen。由于涉及使用速度较慢的存储介质，负载转移会导致显著的延迟增加，因此对于延迟敏感的用例显然不应优先选择此选项。负载转移系统通常用于面向吞吐量的离线批处理等用例。

关于使用多个GPU（对于较大的模型无法避免），将模型分片到多个设备上可以通过利用聚合的内存容量和内存带宽来释放内存压力。

如果选择流水线并行性[23]，模型和KV缓存都将分片到层维度上。如果选择张量并行性[24]（在推理中更为常见），则KV缓存将分片到头维度上。需要注意的是，在这种设置中，MQA变得相当低效：由于无法将单个头分片到多个设备上，KV缓存必须在所有设备上复制，从而失去了MQA的好处。对于实现MQA的模型，另一种选择是将KV缓存分片到批处理大小维度上[25]。

无论哪种情况，以上所有情况都假设单个主机，我们仍然受限于最大多GPU实例的存储容量。据我所知，目前没有推理框架支持多主机模型并行。如果我们能够在多个主机上分片模型和KV缓存，那么可用内存量和能够处理的最大序列长度将变得几乎无限。这正是《Infinite-LLM》论文[26]旨在解决的问题，通过引入新的分布式注意力算法（DistAttention）并适应Ray框架来构建多主机分布式KV缓存管理和调度系统（DistKV-LLM）。

总结

我们了解到选择KV缓存会带来额外的挑战。多头注意力（MHA）模型的KV缓存确实会消耗大量GPU内存，大约每token约1MB，并且很容易比模型权重更大。

鉴于GPU内存有限，KV缓存内存压力促使各种倡议朝不同方向发展：新颖的注意力架构（如MQA、GQA、SWA）、缓存压缩策略（如H2O、Scissorhands、FastGen）、高效的内存管理（如PagedAttention、RadixAttention）、以及量化和存储容量扩展（如负载系统、单主机和多主机模型并行）。

正如下面的内容所述，减少KV缓存大小不仅因为GPU内存有限，而且因为数据移动量实际上是每个自回归步骤的主要延迟因素，因此是整个生成过程的关键。

Reference

[1]: Llama 2: Open Foundation and Fine-Tuned Chat Models (Touvron et al., 2023)

[2]: OPT: Open Pre-trained Transformer Language Models (Zhang et al., 2022)

[3]: Release blog posts for: MPT-7B (May 2023) and MPT-30B (June 2023)

[4]: BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (BigScience, 2023)

[5]: Scaling Laws for Neural Language Models (Kaplan et al., 2020)

[6]: Mistral 7B (Jiang et al., 2023)

[7]: Efficient Streaming Language Models with Attention Sinks (Xiao et al., 2023) + GitHub repository

[8]: H_2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (Zhang et al., 2023) + GitHub repository

[9]: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time (Liu et al. 2023)

[10]: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (Ge et al., 2023)

[11]: Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)

[12]: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)

[13]: PaLM: Scaling Language Modeling with Pathways (Chowdhery et al., 2022)

[14]: The Falcon Series of Open Language Models (Almazrouei et al., 2023)

[15]: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin et al., 2023) + GitHub repository

[16]: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar et al., 2022) + GitHub repository

[17]: LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers et al., 2022) + GitHub repository

[18]: SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao et al., 2022) + GitHub repository

[19]: FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (Sheng et al., 2023) + GitHub repository

[20] Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023) + GitHub repository

[21] vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention (Kwon et al. 2023)

[22] Efficiently Programming Large Language Models using SGLang (Zheng et al., 2023) + Blog post

[23]: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (Huang et al., 2018)

[24]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)

[25]: Efficiently Scaling Transformer Inference (Pope et al., 2022)

[26]: Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Lin et al., 2024)

五、解析模型性能

在上一章中，我们深入探讨了KV缓存的优化。现在我们将转变方向，探索可能影响机器学习模型速度的不同性能瓶颈。本章详细介绍的概念广泛适用于任何机器学习模型，无论是用于训练还是推理，但例子将特别聚焦于LLM推理设置。

在开始之前，我首先强烈推荐阅读这篇博文[1]，本文的很多内容都借鉴于此。

四种性能瓶颈类型

如果你对模型的性能不满意并且准备投入时间进行改进，第一步是确定瓶颈类型。性能瓶颈主要有四类，其中三类与硬件限制相关，一类与软件相关。

让我们从硬件瓶颈开始。每种瓶颈对应于特定的操作模式：

计算受限模式：大部分处理时间，即延迟，花费在执行算术运算上（图1）。与其他模式相比，由于你主要支付的是计算成本，计算受限模式是最经济的模式，因此我们应该朝这个方向努力。

在这里插入图片描述
图1.计算受限过程，计算和数据传输时间分别用黄色和蓝色突出显示

内存带宽受限：大部分处理时间花费在移动数据，如权重矩阵、中间计算等，介于芯片内存和处理器之间（图2）。

在这里插入图片描述
图2.内存带宽受限过程，计算和数据传输时间分别用黄色和蓝色突出显示

通信受限（未在 [1] 中涵盖）：仅当计算和数据分布在多个芯片上时适用。大部分处理时间花费在芯片之间的网络数据传输上（图3）。

图3.通信受限过程，计算、数据传输和网络通信时间分别用黄色、蓝色和绿色突出显示

注意：我使用“芯片”这个词，因为这些概念适用于任何类型的芯片：CPU、GPU、定制硅片（如Google TPU、AWS Neuron Cores等）等。

请注意，现代硬件和框架经过高度优化，计算和数据传输任务通常部分重叠（图4）。为简单起见，在本文中我们将假设它们按顺序发生。

图4.带有重叠数据传输的通信受限过程

最后一个模式称为开销受限模式，与软件引起的限制相关。在这种模式下，大部分处理时间用于调度工作并将其提交给硬件 —— 基本上，我们花费的时间更多用于决定要做什么，而不是在硬件上执行实际操作（图5）。当使用非常灵活的语言（例如Python）或框架（例如PyTorch）时，这种开销受限的情况更为常见，因为这些语言或框架不需要在运行时明确指定所有需要的信息（例如张量数据类型、目标设备、要调用的内核等）。这些缺失的信息必须在运行时推断，对应的CPU周期称为开销。现在加速硬件相对于CPU来说速度非常快，因此在开销影响硬件利用率和成本效率的情况下很可能出现 —— 基本上有时硬件会空闲下来，等待软件提交下一个工作项。

图5.带有计算、数据传输和软件开销时间的开销受限过程，分别用黄色、蓝色和紫色突出显示

执行模型的前向或后向传递涉及运行多个内核执行（例如GPU函数调用）。不同的内核可能不会在同一模式下操作。重要的是确定大部分执行时间花费在哪种模式下。然后，优化的重点将是针对主要瓶颈进行优化，然后是下一个最重要的瓶颈，依此类推。

准确识别瓶颈类型至关重要。每个问题需要不同的解决方案。如果你的诊断有误，可能会浪费时间实施一个甚至起不到帮助作用的优化。

诊断限制因素

这里不会深入细节，但是[1]指出，当处于开销受限模式时，运行时不会与增加的计算或数据传输成比例增加。换句话说，如果你将计算或数据传输容量加倍，但运行时未相应增加，你的程序很可能处于开销受限模式。否则，你就是硬件受限，但区分像计算与内存带宽之间的瓶颈需要访问诸如FLOP计数和数据传输量等指标，即使用分析器。

对于LLM，重要的是要认识到，在训练和推理的预填充阶段，计算资源往往是受限的；然而，在推理解码阶段，大多数硬件则更倾向于受到内存带宽的制约。针对训练过程所设计的优化策略，如采用低精度矩阵乘法等，可能对于减少主要由解码延迟所决定的整体推理延迟，效果并不显著。

基于瓶颈类型进行低延迟优化

让我们看看如何优化每种类型的瓶颈。

如果为计算受限模式：

升级到更强大且更贵的芯片，具有更高的峰值FLOPS。
对于特定操作，如矩阵乘法，利用专门的、更快的核心，例如NVIDIA Tensor Cores。例如，NVIDIA H100 PCIe [2] 使用通用CUDA核心的峰值计算为51 TFLOPS，而使用专用Tensor Cores时达到378 TFLOPS（全精度）。
减少所需操作的数量。在ML模型中，具体来说，这可能意味着使用更少的参数来实现相同的结果。剪枝或知识蒸馏等技术可以帮助实现这一点。
使用较低精度和更快的数据类型进行计算。例如，对于相同的NVIDIA H100 PCIe，8位Tensor Cores的峰值FLOPS（1,513 TFLOPS）是16位峰值FLOPS的两倍（756 TFLOPS），32位峰值FLOPS的四倍（378 TFLOPS）。但这需要对所有输入进行量化（例如权重矩阵和激活值，参见例如LLM.int8() [3] 或SmoothQuant [4] 量化算法）并使用专门的低精度内核。

如果为内存带宽受限模式，可以采取以下优化措施：

升级到更强大和更昂贵的芯片，具有更高的内存带宽。
通过模型压缩技术，如量化或不那么流行的剪枝和知识蒸馏，减少移动的数据量。对于LLMs来说，数据大小问题主要通过仅权重量化技术来解决（例如GTPQ [5] 和AWQ [6] 量化算法），再加上KV-cache量化。
减少内存操作的次数。在GPU上运行任务归结为执行一个内核的有向图（即GPU函数调用）。对于每个内核，需要从内存中获取输入并将输出写入其中。内核融合，即最初分散在多个内核中的操作作为单个内核调用执行，可以减少内存操作的数量。操作融合（图6）可以由编译器自动执行，也可以通过编写自己的内核手动执行（这更难但对于复杂的融合是必要的）。

在Transformer模型的情况下，为注意力层开发高效的融合内核仍然是一个活跃的领域。许多优化过的内核基于流行的FlashAttention算法 [7]。Transformer融合内核库包括FlashAttention、Meta的xFormers和现已废弃的NVIDIA FasterTransformer（已合并到NVIDIA TensorRT-LLM）。

在这里插入图片描述
图6.应用于CNN的水平和垂直层（操作）融合示例，初始状态（上）和最终状态（下）[8]

如果为通信受限模式，可以采取以下措施：

升级到更强大和更昂贵的芯片，具有更高的网络带宽。
通过选择更高效的分区和集体通信策略，减少通信量。例如，[9] 扩展了Transformer模型的流行张量并行布局 [10]，引入了新的张量并行策略，允许通信时间更好地扩展（即防止它们成为瓶颈）用于大芯片数量和/或批处理大小。

例如，[10] 中的张量并行布局保持权重片段静止，而激活片段在芯片之间移动。例如，在预填充阶段和对于非常大的序列批处理，[9] 指出激活可以超过权重。因此，从通信的角度来看，保持激活静止并移动权重片段，如他们的“权重收集”分区策略所示，变得更加有效。

如果为开销受限模式，可以采取以下措施：

通过使用C++等更高效但不那么灵活的语言，以换取更少的开销来减少开销。
将内核分组提交，以将提交开销分摊到多个内核而不是每个内核支付。当需要多次提交同一组短暂生命周期的内核时（例如在迭代工作负载中），这尤其有益。CUDA图（自PyTorch发布1.10版以来已集成）通过提供将代码块产生的所有GPU活动捕获为内核启动的有向图的工具来实现此目的。
提前提取计算图（AOT）为可部署的工件（模型追踪）。例如，PyTorch的torch.jit.trace追踪PyTorch程序并将其打包成可部署的TorchScript程序。

可以通过使用模型编译器进一步优化图。

无论如何，你再次以灵活性换取更少的开销，因为追踪/编译需要参数（如张量大小、类型等）在运行时保持静态和因此保持不变。控制流结构，如if-else，在此过程中通常也会丢失。

对于需要与AOT编译不兼容的灵活性的情况（例如动态张量形状、控制流等），JIT编译器可以在执行模型代码之前动态优化模型代码（尽管没有AOT编译器那么彻底）。例如，PyTorch 2.x具有名为TorchDynamo的JIT编译器。由于你不需要修改PyTorch程序即可使用它，因此你可以在保持Python开发体验的同时获得使用JIT编译器减少开销的好处。

附注：模型优化器和（AOT）编译器之间有什么区别吗？

在我看来，这两个术语的区分有些模糊。以下是概念上如何区分这两个术语。

首先，两者都是提前执行工作。典型的AOT编译器工作流程是：从支持的框架（PyTorch、TensorFlow、MXNet等）中跟踪代码，将计算图提取到中间表示（IR）中，应用硬件无关的优化（代数重写、循环展开、操作融合等），生成优化的图，并最终为目标硬件创建可部署的工件，包括选择最合适的内核、数据移动优化等。AOT模型编译器的示例包括PyTorch的TorchInductor、XLA和Meta的Glow。

模型优化器是一种工具，包括AOT编译，通常针对特定硬件（例如OpenVINO的英特尔硬件，TensorRT和TensorRT-LLM的NVIDIA硬件），并能够执行额外的后训练优化，如量化或剪枝。

到目前为止，我们只关注延迟（处理单个请求所需的时间），但让我们通过深入探讨计算和内存带宽受限模式，重新引入吞吐量（单位时间内可处理的请求数量）。

瓶颈 = f(硬件，算术强度)

有意思的是，处理相同输入的相同算法可以是计算受限或内存带宽受限的，这取决于所使用的硬件。适用的模式由算法的算术强度决定——即每访问的内存字节执行的算术操作数量。

我们希望强度值使我们处于或接近更具成本效益的计算受限模式。正如我们将看到的，更高的强度与更好的吞吐量和成本效率相关。然而，某些强度驱动因素可能会降低延迟。延迟和吞吐量之间的权衡几乎是不可避免的。

设b为每次运行传输到/从内存的数据字节数，p为每次运行执行的浮点运算数（FLOPs）。设BW_mem（以TB/s为单位）为硬件的内存带宽，BW_math（以TFLOPS为单位）为数学带宽，也称为峰值FLOPS。设t_mem为移动数据字节所花费的时间，t_math为执行算术操作所花费的时间。

当算术操作所花费的时间超过数据传输时间时，我们就处于计算受限状态（见图7）。

图7.计算与内存带宽受限模式。计算和数据传输时间分别用黄色和蓝色突出显示

因此，当以下条件成立时，我们处于计算受限状态：

其中A是算法的算术强度，其维度是每字节的FLOP。每传输的数据字节执行的算术操作越多，算术强度越高。

如所示的方程式，要使算法处于计算受限状态，其算术强度必须超过一个硬件相关的峰值FLOPS与内存带宽的比率。相反，如果算术强度低于同一带宽比率，就表示处于内存带宽受限状态（见图8）。

图8.内存带宽/计算受限的边界

让我们来看一些半精度矩阵乘法（即使用Tensor Cores）在NVIDIA硬件上的实际数据（表1）：

表1.NVIDIA数据中心GPU的规格，通常用于训练和/或服务LLMs

这意味着什么？以NVIDIA A10为例，带宽比率为208意味着在该特定硬件上移动一个字节的数据与执行208个FLOP的速度相同。因此，如果在NVIDIA A10上运行的算法每传输的字节不执行至少208个FLOP（或相当于每传输一个半精度数字执行416个FLOP），那么它不可避免地会花费更多时间在数据传输上而不是在计算上，即它是内存带宽受限的。换句话说，算法的算术强度低于硬件带宽比率的情况下，算法是内存带宽受限的。

了解到LLM推理过程中解码阶段的算术强度较低（详见下一篇博文），在大多数硬件上通常表现为内存带宽受限。与NVIDIA H100相比，NVIDIA H200针对这种低强度工作负载具有更有利的带宽比率。这解释了NVIDIA将H200定位为“超级推动生成式AI推理”的硬件设计目标为此内存限制。

现在让我们将算术强度与延迟和吞吐量联系起来：

注意：这里的吞吐量以TFLOPS表示，而不是每秒请求数，但两者成正比。此外，吞吐量以TFLOPS表示突出了它与硬件利用率以及成本效率的联系。为了更明显地表明这种联系，吞吐量更精确地说是每芯片每秒处理的请求数量，芯片每请求的秒数越少，即吞吐量越高，成本效率越大（参见[9]第4节）。

如果我们将算术强度作为X轴，最大可达到的吞吐量作为Y轴，我们就得到了所谓的（朴素）屋顶线模型[12]（图9）。

图9.屋顶线模型

让我们进行一个小思维实验，以更好地理解此图中吞吐量值为最大可达到水平的原因。在计算受限模式下，这显而易见：没有什么能阻止我们利用完整的计算能力，我们只受到硬件峰值容量的限制。在内存带宽受限模式下，我们在1秒内能够获取的最大数据量由硬件的内存带宽BW_mem决定。考虑到算法的算术强度A，我们在1秒内可以实现的最大FLOPs数量因此是BW_mem.A。CED。

增加算法的算术强度会有什么影响？我们可以考虑三种情景（图10）：

在这里插入图片描述
图10.算术强度增加的三种情况

情景1：算术强度的增加微小，不足以摆脱内存带宽受限模式，但会导致吞吐量成比例增加。系统仍然是内存带宽受限的，因此对延迟的影响取决于更高强度如何影响该算法的数据传输。

情景2：算术强度的增加使系统切换到计算受限模式。吞吐量提升至硬件的峰值吞吐量。现在是计算受限的，延迟影响取决于更高强度如何改变该算法的总操作。

情景3：由于已经处于计算受限状态且达到峰值吞吐量，增加的强度不会带来吞吐量增益。延迟影响仍然取决于更高强度如何影响该算法的总计算。

如何具体增加算术强度？这完全取决于算法的具体细节。在下一篇博文中，我们将探讨影响Transformer解码器块算术强度的关键参数。我们将看到，例如提高批处理大小可以增加某些操作的算术强度。

一些已讨论的优化措施还可以提高算术强度，改善吞吐量和资源利用率。对于解码阶段受内存带宽限制的Transformer模型，通过操作融合和数据（权重矩阵，KV缓存）量化主要可以改善算术强度。

到目前为止，我们做出了一个关键假设，即算法实现最优地利用硬件资源。例如，在数据传输时，假设算法实现使用了硬件理论内存带宽的100%。显然，在实践中这并不总是情况（尽管某些实现接近最佳资源使用），那么次优资源使用如何影响分析？

这很简单：上述带宽数字必须由实际实现的数字替换。次优系统位于其自己的屋顶线曲线上，低于最佳屋顶线曲线（图11）。现在有两个改进吞吐量的自由度：增加算术强度和/或改进算法实现以更好地利用硬件资源。

图11.资源利用次优的屋顶线模型

最后，我们通过提供一个实现改进的真实例子来总结。在2.2版本之前，FlashAttention核心的实现在推理解码阶段应用时可能会变得非常次优。先前的数据加载实现使得核心在解码阶段难以有效利用内存带宽。更糟糕的是，随着批处理大小的增加，带宽利用率实际上进一步降低；因此，对于由于内存限制需要较小批次的长序列，性能受到的影响最大。FlashAttention团队通过跨KV缓存序列长度维度并行化数据加载来解决了这个问题，并发布了一个名为FlashDecoding的优化解码阶段核心，显著改善了长序列长度的延迟问题 [13]。

总结

我们了解了影响模型延迟的四种瓶颈类型。识别导致模型延迟的主要因素类型至关重要，因为每种类型都需要特定的缓解策略。

为简化起见，不考虑分布式设置，实际的硬件操作要么是计算受限的，要么是内存带宽受限的。核心的算术强度决定了绑定模式。在较低强度的内存带宽受限模式下，最大可达到的吞吐量与算术强度成线性关系。相比之下，在计算受限模式下，吞吐量受硬件峰值FLOPS限制。根据影响强度的因素，我们可能能够增加强度以提高最大吞吐量，甚至达到计算受限的性能水平。然而，这种强度增益可能会对延迟产生不利影响。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述