4. GPU架构与算子融合

好，咱们今天聊聊GPU架构和算子融合的关系。说实话，很多做模型训练的朋友，一开始都觉得算子融合就是个“黑魔法”——把几个算子拼一起，速度就上去了。但为什么能上去？底层发生了什么？这得从GPU的“脾气”说起。

我个人的习惯是，在动手写融合kernel之前，先花半小时看看GPU的架构图。你想想看，不了解硬件的“性格”，你怎么能写出高效的代码？

GPU计算模型：SIMT与Warp

GPU的核心计算模型叫SIMT——单指令多线程。说白了，就是一条指令下去，一堆线程同时执行。但这里有个关键点：这些线程不是独立的，它们被组织成Warp（在NVIDIA GPU里通常是32个线程一组）。

为什么会这样？因为GPU的设计哲学是“用更多的计算单元换吞吐量”。每个线程其实很“轻量”，切换成本极低。但Warp内的线程必须执行相同的指令——如果遇到分支，那就麻烦了，部分线程得等待。

嗯，这里要注意：不是所有算子都适合融合。如果两个算子的计算模式差异太大（比如一个密集计算，一个稀疏访存），强行融合反而可能让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直接留在寄存器里，延迟几乎为零。

我曾经在一个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

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在做什么？