NLP（十八）：LLM 的推理优化技术纵览

原文：NLP（十八）：LLM 的推理优化技术纵览 - 知乎

目录

收起

一、子图融合（subgraph fusion）

1.1 FasterTransformer by NVIDIA

1.2 DeepSpeed Inference by Microsoft

1.3 MLC LLM by TVM

二、模型压缩（Model Compression）

2.1 稀疏(Sparsity)

2.2 量化(Quantization)

2.3 蒸馏(Distillation)

三、并行化（Parallelism）

3.1 数据并行 (Data Parallelism, DP)

3.2 张量并行(Tensor Parallelism, TP)

3.3 流水线并行(Pipeline Parallelism, PP)

四、Transformer 结构优化

4.1 FlashAttention

4.2 PagedAttention

4.3 FLAT Attention

五、动态批处理（Dynamic Batch, Continuous batch）

5.1 ORCA

5.2 FastServe

5.3 vLLM

5.4 Text Generation Inference

5.5 LMDeploy

六、KV cache 优化

七、投机采样与MoE

7.1 投机采样

7.2 超大模型的 MoE

八、硬件升级

8.1 NVIDIA H100 PCIe

8.2 AMD MI300

8.3 Apple M2 Ultra

参考资料

推理是 LLM 应用的重要一环，在部署服务环节影响重大，本文将讨论主流的 LLM 的推理优化技术。

一、子图融合（subgraph fusion）

图融合技术即通过将多个 OP（算子）合并成一个 OP（算子），来减少Kernel的调用。因为每一个基本 OP 都会对应一次 GPU kernel 的调用，和多次显存读写，这些都会增加大量额外的开销。

1.1 FasterTransformer by NVIDIA

FasterTransformer(FT) 是一个用于实现基于Transformer的神经网络推理的加速引擎。FT框架是用C++/CUDA编写的，依赖于高度优化的 cuBLAS、cuBLASLt 和 cuSPARSELt 库，与 NVIDIA TensorRT 等其他编译器相比，FT 的特点是它支持以分布式方式推理 Transformer 大模型。

图融合是FT 的一个重要特征，将多层神经网络组合成一个单一的神经网络，将使用一个单一的内核进行计算。 这种技术减少了数据传输并增加了数学密度，从而加速了推理阶段的计算。 例如， multi-head attention 块中的所有操作都可以合并到一个内核中。

除此之外，FT还对部分大模型分别支持：

• INT8 低精度量化推理
• Ampere 架构的 GPU 硬件部分支持稀疏化
• Hopper 架构支持 FP8 推理
• Tensor 并行
• Pipeline 并行

1.2 DeepSpeed Inference by Microsoft

对于 Transformer layer，可分为以下4个主要部分：

1. Input Layer-Norm plus Query, Key, and Value GeMMs and their bias adds.
2. Transform plus Attention.
3. Intermediate FF, Layer-Norm, Bias-add, Residual, and Gaussian Error Linear Unit (GELU).
4. Bias-add plus Residual.

如图所示，每一部分可分别进行融合，与未融合相比，以上几个部分的加速比可分别达到 1.5x, 2.9x, 3x, 1.2x 。

除此之外，DeepSpeed Inference 的优化点还有以下几点：

• 多 GPU 的并行优化
• INT8 模型量化
• 推理的 pipeline 方案

更多详细介绍及实践可参考笔者之前的文章：

紫气东来：NLP（十二）：DeepSpeed Inference 在 LLM 推理上的优化探究189 赞同 · 13 评论文章​编辑

1.3 MLC LLM by TVM

之前介绍的推理方案主要是基于GPU的优化，而 MLC LLM 提供了可应用于移动端（例如 iPhone）、消费级电脑端（例如 Mac）和 Web 浏览器的轻设备解决方案。

MLC LLM 的主要工作流基于 Apache TVM Unity，通过扩展 TVM 后端使模型编译更加透明和高效。其中以编译代码转换、融合、内存规划和库卸载（library offloading）为代表的可组合的 ML 编译优化是其中重要的优化特性。

除此之外，MLC LLM 还具有以下特性：

• Dynamic shape：避免了对最大输入长度进行额外填充的需要，并减少了计算量和内存使用量。
• 量化：MLC LLM 利用低位量化来压缩模型权重，并利用 TVM 的 loop-level TensorIR 为不同的压缩编码方案快速定制代码生成。
• 运行时（Runtime）：TVM 编译生成的库能够通过 TVM runtime 在设备的原生环境中运行，TVM runtime 支持 CUDA/Vulkan/Metal 等主流 GPU 驱动以及 C、JavaScript 等语言的绑定。

