LLM Inference Overview

这个网站主要介绍大规模语言模型（LLM, Large Language Model）推理中的一些系统优化技术，涵盖并行化、CUDA 优化、内存优化等方面的内容。以下是网站的主要内容概览：

Transformer 架构: 介绍 Transformer 模型的基本结构和工作原理，这是当前大多数 LLM 的基础架构。
CUDA 优化和 Triton：介绍 GPU 硬件架构和 CUDA 编程，以及如何通过优化 CUDA 代码和使用 Triton 框架来提升 LLM 推理的性能。
并行化技术：探讨在 LLM 推理中常用的并行化策略，如数据并行、模型并行等。
KV Cache 优化：介绍如何高效管理和优化内存使用，以支持大规模模型的推理。
Attention 优化：深入分析 Transformer 模型中的 Attention 机制，并介绍相关的优化技术。

Chapter 1: Attention & Transformers

Attention & Self-Attention

人类在处理信息时会“有选择地关注”最相关的部分，比如看一张图时注意某个局部，读文章时注意关键句子。

基本思想：给输入信息的不同部分分配不同的权重，然后对加权后的信息进行整合。

最常见的是加性注意力或点积注意力，核心是 $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$

$Q$: 当前需要处理的对象
$K$: 输入序列中的所有元素的键
$V$: 输入序列中的所有元素的值
$softmax(QK^{T})$: 表示“Query 与 Key 的相似度”，越相关权重越大

Self-Attention 实际上是 Attention 机制的一种特殊情况，特点是 Q、K、V 都来自同一个序列。

作用：序列中的每个元素都能和序列中的所有其他元素计算相关性，从而获取全局上下文信息。

$$\text{SelfAttention}(X) = text{softmax}\left(\frac{XW_Q(XW_K)^T}{\sqrt{d_k}}\right)XW_V$$ $X$ 来自于同一个序列

Why Transformer use Self-Attention?

传统的序列模型处理方式

在传统的序列处理模型（如 RNN、LSTM 和 GRU）中，模型是按顺序逐个处理序列中的元素（例如单词或字符），并且每个元素的处理依赖于前一个元素的隐藏状态。

:warning: 这种方法在处理长序列时会面临梯度消失或梯度爆炸的问题，导致模型难以捕捉长距离的依赖关系。

Transformer 的计算方式对 GPU 更友好，能更好地利用并行计算资源
- RNN 需要按时间步顺序处理，难以并行化
- Transformer 的 Self-Attention 机制允许同时处理序列中的所有位置，充分利用 GPU 的并行计算能力

Transformers 结构

输入 X 与 attention_mask 的 shape

输入 X 一般形状为[batch_size, seq_len, embedding_dim]
attention_mask 是经过 tokenizer 后返回的 mask 结果，表示哪些样本需要忽略形状一般是[batch_size, num_heads, seq_len]

为什么用 Transformer 中用 layer norm

特性	Batch Norm	Layer Norm	RMSNorm
标准化维度	小批量内各特征维度	每个样本的所有特征维度	每个样本的特征维度的均方根
计算开销	中等	较大	较小
对小批量大小依赖	依赖	不依赖	不依赖
应用场景	CNN、MLP	RNN、Transformer	各类神经网络，尤其在计算效率和稳定性要求高的任务中
正则化效果	有一定正则化效果	无显著正则化效果	无显著正则化效果

列长度的灵活性： Transformer 处理的是序列数据，序列长度可能因输入样本而异。LayerNorm 对每个样本自身的一层神经元的输入进行归一化，与其他样本的序列长度无关，能够很好地处理不同长度的输入序列。而 batch norm 对长度大小不同的 NLP 任务计算的超参数泛化能力差。
并行计算的适应性： Transformer 的多头注意力机制高度并行化，LayerNorm 只需要对单个样本的一层进行计算，不需要等待其他样本的信息，因此更适合并行计算环境。
模型的稳定性： LayerNorm 基于每一层自身的输入进行归一化，能够更好地控制每一层输入的范围和分布，避免梯度消失或梯度爆炸问题。

post-norm 与 pre-norm

原始的 transformer 中使用的是 post-norm，而 llm 中大多使用 pre-norm

norm 位置	优点	缺点
pre-norm	训练稳定：在残差连接之前进行归一化，可以有效缓解梯度消失或爆炸的问题，使深层网络的训练更加稳定。收敛速度快：梯度能够更直接地传递到前面的层，从而加快模型的整体收敛速度。减少调参工作：不需要像 Post-Norm 那样依赖复杂的学习率预热等优化技巧	潜在的表示塌陷问题：靠近输出位置的层可能会变得非常相似，从而对模型的贡献变小，限制了模型的上限。可能削弱层的贡献：由于先进行了归一化，可能会减弱每一层的实际贡献，导致模型的有效深度变浅
post-norm	保留输入特征：更接近原始输入的特征，有助于信息的传递。潜在性能优势：虽然训练不稳定，但有研究暗示其在效果上可能有更大的潜力	训练不稳定：在深层模型中，梯度容易爆炸或消失，对学习率和权重初始化非常敏感，收敛困难。依赖优化技巧：需要使用学习率预热等复杂的优化方法来稳定训练

Chapter 2: Transformer-based LLM

现在的大模型基本是 Transformer-based 的自回归预训练模型，本章我们将以 llama2 以及 GPT-3 为例介绍这种模型的基本结构。

Positional Encoding: RoPE
Attention: Multi-Head Attention or Multi-Query Attention or Grouped-Query Attention or Multi-Head Latent Attention or DeepSeek Sparse Attention
Normalization: LayerNorm or RMSNorm
FFN: MLP or SwiGLU or Mixture of Experts

Positional Encoding

一般分为两大类：绝对位置编码（Absolute Position Encoding）和相对位置编码（Relative Position Encoding）。本节将介绍这两种位置编码的区别及其重要性，并重点解析一种结合了两者优点的创新方法——旋转位置编码（Rotary Position Embedding, RoPE）

Why we need Position Encoding

在自然语言处理领域，Transformer 模型已成为一项革命性的技术。然而，其核心的自注意力机制本身并不具备捕捉序列中单词顺序的能力，即“位置无关性”。为了解决这一问题，位置编码应运而生。

绝对位置编码（Absolute Position Encoding）

给序列中的每个位置一个唯一的向量，把位置信息直接加到 token embedding 上。

固定式：不用学习参数，能推广到比训练时更长的序列在最初的 Transformer 论文《Attention Is All You Need》中，作者使用正弦和余弦函数来生成这些位置编码。其数学表达式如下：

$$ \begin{aligned} PE_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)\ PE_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) \end{aligned} $$

可学习式：模型可以学到适合任务的位置信息。训练时没见过的长序列可能无法泛化 $$ PE_{pos} = \text{Embedding}(pos) $$

相对位置编码（Relative Position Encoding）

相对位置编码是根据单词之间的相对位置关系来计算位置编码。这种编码方式更加灵活，能够捕捉到不同单词之间的相对位置信息，有助于模型更好地理解序列中单词之间的关系。但是也有缺点，计算效率低下，同时大部分相对编码都没有落地可行性。

Rotary Position Embedding (RoPE)

将位置编码与词向量通过旋转矩阵相乘，使得词向量不仅包含词汇的语义信息，还融入了位置信息

$$ (R_mq)^T(R_nk) = q^TR_m^TR_nk = q^TR_{m-n}k $$

给位置为 m 的向量 q 乘上矩阵$R_m$、位置为 n 的向量 k 乘上矩阵$R_n$用变换后的 Q,K 序列做 Attention，Attention 就自动包含相对位置信息

相对位置感知：使用绝对位置编码来达到相对位置编码的效果，RoPE 能够自然地捕捉词汇之间的相对位置关系。
无需额外的计算：位置编码与词向量的结合在计算上是高效的。
适应不同长度的序列：RoPE 可以灵活处理不同长度的输入序列。

目的

我们假设通过下述运算来给 q,k 添加绝对位置信息： [ \tilde{q}_m = f(q, m), \quad \tilde{k}_n = f(k, n) \tag{1} ]

也就是说，我们分别为 q,k 设计操作$f(\cdot,m),f(\cdot,n)$，使得经过该操作后，$\tilde{q}_m,\tilde{k}_n$就带有了位置 m,n 的绝对位置信息。Attention 的核心运算是内积，所以我们希望的内积的结果带有相对位置信息，因此假设存在恒等关系：

