GPU性能-矩阵乘法

• Acronyms
• 2 理论：计算和内存瓶颈分析
  • 2.1 计算量评估
  • 2.2 内存访问量评估（Bytes）
  • 2.3 Arithmetic Intensity
• 3 GPU GEMM优化办法
  • 3.1 Tile the Output
  • 3.2 Tensor Core
  • 3.2 Typical Tile Dimensions in cuBLAS
  • 3.4 tile size的影响
• Tile Quantization
• Wave Quantization
• 总结

Acronyms

GEMMs: General Matrix Multiplications

矩阵乘法在Fully connected layers, Recurrent Layers等计算操作中广为使用，用于计算： $C = \alpha AB + \beta C$，其中A和B是matrix inputs，\alpha 和\beta是scalar inputs，C是一个pre-existing matrix，未来会被计算output覆盖

FC的forward pass中，weight matrix为arguement A，activation为B，\aplha和\bate分别为1，0如果把skip-connection + linear operation加起来的话，\beta也会是1.

2 理论：计算和内存瓶颈分析

2.1 计算量评估

Assume:

A: M * N
B: N * K
C: M * K

则矩阵乘法C=AxB 在FMA中计算量共M*N*K 次

因为每个元素Cij需要累加N次乘法，每次取A的一行和B的一列相乘。而C一共有M*K个元素，所以一共的操作数是M*N*K

每个FMA包含1次乘法1次加法，共 2*M*N*K FLOPS

2.2 内存访问量评估（Bytes）

输入数据：load A （MxK元素）和B（KxN元素） 输出数据：write C（MxN元素） 总字节数：假设数据类型为 FP16，每个元素占2Byte；则总字节数为 2*(M*K+K*N+M*N)

2.3 Arithmetic Intensity

Arithmetic Intensity 
= FLOPS/Bytes
= 2*M*N*K / (2*(M*K+K*N+M*N))
= M*N*K / (M*K+K*N+M*N)
\approx N/2 (When M=N, it equals)

ops:byte V100理论算力=112 TFLOPS 显存带宽 900 GBps

V100 ops: byte=112,000/900=124.4 FLOP/Byte

实际计算中，如果算法Arithmetic Intensity比这个值大，则为math bounded, 否则为memory bounded.

一些例子： 小规模GEMM： 8192 * 128 * 8192 = 128/2 < 112 => Memory Bounded

大规模GEMM： 8192 * 8192 * 8192 = 8192/2 > 112 => Math Bounded

GEMV： M=1或者N=1，Arithmetic Intensity一般小于1，始终内存受限

局限性 上述计算简化的部分：

• 忽略非计算指令（指针运算、循环控制）的影响
• 未考虑GPU的缓存层次，比如L1，共享内存 对数据复用的优化潜力

实际应用的时候的优化方向：

• Memory Bound -> 复用数据/混合精度
• Math Bound -> 对齐维度/开启Tensor Core

Nsight & cuBLAS日志也可以用于验证假设

3 GPU GEMM优化办法

3.1 Tile the Output

GPU伤的GEMM优化采用分块计算策略，即将结果矩阵C话费为多个Tile(Mtile x Ntile)，每个Tile由ThreadBlock负责计算。

计算流程（如Figure1）：

1. 沿K维度分步累加：ThreadBlock逐次家在A和B的子矩阵快，执行乘加操作并累加到输出Tile
2. 数据复用：大尺寸Tile可以提高数据在共享内存中的复用率，减小显存访问次数

3.2 Tensor Core

关键维度对齐：

• 仅“内存中最快变化的维度”需要严格对齐，但是建议所有维度都对齐，以保持最佳性能。
• “内存中最快变化的维度”指：
  • 行优先存储(Row-major)的时候，A[i,j]存储在A[i*N+j]的位置，最快变化的维度为列为j
  • 列优先存储(Column-major)的时候,A[i,j]存储在A[j*M+i]的位置，最快变化的维度的列为i

因为TensorCore的原理，是通过Wrap-Lavel Matrix Multiply Accumulate(WMMA)执行计算，主要利用Shared Memory访存快实现性能优化；所以我们用TensorCore+对齐的矩阵Shape，本质上是希望对齐wrap的shape，减少Bank Conflict，加速wrap tile加载，并且充分利用L2 Cache和Shared Memory。

