4、算子融合技术（下）：Multi-head Attention融合、LayerNorm融合、QKV融合技巧

好，咱们接着聊算子融合。上一讲我们把基础融合讲透了，这一讲直接上硬菜——Transformer模型里的三大融合技巧。说实话，这几年我经手的部署项目里，十个有八个都在跟Transformer打交道。Multi-head Attention、LayerNorm、QKV，这三个家伙要是没处理好，推理速度能差出好几倍。

一、QKV融合：把三个矩阵乘法变成一个

先说说QKV融合。你想想看，标准的Self-Attention里，输入X要分别乘以Wq、Wk、Wv三个权重矩阵。这意味着什么？三次独立的矩阵乘法，三次显存读写，三次kernel launch。我刚开始做BERT部署时，看着profiling结果里那三个并排的MatMul，心里就犯嘀咕：这明明可以一起干啊。

QKV融合的核心思路很简单：把Wq、Wk、Wv拼成一个大矩阵，一次乘法搞定。具体来说：

// 原始做法
Q = X @ Wq   // shape: [B, S, D] @ [D, D] -> [B, S, D]
K = X @ Wk   // shape: [B, S, D] @ [D, D] -> [B, S, D]
V = X @ Wv   // shape: [B, S, D] @ [D, D] -> [B, S, D]

// QKV融合后
W_qkv = concat(Wq, Wk, Wv, axis=1)  // shape: [D, 3*D]
QKV = X @ W_qkv                       // shape: [B, S, 3*D]
Q, K, V = split(QKV, 3, axis=2)      // 切分成三个 [B, S, D]

这样做的好处很明显：

  • 减少kernel launch次数：从3次降到1次。别小看这个，GPU上kernel launch的开销可不小，尤其是小batch时特别明显。
  • 提高计算密度：大矩阵乘法比三个小矩阵乘法更高效，能更好地利用Tensor Core。
  • 减少显存带宽占用：输入X只需要读一次，而不是三次。

实际收益：我在T4上测试过，QKV融合后，单层Transformer的推理时间能减少15%-25%。batch size越小，收益越明显。

小技巧：做QKV融合时，注意权重矩阵的布局。我习惯用NCHW格式，但有些框架用NHWC。融合前一定要确认好，不然split出来的Q、K、V顺序会乱。

二、Multi-head Attention融合：把碎片化操作打包

QKV融合只是第一步。真正的重头戏在Multi-head Attention内部。你想想看，一个标准的MHA里有多少个小操作？

  1. QKV投影（我们已经融合了）
  2. 切分成多个head
  3. 计算attention score（Q @ K^T）
  4. scale + mask + softmax
  5. attention score @ V
  6. 合并多个head
  7. 输出投影

这里面每一步都可能产生中间张量，每一步都需要一次kernel launch。我见过一个项目，MHA里光kernel launch就占了30%的时间。这太浪费了。

MHA融合的思路是：把连续的几个操作合并成一个kernel。常见的融合方式有：

1. QK^T + scale + mask + softmax 融合

这四个操作是串行的，而且都是element-wise或reduction操作。把它们写进一个kernel里，中间结果直接留在寄存器里，不用写回显存。

// 伪代码：融合后的kernel
__global__ void fused_attention_kernel(float* Q, float* K, float* V, float* output, int S, int D) {
    // 每个thread block处理一个query位置
    int row = blockIdx.x;
    
    // 计算Q @ K^T的一行
    float score[S];
    for (int j = 0; j < S; j++) {
        score[j] = dot_product(Q[row], K[j]);
        score[j] /= sqrt(D);  // scale
        score[j] = (mask[row][j]) ? score[j] : -INFINITY;  // mask
    }
    
    // softmax
    float max_val = max(score, S);
    float sum_val = 0;
    for (int j = 0; j < S; j++) {
        score[j] = exp(score[j] - max_val);
        sum_val += score[j];
    }
    for (int j = 0; j < S; j++) {
        score[j] /= sum_val;
    }
    
    // 直接计算score @ V
    for (int k = 0; k < D; k++) {
        float val = 0;
        for (int j = 0; j < S; j++) {
            val += score[j] * V[j][k];
        }
        output[row * D + k] = val;
    }
}