$$ ⟨f(q,m),f(k,n)⟩=g(q,k,m−n)\tag{2} $$

所以我们要求出该恒等式的一个（尽可能简单的）解。求解过程还需要一些初始条件，显然我们可以合理地设 $f(q,0)=q$ 和 $f(k,0)=k$

实现

位置 m 的编码进行解方程，我们得到二维情况下用复数表示的 RoPE：

$$ f(q, m) = R*f(q, m) e^{i \theta f(q,m)} = |q| e^{i(\Theta(q) + m\theta)} = qe^{im\theta} \tag{3} $$

矩阵形式：

$$ f(q, m) = \begin{pmatrix} \cos (m\theta) & -\sin (m\theta) \ \sin (m\theta) & \cos (m\theta) \end{pmatrix} \begin{pmatrix} q_0 \ q_1 \end{pmatrix} \tag{4} $$

由于内积满足线性叠加性，因此任意偶数维的 RoPE，我们都可以表示为二维情形的拼接，即

:warning: 由于$R_m$具有稀疏性，不建议使用 matmul 进行实现，建议使用下面的方式实现：其中$\odot$是逐位对应相乘，即 Numpy、Tensorflow 等计算框架中的 ∗ 运算

Attention Series

在 MHA (Multi Head Attention) 中，每个头有自己单独的 key-value 对；标准的多头注意力机制，h 个 Query、Key 和 Value 矩阵。
在 MQA (Multi Query Attention) 中只会有一组 key-value 对；多查询注意力的一种变体，也是用于自回归解码的一种注意力机制。与 MHA 不同的是，MQA 让所有的头之间共享同一份 Key 和 Value 矩阵，每个头只单独保留了一份 Query 参数，从而大大减少 Key 和 Value 矩阵的参数量。
在 GQA (Grouped Query Attention) 中，会对 attention 进行分组操作，query 被分为 N 组，每个组共享一个 Key 和 Value 矩阵 GQA 将查询头分成 G 组，每个组共享一个 Key 和 Value 矩阵。
- GQA-G 是指具有 G 组的 grouped-query attention。
  
  GQA-1 具有单个组，因此具有单个 Key 和 Value，等效于 MQA。而 GQA-H 具有与头数相等的组，等效于 MHA。

Multi-Head Latent Attention(MLA)

MLA 通过低秩联合压缩技术，减少了推理时的键值（KV）缓存，从而在保持性能的同时显著降低了内存占用

公式

在训练阶段，除了多了一步低秩投影以及只在部分维度加 RoPE 外，MLA 与 Q、K 的计算与 MHA 是基本相同的

RoPE 的$R_m$计算后投影矩阵与位置相关，为了解决这个问题，对 Q 和 K 的低秩投影分为两部分，一部分是原始的投影矩阵，另一部分是与位置相关的投影矩阵
在推理阶段，MLA 的计算可以等效为一个 MQA

投影矩阵吸收

在推理阶段，我们利用

$$ q_t^{(s)}k_i^{(s)T} = (x_tW^{(s)}_q)(c_iW^{(s)}_k)^T = x_t(W^{(s)}_qW^{(s)T}_k)c_i^T $$

将 $W^{(s)}_qW^{(s)T}_k$ 合并起来作为 Q 的投影矩阵，那么 $c_i$ 则取代了原本的 $k_i$，同理，在 $o_t$ 后面我们还有一个投影矩阵，于是 $v^(s)_i=c_iW^{(s)}_v$ 的 $W^{(s)}_v$ 也可以吸收到后面的投影矩阵中去，也就是说此时 KV Cache 只需要存下所有的 $c_i$ 就行，而不至于存下所有的 $k^{(s)}_i、v^{(s)}_i$。注意到 $c_i$跟$^{(s)}$无关，也就是说是所有头共享的，即 MLA 在推理阶段它可以恒等变换为一个 MQA。

RoPE 解耦

为 RoPE 是一个跟位置相关分块对角矩阵$R_m$，满足$R_mR^⊤_n=R_{m−n}$，MLA 加入 RoPE 之后会让固定的投影矩阵与位置相关：

$$ q^{(s)}_i=x_iW^{(s)}_qR_i,k^{(s)}_i=x_iW^{(s)}_kR_i $$

$$ q^{(s)}_tk^{(s)T}_i=(x_tW^{(s)}_q)(c_iW^{(s)}_k)^T = x_t(W^{(s)}_qR_{t-i} W^{(s)T}_k)c_i^T $$

这里的 $W^{(s)}qR{t-i} W^{(s)T}_k$ 就无法合并为一个固定的投影矩阵了（跟位置差 $t−i$ 相关），从而 MLA 的想法无法结合 RoPE 实现。

每个 Attention Head 的 Q、K 新增 $d_r$ 个维度用来添加 RoPE，其中 K 新增的维度每个 Head 共享：

$$ o_t = \big[o_t^{(1)}, o_t^{(2)}, \cdots, o_t^{(h)} ,\big] $$

$$ o_t^{(s)} = Attention!\left(q_t^{(s)}, k_{\le t}^{(s)}, v_{\le t}^{(s)}\right) = \frac{\sum_{i \le t} \exp!\left(q_t^{(s)} k_i^{(s)\top}\right) v_i^{(s)}} {\sum_{i \le t} \exp!\left(q_t^{(s)} k_i^{(s)\top}\right)} $$

$$ q_i^{(s)} = [x_i W_{qc}^{(s)},x_i W_{qr}^{(s)}R_i] ,\quad k_i^{(s)} = [c_i W_{kc}^{(s)},x_i W_{kr}R_i] ,\quad v_i^{(s)} = c_i W_v^{(s)}, \quad c_i = x_i W_c $$

Normalization

LayerNorm: 对某个样本的所有特征维度进行归一化

$$ \text{LayerNorm}(x) = \frac{x - \mu}{\sigma} \cdot \gamma + \beta $$
RMSNorm: 简化版的 LayerNorm，它不减去均值，只基于平方均值 (Root Mean Square) 来归一化 $$ \text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}} \cdot w $$

Why RMSNorm?

Layer-Norm 和 RMS-Norm 在测试集效果上没有明显差异，基本持平
RMS-Norm 的计算效率要更高

Feed Forward Network (FFN)

MLP: 两层全连接网络，中间使用非线性激活函数（如 ReLU 或 SiLU） $$ \text{FFN}(x) = \text{ReLU}(0, xW_1)W_2 $$
SwiGLU: 使用 SiLU 激活函数的变体并增加门控机制 $$ \begin{aligned} \text{SwiGLU}(x) &= (\text{Swish}(xW_1);\odot;xW_2) \ \text{Swish}(x) &= x \cdot \sigma(x) \end{aligned} $$
Mixture of Experts (MoE): 多个专家网络的集合，每个输入样本通过一个路由器选择部分专家进行处理，从而提高模型的表达能力
- Switch Transformer：Top-1 gating，每个 token 只去 1 个专家。
- GShard / GLaM：Top-2 gating，每个 token 走 2 个专家，结果加权。

LLAMA2 模型结构

相较于 Transformer，llama2 使用了 pre-norm 结构
使用 RMSNorm 替代 LayerNorm，不计算样本均值
使用 Rotary Positional Encoding (RoPE)，用绝对编码的方式来实现相对位置编码
使用 GQA (Grouped-Query Attention)，平衡效率和性能
使用 SwiGLU 替代简单的 MLP，激活函数使用 SiLU

$$ \begin{aligned} \text{SwiGLU}(x) &= \text{Swish}(xW_1+b_1);\odot;(xW_2+b_2) \ \text{Swish}(x) &= x \cdot \sigma(x) \end{aligned} $$

参考资料

Chapter 3: Parallelization in LLM

本节主要介绍大模型训练中的并行化技术，涵盖数据并行、模型并行、流水线并行和张量并行等方法。我们将从 Transformer 的参数量、Flops 以及训练占用显存入手，分析为什么需要并行化技术，并介绍这些技术的基本原理。最后我们将以 Transformer-based LLM 为例，对里面 Attention 以及 FFN or MoE 采用的并行化手段进行分析。

Parallelization Concepts

本节主要介绍大模型训练中的并行化技术，涵盖数据并行、模型并行、流水线并行和张量并行等方法。我们将从 Transformer 的参数量、Flops 以及训练占用显存入手，分析为什么需要并行化技术，并介绍这些技术的基本原理。