除了上述3种方案外，其他也支持图融合的方案还包括 NVIDIA TensorRT， Tencent TurboTransformers 等。

二、模型压缩（Model Compression）

模型压缩的基本动机在于当前的模型是冗余的，可以在精度损失很小的情况下实现模型小型化，主要包括3类方法：稀疏(Sparsity)、量化(Quantization)、蒸馏(Distillation)。

2.1 稀疏(Sparsity)

实现稀疏(Sparsity)的一个重要方法是剪枝(Pruning)。剪枝是在保留模型容量的情况下，通过修剪不重要的模型权重或连接来减小模型大小。 它可能需要也可能不需要重新培训。 修剪可以是非结构化的或结构化的。

• 非结构化剪枝允许删除任何权重或连接，因此它不保留原始网络架构。 非结构化剪枝通常不适用于现代硬件，并且不会带来实际的推理加速。
• 结构化剪枝旨在维持某些元素为零的密集矩阵乘法形式。 他们可能需要遵循某些模式限制才能使用硬件内核支持的内容。 当前的主流方法关注结构化剪枝，以实现 Transformer 模型的高稀疏性。

关于剪枝稀疏的基本原理，可参考笔者之前的文章：

紫气东来：NLP（八）：大语言模型的稀疏化技术43 赞同 · 6 评论文章​编辑

除了上文介绍的稀疏方法外，还有其他的稀疏化方法，包括但不限于：

• SparseGPT：该方法的工作原理是将剪枝问题简化为大规模的稀疏回归实例。它基于新的近似稀疏回归求解器，用于解决分层压缩问题，其效率足以在几个小时内使用单个 GPU 在最大的 GPT 模型（175B 参数）上执行。同时，SparseGPT 准确率足够高，不需要任何微调，剪枝后所损耗的准确率也可以忽略不计。
• LLM-Pruner：遵循经典的“重要性估计-剪枝-微调”的策略，能够在有限资源下完成大语言模型的压缩，结果表明即使剪枝 20％ 的参数，压缩后的模型保留了 93.6％ 的性能。
• Wanda: 该方法由两个简单但必不可少的组件构成——剪枝度量和剪枝粒度。剪枝度量用来评估权重的重要性，然后按照剪枝粒度进行裁剪。该方法在 65B 的模型上只需要 5.6 秒就可以完成剪枝，同时达到SparseGPT相近的效果。

以上主要实现了稀疏的方法，那么对于稀疏后的模型如何加速呢？NVIDIA Ampere 架构对与结构化稀疏做了专门的稀疏加速单元，下图展示了结构化稀疏的物理表示：

2:4 结构化稀疏表示

下图展示了稀疏单元GEMM计算与标准GEMM计算的区别（详细解释参见https://arxiv.org/pdf/2104.08378.pdf）

Sparse VS Dense GEMM

2.2 量化(Quantization)

常见量化有两种常见方法：

• 训练后量化（Post-Training Quantization，PTQ）：模型首先经过训练以达到收敛，然后我们将其权重转换为较低的精度，而无需进行更多训练。 与训练相比，实施起来通常相当便宜。
• 量化感知训练（Quantization-Aware Training，QAT）：在预训练或进一步微调期间应用量化。 QAT 能够获得更好的性能，但需要额外的计算资源和对代表性训练数据的访问。

实际上，由于 GPU 内核缺乏对某些类型的矩阵乘法（例如 INT4 x FP16）的支持，理论最优量化策略与硬件内核支持之间的差距，并非以下所有方法都能加速实际推理。

关于量化的基本原理和实现细节，可参考笔者之前的文章：

紫气东来：NLP（十一）：大语言模型的模型量化(INT8/INT4)技术257 赞同 · 15 评论文章​编辑

许多关于 Transformer 模型量化的研究都有相同的观察结果：简单的低精度（例如 8 bit）训练后量化会导致性能显着下降，这主要是由于动态的 activation 和静态的 weight 量化策略无法保持一致。