优先对齐：将 M/N/K 对齐到 16（FP16 为 8，A100 为 64）。

次优选择：若无法完全对齐，选择更小的 2 的幂次（如 4/8 字节）仍可提升性能。

3.2 Typical Tile Dimensions in cuBLAS

1. 分块尺寸选择

• 大分块（如 256×128）：
  • 优点：高数据复用率，减少显存带宽压力，计算效率高（图5）。
  • 缺点：生成的分块数量少，可能导致 GPU 并行度不足（SM 利用率低）。
• 小分块（如 64×64）：
  • 优点：生成更多分块，提升并行度，适合小规模 GEMM。
  • 缺点：数据复用率低，显存带宽压力大，效率下降（图5）。

2. 动态分块策略

• 启发式选择：cuBLAS 根据 GEMM 尺寸自动选择最优分块（如大 GEMM 优先使用 256×128）。

• 基准模式：部分框架（如 PyTorch）在训练前测试所有分块策略，选择最快实现（一次性开销）。

3. GEMM 规模与分块关系

• 大型 GEMM：适合大分块，充分占用 GPU 并行资源，接近峰值算力

• 小型 GEMM：需小分块提升并行度，但效率显著下降

4. 典型分块尺寸

• 优先级排序：
  1. 256×128 / 128×256（最高效）
  2. 128×128
  3. 256×64 / 64×256
  4. 128×64 / 64×128
  5. 64×64（最低效，仅用于极小 GEMM）

3.4 tile size的影响

当M，N都很小，但是K很大的时候，NVIDIA Library也提供了单独对K做Tile的能力。因为K是dot product的方向，所以对K Tile之后的reduction也会成为性能瓶颈；这个功能默认是关闭的状态。

Tile Quantization

当矩阵维度（如输出矩阵的M，N）无法被Thread Block的分块大小整除的时候，导致部分分块的计算区域远小于分块容量，产生无效计算。

例：假设分块大小为128x128，若矩阵维度为192x192，则需要划分为4个分块（2x2布局），但第二个分块并没有利用充分（192-128=64）

影响：

• 无效计算增加：所有分块执行相同量的Math Ops，但部分分块的Tile有效计算占比较低。
• 吞吐骤降：当分块数量不变但有效数据减少的时候，实际FLOPS显著降低。
• 执行时间固定：只要分块数量不变，总时间基本不变

例： 考虑GEMM

M=27648
K=4096
Tile Size 256×128

当N=128->136，Tile Number=1->2但第二个Tile仅含有6.25%有效数据，导致吞吐骤降

优化策略：

• 对齐矩阵维度：确保M，N为分块大小的整数倍（cuBLAS因此默认选择更小的分块大小，如64x64，以减少浪费）
• 动态分块选择：依赖cuBLAS的启发式算法或框架的bencmark模式自动优化

Wave Quantization

由于total Thread Block数量需要和GPU的SM数量对齐，才能尽量让Tail Wave中的SM充分利用，所以我们一般希望总Thread Block数量为Wave的整数倍。

例：

M=2304
Tile Size=256x128
当N=1536, # ThreadBlock = 108
当N=1664, # ThreadBlock = 117，存在一个仅有9个thread的Tail Wave

Tail Wave会导致FLOPS骤降，也会增加总体执行时间。

优化策略：

• 设计FEMM规模，s.t. #Thread Block = #SM的整数倍
• 调整ThreadBlock Size，比如用更小的Thread Block划分方式来增加Thread Block数量，减小Tile Wave浪费

总结

1. 优先对齐维度

• Tile Quantization：确保 M、N 为分块大小的整数倍（如 128×128）。
• Wave Quantization：设计 GEMM 规模使线程块总数接近 SM 数量的整数倍（如 A100 上为 108 的倍数）。

2. Dynamic tile分块选择

• 大型 GEMM 使用大分块（256×128）提升效率，小型 GEMM 使用小分块（64×64）增加并行度。
• 启用框架的“benchmark”自动选择最优分块策略。

3. 硬件适配

• 在 A100 等高端 GPU 上，利用更大对齐粒度（如 FP16 对齐到 64 元素）优化显存带宽。

4. 监控与验证

• 使用性能分析工具（Nsight Systems）观察分块策略与 SM 利用率。
• 结合 cuBLAS 日志验证实际分块选择是否符合预期。