1. 论文基本信息

1.1. 标题

处理 GPU 上 Transformer 推理的重尾输入 (Handling Heavy-tailed Input of Transformer Inference on GPUs)

1.2. 作者

• Jiangsu Du (中山大学)
• Jiazhi Jiang (中山大学)
• Yang You (新加坡国立大学)
• Dan Huang (中山大学)
• Yutong Lu (中山大学)

1.3. 发表期刊/会议

发表于 2022年国际超级计算大会 (International Conference on Supercomputing, ICS '22)。ICS 是高性能计算领域的重要学术会议，通常发表关于并行计算、体系结构和系统软件的原创研究。

1.4. 发表年份

2022年6月16日 (UTC)

1.5. 摘要

Transformer (Transformer-based) 模型在自然语言处理 (NLP) 领域取得了卓越的准确性，并已广泛部署到生产环境中。图形处理单元 (GPU) 作为流行的部署设备，通常采用批处理 (batch processing) 技术来推理 Transformer 模型以实现高硬件性能。然而，由于 NLP 任务的输入序列长度通常是可变的且呈重尾分布 (heavy-tailed distribution)，批处理会带来大量的冗余计算 (redundant computation)，从而损害实际效率。

本文提出了一种统一解决方案，旨在消除 Transformer 模型在 GPU 上推理重尾输入时的大部分冗余计算并获得性能提升。具体而言，该统一解决方案包括针对自注意力 (self-attention) 模块、多层感知器 (Multilayer Perceptron, MLP) 模块和整个 Transformer 模型的三个策略。对于自注意力模块，我们设计了一种细粒度策略 (fine-grained strategy)，通过索引有效的块矩阵乘法 (block matrix multiplication) 来编排自注意力模块中的细粒度并行 (fine-grained parallelism)。对于 MLP 模块，我们采用常用的词累积策略 (word-accumulation strategy)，将批次 (batch) 中的所有序列紧密排列 (densely)。对于整个模型，我们设计了一种块组织策略 (block-organized strategy)，通过以块的粒度组织自注意力模块的数据布局，将细粒度策略和词累积策略连接起来。将我们的解决方案应用于 GLUE 基准测试 (benchmark) 的八个语料库 (corpora)，结果表明自注意力模块的延迟平均降低了 63.9%，Bert-base 模型的延迟平均降低了 28.1%。

1.6. 原文链接

/files/papers/6919d6ad110b75dcc59ae2b8/paper.pdf (此为内部文件链接，若为外部链接则为公开可访问的 URL)

---

2. 整体概括

2.1. 研究背景与动机

近年来，Transformer 模型在自然语言处理 (NLP) 领域取得了主导地位，并在生产环境中广泛部署。为了充分利用现代图形处理单元 (GPU) 的高并行性，推理这些模型时通常采用批处理 (batch processing) 技术，即同时处理多个输入样本以提高硬件利用率和吞吐量。

然而，NLP 任务的一个固有挑战是输入序列的长度是可变的。为了在批处理中处理这些变长序列，通常采用一种朴素的填充策略 (padding strategy)：将所有序列填充 (pad) 到批次中最长序列的长度，或者预设一个最大序列长度 seq_len。这种填充导致了大量的冗余计算。由于实际输入序列长度往往呈重尾分布 (heavy-tailed distribution)，即大多数序列远短于最大长度，冗余计算的量甚至可能超过有效计算量，严重损害了模型推理的实际效率。

现有的优化框架，如 EffectiveTransformer、FasterTransformer 和 TurboTransformer，已经注意到了这个问题。它们通常采用一种称为 EFF-rebuild 的解决方案，通过在 MLP 模块前移除填充并在 MLP 模块后重建填充来消除 MLP 模块中的冗余计算。然而，这种方案有两大缺点：

1. 移除和重建填充需要额外的数据移动 (data movement)，引入了开销。

2. 它无法消除自注意力 (self-attention) 模块中的冗余计算。随着 MLP 模块冗余计算的减少，自注意力模块在总计算时间中的比重变得更大，使其成为更值得优化的瓶颈。

   因此，论文的动机在于提出一个统一的解决方案，能够系统性地解决 Transformer 模型在 GPU 上推理时，变长、重尾输入所导致的自注意力模块和 MLP 模块的冗余计算问题，从而显著提升实际推理性能。

2.2. 核心贡献/主要发现

本文提出了一种统一解决方案，以消除 Transformer 模型在 GPU 上推理重尾输入时的大部分冗余计算，其核心贡献和主要发现如下：

• 揭示重尾分布问题与自注意力模块的优化价值： 论文明确指出 NLP 任务的输入序列长度通常呈重尾分布，并量化分析了自注意力模块中的冗余计算占比，强调了其在 EFF-rebuild 方案下成为主要性能瓶颈，值得重点优化。
• 提出统一解决方案： 论文提出了一个包含三个协同策略的统一解决方案，以消除 Transformer 模型推理中的大部分冗余计算：
  1. 细粒度策略 (fine-grained strategy)： 针对自注意力模块，通过索引有效的块矩阵乘法来编排细粒度并行，从而消除冗余计算。
  2. 词累积策略 (word-accumulation strategy)： 针对 MLP 模块，将批次中的所有序列紧密排列，这是现有框架中常见的有效方法。
  3. 块组织策略 (block-organized strategy)： 针对整个模型，通过以块的粒度组织自注意力模块的数据布局，将细粒度策略和词累积策略联系起来，并减少数据布局切换的开销。
• 定制化 GPU 内核： 提出了实现细粒度策略的三项关键技术：Mini-block Index (MBIX)、Shared Memory Block Transpose (SMBT) 和 Efficient Atomic Operation (EAOP)，并基于这些技术为 $Q \times K^T$ 和 $QK^T \times V$ 函数实现了定制化的 GPU 内核 (kernel)。
• 实验验证与显著性能提升：

  • 在 GLUE 基准测试的八个语料库上进行了广泛评估。

  • 结果显示，与流行的 FasterTransformer 相比，自注意力模块的延迟平均降低了 63.9%。

  • 整个 Bert-base 模型的推理延迟平均降低了 28.1%。

  • 在 QQP 语料库上，自注意力模块的延迟降低高达 87.8%，整个模型延迟降低高达 65%。

    ---

3. 预备知识与相关工作

3.1. 基础概念

理解本文需要以下核心概念：

3.1.1. Transformer 模型 (Transformer Model)

Transformer 模型是一种深度神经网络架构，于2017年由 Google 提出，主要用于处理序列数据，尤其在自然语言处理 (NLP) 领域取得了巨大成功。其核心创新在于完全依赖自注意力机制 (self-attention mechanism)来捕捉输入序列中不同位置之间的依赖关系，而非传统的循环神经网络 (RNN) 或卷积神经网络 (CNN)。这使得 Transformer 能够更好地处理长序列，并支持高度并行化计算。

3.1.2. 自注意力机制 (Self-Attention Mechanism)

自注意力机制是 Transformer 的核心组成部分，它允许模型在处理序列的某个词时，能“关注”到序列中的所有其他词，并根据它们的重要性分配不同的权重。其基本思想是将输入序列中的每个词表示为三个向量：

• 查询 (Query, $Q$)：表示当前词的信息，用于查询其他词。

• 键 (Key, $K$)：表示其他词的信息，用于被查询。

• 值 (Value, $V$)：表示其他词的实际内容，用于加权求和。

  注意力机制的计算通常遵循以下步骤：

1. 计算 Query 与所有 Key 的点积 (dot product)，得到一个注意力分数 (attention score)，表示查询词与每个键词的关联程度。

2. 将注意力分数除以一个缩放因子（通常是 $d_k$ 的平方根，其中 $d_k$ 是 Key 向量的维度），以防止点积过大导致 softmax 函数梯度过小。

3. 对缩放后的分数应用 softmax 函数，将其转换为概率分布，得到注意力权重 (attention weights)。这些权重指示了每个词对当前词的重要性。

4. 将注意力权重与对应的 Value 向量进行加权求和，得到当前词的注意力输出。

   本文重点关注自注意力机制中的两个主要计算步骤：$Q \times K^T$ 和 $QK^T \times V$。 $$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$ 其中：

• $Q$ (Query)、$K$ (Key)、$V$ (Value) 分别是查询、键和值矩阵。这些矩阵通常由输入嵌入通过线性变换生成。
• $Q \in \mathbb{R}^{L \times d_k}$， $K \in \mathbb{R}^{L \times d_k}$， $V \in \mathbb{R}^{L \times d_v}$。
  • $L$ 是序列长度 (sequence length)。
  • $d_k$ 是 Key 向量的维度。
  • $d_v$ 是 Value 向量的维度。
• $K^T$ 是 $K$ 的转置。
• $\sqrt{d_k}$ 是缩放因子。
• $\mathrm{softmax}(\cdot)$ 是 softmax 函数，将分数转换为概率分布。

3.1.3. 多层感知器 (Multilayer Perceptron, MLP)

MLP 模块（在本文中也指代线性函数 linear functions）在 Transformer 模型中扮演着重要角色，主要用于对自注意力机制的输出进行进一步的非线性变换和特征提取。它通常由两个全连接层 (fully connected layers) 构成，中间夹有激活函数 (activation function，如 ReLU 或 GeLU)。在本文中，所有线性函数都被归类为 MLP 模块，包括生成 Query、Key、Value 向量的线性层，以及在自注意力机制之后学习知识的线性层。

3.1.4. GPU 架构 (GPU Architecture)

GPU 是一种高度并行的处理器，由多个流式多处理器 (Streaming Multiprocessors, SMs) 组成。CUDA 是 NVIDIA 提供的并行计算平台和编程模型。

• 线程 (Threads)：GPU 上最小的执行单位。
• Warp：32 个线程的集合，这些线程以单指令多数据 (Single Instruction, Multiple Data, SIMD) 方式并行执行。
• 线程块 (Thread Blocks)：一组 Warp 的集合。同一线程块内的线程可以通过快速的、程序员管理的共享内存 (shared memory) 进行通信。
• 网格 (Grid)：由多个线程块组成，代表一个 GPU kernel 的完整执行。
• 全局内存 (Global Memory)：所有 SM 都可以访问的大容量、高延迟内存。
• 占用率 (Occupancy)：一个 SM 上并发执行的线程块数量。高占用率通常意味着高并行度，有助于隐藏内存和算术操作的延迟，从而带来更高的硬件性能。

3.1.5. 批处理 (Batch Processing)

批处理是一种常见的优化技术，在深度学习中广泛用于 GPU。它将多个独立的输入样本组合成一个批次 (batch)，然后一次性送入模型进行处理。这样做有几个优点：

• 提高硬件利用率： 多个样本并行计算可以更好地利用 GPU 的并行处理能力。
• 减少内核启动开销 (kernel launch overhead)： 每次启动 GPU 内核都有固定开销。批处理可以减少内核启动次数，从而降低总开销。
• 内存访问优化： 批处理可以促进内存访问的局部性，提高缓存命中率。

3.1.6. 填充策略 (Padding Strategy)

由于 NLP 任务的输入序列长度通常是可变的，而 GPU 的批处理要求批次中的所有样本具有相同的形状，因此需要采用填充策略。最常见的填充方法是：

1. 确定批次中最长的序列长度，或预设一个最大序列长度 seq_len。
2. 对于所有短于该长度的序列，在其末尾添加特殊的填充词元 (padding tokens)，直到它们达到相同的长度。 这种策略使得批次中的所有输入矩阵具有统一的维度，从而可以进行高效的矩阵运算。然而，填充词元并不包含实际语义信息，其产生的计算是冗余计算 (redundant computation)。

3.1.7. 重尾分布 (Heavy-tailed Distribution)

重尾分布是一种概率分布，其尾部（即分布的极端值区域）比指数分布或正态分布“更厚”，这意味着极端事件或大值的发生频率相对较高。在 NLP 任务的序列长度分布中，重尾分布意味着绝大多数序列的长度较短，但存在少量序列的长度异常长。当使用填充策略时，这些少数的超长序列决定了批次中的最大长度，导致大量短序列被填充到非常长的长度，从而产生巨额冗余计算。

3.2. 前人工作

本文主要关注 Transformer 推理的性能优化，并与现有的一些知名框架和方法进行了比较。

3.2.1. 现有通用推理框架

• TensorFlow [16] 和 PyTorch [17]： 这些是通用的深度学习框架，虽然可以执行推理，但在充分利用高级硬件进行推理工作负载方面可能效率不高。
• TensorRT [18]： NVIDIA 提供的用于高性能深度学习推理的 SDK，它通过优化模型图、融合层 (layer fusion) 等技术来加速推理。
• TVM [19]： 一个端到端自动优化深度学习编译栈，旨在为各种硬件后端生成高性能代码。

3.2.2. 特定于 Transformer 的推理框架

随着 Transformer 模型的普及，出现了许多专门针对其进行优化的推理框架，这些框架通常通过手动内核融合 (manual kernel fusion) 设计来减少内核启动开销和频繁的内存访问。

• FasterTransformer [11] (NVIDIA)： 这是一个活跃的开源项目，不断更新并收集实际优化方案。它主要通过定制化的 GPU 内核实现算子融合，以及采用批处理和填充策略来提高硬件性能。在新版本中，FasterTransformer 也开始关注大规模 Transformer 模型的分布式推理。

• EffectiveTransformer [10] (字节跳动)： 该框架提出了填充重建方法 (padding-rebuild method) 来解决 MLP 模块中的变长输入问题。它基于早期的 FasterTransformer 版本，并进一步减少了执行时间和内存消耗，特别是在大批次尺寸 (batch size) 情况下。

• TurboTransformer [12] (腾讯)： 该框架提出了一种内存管理策略来减少变长输入问题的内存使用，并首次提出了服务系统的设计。

  这些框架（包括 EffectiveTransformer、FasterTransformer 和 TurboTransformer）通常都采用了**EFF-rebuild 解决方案**来优化 MLP 模块。EFF-rebuild 的核心思想是：

1. 在进入 MLP 模块之前，通过数据移动将批次中所有序列的有效词紧密地排列在一起，形成一个“词累积 (word-accumulation)”的密集矩阵。
2. 对这个密集矩阵执行 MLP 运算，从而避免了填充带来的冗余计算。
3. 在 MLP 运算之后，但在进入自注意力模块之前，进行反向数据移动，重建填充，将密集矩阵恢复到带填充的原始布局，以适应自注意力模块的批处理要求。 这种方法虽然有效减少了 MLP 模块的冗余计算，但引入了额外的数据移动开销，并且未能解决自注意力模块本身的冗余计算问题，因为自注意力模块仍然需要在带填充的矩阵上进行计算。

3.2.3. 其他相关工作

• E.T. [20]： 提出了自注意力模块的内核融合设计。然而，受限于共享内存容量，它只能在有限的序列长度范围内优于 FasterTransformer，并且随着序列长度的增加性能会严重下降。
• 模型大小缩减： 一些工作通过减少模型大小来加速 Transformer 推理，例如硬件感知的剪枝 (pruning) 设计 [20] 或探索自注意力模块的不同结构 [22]。

3.3. 技术演进

Transformer 模型的推理优化是从对通用深度学习框架的优化，逐步演进到针对其特定计算模式（如自注意力）和变长输入挑战的专业优化。

1. 早期阶段： 依赖通用深度学习框架（如 TensorFlow、PyTorch）进行推理。这些框架虽然功能强大，但通常无法完全发挥 GPU 的硬件性能。
2. 通用推理引擎的兴起： TensorRT、TVM 等通用推理框架通过图优化、层融合、自动代码生成等技术，显著提升了深度学习模型的推理性能。
3. Transformer 专用优化： 随着 Transformer 模型的普及，其独特的架构（特别是自注意力机制）和变长输入特性催生了 FasterTransformer、EffectiveTransformer、TurboTransformer 等专用推理框架。这些框架通过手动内核融合、MLP 模块的填充消除（EFF-rebuild 方案）等方法，进一步榨取硬件性能。
4. 细粒度优化： 本文的工作代表了更深层次的优化，它在 EFF-rebuild 的基础上，进一步解决了自注意力模块本身的冗余计算问题，通过细粒度的并行和数据布局优化，旨在消除批处理和填充策略带来的最后一块主要冗余。

3.4. 差异化分析

本文提出的统一解决方案与 FasterTransformer、EffectiveTransformer 和 TurboTransformer 等现有工作的主要区别和创新点在于：

1. 自注意力模块的冗余计算消除： 现有框架（采用 EFF-rebuild 方案）主要解决了 MLP 模块的冗余计算问题，但自注意力模块仍使用批处理和填充策略，导致大量冗余计算。本文首次提出了一种细粒度策略 (fine-grained strategy)，通过定制化 GPU 内核和索引有效的块矩阵乘法，直接在自注意力模块内部消除大部分冗余计算。

2. 统一的数据布局管理： 现有方法需要在 MLP 和自注意力模块之间进行数据布局的频繁切换（密集排列 <-> 批次填充），引入了数据移动开销。本文提出的块组织策略 (block-organized strategy) 通过引入块填充布局 (block padding layout)，以块的粒度组织自注意力模块的数据，旨在更好地匹配 MLP 模块的词累积策略和自注意力模块的细粒度策略，从而减少布局切换的开销和内存需求。

3. 避免原子操作的开销： 在细粒度策略中，针对 $Q \times K^T$ 函数，通过 1-D 分区将原子操作的需要消除；针对 $QK^T \times V$ 函数，通过重新组织偏移量生成和增大写入间隔来缓解原子操作的竞争问题，从而提高性能。

4. 综合性解决方案： 本文的解决方案是一个统一的框架，它不仅仅优化了某个独立模块，而是从整个 Transformer 模型的角度出发，协同优化了 MLP 模块、自注意力模块以及两者之间的数据流。

   简而言之，现有工作主要停留在 MLP 模块的优化和整体的内核融合，而本文则深入到自注意力模块的矩阵乘法内部，通过更精细的粒度来消除冗余，并优化了模块间的数据衔接。

---

4. 方法论

本节将详细阐述论文提出的统一解决方案，包括其针对自注意力模块、MLP 模块和整个模型的三个核心策略。

4.1. 统一解决方案概述

本文提出的统一解决方案旨在消除 Transformer 模型在 GPU 上推理时，变长重尾输入所导致的冗余计算。该方案由三个协同工作的策略组成：

1. 细粒度策略 (Fine-grained Strategy)：针对自注意力模块，通过编排细粒度并行来消除冗余计算。

2. 词累积策略 (Word-accumulation Strategy)：针对 MLP 模块，这是现有的、将序列紧密排列以消除冗余计算的通用方法。

3. 块组织策略 (Block-organized Strategy)：针对整个模型，将上述两种策略通过优化的数据布局和切换机制连接起来，以减少开销。

   以下是原文 Figure 6 的示意图，展示了我们统一解决方案的整体结构：

   该图像是示意图，展示了我们提出的统一解决方案的整体结构。图中包含了针对自注意力模块、MLP模块和整个变换器模型的策略，其中涉及到的关键步骤包括基于序列长度生成索引数组以及运用 $Q imes K^T imes V$ 计算的细粒度策略和词维度策略。

图示中，Fine-grained Strategy 通过基于序列长度生成索引数组，并利用 $Q \times K^T$ 和 $QK^T \times V$ 的细粒度计算来移除冗余。Word-accumulation Strategy 则关注 MLP 模块的优化，通过紧密排列序列来提高效率。Block-organized Strategy 作为桥梁，在整个模型层面组织数据布局，以协同两种策略并减少布局切换开销。

4.2. 细粒度策略 (Fine-grained Strategy)

细粒度策略是本文针对自注意力模块（特别是 $Q \times K^T$ 和 $QK^T \times V$ 函数）提出的核心创新。它基于块矩阵乘法 (block matrix multiplication) 的思想，将大的矩阵乘法分解为多个小块的矩阵乘法，并只调度执行有效的迷你块 (mini-block) 运算，从而避免冗余计算。

块矩阵乘法 (block matrix multiplication) 的数学表示如下： $$\begin{array}{c} { \begin{array} { r l } { AB = { \binom { A _ { 1 1 } } { A _ { 2 1 } } } } & { A _ { 1 2 } } \\ { A _ { 2 1 } } & { A _ { 2 2 } } \end{array} } { \left[ \begin{array} { l l } { B _ { 1 1 } } & { B _ { 1 2 } } \\ { B _ { 2 1 } } & { B _ { 2 2 } } \end{array} \right] } \\ { = { \binom { A _ { 1 1 } B _ { 1 1 } + A _ { 1 2 } B _ { 2 1 } } { A _ { 2 1 } B _ { 1 1 } + A _ { 2 2 } B _ { 2 1 } } } } & { A _ { 1 1 } B _ { 1 2 } + A _ { 1 2 } B _ { 2 2 } } \end{array}$$ 其中：

• $A$ 和 $B$ 是可以被划分为块矩阵的大矩阵。

• $A_{ij}$ 和 $B_{ij}$ 是 $A$ 和 $B$ 的子块矩阵。

• 矩阵乘法 AB 的结果也可以表示为子块矩阵的乘法和加法组合。

  细粒度策略将自注意力模块中的 $Q \times K^T$ 和 $QK^T \times V$ 函数的矩阵乘法分解为多个迷你块矩阵乘法。这些迷你块的大小和行为完全相同。该策略的实现依赖于三项关键技术：

1. 迷你块索引 (Mini-block Index, MBIX)

2. 共享内存块转置 (Shared Memory Block Transpose, SMBT)

3. 高效原子操作 (Efficient Atomic Operation, EAOP)

   通过这些技术，本文实现了针对 $Q \times K^T$ 和 $QK^T \times V$ 函数的定制化 GPU 内核。

4.2.1. 迷你块索引 (Mini-block Index, MBIX)

MBIX 的核心思想是为有效的迷你块矩阵乘法构建索引数组，以指导 GPU 只执行必要的计算。

对于 $Q \times K^T$ 函数： 如图 (原文 Figure 7) 所示，MBIX 构建了两个索引数组来记录有效查询 (query) 和键 (key) 迷你块的执行顺序（头部偏移量 head offsets）。这些偏移量根据输入序列的实际长度 val_len 计算并存储。此外，为了正确写入结果矩阵，MBIX 还构建了结果矩阵的偏移量数组，记录 $Q K^T$ 的哪个迷你块应该写入计算结果。

该图像是示意图，展示了 $Q \times K^T$ 计算中的二维分区过程。左侧展示了参与计算的矩阵分区，中央部分为矩阵乘法运算，右侧为结果矩阵及相应的偏移量，包括 $Q\_offset$、$K\_offset$ 和 $QK\_offset$。该示意图用于阐释在自注意力模块中进行有效计算的过程。

例如，如果迷你块大小设置为 $2 \times 2$，对于一个三词序列，q_offset 数组可能记录 [0, 2, 0, ..., 10]，k_offset 数组记录 [0, 2, 8, ..., 10]，然后可以计算出输出矩阵中对应的偏移量（例如 [0, 0, 2, ..., 14]）。GPU 根据这些偏移量，仅访问和处理有效的迷你块对，从而避免了大部分冗余计算。

对于 $QK^T \times V$ 函数： 如图 (原文 Figure 8) 所示，MBIX 也采用类似的思想，对 $QK^T \times V$ 函数进行 2-D 分区，构建相应的索引数组来指导有效迷你块的计算和结果写入。

MBIX 的开销与优化： MBIX 虽然能消除大量冗余计算，但引入了新的开销，例如索引数组的构建时间、不规则数据访问和额外的内存占用。迷你块的大小是一个关键参数，用于权衡冗余计算与新引入开销。

• 迷你块过小： 索引数组会非常大，导致构建时间长，内存消耗多，并可能带来较差的空间局部性 (spatial locality)。
• 迷你块过大： 虽然索引开销小，但适应不规则数据结构的能力差，可能留下更多冗余计算。 本文经验性地选择 $32 \times 32$ 作为 2-D 分区的迷你块大小。

4.2.2. 共享内存块转置 (Shared Memory Block Transpose, SMBT)

为了减少内存访问开销，SMBT 利用 GPU 的共享内存 (shared memory)。在处理一对迷你块时，数据首先从全局内存 (global memory) 预取到共享内存中。

如图 (原文 Figure 9) 所示：

该图像是示意图，左侧展示了直接复制的数组布局，右侧展示了转置复制的数组布局。红线指示了复制的方向。

• 左侧部分展示了共享内存块转置：迷你块数据从全局内存复制到共享内存时，进行转置操作。

• 右侧部分展示了直接复制到共享内存可能导致的问题：当 MBIX 中每个线程负责两个向量的点积时，直接复制可能导致同一 warp 中的多个线程访问同一个内存 Bank，从而产生 Bank 冲突 (bank conflict)，严重降低内存带宽效率。

  SMBT 的目的是通过调整共享内存中的数据布局来解决 Bank 冲突。通过在数据从全局内存迁移到共享内存时进行转置，可以确保同一 warp 中的线程访问不同的内存 Bank，从而消除 Bank 冲突，提高共享内存的访问效率。

4.2.3. 高效原子操作 (Efficient Atomic Operation, EAOP)

在 2-D 分区中，多个迷你块对的结果可能需要写入到同一个位置，因此需要原子操作 (atomic operation) 来保证正确性。然而，原子操作的竞争问题会降低性能。EAOP 技术旨在减少这种竞争开销。

对于 $Q \times K^T$ 函数： 当查询 (query) 和键 (key) 的形状相同时，$Q \times K^T$ 函数的 MBIX 可以退化为更特殊的 1-D 分区。如图 (原文 Figure 10) 所示，如果块大小设置为 $1 \times 4$（即 1-D 分区），输出矩阵中的一个位置只对应一对迷你块的计算。

该图像是一个示意图，展示了 $Q imes K^T$ 操作的 1-D 划分。左侧矩阵和中间绿色矩阵的乘法结果映射到右侧矩阵，同时下方给出了 Q、K 和 QK 的偏移量。

在这种情况下，可以避免原子操作。因此，本文经验性地选择 $64 \times 32$ 作为实现中的块大小，以利用 1-D 分区带来的原子操作消除优势。

对于 $QK^T \times V$ 函数： 由于 $QK^T \times V$ 函数的结构不同，通常仍需要原子操作。EAOP 通过以下方式减少其开销：

• 重组偏移量生成： 调整偏移量生成的顺序。例如，在 Figure 8 中，第一行的蓝色矩阵的迷你块与第一列的黄色矩阵的迷你块相乘，并写入相同位置。通过调整偏移量的顺序（例如 [0,12,2,14,0,12,2,14] 和 [0,0,8,8,2,2,10,10]），可以增大两个依赖值写入之间的间隔。
• 增大写入间隔： 通过在其他维度（如批次维度）插入循环，可以进一步增大依赖值写入的间隔，从而更好地缓解竞争问题，提高原子操作的效率。

4.3. 词累积策略 (Word-accumulation Strategy)

词累积策略是针对 MLP 模块的优化，它不是本文首创，而是现有框架（如 EffectiveTransformer、FasterTransformer 和 TurboTransformer）中普遍采用的有效方法。

核心思想： 如图 (原文 Figure 11 的 Dense Arrangement 部分) 所示，该策略将批次中所有序列的有效词紧密地排列在一起，形成一个没有填充的密集矩阵。

该图像是示意图，展示了不同的数据排列方式，包括密集排列（Dense Arrangement）、块排列（Block Padding）和批处理填充（Batch Padding）。这些布局方法旨在优化输入序列的处理效率，减少冗余计算。

• 消除冗余： 通过移除序列间的填充，MLP 模块的计算只作用于有效词，从而彻底消除了填充带来的冗余计算。
• 利用 GEMM 库： 这种密集排列的矩阵可以直接利用现有的通用矩阵乘法 (General Matrix Multiplication, GEMM) 库（如 cuBLAS 中的 cublasLtMatmul 和 cublasGemmEx）进行高效计算。

与自注意力模块的连接问题： 虽然词累积策略对 MLP 模块非常有效，但它产生了与自注意力模块所需的批次填充布局不兼容的数据布局。传统解决方案需要在这两种布局之间进行数据移动和转换（即 EFF-rebuild 方案中的移除填充和重建填充）。

4.4. 块组织策略 (Block-organized Strategy)

块组织策略旨在更好地连接细粒度策略（针对自注意力模块）和词累积策略（针对 MLP 模块），减少数据布局切换的开销，并从整个模型的角度统一冗余计算消除方法。

4.4.1. 块填充数据布局 (Block Padding Layout)

如图 (原文 Figure 11) 所示，本文提出了一种新的数据布局：块填充 (Block Padding)。

该图像是示意图，展示了不同的数据排列方式，包括密集排列（Dense Arrangement）、块排列（Block Padding）和批处理填充（Batch Padding）。这些布局方法旨在优化输入序列的处理效率，减少冗余计算。

• 动机： 由于细粒度策略只计算有效的迷你块，传统的批次填充 (Batch Padding) 布局中许多填充区域将不再参与计算，变得低效。
• 设计： 块填充布局只在块的粒度上进行必要的填充。细粒度策略将有效值分割成块，如果有效值的长度不能被块大小整除，则边缘的块需要填充。因此，块填充只填充边缘的块，从而比批次填充大大减少了填充区域。
• 内存效率： 相比于批次填充，块填充以更密集的方式排列数据，可以减少峰值内存需求。

4.4.2. 定制化布局切换 (Customized Layout Switch)

尽管块填充布局比批次填充布局更密集，但它与词累积策略的密集排列布局仍然不同，因此仍需要数据移动进行兼容。块组织策略设计了定制化的布局切换方法来处理这种转换：

• 减少数据移动开销： 由于块填充的冗余填充区域远少于批次填充，从块填充布局切换到密集排列布局（以及反向操作）所需的数据移动量也大大减少。

• 借鉴 FasterTransformer： 布局切换函数通常与转置操作融合，以减少内核启动和数据移动开销。本文主要参考了 FasterTransformer 的方法。

• 块粒度索引： 本文的布局切换逻辑通过构建一个索引数组来记录块填充矩阵的目标偏移量。与 FasterTransformer 中以词的粒度构建索引不同，本文以块的粒度 (grain of block) 构建索引数组，使得索引数组更小，并且每个 warp 负责移动一个块，从而提高效率。

• 偏移量生成函数： 块组织策略还需要定制偏移量生成函数。这通常是一个串行累积工作，实现相对简单。

  通过块组织策略，整个模型的数据流可以更有效地在 MLP 模块（使用词累积策略的密集排列）和自注意力模块（使用细粒度策略和块填充布局）之间进行切换，同时最小化了数据移动和冗余计算的开销。

---

5. 实验设置

本节详细介绍论文的实验设置，包括数据集、评估指标和对比基线。

5.1. 数据集

实验使用了来自流行 GLUE 基准测试 (General Language Understanding Evaluation benchmark) 的八个语料库 (corpora)。GLUE 是一个广泛用于自然语言理解任务的基准测试集。

以下是原文 Table 3 中关于这八个语料库的统计结果：

数据集特点和选择理由：

• 来源： GLUE benchmark 是 NLP 领域广泛接受和使用的标准数据集集合，确保了实验结果的普遍性和可比性。

• 规模与多样性： 包含了不同规模（如 CoLA 的 8544 条到 MNLI 的 391168 条）和不同平均序列长度 (Mean) 的语料库，这使得实验能够全面评估解决方案在各种实际场景下的性能。

• 重尾分布： 论文在 Figure 2 中展示了这些语料库的序列长度均遵循重尾分布，证实了其研究动机的合理性。

  该图像是图表，展示了GLUE基准中各语料的长度统计，X轴为序列长度区间，Y轴为百分比，各条线代表不同数据集的分布情况。

• 最大序列长度 (Seq_Len)： 论文为每个语料库设定了 Bert 模型能够处理的最大序列长度 seq_len。这个长度是根据语料库的统计结果（去除极端异常值）来确定的，它决定了填充策略下的计算量，也是冗余计算的主要来源。例如，CoLA 的 seq_len 为 256，SST-2 和 WNLI 为 512，最长的 QQP、QNLI 和 RTE 为 1536。

5.2. 评估指标

论文主要使用延迟 (Latency) 作为性能评估指标。

5.2.1. 延迟 (Latency)

• 概念定义： 延迟是指从发出请求到收到响应之间经过的时间。在模型推理中，它通常指处理一个或一批输入样本所需的总时间。较低的延迟表示模型推理速度更快，响应更及时。
• 数学公式： 延迟通常直接以时间单位（如毫秒 ms）来测量，没有一个统一的数学公式。它通过计时器测量代码执行时间得到。
• 符号解释： 无特定符号。 论文中还提到了延迟降低 (Latency Reduction)，其计算公式为： $$\text{Latency Reduction} = \frac{\text{Baseline Latency} - \text{Our Solution Latency}}{\text{Baseline Latency}} \times 100\%$$ 其中：
• $\text{Baseline Latency}$ 是基线方法测得的延迟。
• $\text{Our Solution Latency}$ 是本文提出的解决方案测得的延迟。

5.3. 对比基线

论文针对不同的评估目标，选择了不同的基线模型和实现进行比较。

5.3.1. 细粒度策略的模块评估

• 目标： 评估自注意力模块（$Q \times K^T$ 和 $QK^T \times V$ 函数）的延迟。
• 基线： 使用 cuBLAS 库中的 cublasGemmStridedBatchedEx 内核。
  • 代表性： FasterTransformer 和 TurboTransformer 等流行框架都采用此 cuBLAS 内核来实现自注意力模块。
  • 配置： 论文选择了 FasterTransformer 中使用的 CUBLAS_GEMM_ALGO0 用于 $Q \times K^T$ 函数，CUBLAS_GEMM_ALGO1 用于 $QK^T \times V$ 函数。

5.3.2. 块组织策略的模块评估

• 目标： 评估数据布局切换（密集排列 <-> 块填充布局）的性能。
• 基线：
  • Fast： FasterTransformer 中实现的密集排列 <-> 批次填充布局切换方法。该方法通过构建有效词的前缀和数组来减少冗余数据移动。
  • Turbo： TurboTransformer 中实现的密集排列 <-> 批次填充布局切换方法。该方法在内核函数中计算偏移量，并通过条件语句判断是否移动值。

5.3.3. 统一解决方案的整体评估

• 目标： 评估整个 Transformer 模型推理的延迟。
• 基线： 采用 FasterTransformer 中实现的 EFF-rebuild 解决方案。
  • 特点： 该基线已通过词累积策略消除了 MLP 模块的冗余计算。
  • 本文解决方案的集成： 在此基线基础上，本文的统一解决方案进一步替换了 cuBLAS 内核为本文的细粒度内核，开发了新的偏移量生成函数，在自注意力模块中使用了块组织填充，并采用了新的数据布局切换方法。

5.3.4. 硬件与软件环境

• GPU： Tesla V100 GPU (16GB)

• CPU： Intel Xeon Silver 4208 CPU

• 操作系统： Ubuntu 18.04

• 编译器： GCC 7.5.0

• CUDA 版本： CUDA 11.3

• 模型配置： Bert-base 和 Bert-large，但主要结果以 Bert-base 为主。

• 实验重复： 每个测试用例重复 3 次，取平均结果以确保可靠性和有效性。

  ---

6. 实验结果与分析

本节将详细分析论文的实验结果，包括细粒度策略、块组织策略的模块评估以及统一解决方案的整体性能评估。

6.1. Bert 模型推理时间分解

论文首先给出了 Bert 模型推理时间的分解，以展示 MLP 和自注意力模块在不同优化策略下的耗时占比。这里使用了 Bert-base 配置和 SST-2 语料库，最大序列长度 seq_len 为 512，批次大小 batch_size 为 8。

以下是原文 Table 2 的结果：

其中：

• $GEMM(3)$：MLP 模块中用于生成 query、key 和 value 的三个线性函数。
• $GEMM(2)$：MLP 模块中用于从自注意力机制学习知识的两个线性函数。
• BatchedGEMM：自注意力模块中的批处理 GEMM（即 $Q \times K^T$ 和 $QK^T \times V$ 函数）。
• Padding 行表示使用批处理和填充策略的总耗时。
• No Padding 行表示使用 EFF-rebuild 方法（即 MLP 模块消除填充）后的耗时。

分析：

1. 初始填充策略下： MLP 模块 ($GEMM(3)$ + $GEMM(2)$) 占据了绝大部分时间 (23.6% + 58.1% = 81.7%)，而自注意力模块 (BatchedGEMM) 仅占 18.2%。这意味着在未优化 MLP 的情况下，优化自注意力模块带来的整体收益有限。
2. EFF-rebuild 方案下 (No Padding 行)： MLP 模块的计算时间大幅减少 ($GEMM(3)$ 从 20.4 降至 3.24，$GEMM(2)$ 从 50.16 降至 6.84)。此时，自注意力模块 (BatchedGEMM) 的耗时 (15.26) 变得突出，占据了总计算时间的 60.2%。这有力地证明了在 EFF-rebuild 方案之后，消除自注意力模块中的冗余计算变得非常有前景和必要。

6.2. 细粒度策略的模块评估

本部分评估细粒度策略在处理 8 个语料库时，自注意力模块（$Q \times K^T$ 和 $QK^T \times V$ 函数）的延迟。基线是使用 cuBLAS 的 cublasGemmStridedBatchedEx 内核。

以下是原文 Figure 12 的结果：

该图像是条形图，展示了在Bert-large配置下，计算 $Q \times K^T$ 和 $QK^T \times V$ 函数的平均延迟。图中包含四个自然语言处理任务（CoLA、SST-2、WNLI、SST-B）的数据，分别显示了不同序列长度（4、8、16、32）下，使用 cublas 和 FineGrained 两种策略的延迟对比。可以看出，FineGrained策略在多个任务中显著降低了延迟。

图表展示了 Bert-base (a) 和 Bert-large (b) 配置下，cuBLAS 策略和 FineGrained 策略的平均延迟。

分析：

1. 偏移量生成开销： 细粒度策略引入的索引数组构建时间被排除在延迟统计之外，因为它只占总执行时间的一小部分（约 5%），并且可以与词嵌入模块并行或隐藏。
2. cuBLAS 策略： 对于 cuBLAS 实现，具有相同 seq_len 的语料库表现出非常相似的延迟。当 seq_len 翻倍时，延迟会急剧增加，因为自注意力模块的计算量与 seq_len 呈平方关系。
3. 细粒度策略的优势：
   • 细粒度内核的延迟主要取决于语料库的平均序列长度 (mean length)，而不是最大序列长度 seq_len。这意味着它能有效避免填充带来的冗余计算。
   • 在大多数语料库上，细粒度策略显著降低了延迟。例如，在 QQP 语料库上，延迟降低高达 87.8%。
   • 在 8 个 GLUE 语料库上，自注意力模块的延迟平均降低了 63.9%。
4. 负优化情况： 对于 WNLI 语料库，在某些情况下（例如 Bert-base 配置，批次大小 4），细粒度内核的延迟甚至可能超过 cuBLAS 内核，表现出负优化。这可能是因为 WNLI 的平均序列长度相对较大 (148.5)，使得冗余计算减少带来的收益不足以弥补细粒度调度和数据访问引入的额外开销。
5. 模型规模和批次大小的影响：
   • 从 Figure 12(a) 到 Figure 12(b)，Bert-large 配置下的延迟增加，但细粒度策略与 cuBLAS 策略之间的性能比保持稳定，表明该优化方法与模型头数（Bert-large 主要增加头数）关系不大。
   • 增加批次大小通常会提高两种策略的性能，因为可以更好地饱和硬件。但随着批次大小的进一步增加，性能提升的趋势逐渐减弱。

6.3. 块组织策略的模块评估

本部分评估块组织策略中的数据布局切换性能。Ours 表示本文实现的从密集排列到块填充布局的切换。Fast 表示 FasterTransformer 中实现的密集排列到批次填充布局的切换。Turbo 表示 TurboTransformer 中实现的密集排列到批次填充布局的切换。

以下是原文 Figure 13 的结果：

该图像是一个图表，展示了不同策略（Ours、Fast 和 Turbo）在多个数据集（CoLA、SST-2、WNLI、SST-B）上所需时间的比较。横轴为序列长度，纵轴为时间（毫秒）。每个图显示了在不同批大小（4、8、16、32）下的性能表现。

图表展示了在不同批次大小（4、8、16、32）下，三种布局切换策略在四个语料库（CoLA、SST-2、WNLI、SST-B）上的耗时。

分析：

1. TurboTransformer 的性能：
   • TurboTransformer 在几乎所有情况下表现最差。
   • 原因在于其在内核函数中计算偏移量，并包含条件语句来判断值是否冗余，导致性能主要取决于预设的最大序列长度 seq_len，而不是实际的平均序列长度。因此，具有相同 seq_len 的语料库表现出非常相似的切换时间。
   • 但在某些情况下（例如 CoLA 语料库，批次大小为 4），Turbo 表现最佳。这归因于：Turbo 启动了更多的 GPU warp，能在较小工作负载下饱和设备；且它没有索引构建开销。
2. Ours 和 Fast 的性能：
   • Ours 和 Fast 的性能主要取决于语料库的平均长度 (average length)，平均长度越长，所需时间越多。这是因为这两种方法都处理较少的冗余值，使性能与有效数据量成比例。
   • 本文方法 (Ours) 在数据布局切换方面略优于 FasterTransformer (Fast)。FasterTransformer 的实现中，每个 GPU warp 只负责移动一个有效词，导致需要更详细的前缀和索引构建和更多的串行计算。而本文方法以块粒度移动数据，减少了索引开销。
3. 批次大小的影响：
   • 本文方法 (Ours) 的切换时间随着批次大小的增加而缓慢增长。当工作负载较小时（例如批次大小从 16 增加到 32，对于 CoLA 和 SST-2），时间增加倍数较小（1.05x 和 1.16x）。但对于更大的工作负载（MNLI、QNLI、RTE），时间增加倍数更大（1.36x、1.4x、1.48x）。当批次大小持续增加（例如从 512 到 1024）时，时间开始翻倍。这表明在小工作负载下，GPU warp 数量可能不足以饱和设备。

6.4. 统一解决方案的整体评估

本部分评估本文提出的统一解决方案在整个 Bert-base 模型推理中的延迟。Fast 表示使用 FasterTransformer 的 EFF-rebuild 解决方案（已消除 MLP 模块冗余计算）作为基线。Ours 表示应用了本文统一解决方案后的性能。

以下是原文 Figure 14 的结果：

该图像是图表，展示了在不同数据集下，快速（Fast）和我们的方法（Ours）在BERT-base模型推理中的平均延迟（Latency）。图中包括了四个数据集：CoLA、SST-2、WNLI和SST-B，各使用不同的批次大小（4、8、16和32）进行比较。

图表展示了在不同批次大小（4、8、16、32）下，Fast 策略和 Ours 策略在四个语料库（CoLA、SST-2、WNLI、SST-B）上的平均延迟。

分析：

1. 基线延迟的依赖性： 与细粒度策略的模块评估不同，整体基线 (Fast) 的延迟不仅取决于 seq_len，还取决于语料库的平均长度，因为平均长度也会影响 MLP 模块的计算量。
2. 统一解决方案的整体性能提升：
   • 本文的统一解决方案 (Ours) 在大多数语料库上都取得了延迟改进。
   • 性能提升与序列集中在短长度区域的程度呈正相关（即平均序列长度与 seq_len 的差距越大，提升越明显）。
   • 在 QQP 语料库上，整体延迟降低高达 65%。
   • 在 8 个 GLUE 语料库上，本文的统一解决方案平均实现了 28.1% 的延迟降低，相较于 FasterTransformer 的基线实现。
3. 负优化情况： 对于 WNLI 语料库，由于其平均序列长度相对较大，出现了负优化（延迟反而增加）。这与细粒度策略模块评估中的观察一致，表明当冗余计算量相对较小时，细粒度策略引入的开销可能超过其带来的收益。
4. seq_len 的影响：
   • 本文的统一解决方案在 seq_len 较大的语料库上表现更好。例如，在细粒度策略评估中，QQP 和 QNLI 实现了超过 70% 的延迟降低，在整体评估中也达到了 50% 以上。
   • 原因在于：线性函数（MLP）的计算量与序列长度呈线性关系，而自注意力模块（$Q \times K^T$ 和 $QK^T \times V$）的计算量与序列长度呈平方关系。因此，在 seq_len 较大时，自注意力模块在总计算中的比重更大，使得其优化效果更为显著。
5. 硬件效率考量： 论文提到，细粒度策略的硬件效率并不高，仅为 cublasGemmStridedBatchedEx 的约 1/4。这意味着其主要性能提升来源于冗余计算的消除，而非纯粹的硬件利用率提升。未来可能仍有空间进一步优化细粒度内核的实现效率。

6.5. 总结

实验结果有力地证明了本文提出的统一解决方案的有效性。

• 细粒度策略通过消除自注意力模块中大量的冗余计算，显著提高了其性能，平均实现了 63.9% 的延迟降低。

• 块组织策略通过优化的数据布局和切换机制，减少了整个模型的数据移动开销。

• 综合应用这两种策略，统一解决方案在 Bert-base 模型推理中实现了平均 28.1% 的延迟降低，并在 QQP 等语料库上达到了高达 65% 的延迟降低。这表明，在处理重尾分布的变长输入时，针对自注意力模块的细粒度优化是至关重要的，尤其是在 MLP 模块已被优化的前提下。尽管在某些平均序列长度较大的情况下可能出现负优化，但对于大多数 NLP 任务的实际场景，该方案仍能带来显著的性能提升。

  ---

7. 总结与思考

7.1. 结论总结

本文深入探讨了 Transformer 模型在 GPU 上推理时，由于 NLP 任务输入序列长度的可变性及重尾分布特性，所导致的批处理中的大量冗余计算问题。为了解决这一挑战，论文提出了一套统一解决方案 (unified solution)。

该方案包含三大核心策略：

1. 细粒度策略 (Fine-grained Strategy)： 针对自注意力模块，通过 Mini-block Index (MBIX)、Shared Memory Block Transpose (SMBT) 和 Efficient Atomic Operation (EAOP) 技术，设计定制化 GPU 内核，以块矩阵乘法的方式仅计算有效的迷你块，从而消除 $Q \times K^T$ 和 $QK^T \times V$ 函数中的大部分冗余。

2. 词累积策略 (Word-accumulation Strategy)： 沿用现有技术，在 MLP 模块中将输入序列紧密排列，避免了该模块的冗余计算。

3. 块组织策略 (Block-organized Strategy)： 通过引入块填充布局 (block padding layout)，以块的粒度组织自注意力模块的数据，并设计定制化的布局切换方法，有效地连接了前两种策略，减少了数据布局转换的开销和内存需求。

   实验结果表明，在 GLUE 基准测试的八个语料库上，本文的解决方案取得了显著的性能提升：

• 自注意力模块的延迟平均降低了 63.9%。

• 整个 Bert-base 模型的推理延迟平均降低了 28.1%，在 QQP 语料库上甚至达到了 65% 的延迟降低。

  这些成果证明，通过对自注意力模块进行细粒度优化，并结合整体模型的数据布局管理，可以有效解决 Transformer 模型在 GPU 推理中的重尾输入问题，显著提升实际性能。

7.2. 局限性与未来工作

论文作者指出了以下局限性：

• 细粒度策略的硬件效率： 论文坦言，其细粒度策略的硬件效率并不高，与 cuBLAS 相比，大约只有 1/4 的硬件性能。这意味着当前的性能提升主要来源于冗余计算的消除，而非纯粹的底层硬件利用率优化。

• 负优化情况： 在某些平均序列长度相对较大的语料库（如 WNLI）上，细粒度策略可能出现负优化，即其引入的调度和数据访问开销抵消了冗余计算消除带来的收益。

  基于这些局限性，论文提出了一些未来可能的研究方向：

• 进一步优化细粒度内核的实现： 鉴于目前细粒度策略的硬件效率不高，未来可以深入研究定制化内核的底层实现，例如利用更多的 GPU 硬件特性（如 Tensor Cores），或更精细的线程块调度和内存访问模式，以提高其硬件利用率。

• 自适应策略： 针对不同输入序列长度分布的语料库，开发自适应的优化策略。例如，对于平均序列长度较长的语料库，可能需要权衡细粒度策略的开销，甚至考虑混合使用填充和非填充方法。

• 与其他优化技术的结合： 将本文的解决方案与模型剪枝 (pruning)、量化 (quantization) 等其他 Transformer 优化技术相结合，以探索更大的性能提升空间。

• 针对不同硬件平台的优化： 探索如何在其他类型的 GPU（例如消费级 GPU）或专用加速器上实现和优化这些策略。

7.3. 个人启发与批判

7.3.1. 个人启发

1. 问题定义的精准性： 论文从实际部署中遇到的“重尾分布”这一关键问题出发，并深入分析了其如何导致“冗余计算”这个性能瓶颈。这种从实际痛点出发，并进行量化分析来确定优化重点的方法，对指导研究方向非常有启发性。很多时候，表面上的性能瓶颈可能只是深层问题的一个表象。
2. 分层优化思想： 论文没有简单地“发明”一个全新的 Transformer，而是针对其核心组件（MLP、自注意力）进行了分层、精细的优化。它首先利用了现有针对 MLP 的有效策略，然后将优化重点转移到之前被忽略的自注意力模块。这种模块化、循序渐进的优化思路，在复杂系统优化中非常实用。
3. 细粒度控制的重要性： 尽管 cuBLAS 等库提供了高度优化的通用 GEMM，但面对特定模式（如变长输入、稀疏计算），定制化、细粒度的内核仍能带来巨大收益。这提醒我们，在追求极致性能时，理解底层硬件和计算模式，并进行精细的软件-硬件协同设计是不可或缺的。
4. 数据布局的决定性作用： 块组织策略强调了数据布局在不同计算模块之间衔接的重要性。优化数据布局不仅能减少数据移动，还能更好地适配计算模式，是高性能计算中的一个核心考量。

7.3.2. 批判

1. 负优化场景的深入分析不足： 论文提到了 WNLI 语料库等可能出现负优化的情况，并归因于平均序列长度相对较大。然而，可以进一步深入分析导致负优化的具体原因，例如：

   • 细粒度策略的内核启动开销、索引查找开销和不规则内存访问模式，在冗余度不高时，是否确实超过了节省的计算量？
   • 是否存在一个临界点，当平均序列长度与最大序列长度之比高于某个阈值时，细粒度策略的优势就不再明显，甚至转为劣势？
   • 对于这些场景，是否有混合策略，例如根据输入批次的长度分布动态选择使用细粒度策略还是 cuBLAS？

2. 硬件效率的提升空间： 论文承认其细粒度内核的硬件效率仅为 cuBLAS 的 1/4，这暗示了其优化主要依赖于消除冗余而非提升计算密度。未来的工作可以更深入地探索如何同时提高细粒度内核的硬件效率，例如通过更复杂的 CUDA 编程技巧、Tensor Core 利用、或更优的线程块/warp 调度策略。如果能够将硬件效率也提升上来，那么即使是平均序列长度较大的场景也能获得收益。

3. 泛化性与维护成本： 定制化内核通常意味着更高的开发和维护成本。虽然本文的解决方案在 GLUE 基准上表现出色，但其在其他 Transformer 架构（如 GPT 系列）或不同 NLP 任务上的泛化性如何？这些定制化内核是否易于集成到不同的深度学习框架中？这些都是实际部署中需要考虑的问题。

4. 块大小选择的经验性： 论文中提到了经验性地选择 $32 \times 32$ 或 $64 \times 32$ 作为块大小。虽然这在验证概念时是合理的，但对于实际生产环境，可能需要更系统的自动化搜索或基于模型-硬件特性的自适应块大小选择机制。

   总的来说，这篇论文提供了一个非常扎实的针对 Transformer 推理中重尾输入问题的解决方案，特别是在自注意力模块的优化上提供了新的思路。它在强调了冗余计算是主要瓶颈的同时，也为未来在硬件效率和自适应优化方面的工作指明了方向。