为了不损失精度而提高性能，可以考虑 WeightOnly 量化技术，即只把 Weight 量化成 int8 格式，以降低访存压力。到实际 Kernel 内部再 Dequantize 回 fp16，进行矩阵乘计算。这种方法在 BS 较小是比较有效(因为此时的瓶颈在IO)，BS 较大时(因为此时的瓶颈在计算)效果变差。

WeightOnly 量化的典型案例是 AWQ: Activation-aware Weight Quantization，即只对 weight 进行量化以实现压缩和加速的效果。

2.3 蒸馏(Distillation)

知识蒸馏是一种构建更小、更便宜的模型（“student 模型”）的直接方法，通过从预先训练的昂贵模型中转移技能来加速推理（“ teacher 模型”）融入 student。 除了与 teacher 匹配的输出空间以构建适当的学习目标之外，对于如何构建 student 架构没有太多限制。

知识蒸馏基本框架

给定数据集，训练 student 模型通过蒸馏损失来模仿 teacher 的输出。 通常神经网络有一个softmax层； 例如，LLM 输出 token 的概率分布。 将 softmax 之前的 logits 层表示为 �� 和 �� , 分别表示 teacher 和 student 模型。 蒸馏损失最小化两个 softmax 输出之间的差异（温度 � ）。 当标签 � 已知，可以将其与student 的 logits 之间计算交叉熵，最后将两个损失相加，如下：

�KD=�distll (softmax⁡(��,�),softmax⁡(��,�))+��CE(�,��)

在 Transformer 中一个典型案例是DistilBERT，模型参数减少 40%，速度提升71%。在大模型时代，蒸馏可以与量化、剪枝或稀疏化技术相结合，其中 teacher 模型是原始的全精度密集模型，而 student 模型则经过量化、剪枝或修剪以具有更高的稀疏级别，以实现模型的小型化。

三、并行化（Parallelism）

当前的推理的并行化技术主要体现在3个维度上，即 3D Parallelism:

• Data Parallelism(DP)
• Tensor Parallelism(TP)
• Pipeline Parallelism(PP)

3D Parallelism 的3个维度

3.1 数据并行 (Data Parallelism, DP)

在推理中，DP 主要是增加设备数来增加系统整体 Throughput，其中最经典的即DeepSpeed的Zero系列

另外 FSDP 也比较高效和易用

3.2 张量并行(Tensor Parallelism, TP)

在推理中，TP 主要是横向增加设备数通过并行计算来减少 latency，其实现原理及细节可参考笔者之前的文章

紫气东来：NLP（六）：GPT 的张量并行化（tensor parallelism）方案47 赞同 · 11 评论文章​编辑

当前也有一些方便易用的 TP 方案，如 BlackSamorez/tensor_parallel ，使用起来非常简单：

import transformers
import tensor_parallel as tp
tokenizer = transformers.AutoTokenizer.from_pretrained("facebook/opt-13b")
model = transformers.AutoModelForCausalLM.from_pretrained("facebook/opt-13b")  # use opt-125m for testing

model = tp.tensor_parallel(model, ["cuda:0", "cuda:1"])  # <- each GPU has half the weights

inputs = tokenizer("A cat sat", return_tensors="pt")["input_ids"].to("cuda:0")
outputs = model.generate(inputs, num_beams=5)
print(tokenizer.decode(outputs[0])) # A cat sat on my lap for a few minutes ...

model(input_ids=inputs, labels=inputs).loss.backward()  # training works as usual

当前主流的推理框架都支持 TP 的方式，包括但不限于：

• Megatron-LM
• FasterTransformer
• DeepSpeed Inference
• vLLM
• Text Generation Inference
• ParallelFormers
• ColossalAI
• FlexFlow
• LiBai
• AlpaServe

3.3 流水线并行(Pipeline Parallelism, PP)

在推理中，PP 主要是纵向增加设备数通过并行计算来支持更大模型，同时提高设备利用率。

通常来说，PP 需要与 TP 结合以支持更大模型，并实现最佳效果

四、Transformer 结构优化

该类方法主要通过优化 Transformer 的结构以实现推理性能的提升。

4.1 FlashAttention

该部分的实现细节可参考笔者之前的文章，在此不予赘述

紫气东来：NLP（十七）：从 FlashAttention 到 PagedAttention, 如何进一步优化 Attention 性能750 赞同 · 43 评论文章​编辑

FlashAttention-v2 在原基础上做了改进，使其在算法、并行化和工作分区等方面都有了显著改进，对大模型的适用性也更强。在A100 上性能数据如下：

