GPU 算子优化中的数据复用

在 GPU 高性能算子开发中，我们经常会听到这样一句话：

尽可能提高数据复用，减少对 Global Memory 的访问。

这句话本身没有问题。Global Memory 的访问延迟高、带宽有限，而 Shared Memory 和寄存器都位于片上，访问速度通常快得多。将数据加载到片上后重复使用，是 GEMM、卷积、Attention 和 MoE 等算子优化的核心手段。

然而，在实际优化中，我们也经常遇到一个看似矛盾的现象：

增大 Tile、增加 Shared Memory 缓冲、提高数据复用率之后，算子反而变慢了。

其中一个重要原因就是：为了提高数据复用，单个线程块消耗了更多 Shared Memory 和寄存器，导致 Occupancy 下降，GPU 隐藏延迟的能力变弱。

本文将详细解释这条因果链。

1. 高性能算子优化的四个支柱

高性能 GPU 算子的开发，本质上是在多个因素之间进行平衡。

常见的四个维度包括：

并行度：让足够多的线程、Warp 或 Wave 同时运行，填满计算资源。
分支发散：减少同一个 Warp 内不同线程走不同分支造成的串行执行。
访存合并：让相邻线程访问相邻地址，最大化有效显存带宽。
数据复用：让从显存加载的数据在片上被多次使用，减少重复访存。

这几个方向并不是相互独立的。

例如，增大 Tile 可以提高数据复用，但也会增加 Shared Memory 和寄存器占用；资源占用增加后，同一个 SM 上能够同时驻留的线程块减少，Occupancy 就可能下降。

因此，GPU 优化不是简单地让每一个指标都达到最大，而是在它们之间寻找平衡点。

2. 什么是数据复用？

以矩阵乘法为例：

[
C=A\times B
]

矩阵乘法中的每个输入元素通常都会参与多次乘加运算。

例如，矩阵 (A) 中的一个元素，会与矩阵 (B) 中一行或一列的多个元素进行运算。

如果每次计算都重新从 Global Memory 读取数据，就会产生大量重复访存。

一种常见的优化方式是分块计算：

Global Memory
      ↓
将 A_tile 和 B_tile 加载到 Shared Memory
      ↓
多个线程重复使用 A_tile 和 B_tile
      ↓
计算结果保存在寄存器中
      ↓
写回 Global Memory

假设某个数据需要参与 128 次乘法。

不做复用时：

从 Global Memory 读取 128 次

使用 Shared Memory 复用后：

从 Global Memory 读取 1 次
从 Shared Memory 使用 128 次

这会显著提高算术强度：

[
\text

\frac{\text{计算操作数量}}{\text{Global Memory 访问字节数}}
]

算术强度越高，算子越有机会从访存受限转变为计算受限。

3. 数据复用为什么会消耗更多片上资源？

数据复用不会凭空发生。

如果希望一个数据被重复使用，就必须先把它保存在某个片上位置。GPU 中常见的位置有：

Shared Memory
Register
Cache

在算子开发中，开发者最直接控制的通常是 Shared Memory 和寄存器。

因此，数据复用本质上是在做一种资源交换：

[
\boxed{
\text{消耗更多片上存储资源}
\quad\Longleftrightarrow\quad
\text{减少片外显存访问}
}
]

需要注意的是：

数据复用本身不会直接导致 Occupancy 下降。

真正导致 Occupancy 下降的，是为了提高数据复用而采取的具体手段，例如：

增大计算 Tile
增加每个线程负责的输出数量
使用 Double Buffering
使用 Triple Buffering
增加流水线 Stage
保存更多中间结果
增加寄存器中的累加器数量

完整的因果链可以写成：

[
\text{提高数据复用}
\rightarrow
\text{增大 Tile 或增加 Buffer}
\rightarrow
\text{Shared Memory、寄存器消耗增加}
\rightarrow
\text{每个 SM 可驻留线程块减少}
\rightarrow
\text{Occupancy 下降}
]

4. Shared Memory 为什么会降低 Occupancy？

每个 SM 拥有的 Shared Memory 容量是有限的。

假设某个 SM 一共有：

128 KB Shared Memory

4.1 每个线程块使用 16 KB

一个线程块使用 16 KB Shared Memory，那么理论上最多可以驻留：