下一节我们将以 Transformer-based LLM 为例，对里面 Attention 以及 FFN or MoE 采用的并行化手段进行分析。

Transformer 模型的参数量、计算量、显存占用 1

参数量

每个 transformer 层的参数量为 $$ \text{Params per layer} = 12h^2 + 13h $$
当隐藏维度 $h$ 较大时，可以忽略一次项，模型参数量近似为 $$ \text{Params per layer} \approx 12lh^2 $$

计算量(FLOPs)

l 层的 Transformer 一次训练迭代计算量： $$ \text{FLOPs per iteration} = l\cdot(24hsh^2+4bs^2h)+2bshV $$
- $l$: Transformer 层数
- $h$: 隐藏层维度
- $s$: 序列长度
- $b$: 批量大小
- $V$: 词汇表大小
当 $h$ 较大且$s \ll h$时，可以忽略低阶项，我们可以近似认为：在一次前向传递中，对于每个 token，每个模型参数，需要进行 2 次浮点数运算，即一次乘法法运算和一次加法运算
一次训练迭代包含了前向传递和后向传递，后向传递的计算量是前向传递的 2 倍。

一次训练迭代中，对于每个 token，每个模型参数，需要进行 6 次浮点数运算

显存占用

参数训练时的显存
- 训练时需要存储模型参数、梯度和优化器状态
- 每个可训练模型参数都会对应 1 个梯度，并对应 2 个优化器状态（Adam 优化器梯度的一阶动量和二阶动量），使用混合精度训练 $$ \text{Total Memory for Parameters} = 20 \times \text{Params num} $$
中间激活值的显存
- 对于 l 层 transformer 模型，中间激活占用的显存大小可以近似为 $$ \text{Activation Memory} = (34bsh+5hs^2a)\cdot l $$

通常会尝试减小批次大小来避免显存不足的问题，这种方式减少的其实是中间激活占用的显存，而不是模型参数、梯度和优化器的显存

KV Cache 的显存
- 设输入序列的长度为 $s$ ，输出序列的长度为 $n$ ，以 float16 来保存 KV cache，那么 KV cache 的峰值显存占用大小为 $$ b(s+n)lh \times 2 \times 2= 4blh(s+n) $$ 这里第一个 2 表示 K/V cache，第二个 2 表示 float16 占 2 个 bytes。

Why Parallelism?

Computation Bottleneck: LLM 的参数量巨大，以 GPT-3 为例，1750 亿参数，对 GPU 计算能力要求极高，单个 GPU 无法满足。(TP, PP, SP)

FLOPs 近似 $6 * 1750 * 10 ^9$FLOPs
Memory Bottleneck: LLM 的模型参数和中间激活值占用大量内存，单个 GPU 的显存有限，无法容纳整个模型。(TP, PP, ZeRO, SP)

参数训练时的显存近似 $20 * 1750 * 10 ^9$ bytes 远大于目前单个 GPU 的显存容量 中间激活值的显存占用取决于批次大小、序列长度和模型层数，通常也非常大。

Parallelism Overview 2

Before Parallelism

:smile:在分布式训练中，GPU 之间的通信必不可少，所以我们需要知道 NCCL 的基本通信原语³

Minimum Spanning Tree: HPC 通信常用，最小化通信 ROUND 数
Ring Algorithm: GPU 通信通常有极高的带宽，通信启动带来的延迟并不是主要 bottleneck

AllReduce: 将所有 GPU 的数据进行汇总，计算平均值，并将结果广播回所有 GPU，常用于梯度更新 Broadcast: 将一个 GPU 的数据发送到所有其他 GPU，常用于模型参数初始化

Reduce: 将所有 GPU 的数据汇总到一个 GPU 上，常用于收集结果 reduce

AllGather: 将所有 GPU 的数据收集到每个 GPU 上，常用于收集中间结果

ReduceScatter: 将数据先进行 Reduce 操作，然后将结果分散到各个 GPU 上，常用于分布式计算中的中间步骤

其中 AllReduce 可以看作是 ReduceScatter 和 AllGather 的组合

ZeRO-3 的优化基于此

Data Parallelism

Parameter Server: workers 存有副本保证容错，PS 收集梯度更新并广播保证一致性
All-reduce: 每个 GPU 计算梯度，然后通过 All-reduce 汇总并更新参数，没有容错，但是简单(工业界喜欢简单有效的方案)
DDP in PyTorch⁴

Model Parallelism

模型并行原因是因为模型权重过大，单个 GPU 无法容纳或者模型中间激活值过大不能在单个 GPU 上进行完整的一次训练迭代

一般来说，模型并行可以分为张量并行和流水线并行，以及对 MoE 模型的专家并行。

TP：通信载量很大，一般在小于 8 张 GPU 的场景下使用
PP：通信载量较小，模型分割比较独立

Pipeline Parallelism

把模型切分多个 stage，每个 device 负责一个 stage 的计算，类似 CPU 的 pipeline，问题变成了如何让 pipeline 的 bubbles 尽可能少

实际上 llm 推理和训练的 pipeline 并不完全一样，推理的 pipeline 更简单，因为不需要反向传播，而训练的话需要考虑反向传播和梯度更新。反向传播和梯度更新对前向传播有依赖

Device Placement

只对特殊的有分支的模型有效

Synchronous Pipeline Parallelism