FlashDecoding 是在 FlashAttention 的基础上针对 inference 的优化主要分为三步：

1. 长文本下将KV分成更小且方便并行的chunk
2. 对每个chunk的KV，Q和他们进行之前一样的FlashAttention获取这个chunk的结果
3. 对每个chunk的结果进行reduce

4.2 PagedAttention

可参考

紫气东来：NLP（十七）：从 FlashAttention 到 PagedAttention, 如何进一步优化 Attention 性能750 赞同 · 43 评论文章​编辑

4.3 FLAT Attention

FLAT-Attention 与 FlashAttention 采取不同的路线来解决同一问题。 提出的解决方案有所不同，但关键思想是相同的（tiling 和 scheudling）。下面主要讨论二者不同之处：

4.3.1 Tiling 策略比较

FlashAttention 使用块平铺和权重固定。 FLAT-Attention 使用行平铺（行粒度）和输出固定。

4.3.2 Scheduling 策略(数据流)比较

FlashAttention 的 Scheduling 过程

FLAT-Attention 的 Scheduling 过程

五、动态批处理（Dynamic Batch, Continuous batch）

该类方法主要是针对多 Batch 的场景，通过对 Batch 的时序优化，以达到去除 padding、提高吞吐和设备利用率。传统的 Batch 处理方法是静态的，因为Batch size 的大小在推理完成之前保持不变。

如下图所示，使用静态 Batch 完成四个序列。 在第一次迭代（左）中，每个序列从prompt（黄色）生成一个token（蓝色）。 经过几次迭代（右）后，每个完成的序列都有不同的大小，因为每个序列在不同的迭代中发出其序列结束标记（红色）。 可见序列 3 在两次迭代后就已经结束，但仍然需要等待 Batch 中的最后一个序列完成生成（在本例中，序列 2 在六次迭代后）才能统一输出，这意味着 GPU 未被充分利用。

静态 Batch 的推理情况

下面我们来研究 Dynamic Batch 是如何优化这一过程的。

5.1 ORCA

Orca 不是等到 Batch 中的所有序列完成生成，而是实现 iteration 级调度，其中Batch size由每次迭代确定。 结果是，一旦 Batch 中的序列完成生成，就可以在其位置插入新序列，从而比静态 Batch 产生更高的 GPU 利用率。

下图展示了使用 Dynamic Batch 完成七个序列的过程。 左侧显示单次迭代后的批次，右侧显示多次迭代后的 Batch 。 一旦序列发出序列结束标记，就在其位置插入一个新序列（即序列 S5、S6 和 S7）。 这样可以实现更高的 GPU 利用率，因为 GPU 不会等待所有序列完成才开始新的序列。

结果显示在延时不变的情况下，其相对于FasterTransformer 可获得 36.9 倍的吞吐提升。

5.2 FastServe

ORCA 使用first-come-first-served (FCFS) 处理推理作业, 计划任务持续运行直至完成。 由于 GPU 内存容量有限以及推理对延时敏感，无法通过任意数量的传入函数来增加处理，由此可能会导致队列阻塞。

FastServe 使用 preemptive scheduling，通过新颖的跳跃连接 Multi-Level Feedback Queue 程序来最小化延时。 基于 LLM 推理的长度无法确定，调度程序利用输入长度信息来分配适当的初始值每个到达作业要加入的队列。 较高优先级队列跳过加入的队列以减少降级。 设计高效的GPU内存管理机制主动下载和上传 GPU 内存和主机内存之间的中间状态，以进行 LLM 推理。

实验表明，该方法比ORCA有明显的性能提升

5.3 vLLM

vLLM 的核心是 PagedAttention，其灵感来自传统操作系统概念，例如分页和虚拟内存。 它们通过在固定大小的“页面”或块中分配内存，允许 KV 缓存变得不连续。 然后可以重写 attention 机制以对块对齐的输入进行操作，从而允许在非连续的内存范围上执行 attention 。

这意味着 cache 分配可以 just-in-time，而不是 ahead-of-time：当启动一个新的生成任务时，框架不需要分配大小为 Maximum_context_length 的连续 cache。 每次迭代，调度程序都可以决定特定生成任务是否需要更多空间，并动态分配，而不会降低 PagedAttention 的性能。 这并不能保证内存的完美利用（浪费现在限制在 4% 以下，仅在最后一个块中），但它明显改善了当今业界广泛使用的提前分配方案的浪费 。