[
\frac{128\text{ KB}}{16\text{ KB}}
=8
]

即 8 个线程块。

4.2 每个线程块使用 48 KB

如果增大 Tile 后，一个线程块需要 48 KB：

[
\left\lfloor
\frac{128\text{ KB}}{48\text{ KB}}
\right\rfloor
=2
]

虽然还剩下 32 KB，但这 32 KB 不足以再容纳一个需要 48 KB 的线程块，因此只能驻留两个线程块。

4.3 使用 Double Buffering

假设一个 Tile 本来需要 48 KB Shared Memory。

为了让数据加载和计算重叠，可以使用 Double Buffering：

Buffer 0：当前正在参与计算
Buffer 1：正在加载下一批数据

此时一个线程块需要：

[
48\text{ KB}\times 2=96\text{ KB}
]

一个 128 KB Shared Memory 的 SM 最多只能驻留：

[
\left\lfloor
\frac{128\text{ KB}}{96\text{ KB}}
\right\rfloor
=1
]

于是，该 SM 上同时只能存在一个线程块。

虽然 Double Buffering 提高了单个线程块内部的流水效率，但同时降低了整个 SM 上线程块级别的并发度。

5. 为什么提高复用通常也会增加寄存器压力？

Shared Memory 增加并不必然意味着寄存器也一定增加。

但是在实际算子优化中，这两者经常同时增长，因为它们通常来自同一种优化策略：增大单个线程块或单个线程处理的数据规模。

5.1 更大的输出 Tile 需要更多累加器

在矩阵乘法中，输出结果一般先保存在寄存器中累加，最后再写回显存。

假设一个线程原本负责计算一个 (4\times4) 的输出块，需要保存：

[
4\times4=16
]

个累加结果。

如果扩大到 (8\times8)，则需要保存：

[
8\times8=64
]

个累加结果。

累加器通常会占用寄存器。

因此，仅输出累加器的数量就从 16 个增加到了 64 个。

除此之外，线程还需要寄存器保存：

A 的计算片段
B 的计算片段
地址和指针
循环变量
Tile 索引
Pipeline 状态
Mask
中间计算结果

所以，增大 Tile 往往会显著增加每线程寄存器数量。

5.2 多级流水线增加变量存活时间

假设使用 Triple Buffering：

Stage 0：正在计算
Stage 1：已经加载完成，等待计算
Stage 2：正在加载下一批数据

Shared Memory 中需要同时保存三份 Tile。

与此同时，寄存器中也可能需要保存：

当前 Stage 的地址
下一 Stage 的地址
多组数据指针
多组 A/B Fragment
当前 Buffer 编号
异步拷贝状态
同步状态

而且，这些变量需要跨越更长的代码范围保持有效。

在编译器中，这被称为变量的活跃区间变长。

变量活跃期间，寄存器不能被其他变量复用，因此：

[
\text{流水级数增加}
\rightarrow
\text{变量活跃区间变长}
\rightarrow
\text{寄存器压力增加}
]

6. 寄存器压力为什么会降低 Occupancy？

寄存器通常按线程分配，但一个 SM 上的寄存器总量是有限的。

假设某个 SM 有：

65536 个 32-bit 寄存器

一个线程块包含：

256 个线程

6.1 每线程使用 32 个寄存器

一个线程块需要：

[
256\times32=8192
]

个寄存器。

理论上最多可驻留：

[
\frac{65536}{8192}=8
]

个线程块。

6.2 每线程使用 64 个寄存器

一个线程块需要：

[
256\times64=16384
]

个寄存器。

最多驻留：

[
\frac{65536}{16384}=4
]

个线程块。

6.3 每线程使用 128 个寄存器

一个线程块需要：

[
256\times128=32768
]

个寄存器。

最多只能驻留：

[
\frac{65536}{32768}=2
]

个线程块。

因此，每线程寄存器数量翻倍时，可驻留线程块数量可能直接减半。

6.4 Occupancy 通常是阶梯式变化

实际硬件通常不会按单个寄存器精确分配，而是按固定粒度分配。

因此，Occupancy 往往不是连续下降，而是阶梯式变化。

例如：

每线程使用 64 个寄存器：可驻留 4 个线程块
每线程使用 65 个寄存器：可能只能驻留 3 个线程块