Idea: 尽可能减少 Pipeline Bubbles
Pros: 保证收敛语义，训练过程和在单张卡上相同
Cons:
- compute: 流水线气泡
- compute: 为了消除流水线气泡，我们必须切分 input 变成更小的 micro-batches，但是 micro-batch 太小会影响硬件计算的效率($AI \space = \space #ops \space/ \space #bytes$)
  
  算法上其实对 global batch size 也有限制

GPipe 5

把 input 切分成多个 micro-batches，把 micro-batches 依次送入 pipeline。梯度计算如下：

$$ \nabla L_{\theta}(x) = \frac{1}{N} \sum_{i=1}^{N} \nabla L_{\theta}(x_i) $$

Memory Usage 问题： 存在一个不停增长内存的阶段直到峰值，因为每个 micro-batch 的中间激活值都需要存储直到开始反向传播

1F1B & Interleaved 1F1B 6

1F1B: 并没有降低 GPipe 的延迟，只是改变了调度的顺序，尽可能优先地执行反向传播
Interleaved 1F1B: 切分成更细粒度的 stage，每个 GPU 负责多个不邻近的 stage

有 16 transformer layers （L1 to L16），他们被均匀地分配到四个 devices 上 Device 1: [L1, L2], [L9, L10] Device 2: [L3, L4], [L11, L12] Device 3: [L5, L6], [L13, L14] Device 4: [L7, L8], [L15, L16]

这样的话，一个数据在执行了 device 4 的 L7, L8 之后，可以转到 device 1 继续 L9 L10；这也不耽误 device 4 继续接受来自 forward/backward 的计算任务了。

TeraPipe 7

GPipe 等 PP 在大规模扩展时面临的“流水线气泡”和“内存限制”两大瓶颈。

token-level 并行：专门针对 autoregressive 系列模型设计的 pipeline 并行，一个设备每处理完一小块数据(token-level)，就立刻将其发送给下一个设备

减小 activation 占用的显存： “激活重计算”（gradient Checkpointing） ，即在正向传播时只保存少量关键的激活值，在反向传播需要时再重新计算其他的。但这会增加计算开销，拖慢反向传播的速度。TeraPipe 对此进行了优化：

传统方式：在反向传播的主计算流中同步地进行重计算，导致主流程需要等待重计算完成。
TeraPipe 方式：将激活值的重计算任务放到一个**独立的、并行的计算流（Asynchronous Stream）**中执行。这样，主计算流在进行反向梯度计算的同时，另一个计算流在“后台”悄悄地准备好下一步需要的激活值。

Chimera 8

提出双向流水线 (bi-directional pipeline)，即同时在 forward 阶段和 backward 阶段并行运行不同 micro-batch(DeepSeek 使用类似的变体)

Zero Bubble Pipeline Parallelism 9

Motivation: backward 需要 forward 的两倍的时间开销，forward 只有一个 wx 乘法；而 backward 里面有两个乘法，将 backward 拆成对 Input 的 backward 和对 weight 的 backward 两个部分

当前 layer 的 backward 只依赖对上一个 layer 的 input 的 backward(B)，不依赖 weights 的 backward(W)
ZB-H1：device 1 的峰值内存与 1F1B 相比没有变化，其他 device 相对增加，但是 bubble 减少了 2/3
ZB-H2：device 1 的峰值增加到 2 * 1F1B，bubbles 接近为 0
- :warning: 这里实际上有一个问题，optimizer 一般需要全局同步，但是 ZB-H2 里每个 device 都有自己的 optimizer state，需要额外的机制来实现
- optimizer sync: 为了进行梯度裁剪，需要计算全局梯度范数；在混合精度训练中，需要对所有参数进行全局 NaN/INF 检查。
  
  当梯度的“整体大小”超过某个阈值时，就把它按比例缩小。
- 实际上触发梯度裁剪和 NaN/INF 检查的频率并不高，所以异步进行每个阶段的 gradients 检查
  
  这里在每个流水线阶段执行优化器步骤之前，该阶段会从前一个阶段接收一个部分归约（partially reduced）的全局状态，将其与本阶段的本地状态（local state）合并后，再传递给下一个阶段。如果最后发现需要裁剪梯度或者发现 NaN/INF，进行 rollback，丢弃此次更新即可。
ZB-V: ZB-H2 结合 Interleaved 1F1B，通过这种方式将集群峰值内存降低到 1F1B 的水平，同时保持接近零的流水线气泡
Comparison of different PP methods

DualPipe(DeepSeek-V3 Used)10

流水线稳定的标准：
- 双向管道都已“灌满”：
  
  第一个流水线（F0）的第一个前向块（黄色的 0）必须已经走完了全程，从 Device 0 到达了 Device 7。第二个流水线（F1）的第一个前向块（橙色的 10）也必须已经走完了全程，从 Device 7 到达了 Device 0。
- 双向反向传播都已开始：
  
  当 F0 的块 0 到达 Device 7 后，它的反向传播 B0（绿色的 0）就可以从 Device 7 开始，并向着 Device 0 的方向回传。当 F1 的块 10 到达 Device 0 后，它的反向传播 B1（比如 T=8 时 Device 4 上的橙色块 10）就可以从 Device 0 开始，并向着 Device 7 的方向回传。
Overlap 计算和通信¹¹
完整的 DualPipe 调度
:warning: 下面的图里，有两个 X，代表 bubble，即，device 2 需要等待 device 3 传递过来 14+0 的执行结果之后，才可以自己执行。

Asynchronous Pipeline Parallelism

Idea: 在反向传播完成之前开始下一阶段的前向传播
Pros: No pipeline bubbles
Cons:
- compute: 打破了同步训练语义(forward 和 backward 必须匹配同一份权重)，训练过程中可能有滞后 gradient
- memory: 算法会存储多个版本的 weights 来保证一致性

AMPNet¹²: 完全异步。每个设备在空闲时就执行前向传播，并在每次反向传播后更新权重。

Techniques:
- Task Queue: 每个计算单元维护一个本地队列，forward 和 backward 都被表示为任务。任务只有在依赖的数据 (activation 或梯度) 到达时，才会被触发
- AMPNet 引入了依赖计数 (dependency counter)，当所有依赖满足时，任务才能被调度执行。
Pros: No pipeline bubbles
Cons:
- :sob:只在本地更新权重，没有全局同步。小数据集上可以很好的泛化，但是大数据集上很难
- activation 缓存造成 memory overhead(forward 结果需要存储直到 backward)

Pipedream¹³: 每个 batch 存储多个权重版本(类似 MVCC)

Techniques:
- Weight Versioning:
  1. forward 时保存一份当前参数快照 (stash)。
  2. backward 时必须使用对应 forward 时的那份参数（确保一致性）。
  3. 更新是异步的，但保证 forward/backward 用的是同一版本的权重。
- 采用 1F1B (one forward, one backward) pipeline 调度。
:skull: No Memory saving compared to single device case, which is opposite to the purpose of model parrallelism

Pipedream-2BW: 只保留 2 个版本的权重

Techniques:
- 每个 stage 只保留两个版本的权重（current 和 previous）。不需要为每个微批次 stash weight，大幅降低显存占用。
- forward 总是用前一版本（stalled by 1 update）
- 调度：改进了 1F1B 调度，使得在 forward/backward 之间能更高效复用 weight。

RL-Based Partitioning Algorithm

将计算图分 stage 问题变成一个强化学习问题

Optimization-Based Partitioning Algorithm

讲求解划分问题变成一个整数线性规划问题

Tensor Parallelism

Parallelize One Operator

Manual: 枚举所有可能的划分方式，选择在你环境限制下最优的划分方式

Re-partition in Tensor Parallelism

不同的 op 并行策略要求 tensor 的形状不同，通信开销也不同

Parallelize All Operators in a Computation Graph

Problem Definition:

给定一个计算图，选择一个划分策略，使得 Node costs(计算开销+通信开销) + Edge costs(re-partition 通信开销) 最小

Model-specific Intra-op Parallel Strategies

AlexNet

Megaton-LM¹⁴

f 和 g 是共轭的。 f 在正向传播中是恒等算子，在反向传播中是所有 reduce 操作，而 g 在正向传播中是所有 reduce 操作，在反向传播中是恒等操作

Gshard-MoE¹⁵

除了 MoE 层分布式在多个设备上，其他层都 copy 到每个设备上
一次 forward 和 backward 分别需要 2 次 all-to-all 操作

路由分发 token 到不同的专家，将路由到自己专家的输入 gather

Systems for Intra-op Parallelism

ZeRO Optimizer¹⁶: 解决数据并行中内存瓶颈，将 gradients, optimizer states and model weights 分区并分发到不同设备上

这里 Optimizer States 使用了 Adam 优化器，需要为每个参数额外存储一阶动量 + 二阶动量 + 参数副本(fp32 copy)，ZeRO 论文中使用了 FP16 + FP32 混合精度训练，所以会存储两倍的 Optimizer States

ZeRO 2: 不仅将 optimizer states 分布到不同设备上，还将 gradients 分布到不同设备上

ZeRO 3: 将模型参数也分布到不同设备上，但是增加了通信开销

Automatic Parallelization Methods

Alpa¹⁷

Summary for Parallelism Concepts

Why we do not use automatic parallelization in LLM?

因为 LLM 现在的架构基本是 Transformer，Transformer 的计算图比较简单，手动设计的并行策略已经足够好，并且自动并行化方法的搜索空间和计算开销都比较大，收益不明显。

Communication in Different Parallelism

² | Parallelism Techniques | Communication Pattern | Data Volume Per Transfer | Total Transfer | Total Volume | Data Traffic | | :--------------------- | :-------------------- | :----------------------- | :------------- | :----------- | :----------- | | DP | AllReduce | 711.75 MB | 64 | 44.48 GB | 1.34% | | PP | P2P | 192 MB | 26 | 4.875 GB | 0.14% | | TP | AllReduce | 360 MB | 4992 | 1775 GB | 52.9% | | SP | AllGather | 180/360 MB | 4992/1664 | 1462.5 GB | 44.08% | | EP | AlltoAll | 10.5 MB | 4992 | 51.19 GB | 1.54% |

数据并行（Data Parallelism, DP）：
- 为了同步梯度，GPU 之间需要进行 AllReduce 通信操作；
- DP 的流量虽然只占不到 2%，但需长距离传输，同时可与计算部分掩盖，需控制通信开销占比。
流水线并行（Pipeline Parallelism, PP）：
- 涉及用于跨层传输参数的 P2P 通信；
- 通信量小、通信次数少，可通过流水掩盖，对网络诉求较低。
张量并行（Tensor Parallelism, TP）：
- 张量并行所需的通信量最大，且不可被计算掩盖，因此需在同一个服务器内使用张量并行。
序列并行（Sequence Parallelism, SP）：
- 序列并行所需的通信量也很大，，因此需在同一个服务器内使用序列并行。
专家并行（Experts Parallelism, EP）：
- 例如将需要给专家 1 计算的数据收集起来放在专家 1 处、将需要给专家 2 计算的数据收集起来放在专家 2 处，因此需要使用 All-to-All 通信；
- 单次通信量小，但通信次数多。

参考资料

分析 transformer 模型的参数量、计算量、中间激活、KV cache

NCCL 官方文档

⁹

Zero Bubble Pipeline Parallelism

¹⁰

Deepseek-v3 技术报告-图的逐步解析-4-DualPipe

¹¹

DeepSeek-V3 Technical Report

UB-Mesh: a Hierarchically Localized nD-FullMesh Datacenter Network Architecture

⁴

[PyTorch Distributed: Experiences on Accelerating
Data Parallel Training](https://arxiv.org/pdf/2006.15704)

⁵

GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

¹⁵

GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

¹⁴

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

¹⁶

ZeRO: Memory Optimization Towards Training A Trillion Parameter Models

⁶

[Efficient Large-Scale Language Model Training on GPU Clusters
Using Megatron-LM](https://arxiv.org/pdf/2104.04473)

⁷

TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models

⁸

Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines

¹²

AMPNet: Asynchronous Model-Parallel Training for Dynamic Neural Networks

¹³

PipeDream: Fast and Efficient Pipeline Parallel DNN Training

¹⁷

Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning

附录

DDP in PyTorch⁴

Original Solution

从相同的模型状态开始

在 DDP 构建时将模型状态从一个进程广播到所有其他进程来实现
在每次迭代中消耗相同的梯度来保证正确性。

一个简单的解决方案可以在本地反向传播之后和更新本地参数之前插入一个梯度同步阶段。

DDP 可以注册 autograd 钩子以在每个反向传播后触发计算。当触发时，每个钩子会扫描所有本地模型参数，并从每个参数中检索梯度张量。

然后，它使用 AllReduce 集体通信调用来计算所有进程中每个参数的平均梯度，并将结果写回梯度张量。

实际上这种方案导致了性能问题：

集群通信在小 tensor 上效率低下
因为将梯度计算和同步分开，无法重叠计算和通信

Improved Technique

Gradient Bucketing: DDP 可以通过等待一段时间并将多个梯度合并到一个 AllReduce 操作中，实现更高的吞吐量和更低的延迟。为了重叠计算和通信，DDP 不应在一次 AllReduce 中通信所有梯度，否则在计算完成之前无法开始通信
Overlap Computation with Communication: 使用桶划分，DDP 只需要等待同一桶中的所有内容都准备好后才开始通信。DDP 为每个梯度累加器注册了一个自动微分钩子。钩子在相应的累加器更新梯度后触发，并将检查它所属的桶。如果同一桶中所有梯度的钩子都已触发，最后一个钩子将触发该桶的异步 AllReduce。

:skull:Attentions

所有进程必须使用相同的桶排序顺序，并且没有进程可以在启动桶 i+1 之前启动桶 i 的 AllReduce

（a）显示了一个例子，其中两个垂直轴代表时间，虚线表示梯度何时准备好。在进程 1 中，四个梯度按顺序计算，但在进程 2 中，梯度$g_2$是在$g_3$和$g_4$之后计算的。在这种情况下，如果所有进程在准备好后立即进行 AllReduce，AllReduce 内容将不匹配($g_2$还没有计算完成，结果不正确)

PyTorch v1.5.0 通过使用 model.parameters()的逆序作为桶排序顺序来解决这个问题，假设层很可能按照与它们在正向传递中调用的相同顺序注册。因此，反向顺序应大致代表反向传播中的梯度计算顺序
对于跳过的梯度，DDP 可以通过在正向传播结束时主动将它们标记为就绪来避免等待其余参数梯度

由于梯度到桶的映射是在构建时确定的，那些缺失的梯度会导致一些桶从未看到最终的 autograd 钩子，从而未能将桶标记为就绪。因此，反向传播可能会挂起

（b）显示了一个例子，其中一个迭代中跳过了对应梯度$g_3$的参数，导致桶 2 永远不会准备好

DistributedDataParallel Algorithm

Alpa¹⁷

Techniques

分层抽象 (Hierarchical Abstraction)：把大问题拆分成两层：
- Inter-Operator Parallelism（算子级别并行） → 决定每个 layer/module 怎么切（DP, PP, TP）。
- Intra-Operator Parallelism（算子内并行） → 针对矩阵乘法、卷积等操作进一步做张量切分。
两阶段搜索 (Two-level search)：
- Stage 1: 通过动态规划（DP）+ cost model，在 operator graph 层面找到最优划分。
- Stage 2: 在每个 operator 内部，用 ILP / search 确定最佳 intra-op parallelism。
编译器支持 (Compiler-based)：
- Alpa 构建在 XLA 上，可以自动生成 SPMD (single program multiple data) 程序。
- 通过 parallel execution plan → lowered to XLA SPMD → 映射到 GPU/TPU。

搜索空间

对于 computation block，Alpa 允许 DP PP TP 以及他们的组合

Search Methods

Inter-Operator 层级：把模型分成多个 stage 分配给不同的 mesh。每个 segment 选择一个并行方式（DP / PP / TP）。动态规划搜索 segment 划分，利用 cost model 评估执行的延迟。
Intra-Operator 层级：对每个算子，枚举可能的 tensor 切分方式。用 ILP 找到代价最低的方案。

Parallelism in Transformer-Based LLM

DP: LLM 训练常结合 ZeRO 使用，减少 optimizer/gradient/weight 的冗余。
PP: DeepSeek 使用了 Dualpipe，大多数都会使用 1F1B
TP: Megaton-LM 的 attention 算子并行和 matmul 并行
Expert Parallelism: MoE 层的专家并行，GShard-MoE
- :laughing:Inference 越来越常见
SP(Sequence Parallelism)/CP(context parallelism): 超长上下文 LLM 训练使用

接下来我们将以 Transformer-based LLM 为例，分析 Attention 和 FFN or MoE 采用的并行化手段(Training & Inference)

MLP Parallelism

MLP 的并行运算，包括层内切分、层间切分以及参数冗余消除，并行方式包括了数据并行（DP）、序列并行(SP)、张量并行（TP）、层并行（PP）、参数冗余消除（Zero）

层内并行

DP: batch size 切分
SP: seq len 切分
TP: hidden size 切分 $$ activation size:[batch_size, seq_len, hidden_size]$$

SP 和 TP 通常一起使用，因为 TP 仅仅切分了 weights，而 input 没有切分，GPU 上仍然保存完整的输入计算的 activation，而 SP 则切分了 activation

对于 MLP 来说，SP 是比较容易的，

Example megatron-v3 中使用 SP+TP 来减小 activation 的大小¹，TP 的计算过程类似 MatMul Parallelism 的描述

MLP TP 在 A x B x C 矩阵乘法场景，B 矩阵列切， C 矩阵行切。计算量减半，多了一次集群通信（allReduce）中间值的存储大小减半，Input, Weight 减半
因为 GeLU 非线性，必须使用 allGather 来收集 activation¹

层间并行

Pipeline Parallelism，Expert Parallelism，Attention-FFN Decoupling

现在的普遍做法是 AFD(Attention-FFN Decoupling)，这是一种层间异构并行策略，把 Transformer 里的 Attention 和 MoE-FFN 拆开到不同设备上执行，以适配它们算力/显存需求的差异，从而提高推理吞吐和硬件利用率，MoE 的不同专家也分配到不同的 GPU 上²³

Attention 是 Matmul + softmax + activation，通常是访存瓶颈（memory access bound）
FFN 是占用显存大，算力要求高，通常表现为计算瓶颈（compute bound）
ByteDance 的 AFD 方案 overview
- 将一批请求拆分成 m 个微批次 (micro-batches)，在注意力节点和专家节点之间创建一个“乒乓流水线”。这些节点执行微批次的前向传播，并在每个 MoE 层中交换两次中间结果。这种设置使得前向计算能够覆盖通信开销

冗余参数消除

一般使用 ZeRO，将参数分散存储

模型参数较小时，一般选择参数广播
参数较大时，将每一层的权重 n 等分，每个 GPU 设备上面存一份，当需计算时将其 allgather 回来。

Attention Parallelism

DP: batch size
TP: heads & $d_k$
SP/CP(context parallelism): seq_len $$attention \space size: [batch_size, heads, seq_len, d_k]$$

Attention 的层间并行、冗余参数消除方式与线性层的方式一致，层内并行的主要差异是 TP 和 SP

Sequence Parallelsim 4⁵

Motivation

self-attention 的内存需求是输入长度（sequence length）的 2 次方。其复杂度为 $O(n^2)$ ，其中，n 是序列长度。换言之，长序列数据将增加中间 activation 内存使用量，从而限制设备的训练能力。

这里出现了三种切分方式：

只切分 Q 的序列 Q 的切分后的尺寸为$[bs, heads, seq_len/SP, head_dim]$，按照 attention 计算：
1. 求解 score，Q x K 相当于左矩阵行切，score 尺寸：$[bs, heads, seq_len/SP, seq_len]$
2. softmax 求解的是最后一个维度，计算元素值相同，得到 attention_weights，
3. attention_weights 与 V 进行矩阵乘，还是左矩阵行切运算，元素值相同，计算得到 O 的分块结果
4. 将计算的 O 进行 allgather，结果相等。

:warning: 虽然计算上可行，但每个 GPU 都需要一份完整的 K 和 V，没有节省下长序列带来的巨大 KV Cache 内存开销，这违背了 SP 的初衷。

只切分 K 的序列 score 尺寸是$[bs, heads, seq_len, seq_len/SP]$，进一步计算 softmax，由于最后一个维度的数据只有之前的一半长度，而 softmax 的计算跟整个序列相关，直接拼接会导致结果不相等。所以，单独切 K 序列后拼接，结果不等
只切分 V 序列得到 attention_weights 尺寸完整，计算 attention_weights x V，因为 V 矩阵被行切，所以 attention_weights 需要列切(Matmul parallelism 的图)，最后 allreduce 能够获得完整结果

:warning: softmax 计算需要完整的 score 矩阵，内存瓶颈没有解决

当前常用的方式：QKV 按照相同方式切分，然后对 softmax 修正：

数据协同切分：每个切分 Qi 要与所有的 Ki、Vi 进行一次计算，得到 Oi，尺寸均为$[bs, heads, seq_len/N, head_dim]$，$N$是 sp 的并行度
分块计算与通信：每个$Rank_i$ 的目标是计算出它所负责的输出 $O_i$。要计算 $O_i$，$Q_i$ 必须和所有的 $K_j$ 与 $V_j$（其中 $j=0, 1, ..., N-1$）进行交互。这通过通信实现（例如，All-to-All 或者 Ring AllGather）。
在线 Softmax 修正：在与每个 $K_j$, $V_j$ 块交互计算时，不能直接计算局部的 softmax 然后相加。必须使用一种Online 的算法来迭代地更新 softmax 的结果，从而保证最终结果与全局计算完全一致。

prefill

推理的 prefill 阶段，Q 的序列长度与 KV 保持一致，开启 SP 后 GPU 之间需要交换 KV 值与 Q 进行运算

每个 rank 的 Q 与 KV 匹配计算完后获得三个输出值，然后进行结果修正得到$[O_{X0}, O_{X1}, O_{X2}]$，X 值为 rank 序号。最后每个 rank 将自己的分块结果进行聚合（加法）运算得到结果 $O_X$⁶

可以选择 pass Q 或者 pass KV

类似分布式的 flash attention⁷

decode

在 decode 阶段 Q 的长度为 1，若对应的 KV 值分散在各个 GPU 设备中，可以将 Q 复制 N 份与切片的 KV 值进行 attention 计算后，最后将结果 gather 到一个设备上再计算修正结果。

这个方式意味着 GPU0 需要完成最大值、修正运算、求和运算，GPU0 可能成为瓶颈。一种 Tree attention 算法对过程做了改进，提升了通信效率。就是把 gather 换成多次 allreduce，这种方式在跨节点场景中优势明显⁸。

Tensor Parallelism

Attention 的张量并行分为两个部分:QKV 运算、线性投影运算。

QKV 运算的 TP 切分一般针对 heads 维度，相当于矩阵的并行的批量运算¹

attention 按照 heads 进行切分
linear 按照 hidden_size / heads 进行切分

DeepSeek-V3 9

DeepSeek 一共有 61 层，但并非每层都采用 MoE 结构，前 3 层依然是只有 1 个 MLP 的 Dense 层

Training Parallelism

$16PP \times 64EP \times 2DP(ZeRO) = 2048$ GPUs

DualPipe with Overlap
- 可以让 compute 和 communication 重叠，在整个双向管道都充满数据的情况下
MoE Expert Parallelism
Not Tensor Parallelism: 主要原因是 DeepSeek 没有 H100GPU 和 NVLink 的支持，Compute 能力和通信能力不足以支撑 TP

Inference Parallelism

Prefill

Compute:
- Attention: TP4 + SP + DP8 ⁹
  
  TP4 是为了降低通信开销
- FFN: EP32 + TP
  
  Dense MLP（浅层）部分：TP1（无张量并行）
Scheduler: 同时处理两个 micro-batch；
- 在一个 micro-batch 做 Attention + MoE 计算时，
- 另一个 micro-batch 正在进行 token 的 dispatch/combine（all-to-all 通信）。
- 用流水线（pipeline）方式在 micro-batch 级别做并行计算与通信重叠

Decode

Compute:

模块并行策略说明

Attention 部分 TP4 + SP + DP80 Tensor Parallel 4 路 + Sequence Parallel + Data Parallel 80 路

MoE 部分 EP320 Expert Parallel 320 路，每张 GPU 承载一个专家

每张 GPU 上只有 1 个专家；320 个专家覆盖所有 GPU；另有 64 张 GPU 专门放置冗余专家和共享专家。
Scheduler: 和 prefilling 阶段类似，他们尝试并行处理两个微批（micro-batch）；在解码阶段，Attention 占用的计算时间更多；因此采用新的重叠方式：

阶段 micro-batch A micro-batch B

时间片 1 Attention Dispatch + MoE + Combine

时间片 2 Dispatch + MoE + Combine Attention

模块	并行策略	说明
Attention 部分	TP4 + SP + DP80	Tensor Parallel 4 路 + Sequence Parallel + Data Parallel 80 路
MoE 部分	EP320	Expert Parallel 320 路，每张 GPU 承载一个专家

阶段	micro-batch A	micro-batch B
时间片 1	Attention	Dispatch + MoE + Combine
时间片 2	Dispatch + MoE + Combine	Attention

参考资料

Reducing Activation Recomputation in Large Transformer Models

[MegaScale-Infer: Serving Mixture-of-Experts at Scale
with Disaggregated Expert Parallelism](https://arxiv.org/pdf/2504.02263)

LLM 推理提速：Attention 与 FFN 分离(AFD)方案解析

⁴

Sequence Parallelism: Long Sequence Training from System Perspective

⁵

[并行训练]Context Parallelism 的原理与代码浅析

⁶

ring attention + flash attention：超长上下文之路

⁸

Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters

⁹

DeepSeek-V3 Technical Report

⁷

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Chapter 4: CUDA Optimization for LLM Inference

Overview

在大规模语言模型（LLM）推理中，优化 CUDA 代码对于提升性能和效率至关重要。本文档介绍了一些关键的 CUDA 优化技术，帮助开发者更好地利用 GPU 资源进行 LLM 推理。这里我们将介绍 Transformer-based 自回归预训练模型推理用到的算子的优化方法。这些优化的方法也适用于其他算子。

我们主要关注以下几个方面：

GPU 硬件架构(Hirerarchy Memory, SM, Warp 等)
CUDA 计算模型(Thread, Thread Block, Grid 等)
CUDA Kernel 性能调优(Nsight Compute, Occupancy)
CUDA 常用优化技巧(Double buffering, Memory Coalescing, Overcoming Bank Conflict 等)
Transformer 内部算子优化
- Matrix Multiplication 优化
- Softmax 优化
- LayerNorm 优化
- Self-Attention 优化
- Flash Attention 优化

GPU 硬件架构

GPU 以 Throughput 为设计目标，和 CPU 有很大的不同。

GPU 中虽有缓存结构但是数量少。因为要减少指令访问缓存的次数。
GPU 中控制单元非常简单。控制单元中没有分支预测机制和数据转发机制，对于复杂的指令运算就会比较慢。
GPU 的运算单元 (Core) 非常多，采用长延时流水线以实现高吞吐量。每一行的运算单元的控制器只有一个，意味着每一行的运算单元使用的指令是相同的，不同的是它们的数据内容。那么这种整齐划一的运算方式使得 GPU 对于那些控制简单但运算高效的指令的效率显著增加。¹

Streaming Multiprocessor (SM)2

Functional Units (FP32, INT32, SF, LS)
Data Storage
- Registers
- Constant Cache (Const.)
- Shared Memory
Warp Contexts (Warp Ctxs)
- One for each warp assigned to SM.
- Holds PC (addr of next insn to execute)
- Execution state
Warp Scheduler
- round-robin policy

Functional Units 2

CPU 可以通过重新调度指令来避免 stalls
- 指令重排序: 将无关指令插入到有依赖关系的指令之间
- 数据前递: 硬件层面的优化，允许数据在流水线中提前传递
GPU 可以通过线程切换（thread switching / warp scheduling）来规避延迟带来的停顿。

GPU 有成百上千个线程（或 warp），当某个线程在等待功能单元结果时，GPU 可以直接调度另一个线程执行。这样，延迟不会直接影响整体吞吐量

Warp Contexts 2

warp contexts 的数量决定了 SM 上能同时并发的 block 数量 warp_contexts

Warp Scheduler 2

硬件来决定哪个 wrap 将下一个执行
Branches and Warp Divergence
1. NVIDIA 7.0 之前，当 warp 遇到分支时，必须等到最外层的重新收敛点才能继续执行
```
if (condition) {
// Branch A: 一些线程执行这里
instruction_a1();
instruction_a2();
} else {
// Branch B: 另一些线程执行这里
instruction_b1();
instruction_b2();
instruction_b3(); // 更多指令
}
// 重新收敛点: 所有线程在这里汇合
instruction_after_branch();
```
1. 替代方案：调度器可以选择任意具有相同 PC 值的线程子集来执行，不需要等待最外层的收敛点
```
时刻1: 调度器选择8个在Branch A的线程执行
时刻2: 调度器选择16个在Branch B的线程执行
时刻3: 调度器选择剩余的线程继续执行
...
```
1. 问题
- 仍需要为所有未 masked 的线程解码相同指令解码复杂度没有降低
- 需要更复杂的调度逻辑，要追踪每个线程的状态，要决定最佳的线程子集组合
- 每次都要计算最优的线程调度组合，增加了调度开销

GPU Memory Hierarchy 3

Global Memory: 全局内存的主要角色是为核函数提供数据，并在主机与设备及设备与设备之间传递数据
Constant Memory: 特殊的常量内存缓存（constant cache）进行缓存读取，常量内存为只读内存
Texture Memory & Surface Memory: 纹理内存和表面内存类似于常量内存，也是一种具有缓存的全局内存，有相同的可见范围和生命周期，而且一般仅可读(表面内存也可写)
Register: 寄存器的速度是访问中最快的，但是它的容量较小。
Local Memory: 局部内存是每个线程私有的内存空间，但从硬件来看，局部内存只是全局内存的一部分。所以，局部内存的延迟也很高
Shared Memory: 共享内存是每个线程块（block）内所有线程共享的内存空间。共享内存的访问延迟远低于全局内存
L1 / L2 Cache: L1 缓存是每个 SM 独有的，而 L2 缓存是所有 SM 共享的

{ width="600px" }

:laughing:Summary

每个 thread 都有自己的一份 register 和 local memory 的空间
同一个 block 中的每个 thread 则有共享的一份 share memory
所有的 thread (包括不同 block 的 thread) 都共享一份 global memory

不同的 grid 则有各自的 global memory。

层级	类型	特性
Registers	寄存器	每线程私有，延迟最低
Shared Memory	共享内存	每 block 共享，低延迟 (~100x DRAM 快)
L1 / L2 Cache	Cache	SM 局部 / GPU 全局缓存
Global Memory	全局内存	所有 SM 共享，延迟高 (~400-600 cycles)
Constant / Texture Memory	只读缓存	适合广播数据，缓存优化访存

Warp

程序员为单个线程编写代码，但硬件层面会将线程组织成固定大小（32 个）的束，称为 Warp 。Warp 是 SM 上调度和执行的真正基本单位

SIMT 执行：一个 Warp 中的所有 32 个线程在同一时刻执行相同的指令，但处理不同的数据。
Warp 分化 (Warp Divergence)：如果一个 Warp 内的线程因条件判断而走向了不同的代码路径，硬件必须串行化执行这些路径

它会首先为走向 if 分支的线程执行代码，屏蔽掉其他线程；然后再为走向 else 分支的线程执行代码。这种分化会严重降低性能，因为在每个分支的执行过程中，总有一部分线程处于空闲状态 。
延迟隐藏 (Latency Hiding)：当一个 Warp 停顿（例如，等待从全局内存读取数据）时，SM 上的 Warp 调度器会立即切换到另一个“准备就绪”的 Warp 来执行。

CUDA 计算模型 3

CUDA 将计算任务组织成一个三级层次结构。这是一个由程序员创建的、用于组织问题的逻辑层次，而非硬件的直接体现。

线程 (Thread)：最基本的执行单位。单个线程执行一个 Kernel 函数的实例。每个线程在其所属的线程块内拥有一个唯一的 ID(threadIdx)
线程块(Thread Block)：一组可以相互协作的线程（在现代架构上最多 1024 个）。一个块内的所有线程可以通过高速的片上共享内存共享数据，并能通过__syncthreads()来协调它们的执行。每个线程块在其所属的 Grid 内也拥有一个唯一的 ID(blockIdx)
网格 (Grid)：为执行单个 Kernel 而启动的所有线程块的集合。一个 Grid 内的所有线程都可以访问同一个全局内存空间。Grid 内的线程块被假定为独立执行，且执行顺序不确定；它们之间没有直接的同步机制。

线程块和网格可以被组织成一维、二维或三维的结构，这为将计算任务映射到向量、矩阵、体数据等数据结构上提供了极大的便利。 { width="600px" }

Kernel 执行流程

当一个 Kernel 被启动时，由硬件调度器将 Grid 中的所有线程块分配到 GPU 上可用的 SM 中
一个 Thread Block 会被完整地分配给一个 SM，并在其上完整地执行。在其生命周期内，它不会被迁移到其他 SM
一个 SM 可以并发地执行多个 Thread Block，前提是它拥有足够的资源（如寄存器、共享内存）来容纳这些线程块
在 SM 内部，一个 Thread Block 的所有线程被划分为若干个 Warp（每组 32 个线程）。这些 Warp 才是被 SM 的 Warp 调度器实际调度执行的单元。

这个映射关系是层级化的：Grid -> GPU，Block -> SM，Thread -> Warp -> CUDA核心

CUDA Kernel 性能调优

CUDA 常用优化技巧

Maximize Compiler Computation 2

Unroll Loops

展开循环（loop unrolling），让循环体重复展开多次，减少循环控制开销（比如 i++、i<N 的判断），提高 GPU 的吞吐量。
Write code using compile-time constants (not same as constant registers)

在代码里用编译期已知的常量来做索引、循环次数、数组大小等，而不是依赖 GPU 的常量寄存器

Coalescing Memory Access

当一个 Warp 中的所有 32 个线程访问全局内存中的连续位置时，硬件可以将这 32 个小的请求“合并”成一个单一、大型、高效的内存事务
Memory Access Patterns:
- 合并访问（理想）：Warp 中的线程 i 访问内存地址 base + i。这在处理按行主序存储的矩阵的行时非常常见
- 跨步访问（问题）：线程 i 访问 base + i * stride。如果步长（stride）很大，这将导致许多独立的、低效的内存事务。这在访问按行主序存储的矩阵的列时很常见
- 非对齐访问：Warp 访问的起始地址未与内存事务的大小对齐

Avoid Bank Conflicts in Shared Memory

:warning: Shared memory is organized into 32 banks. Each bank is a slice of SRAM that can load or store 4 bytes of data every cycle. smem

当一个 Warp 中的所有 32 个线程访问全局内存中的连续位置时，硬件可以将这 32 个小的请求 “合并”成一个单一、大型、高效的内存事务
当同一个 Warp 中的两个或更多线程试图访问位于同一个内存银行中的不同地址时，就会发生银行冲突。此时，这些访问会被串行化处理，从而降低了共享内存的有效带宽

Solutions 4

Padding: 在数据结构中插入填充元素，以改变数据在内存中的布局，避免多个线程访问同一银行
- 可能降低 SM 的 occupancy
- 可能地址访问不对齐，无法使用向量化访问
swizzling: 重新组织数据的存储方式，使得并行访问时更少冲突(更常用):rocket:
- 某些 swizzling 方法在从 shared memory 读数据到 register 时不能进行 float4 的合并读取
- 逻辑位置表示元素在矩阵中的逻辑坐标。
- 物理位置表示其对应元素在实际存储数据的 shared memory 中的位置坐标。
  
  当我们说读取矩阵的第 2 行第 3 列的元素，(2,3)就表示逻辑位置，而真正读取数据的时候，我们需要从实际存储数(2,1)的 shared memory 中对应的位置

:warning: 广播 (Broadcast): 如果一个 Warp 中的所有线程访问同一个 bank 中的完全相同的地址，这是一种广播操作，不会产生冲突

Matmul Bank Conflict Avoidance Example 5

const int warp_id = tid / 32;
const int lane_id = tid % 32;
const int a_tile_index =  warp_id / 2 * 16 + lane_id / 8 * 4
const int b_tile_index =  warp_id % 2 * 32 + lane_id % 8 * 4;

使用的是 4×2 的 warp 布局, warp 中的每个线程按照 4×8 进行排布, 每个 warp 对应 16×32 的数据

每个 wrap 32 个线程只获取 As 的 4*4=16 个数据
每个 wrap 32 个线程只获取 Bs 的 8*4=32 个数据
shared memory 中有 32 个 bank，所以不会产生 bank conflict

Double Buffering 5

在共享内存中分配两个缓冲区: 当 SM 正在对缓冲区 1 中的数据进行计算时，硬件可以异步地将下一块数据从全局内存预取到缓冲区 2 中。一旦计算完成，两个缓冲区的角色互换。这种方式有效地将全局内存的访问延迟隐藏在了计算的背后 double-buffering

if (doubleBufferIdx == 0) {
  // load first (B0)
  db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
      N, K, A, B, As, Bs, innerRowA, innerColA, innerRowB, innerColB);
}
__syncthreads();

// outer-most loop over block tiles
for (uint bkIdx = 0; bkIdx < K; bkIdx += 2 * BK) {
  if (doubleBufferIdx == 0) {
    // process current (B0)
    db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
                        TN>(regM, regN, threadResults, As, Bs, warpRow,
                            warpCol, threadRowInWarp, threadColInWarp);
    __syncthreads();

    // process current+1 (B1)
    if (bkIdx + BK < K) {
      db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN,
                          TM, TN>(regM, regN, threadResults, As + (BM * BK),
                                  Bs + (BK * BN), warpRow, warpCol,
                                  threadRowInWarp, threadColInWarp);
    }
    __syncthreads();

    // load current + 2 (B0)
    if (bkIdx + 2 * BK < K) {
      db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
          N, K, A + 2 * BK, B + 2 * BK * N, As, Bs, innerRowA, innerColA,
          innerRowB, innerColB);
    }
  } else {
    // load current + 1 (B1)
    if (bkIdx + BK < K) {
      db::loadFromGmem<BM, BN, BK, rowStrideA, rowStrideB>(
          N, K, A + BK, B + BK * N, As + (BM * BK), Bs + (BK * BN), innerRowA,
          innerColA, innerRowB, innerColB);
    }
    __syncthreads();

    // process current (B0)
    db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN, TM,
                        TN>(regM, regN, threadResults, As, Bs, warpRow,
                            warpCol, threadRowInWarp, threadColInWarp);
    __syncthreads();

    // process current+1 (B1)
    if (bkIdx + BK < K) {
      db::processFromSmem<BM, BN, BK, WM, WN, WMITER, WNITER, WSUBM, WSUBN,
                          TM, TN>(regM, regN, threadResults, As + (BM * BK),
                                  Bs + (BK * BN), warpRow, warpCol,
                                  threadRowInWarp, threadColInWarp);
    }
  }

  A += 2 * BK;     // move BK columns to right
  B += 2 * BK * N; // move BK rows down
  __syncthreads();
}

Tile 6

将原本一行 × 一列的计算进行分块，每次只计算一块
一次性从全局内存中加载一小块 A (BM x BK) 和一小块 B (BK x BN) 到共享内存中
一个线程块内的所有线程就可以在共享内存上快速地进行大量的计算，以完成对应的一小块 C (BM x BN) 的计算
每个线程不再是只计算 C 块中的一个元素，而是负责计算一个更小的结果网格（图中是 2x2）。这样做可以进一步提升数据复用率和计算效率

Wrap Tile 6

线程块分片(blocktiling): 不同的线程块可以在不同的 SM 上并行执行.
warp 分片(warptiling): 不同的 warp 可以在不同的 warp 调度器上并行执行, 也可以在同一个 warp 调度器上并发执行.
线程分片(threadtiling): (数量非常有限的)指令可以在相同的 CUDA 内核上并行执行(=指令级并行, instruction-level parallelism, 即 ILP).

for (uint bkIdx = 0; bkIdx < K; bkIdx += BK) {
  // load from gmem
  for (uint offset = 0; offset + rowStrideA <= BM; offset += rowStrideA) {
    const float4 tmp = reinterpret_cast<const float4 *>(
        &A[(innerRowA + offset) * K + innerColA * 4])[0];
    // float4 tmp;
    // asm("ld.global.nc.v4.f32 {%0, %1, %2, %3}, [%4];"
    //     : "=f"(tmp.x), "=f"(tmp.y), "=f"(tmp.z), "=f"(tmp.w)
    //     : "l"(&A[(innerRowA + offset) * K + innerColA * 4]));
    As[(innerColA * 4 + 0) * BM + innerRowA + offset] = tmp.x;
    As[(innerColA * 4 + 1) * BM + innerRowA + offset] = tmp.y;
    As[(innerColA * 4 + 2) * BM + innerRowA + offset] = tmp.z;
    As[(innerColA * 4 + 3) * BM + innerRowA + offset] = tmp.w;
  }

  for (uint offset = 0; offset + rowStrideB <= BK; offset += rowStrideB) {
    reinterpret_cast<float4 *>(
        &Bs[(innerRowB + offset) * BN + innerColB * 4])[0] =
        reinterpret_cast<const float4 *>(
            &B[(innerRowB + offset) * N + innerColB * 4])[0];
    // asm("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];"
    //     : "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 0]),
    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 1]),
    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 2]),
    //       "=f"(Bs[(innerRowB + offset) * BN + innerColB * 4 + 3])
    //     : "l"(&B[(innerRowB + offset) * N + innerColB * 4]));
  }
  __syncthreads();
  // dotIdx loops over contents of SMEM
  for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) {
    // 实际上是修改 warp 的布局，变成 2x2 的布局
    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
      for (uint i = 0; i < TM; ++i) {
        regM[wSubRowIdx * TM + i] =
            As[(dotIdx * BM) + warpRow * WM + wSubRowIdx * WSUBM +
              threadRowInWarp * TM + i];
      }
    }
    for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
      for (uint i = 0; i < TN; ++i) {
        regN[wSubColIdx * TN + i] =
            Bs[(dotIdx * BN) + warpCol * WN + wSubColIdx * WSUBN +
              threadColInWarp * TN + i];
      }
    }

    // 每次每个 warp 加载 32 * 16 个数据，不会产生 bank conflict，需要迭代 4 次
    // 每个 warp 多次迭代的数据是连续的一块, 不同 warp 同一次迭代的数据则是分散的多块
    for (uint wSubRowIdx = 0; wSubRowIdx < WMITER; ++wSubRowIdx) {
      for (uint wSubColIdx = 0; wSubColIdx < WNITER; ++wSubColIdx) {
        // calculate per-thread results with register-cache locality
        for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
          for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
            threadResults[(wSubRowIdx * TM + resIdxM) * (WNITER * TN) +
                          (wSubColIdx * TN) + resIdxN] +=
                regM[wSubRowIdx * TM + resIdxM] *
                regN[wSubColIdx * TN + resIdxN];
          }
        }
      }
    }
  }
  A += BK;     // move BK columns to right
  B += BK * N; // move BK rows down
  __syncthreads();
}

参考资料

CUDA（一）：CUDA 编程基础

CUDA（二）：GPU 的内存体系及其优化指南

⁶

[CUDA 学习笔记] 如何优化 CUDA 矩阵乘内核以获得类似 cuBLAS 的性能: 工作日志

EE 7722, Lecture Slides: NVIDIA GPU Microarchitecture

⁴

CUDA shared memory 避免 bank conflict 的 swizzling 机制解析

⁵

GEMM 优化: 双缓冲 (Prefetch) 和 Bank Conflict 解决

LLM Study