总而言之，PagedAttention + vLLM 可节省大量内存，因为大多数序列不会消耗整个上下文窗口。 这些内存节省直接转化为更高的 Batch 大小，这意味着更高的吞吐量和更便宜的服务。

实验表明，该方法相比于静态 Batch 与其他动态 Batch 的方法吞吐性能提升明显。

5.4 Text Generation Inference

TGI 是 HuggingFace 开发的基于 Rust, Python 和 gRPC 的推理服务工具，其基本框架如下：

关于 TGI 的用法，可参考笔者的文章，同时对比了和 vLLM 和 FasterTransformer 的性能。

紫气东来：小记：主流推理框架在Llama 2 的上性能比较141 赞同 · 33 评论文章​编辑

5.5 LMDeploy

LMDeploy 是由 MMRazor 和 MMDeploy 团队开发的用于压缩、部署 LLM 服务的工具包。 它具有以下核心特点：

• TurboMind：基于FasterTransformer 的高效推理引擎。
• 交互推理：通过缓存多轮对话过程中的 k/v，记住对话历史，以避免对历史会话的重复处理。
• 多GPU模型部署和量化
• Dynamic Batch

六、KV cache 优化

关于 KV cache 的典型优化方法及原理，可参考笔者的文章

紫气东来：NLP（二十）：漫谈 KV Cache 优化方法，深度理解 StreamingLLM260 赞同 · 8 评论文章​编辑

七、投机采样与MoE

7.1 投机采样

该类方法的典型代表是 SpecInfer 和 Google 投机采样，下面以 SpecInfer 为例说明其原理。

SpecInfer 的核心思路为：通过一些列小模型 SSM（Small Speculative Model）联合预测 LLM 输出，并将多个模型的预测输出组织为 Token 树，树中每个分支表示一个候选 Token 序列。然后 LLM 使用基于树的并行解码（tree-based parallel decoding）机制来并行的验证 Token 树中所有 Token 的正确性。SpecInfer 使用 LLM 作为 Token 树验证器而非增量解码器，这显著降低了生成式 LLM 的端到端延迟和计算量，同时可以保持模型的质量。

下图 (a) 为传统的增量解码方式，每次解码 1 个 Token, (b) 中对比了传统增量解码方式和投机推理解码方式差异。(c) 即为 SpecInfer 方案，下面将主要介绍该方案。

SpecInfer 中的一个核心为预测器（speculator）的设计和实现。

• 一方面，更准确的预测器能生成更长匹配长度的预测 Token 序列，这将帮助减少所需的验证步数。
• 另一方面，句子中某些短语更容易推测，而有些短语就很有挑战性，使用固定的配置（比如搜索时 beam search 的宽度和深度）会导致性能不足，因为非常小的预测 Token 窗口可能导致错过匹配更长的 Token 序列，而更大的预测窗口也可能导致生成很多无效的 Token。

SpecInfer 采用了两个关键技术来解决这一挑战：

1. 为了提升预测 Token 树的性能，引入了 collective boost-tuning，一种新的微调技术，其通过自适应 boosting 来实现一组 SSM 的聚合预测结果和 LLM 输出的对齐。
2. 为了解决不同预测之间的动态性，提出了一个可学习的推测调度器，其学习给定一个输入 Token 序列来找到一组 SSM 和其最优的推测配置。

Collective Boost-Tuning

SpecInfer 使用一种无监督的方法来联合微调多个 SSM，通过自适应增强技术来使它们的输出与 LLM 的输出对齐。如下图所示，SSM 的作用是用来预测 LLM 将要生成的 Token，因此 SpecInfer 使用通用文本数据集以完全无监督的方式将 SSM 的聚合输出与 LLM 自适应对齐。其中，作者将文本语料库转成 prompt 样本集合，然后使用 LLM 来为每一个 prompt 生成 Token 序列。

微调的过程中，SpecInfer 一次将一个 SSM 微调到最充分，然后从训练样本中去除当前 SSM 和 LLM 生成完全一致的样本；然后使用剩余样本来最大化的微调下一个 SSM。对 SSM 池中的所有 SSM 重复这一过程，以此 SpecInfer 将获得一组不同的 SSM，其聚合输出在很大程度上保持与 LLM 输出重叠。

Learning-based Speculative Scheduler