虽然只增加了一个寄存器，但因为资源分配需要向上取整，最终可能跨过某个资源阈值。

这也是算子调优中经常出现性能突变的原因之一。

7. Occupancy 下降为什么可能让速度变慢？

Occupancy 的主要作用，并不是简单增加“同时执行的指令数量”，而是为 GPU 提供足够多的 Warp 或 Wave，用来隐藏各种延迟。

常见延迟包括：

Global Memory 访问延迟
Shared Memory 访问延迟
指令数据依赖
MMA Pipeline 延迟
Barrier 同步等待
Cache Miss

以 Wave 为例，假设 Wave 0 发出了一个 Global Memory 请求：

Wave 0：等待 Global Memory 返回数据
Wave 1：执行 MMA
Wave 2：读取 Shared Memory
Wave 3：计算下一组地址

当 Wave 0 等待时，调度器可以选择其他 Ready Wave 执行。

这样，虽然单个 Wave 正在等待，但整个计算单元仍然可以保持忙碌。

7.1 Occupancy 较高时

Wave 0：等待访存
Wave 1：执行 MMA
Wave 2：执行 Load
Wave 3：执行整数运算
Wave 4：准备就绪
Wave 5：准备就绪

调度器通常能够找到可执行的 Wave，因此计算资源不容易空闲。

7.2 Occupancy 太低时

假设一个 SM 只能驻留一个线程块，并且该线程块只有少量 Wave：

Wave 0：等待 Global Memory
Wave 1：等待 Shared Memory
Wave 2：等待 Barrier
Wave 3：等待上一条 MMA 的结果

此时所有 Wave 都处于等待状态，调度器找不到可以发射的指令。

于是硬件流水线中就会出现空闲周期，也就是所谓的 Bubble。

时间 →

MMA：████████░░░░░░████████
MTE：░░████████░░░░░░██████
              ↑
            Bubble

完整过程可以概括为：

[
\text{Occupancy 过低}
\rightarrow
\text{Ready Wave 数量不足}
\rightarrow
\text{无法隐藏延迟}
\rightarrow
\text{执行单元出现 Bubble}
\rightarrow
\text{性能下降}
]

8. 一个典型例子：Tile 越大，为什么可能先快后慢？

考虑三个不同的 GEMM Tile 方案。

8.1 方案 A：小 Tile

Tile：64 × 64
Shared Memory：16 KB
每线程寄存器：32
每 SM 驻留线程块：8

优点：

Occupancy 高
Ready Wave 数量多
延迟隐藏能力强
单个线程块资源占用较低

缺点：

数据复用率较低
Global Memory 流量较大
地址计算和循环控制开销较高
MMA 计算密度不足

可能出现：

Occupancy：高
AP MMA Duty：35%
显存流量：高

此时虽然并发度很高，但每个线程块做的有效计算太少。

8.2 方案 B：中等 Tile

Tile：128 × 128
Shared Memory：48 KB
每线程寄存器：64
每 SM 驻留线程块：2

优点：

数据复用明显提高
Global Memory 访问减少
MMA 指令占比增加
地址计算开销被更多计算摊薄
仍然有足够 Wave 隐藏延迟

可能出现：

Occupancy：中等
AP MMA Duty：65%
显存流量：明显下降

虽然 Occupancy 比方案 A 低，但整体性能可能明显更高。

8.3 方案 C：超大 Tile 和 Triple Buffering

Tile：256 × 256
Shared Memory：120 KB
每线程寄存器：128
每 SM 驻留线程块：1

优点：

数据复用率非常高
Global Memory 流量进一步下降
理论计算密度很高

缺点：

每个 SM 只有一个线程块
Ready Wave 数量不足
Barrier 期间没有其他线程块可调度
指令依赖链变长
边界 Tile 浪费更严重
可能发生寄存器 Spill

可能出现：

Occupancy：很低
AP MMA Duty：48%
显存流量：更低
总运行时间：反而增加

因此，性能通常不是随着 Tile 增大而一直上升，而是存在一个最优点：

性能
  ^
  |             ● 最优点
  |          ／   ＼
  |       ／         ＼
  |    ／               ＼
  +----------------------------> Tile 大小

9. 寄存器压力过高还可能导致 Spill

寄存器压力除了降低 Occupancy，还可能引发更严重的问题：寄存器溢出。

假设一个线程理论上需要：