注意：这种融合对显存带宽的优化非常明显，但实现起来有坑。我曾经在A100上踩过一个雷——当序列长度S超过1024时，score数组放不进shared memory，导致性能反而下降。后来我改用tiling策略，把长序列切分成多个block处理，才解决了这个问题。

2. 多head并行计算

不同head的计算是独立的，天然适合并行。我一般把head维度映射到block或warp上，让所有head同时计算。这样能充分利用GPU的并行能力。

三、LayerNorm融合：把归一化塞进前一个kernel

LayerNorm在Transformer里无处不在。每个子层后面都跟着一个LayerNorm。但LayerNorm本身的计算量不大，主要开销在显存读写上——它需要读输入、计算均值和方差、再做归一化。

LayerNorm融合的思路有两种：

1. 与前面的MatMul融合

比如在FFN里，做完GELU激活后紧接着就是LayerNorm。我们可以把LayerNorm的计算塞进GELU的kernel里，省掉一次中间结果的读写。

// 融合前
h = gelu(X @ W1)      // 写回显存
y = layernorm(h)      // 从显存读h

// 融合后
// 在gelu kernel里直接做layernorm
__global__ void fused_gelu_layernorm(float* X, float* W, float* gamma, float* beta, float* Y, int M, int N) {
    // 先计算gelu
    float val = gelu(dot_product(X[row], W[col]));
    // 这里不写回显存，直接做layernorm
    // ... 计算均值和方差
    // ... 归一化
    Y[row * N + col] = gamma[col] * normalized_val + beta[col];
}

2. 与后面的Add融合

Transformer里有个残差连接：output = layernorm(x + sublayer(x))。这里Add和LayerNorm可以融合成一个kernel。我习惯叫它「Add+LayerNorm融合」。

我的经验：LayerNorm融合的收益在batch size较小时特别明显。有一次我在一个对话模型上做优化，batch size=1时，LayerNorm融合让端到端推理快了8%。但batch size=64时，收益就降到2%了。原因很简单：大batch时计算瓶颈在MatMul上，小batch时显存带宽才是瓶颈。

四、融合策略的选择与权衡

说了这么多融合技巧，你可能会问：是不是融合得越多越好？

嗯，还真不是。我见过有人把整个Transformer层写成一个超大的kernel，结果编译时间长达半小时，而且因为寄存器压力太大，实际性能反而下降了。

我的建议是：

融合策略 适用场景 收益 风险
QKV融合 几乎所有Transformer 高，稳定 低
QK^T+softmax+PV融合 短序列（S < 512） 高 中等（寄存器压力）
LayerNorm+Add融合 小batch场景 中等 低
全层融合 极少使用 不确定 高（编译时间长，调试困难）

核心原则：融合的目标是减少显存带宽瓶颈，而不是减少计算量。如果融合后寄存器压力过大导致spilling，那还不如不融合。我一般用nvcc的--ptxas-options=-v参数检查寄存器使用情况，超过64个寄存器就要小心了。

五、实战中的避坑指南

最后分享几个我踩过的坑：

  • 精度问题：融合后的kernel里，中间计算精度要统一。我遇到过FP16下QK^T累加时精度不够，导致attention score分布异常。后来在累加时用了FP32累加器才解决。
  • 动态shape：很多融合kernel需要知道序列长度S。如果S是动态的，就得用模板或动态编译。我习惯用CUDA的JIT编译，在运行时根据实际S生成最优kernel。
  • 调试困难：融合后的kernel很难调试。我的做法是先写一个naive版本做golden，然后逐块替换成融合版本，每一步都验证精度。

好了，算子融合这部分就讲到这里。下一讲我们聊聊量化，那又是一个大话题。记住，融合不是目的，性能才是。别为了融合而融合，多看看profiling结果，让数据说话。

← 上一章 📖 返回目录 下一章 →