为了找到在每个 decoding step 中启动多个 SSM 的最佳配置，作者设计了一个可学习的推测调度器，其学习给定一个输入 Token 序列，决定使用哪些 SSM 以及这些 SSM 的最优推测配置。

调度器包含一个匹配长度预测器和一个代价模型。匹配长度预测器将 LLM 的最后一个隐藏层的最后一个 Token embedding 作为输入，然后输出一个连续数字矢量，每个矢量表示特定推测配置下的预期匹配长度。SpecInfer 使用 3 层 MLP 来构建匹配长度预测器模型，并使用每个 SSM beam seach 的宽度、深度构成配置空间，比如宽度为 [1, 2, 4]，深度为 [1, 2, 4, 8, 16]，可以构成一个 15 维的搜索空间。因此预测器输出一个由 15 个数字组成的矢量，每个数字都表示一个 SSM 对应的推测配置下的预测匹配长度。长度预测器也是在公开可用的开源数据集上训练而来。需要说明的是，得到预测器的输入 Token embedding 并不涉及额外的成本，因为其已经包含在 SpecInfer 的验证器（LLM 推理验证）中。

Token Tree Verifier

这个部分介绍 Token 树验证，其输入为由推测器（speculator）生成的 Token 树，然后使用 LLM 来验证 Token 序列的正确性。

Token 树：SpecInfer 使用 Token 树来存储可学习推测器生成的 Token 结果。Token 树是一个树结构，其中的每一个节点 u 表示一个 Token �� ， �� 表示该节点对应的树中的父节点。对于树中的每一个节点 u，将其沿父节点遍历即可得到对应的 Token 序列 �� 。SpecInfer 从多个 SSM 会得到多个 Token 序列，然后会将它们合并成一个 Token 树。

SpecInfer 设计的核心想法是通过一个 LLM decoding step 即可同时验证 Token 树中的所有 Token 序列。Token 树验证允许 SpecInfer 有机会同时解码多个 Token（而不是增量解码方式中的单个 Token），从而减少对 LLM 参数的访问。在 Token 树验证中必须解决的一个挑战是高效地计算 Token 树中所有序列的注意力得分。SpecInfer 使用树注意力机制来实现快速且廉价的计算，同时也有许多其他系统级的优化。

Tree-based Parallel Decoding

对于 SpecInfer 来说，由传统的序列方式变成 Token 树的形式，不同的分支序列中可能依赖不同的 key 和 value，比如上图左上（Speculated Token Tree）中的子序列（t2, t3, t4, t5）和（t2, t3, t8, t9）在第 3 和第 4 个位置的 key、value 不同。

在 SpecInfer V1 中作者介绍了如下的三种方案：

1. 图左中（Sequence-based Decoding）：最简单、直观的方案为并行处理 3 个子序列（一个子序列可以经过一次 LLM decoding step 完成），此方式会导致极大的浪费，比如每个子序列都计算并保留了 t2 的 key 和 value。
2. 图右中（Token-based Decoding）：为了避免重复计算和存储 key、value cache，可以使用深度优先搜索机制，每次只计算一个 Token，直到把所有 Token 计算完成。这种方式也存在比较大的缺陷，每个 Token 单独计算，导致（t3, t4, t5, t6, t7, t8, t9）需要 7 个独立的 kernel 调用，而且不能并行运行，退化成传统增量生成模式。因为部分 Token 还要被丢弃，所以最终的时延反而比传统增量生成方式还长。
3. 图下（Tree-based Decoding）：SpecInfer 提出了 Tree-based Decoding 算法，其在 Token-based Decoding 的基础上将部分 Token 分组合并，如图所示，（t3, t4, t5, t6, t7, t8, t9）可以分为三组:（t3, t4, t5）、（t6, t7）和（t8, t9），这三组就可以在 3 个 kernel 里面完成，但是需要注意（t6, t7）依赖（t3, t4），（t8, t9）依赖（t3），因此需要保证（t3, t4, t5）对应的 kernel 首先计算。基于此方式，Token-based Decoding 需要的 7 个 kernel 变成 3 个 kernel。

在 SpecInfer V2 中，作者重新更新了 Tree-based Parallel Decoding 部分，将上一版本的 Tree-based Parallel Decoding 改成 Sequence-based Parallel Decoding，而将直接依赖 Topolopy-aware causal mask 的方案改成 Tree-based Parallel Decoding，相应的，经过一次 Kernel 即可完成计算。