160 个寄存器

但编译器最终只能分配：

128 个寄存器

剩余变量可能被 Spill 到 Local Memory。

需要 160 个寄存器
        ↓
只能分配 128 个
        ↓
32 个寄存器对应的数据进入 Local Memory

Local Memory 的名字容易让人误以为它是片上的高速私有存储。

实际上，在许多 GPU 架构中，Local Memory 物理上位于显存中，只是每个线程拥有独立的地址空间，并可能经过缓存访问。

于是原本期望的数据路径是：

Global Memory
      ↓
Shared Memory
      ↓
Register
      ↓
重复使用

发生 Spill 后可能变成：

Global Memory
      ↓
Shared Memory
      ↓
Register 不足
      ↓
写入 Local Memory
      ↓
再次从 Local Memory 读取

这样会增加额外的显存读写，甚至抵消数据复用带来的收益。

因此，在 Profiler 或编译日志中看到 Local Memory 使用量明显增加时，需要特别警惕寄存器 Spill。

10. Shared Memory 复用还可能带来哪些额外成本？

除了资源占用和 Occupancy 下降，提高 Shared Memory 复用还可能引入其他成本。

10.1 Barrier 同步开销

线程块中的线程共同使用 Shared Memory 时，通常需要同步。

load_to_shared();
barrier();

compute();

barrier();

如果某些 Wave 先完成加载，它们必须等待其他 Wave：

Wave 0：到达 Barrier，等待
Wave 1：到达 Barrier，等待
Wave 2：仍在搬运数据
Wave 3：仍在计算地址

等待期间，这些 Wave 无法继续执行。

如果一个 SM 上还有其他线程块，调度器可以切换到其他线程块。

但如果高 Shared Memory 使用量导致 SM 上只驻留一个线程块，那么 Barrier 等待就更难被隐藏。

10.2 Shared Memory Bank Conflict

将数据放入 Shared Memory，并不代表访问一定高效。

Shared Memory 通常被划分为多个 Bank。

如果多个线程访问同一个 Bank 中的不同地址，访问可能被拆分成多次：

无冲突：一次访问完成
2 路冲突：拆成两次完成
4 路冲突：拆成四次完成

这会降低 Shared Memory 的有效带宽。

因此，高性能 GEMM 中经常需要设计：

Padding
Swizzle
Interleave
向量化读取
Shared Memory Layout

它们的目的之一，就是减少 Bank Conflict。

10.3 边界 Tile 浪费

较大的 Tile 还可能增加边界浪费。

假设：

M = 300
Tile_M = 256

需要两个线程块：

Tile 0：处理 256 行
Tile 1：只处理剩余 44 行

第二个线程块中的大量线程会被 Mask 掉。

虽然这些线程不产生有效结果，但线程块仍然可能占用完整的：

Shared Memory
寄存器
Wave Slot
Block Slot

因此，Tile 越大，在不规则形状或小矩阵场景下越可能出现资源浪费。

11. Occupancy 不是越高越好

一个常见误区是：

Occupancy 100% = 性能最好
Occupancy 50% = 性能只有一半

这种理解是不正确的。

Occupancy 是隐藏延迟的手段，而不是最终优化目标。

11.1 内存受限算子

如果 Kernel 中存在大量 Global Memory 访问，那么通常需要更多 Wave 来隐藏访存延迟。

此时从较低 Occupancy 提升到中等 Occupancy，可能明显提高性能。

例如：

Occupancy：25% → 50%

可能显著减少等待显存时的 Bubble。

11.2 计算密集型 GEMM

对于优化良好的 GEMM，如果：

数据已经位于 Shared Memory 或寄存器
MMA 流水线安排合理
单线程指令级并行度较高
单个线程块已有足够 Wave
Global Memory 延迟已被异步流水覆盖

那么 50% 左右的 Occupancy 可能已经足够。

此时继续追求 100% Occupancy，可能需要：

缩小 Tile
减少累加器
降低寄存器使用
减少 Shared Memory
减少 Pipeline Stage

这些操作虽然提高了 Occupancy，却可能降低数据复用和计算密度，最终反而变慢。

因此，更准确的目标是：

获得足以隐藏延迟的 Occupancy，而不是追求最大的 Occupancy。

12. 一个 SM 能驻留多少线程块，由多个因素共同决定

