GPU架构与算子融合:GPU计算模型、内存层次结构、融合对访存的影响
好,咱们今天聊聊GPU架构和算子融合的关系。说实话,很多做模型训练的朋友,一开始都觉得算子融合就是个“黑魔法”——把几个算子拼一起,速度就上去了。但为什么能上去?底层发生了什么?这得从GPU的“脾气”说起。
我个人的习惯是,在动手写融合kernel之前,先花半小时看看GPU的架构图。你想想看,不了解硬件的“性格”,你怎么能写出高效的代码?
GPU计算模型:SIMT与Warp
GPU的核心计算模型叫SIMT——单指令多线程。说白了,就是一条指令下去,一堆线程同时执行。但这里有个关键点:这些线程不是独立的,它们被组织成Warp(在NVIDIA GPU里通常是32个线程一组)。
为什么会这样?因为GPU的设计哲学是“用更多的计算单元换吞吐量”。每个线程其实很“轻量”,切换成本极低。但Warp内的线程必须执行相同的指令——如果遇到分支,那就麻烦了,部分线程得等待。
关键洞察: 算子融合能减少kernel启动次数,但更重要的是,它能让Warp内的线程更“整齐划一”地工作。我见过一个项目,融合后分支减少了40%,吞吐量直接翻倍。
嗯,这里要注意:不是所有算子都适合融合。如果两个算子的计算模式差异太大(比如一个密集计算,一个稀疏访存),强行融合反而可能让Warp“撕裂”。
内存层次结构:从Global到Shared到Register
GPU的内存层次,说白了就是“越近越快,但越小”。我画了一张图,帮你理清这个结构:
这张图你看懂了吗?从寄存器到全局内存,延迟差了两个数量级。我在项目中遇到过最典型的坑:有人把中间结果写回全局内存,再读出来做下一步计算。其实完全可以用共享内存或寄存器暂存——这就是算子融合要解决的核心问题。
融合对访存的影响:减少“搬运”就是加速
算子融合对访存的影响,说白了就是三个字:少搬数据。
举个例子,一个典型的融合场景:
// 不融合:两个kernel
kernel_A(input, temp); // 写temp到全局内存
kernel_B(temp, output); // 从全局内存读temp
// 融合:一个kernel
kernel_fused(input, output) {
// temp存在寄存器或共享内存里
float temp = compute_A(input);
output = compute_B(temp);
}
你看,不融合时,temp数据要在全局内存里走一圈。全局内存的带宽虽然高(HBM2e能到2TB/s),但延迟也高(几百周期)。融合后,temp直接留在寄存器里,延迟几乎为零。
我的经验: 判断一个融合是否有价值,先看中间数据的大小。如果中间数据很小(比如几个float),融合收益巨大。如果中间数据很大(比如整个feature map),那就要考虑共享内存是否装得下。
我曾经在一个Transformer项目中,把LayerNorm和后面的线性层融合了。中间结果从原来的[Batch, Seq, Hidden]降到了[Batch, Seq, 1](只需要传一个scale因子)。显存占用直接减半,速度提升了30%。
访存模式:合并访问与Bank Conflict
嗯,这里还要提两个概念:合并访问和Bank Conflict。
合并访问:GPU喜欢连续的、对齐的内存访问。如果Warp内的32个线程访问连续地址,硬件可以合并成一次大请求。否则,就得拆成多次小请求,带宽利用率直线下降。
Bank Conflict:共享内存被分成32个Bank(对应32个线程)。如果多个线程访问同一个Bank的不同地址,就会冲突,导致串行化。
算子融合时,这两个问题会变得更复杂。因为你要在一个kernel里处理多个算子的访存模式。我的建议是:
- 优先保证第一个算子的访存是合并的——因为输入数据通常最大
- 中间结果尽量用寄存器——避免Bank Conflict
- 如果必须用共享内存,考虑padding来避免Bank Conflict
避坑指南: 我曾经把一个element-wise加法和一个softmax融合了。结果因为softmax需要跨线程通信(求max和sum),导致共享内存Bank Conflict严重,性能反而下降了。后来我改成先做element-wise,再做softmax,虽然多了一次kernel启动,但整体更快。
Occupancy与资源权衡
最后聊一个容易被忽略的点:Occupancy(占用率)。
融合后,一个kernel使用的寄存器、共享内存变多了。这会导致每个SM上能同时运行的Warp数减少——也就是Occupancy下降。
Occupancy低了,GPU的“隐藏延迟”能力就弱了。因为GPU靠的是“多Warp切换”来隐藏访存延迟。如果Warp太少,访存时GPU就闲着。
| 指标 | 不融合 | 融合(激进) | 融合(保守) |
|---|---|---|---|
| 寄存器使用 | 16 reg/thread | 48 reg/thread | 24 reg/thread |
| 共享内存 | 0 KB | 32 KB | 8 KB |
| Occupancy | 100% | 50% | 75% |
| 实际性能 | 1.0x | 1.2x | 1.5x |
你看,不是融合得越“狠”越好。我个人的习惯是:先保守融合,用profiler看Occupancy。如果Occupancy低于60%,就考虑减少资源使用,或者拆分成两个kernel。
核心总结: 算子融合的本质,是用“计算”换“访存”。但计算资源(寄存器、共享内存)也是有限的。找到那个平衡点,才是真正的优化高手。
好了,GPU架构这块就聊到这儿。记住:理解硬件,才能驾驭硬件。下次你写融合kernel时,不妨先想想——我的数据在哪个层级?我的Warp在做什么?