Verification

对于一个给定的预测 Token 树 N，SpecInfer 使用 tree attention 机制来计算 N 中每个节点 u 的输出。该方案的核心优势是只访问 LLM 参数一次即可并行验证所有的 Token，对应会为每个节点 u 生成一个新的 Token。

SpecInfer 使用的投机推理算法以及 Token 树验证算法，其中全流程为：

1. SpecInfer 将当前 Token 序列 S 作为输入，并生成候选的 Token 树 N。
2. SpecInfer 对 LLM 的使用和现有的系统不太一样：
   1. LLM 将 Token 树作为输入，并为每个节点 u∈N 生成一个 Token O(u)。TreeParallelDecode 函数能在 LLM 的一个 decoding 步骤同时生成 O 中的所有 Token。
   2. Verify 根据 LLM 的输出 O 检查候选的 Token 树 N，并生成一系列经过验证的 Token V。
   3. 将 Token 序列 V 附加到 Token 序列 S 中

其中的 Verify 过程为：从 Token 树 N 的根结点开始。给定节点 u，如果子节点 v 中的 LLM 新生成 Token 和推测 Token 相同，则接受节点 v，并继续验证节点 v 的下一个子节点。如果 u 的子节点中没有任何一个被接受，则在 u 节点终止。

算法开销

SpecInfer 以内存和计算开销为代价加速生成式 LLM 推理。本节分析这些开销，并表明它们通常比使用增量解码方式执行 LLM 推理的内存和计算成本小一两个数量级。

i) 内存开销（Memory overhead）

SpecInfer 的内存开销主要来自两个方面：

1. 相比传统增量解码方式，SpecInfer 需要额外保存一个或多个小模型。作者评估使用比 LLM 小 40-100 倍的模型即可很大程度上提升效果（SpecInfer V2 中为 100-1000倍），因此这些模型只增加了 1-2% 的内存占用。
2. 相比传统增量解码方式一次只处理一个 Token，SpecInfer 的 Token 树验证阶段需要额外存储 Token 树中所有 Token 的 K 和 V。对于长序列生成任务来说，这些存储相比整个序列的 KV cache 要小的多，几乎可以忽略。比如 GPT-4 支持 32K 的 Token，而作者评估 SpecInfer 中的 Token 树只需 32 或 64 个 Token 即可。（对于现在常用的 2K、4K 序列长度模型来说，这个代价也不大，差不多为 1-3%）。

ii) 计算开销（Computation overhead）

与内存开销类似，SpecInfer 的计算开销也主要来自两个方面：

1. 以增量解码方式运行多个 SSM 模型来生成候选 Token 树序列，SpecInfer 可以并行的运行这些 SSM 模型来尽量降低时延。作者评估显示，以增量方式运行 SSM 生成比 LLM 快 3.7 倍。
2. SpecInfer 在 Token 树验证阶段会计算所有 Token 的输出，而这里的很多 Token 最终会被丢弃，因此这其中存在很多无效计算。而 LLM 的增量解码机制都会使用 KV cache，这会占用很多的内存，因此无法使用很大的并发（batch size），也导致其存在明显的 IO bound，而 SpecInfer 的 Token 树验证阶段虽然增加了一定计算，但是没有明显增加 IO，相当于使用了未被充分利用的计算资源，因此也基本可以忽略。

7.2 超大模型的 MoE

随着模型越来越大，训练和推理成本会平方级增长，为了平衡成本和规模，即希望在较小成本下实现超大模型的效果，通常会采用 MOE（Mixture-of-Experts）技术。

MoE 层即将大模型拆分成多个小模型(专家，expert)， 每轮迭代根据样本决定激活一部分专家用于计算，达到了节省计算资源的效果；并引入可训练并确保稀疏性的门( gate )机制，以保证计算能力的优化。

与密集模型不同，MoE 将模型的某一层扩展为多个具有相同结构的专家网络( expert )，并由门( gate )网络决定激活哪些 expert 用于计算，从而实现超大规模稀疏模型的训练。MoE 细节如下图所示，中间层具有 n 个 expert 的 MoE 结构，并引入 Gating network 和 Top_k 机制。计算过程如下公式：

���(�)=∑�=1�(�(�)���(�))�(�)=TopK⁡(softmax⁡(��(�)+�))