线程块驻留数量并不只由 Shared Memory 决定。

通常需要同时满足以下限制：

[
N_{\text{resident blocks}}

\min
\begin{cases}
N_{\text{shared memory limit}}
N_{\text{register limit}}
N_{\text{thread limit}}
N_{\text{warp or wave limit}}
N_{\text{hardware block limit}}
\end{cases}
]

例如：

Shared Memory 最多允许：2 Blocks
寄存器最多允许：3 Blocks
线程数量最多允许：4 Blocks
硬件上限允许：8 Blocks

最终驻留数量为：

[
\min(2,3,4,8)=2
]

此时 Shared Memory 是 Occupancy 的主要限制因素。

如果减小 Shared Memory 后变成：

Shared Memory 最多允许：4 Blocks
寄存器最多允许：2 Blocks

最终仍然只能驻留两个线程块。

此时继续减少 Shared Memory 已经无法提高 Occupancy，因为瓶颈已经转移到了寄存器。

13. 在 TileLang 和 C500 算子优化中如何理解？

在 TileLang 算子中，经常可以调整以下参数：

block_token
block_dhidden
block_dexpert
threads
num_stages
每线程输出数量
Shared Memory Layout
向量化宽度

这些参数本质上都在影响计算粒度、数据复用和资源占用。

参数变化	可能收益	可能代价
增大 `block_token`	提高 Token 维度复用	Shared Memory 和寄存器增加
增大 `block_dhidden`	提高连续访存和计算密度	Tile 变大，边界浪费增加
增大 `block_dexpert`	提高 K 维数据复用	单线程块资源占用增加
增加 `num_stages`	让加载与计算重叠	Shared Memory 成倍增加
增加每线程输出数量	提高寄存器数据复用	累加器寄存器增加
增加线程数	提高并行搬运能力	单 Block 总寄存器消耗增加

例如，将：

num_stages = 1

修改为：

num_stages = 2

理论上可以实现：

当前 Tile 执行 MMA
下一 Tile 同时进行数据加载

但如果 Shared Memory 使用量从 50 KB 增加到 100 KB，使得每个 AP 或 SM 上的驻留线程块数量从两个下降到一个，那么需要比较两方面效果。

收益：

单个线程块内部的加载与计算重叠更充分

代价：

没有第二个线程块帮助隐藏 Barrier、访存和数据依赖延迟

最终是否提速，不能只根据 num_stages 判断，必须通过实际 Profile 和计时确认。

14. Profiler 中应该联合观察哪些指标？

只看 Occupancy 无法判断一个优化是否成功。

更合理的方法是联合观察多个指标。

14.1 AP MMA Duty

用于观察 MMA 计算单元是否持续工作。

如果数据复用优化后：

AP MMA Duty：40% → 65%

通常说明 MMA 单元得到了更充分的供数。

14.2 AP MTE Duty

用于观察数据搬运单元的繁忙程度。

如果 AP MTE 长时间满载，而 MMA Duty 较低，可能说明算子仍然受数据搬运限制。

14.3 Active Waves 和 Occupancy

用于观察每个 AP 或 SM 上是否有足够多的活跃 Wave。

Occupancy 很低且 MMA Duty 也很低时，可能存在 Ready Wave 不足的问题。

14.4 Shared Memory 使用量

观察每个线程块消耗的 Shared Memory，以及它是否限制驻留线程块数量。

14.5 每线程寄存器数量

观察寄存器是否成为 Occupancy 的限制因素。

即使 Shared Memory 很少，寄存器过多也可能导致驻留 Block 数量下降。

14.6 Local Memory 和 Spill

如果 Local Memory 访问增加，需要判断是否发生了寄存器 Spill。

Spill 通常会带来额外的显存访问和延迟。

14.7 L2 Hit Rate 和显存带宽

数据复用优化后，应该检查：

L2 命中率是否提高
Global Memory 流量是否减少
有效带宽是否改善

否则可能只是增加了片上资源占用，却没有真正减少片外访存。

15. 如何判断 Occupancy 下降是好事还是坏事？

考虑下面两组情况。

情况一：Occupancy 降低，但整体性能提高

修改前：
Occupancy = 75%
AP MMA Duty = 40%
L2 流量较高