上述第 1 个公式表示了包含 n 个专家的 MoE 层的计算过程。具体来讲，首先对样本 x 进行门控计算， W 表示权重矩阵；然后，由 Softmax 处理后获得样本 x 被分配到各个 expert 的权重；然后，只取前 k (通常取 1 或者 2）个最大权重；最终，整个 MoE Layer 的计算结果就是选中的 k 个专家网络输出的加权和。

典型的 MoE 的方案包括 GShard, Switch-Transformer, GLaM, FastMoE, DeepSpeed-MoE 等方案。

FastMoE 支持在同一个 worker 上运行多个 experts，从而减少模型研究者在探索更多 experts 数量时所需的硬件资源。当 experts 数量较多时，FastMoE 针对传统的两层 MLP 全连接网络（即 Transformer 中的 FFN 网络）使用了更精细的并行策略，从而使得 Transformer 模型中 MLP 部分的运算速度相比朴素的实现较大的加速。

八、硬件升级

以上主要介绍了在算法和模型层面的优化方法，除此之外，升级硬件系统可以进一步提升整体性能，下面将介绍几种可用于(和潜在的)推理加速的硬件产品。

8.1 NVIDIA H100 PCIe

NVIDIA H100 核心架构与 Ampere 相似，数学运算部分布置在144组CUDA上，最高可拥有18432个FP32(单精度)、9216个FP64(双精度)CUDA核心，辅以576个第四代Tensor核心。H100核心采用台积电的N4工艺制造，内建800亿个晶体管，核心面积仅有814m㎡。其与A100 主要参数对比如下：

在性能方面，H100 较 A100 也有明显提升，其部分数据如下所示。

8.2 AMD MI300

AMD MI300 处理器集成了24个Zen 4架构CPU核心，以及CDNA 3架构GPU核心，周围还有着8颗HBM3高速缓存，容量高达128GB，总计拥有1460亿个晶体管。与上一代 MI250相比，MI300进行AI运算的速度将提高至8倍，能效方面也将提升5倍。

目前未找到公开的在 LLM 方面的推理性能数据。

8.3 Apple M2 Ultra

M2 Ultra 采用第二代 5 纳米工艺制造，并使用 Apple 突破性的 UltraFusion 技术连接两个 M2 Max 芯片的芯片，使性能提高一倍。 M2 Ultra 由 1340 亿个晶体管组成，比 M1 Ultra 多了 200 亿个。 其统一内存架构支持突破性的192GB内存容量，比M1 Ultra多出50%，并具有800GB/s的内存带宽，是M2 Max的两倍。 M2 Ultra 配备更强大的 CPU（比 M1 Ultra 快 20%）、更大的 GPU（快 30%）以及神经引擎（快 40%）。

目前未找到公开的在 LLM 方面的推理性能数据。

参考资料

[1] Large Transformer Model Inference Optimization

[2] Efficient Inference on a Single GPU

[3] How continuous batching enables 23x throughput in LLM inference while reducing p50 latency | Anyscale

[4] https://www.coreweave.com/blog/serving-inference-for-llms-nvidia-triton-inference-server-eleuther-ai

[5] https://medium.com/mobius-labs/accelerating-large-language-models-strategies-for-enhancing-your-ai-inference-speed

[6] LLM-Pruner: On the Structural Pruning of Large Language Models

[7] https://jonathan-hui.medium.com/ai-chips-a100-gpu-with-nvidia-ampere-architecture-3034ed685e6e

[8] arxiv.org/pdf/2302.14017.pdf

[9] 7 ways to speed up inference of your hosted LLMs. «In the future, every 1% speedup on LLM… | by Sergei Savvov | Jun, 2023 | Medium | Medium

[10] https://arxiv.org/pdf/2104.08378.pdf

[11] Introduction to Weight Quantization | Towards Data Science

[12] Accelerating Large Language Models via Low-Bit Quantization | NVIDIA On-Demand

[13] 认真读读WeightOnly - 知乎 (zhihu.com)

[14] https://hackmd.io/@felixkao/HkZaooPD3

[15] SpecInfer: Accelerating Generative Large Language Model Serving with Speculative Inference and Token Tree Verification

[16] Fast Inference from Transformers via Speculative Decoding

[17] Accelerating large language model decoding with speculative sampling

钟鼎山林都是梦，人间宠辱休惊。 —— 辛弃疾 《临江仙》

编辑于 2024-01-01 16:40・IP 属地美国