修改后：
Occupancy = 50%
AP MMA Duty = 65%
L2 流量下降
运行时间下降

虽然 Occupancy 下降了，但数据复用让 MMA 单元更加繁忙，减少了显存访问，因此这是一次成功的优化。

情况二：Occupancy 降得过低，性能恶化

修改前：
Occupancy = 50%
AP MMA Duty = 65%

修改后：
Occupancy = 12.5%
AP MMA Duty = 38%
Shared Memory 接近上限
运行时间增加

这通常说明：

Tile 或 Buffer 过大
驻留 Wave 数量不足
无法隐藏访存和 Barrier 延迟
AP 上出现大量 Bubble

此时数据复用做得过头了。

16. 总结

数据复用是在用片上资源换取更少的片外访存：

[
\boxed{
\text{更多 Shared Memory 和 Register}
\quad\Longleftrightarrow\quad
\text{更少 Global Memory 访问}
}
]

但单个线程块占用的片上资源越多，同一个 SM 或 AP 上能够同时驻留的线程块越少：

[
\boxed{
\text{单个线程块计算能力更强}
\quad\Longleftrightarrow\quad
\text{线程块级并发度更低}
}
]

当 Occupancy 降得过低时，GPU 中可供调度的 Ready Wave 数量不足，访存、同步和指令依赖延迟无法被隐藏，最终会产生 Bubble。

因此，高性能算子优化的目标并不是：

数据复用越高越好

也不是：

Occupancy 越高越好

而是：

在不过度牺牲并发度的前提下，提高数据复用；同时保证有足够多的 Ready Wave，让 MMA、访存和流水线持续运行。

最终需要在以下因素之间找到平衡：

Global Memory 流量
Shared Memory 容量
寄存器压力
Occupancy
Pipeline Stage
Tile 大小
MMA 利用率
Barrier 和依赖延迟
边界浪费
Bank Conflict

这正是 GPU 高性能算子开发中所谓的 Balancing Act：平衡的艺术。

posted @ 2026-07-12 14:31 White_Swan 阅读(3) 评论(0) 收藏举报

刷新页面返回顶部

GPU 算子优化中的数据复用

GPU 算子优化中的数据复用

1. 高性能算子优化的四个支柱

2. 什么是数据复用？

[ \text

3. 数据复用为什么会消耗更多片上资源？

4. Shared Memory 为什么会降低 Occupancy？

4.1 每个线程块使用 16 KB

4.2 每个线程块使用 48 KB

4.3 使用 Double Buffering

5. 为什么提高复用通常也会增加寄存器压力？

5.1 更大的输出 Tile 需要更多累加器

5.2 多级流水线增加变量存活时间

6. 寄存器压力为什么会降低 Occupancy？

6.1 每线程使用 32 个寄存器

6.2 每线程使用 64 个寄存器

6.3 每线程使用 128 个寄存器

6.4 Occupancy 通常是阶梯式变化

7. Occupancy 下降为什么可能让速度变慢？

7.1 Occupancy 较高时

7.2 Occupancy 太低时

8. 一个典型例子：Tile 越大，为什么可能先快后慢？

8.1 方案 A：小 Tile

8.2 方案 B：中等 Tile

8.3 方案 C：超大 Tile 和 Triple Buffering

9. 寄存器压力过高还可能导致 Spill

10. Shared Memory 复用还可能带来哪些额外成本？

10.1 Barrier 同步开销

10.2 Shared Memory Bank Conflict

10.3 边界 Tile 浪费

11. Occupancy 不是越高越好

11.1 内存受限算子

11.2 计算密集型 GEMM

12. 一个 SM 能驻留多少线程块，由多个因素共同决定

[ N_{\text{resident blocks}}

13. 在 TileLang 和 C500 算子优化中如何理解？

14. Profiler 中应该联合观察哪些指标？

14.1 AP MMA Duty

14.2 AP MTE Duty

14.3 Active Waves 和 Occupancy

14.4 Shared Memory 使用量

14.5 每线程寄存器数量

14.6 Local Memory 和 Spill

14.7 L2 Hit Rate 和显存带宽

15. 如何判断 Occupancy 下降是好事还是坏事？

情况一：Occupancy 降低，但整体性能提高

情况二：Occupancy 降得过低，性能恶化

16. 总结

公告

[
\text

[
N_{\text{resident blocks}}