想象你是这家 GPU 大饭店的厨房总管。食材仓库(HBM)在楼下,备菜案板(SRAM)就在手边,而炉台上的每位厨师都算得飞快。可老厨师有个坏习惯:每做好一步,都要把半成品端回楼下仓库,下一步再跑下去端上来。炒一盘 Linear 配 GELU,要跑三四趟;调 Attention 的三款酱料 Q、K、V,更是来回折腾。菜没出几道,腿先跑细了——整个后厨的瓶颈,从来不在火候,而在跑腿。 Triton 算子融合,就是给后厨立的一条新规矩:能留在案板上做完的,绝不下楼。把 Linear 和 GELU 放进同一个 kernel 一口锅炒,中间结果留在 SRAM;把 Q、K、V 三次投影合成一次大矩阵乘法,调好酱料再分装。省掉的,正是 GPU 上最贵的东西——访存带宽。 这篇文章会从“为什么要融合”讲起,带你认识 Triton 的块级编程思维,亲手实现 Linear+GELU 融合,再攻下 Attention 的 QKV 投影融合与在线 Softmax,最后给出工程验收清单。备好围裙,我们开始。

岗位关联

算子融合是 AI 编译与性能优化的核心杠杆:理解它,你才能把同样的 GPU 资源榨出成倍的吞吐,是算子开发工程师从“能跑”迈向“跑得快”的必修课。

目录
  1. 可视化演示
  2. 融合背景动机
  3. Triton核心概念
  4. Linear算子剖析
  5. GELU激活融合
  6. QKV投影融合
  7. 注意力内核融合
  8. 工程实践验收
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

本节配有交互动效(进入视口后自动加载,约 3KB)。

算子融合前后对比:访存开销变化 \u2717 未融合 (3个独立kernel)\u2193 6次访存\u2713 融合后 (1个kernel)\u2717 未融合 (3个独立kernel)\u2717 未融合 (3个独立kernel)\u2193 6次访存\u2713 融合后 (1个kernel)\u2717 未融合 (3个独立kernel) \u2717 未融合 (3个独立kernel)\u2193 6次访存\u2713 融合后 (1个kernel)\u2717 未融合 (3个独立kern...

融合背景动机

老厨师做青椒炒肉,流程是这样的:下楼取肉,切好,端回仓库;下楼取青椒,切好,端回仓库;再把两份半成品端上来,下锅。每一步“端盘子”都要穿过又长又挤的走廊。GPU 里的走廊就是 HBM 总线,带宽再大,也架不住来回倒腾。

计算单元在芯片上飞速运转,但数据从 HBM 搬进寄存器,走的是片外总线,延迟和能耗是片上 SRAM 的几十倍。如果两个算子之间存在大块中间结果,每写一次 HBM 再读一次,就等于白跑两趟仓库。

把时间开销写成公式,就一目了然:

所以融合的本质不是少算,而是少搬。Linear+GELU 融合后,中间结果从 HBM 的写读两次,变成 SRAM 上的原地变换;Attention 融合后,QKV 三次投影合并为一次。跑腿少了,上菜自然就快。

Tmove=DB\displaystyle T_{\text{move}} = \frac{D}{B}
  1. T — 搬运总耗时这是我们要优化的目标。融合前后,火候(计算量)几乎不变,变的只有 T。
    T=完成数据搬运所需的时间(秒)T = \text{完成数据搬运所需的时间(秒)}
  2. D — 搬运数据量所有中间结果进出 HBM 的字节总数。融合把这个分子压到最小。
    D=所有读写 HBM 的字节数之和D = \text{所有读写 HBM 的字节数之和}
  3. B — 访存带宽走廊每秒能通过的人数,由硬件决定。我们能改变的是 D,而不是 B。
    B=HBM 带宽(字节/秒)B = \text{HBM 带宽(字节/秒)}

Triton核心概念

新主厨 Triton 拿到一本《块级菜谱》:他不用操心整个饭店有几张桌子,只需写清楚“一桌菜”怎么做——比如一桌四人的分量怎么切配、怎么下锅。写好后,饭店自动把这套流程复制到每一桌。这就是块级编程(Block-Level Programming)。

在 Triton 里,我们自己写的这个“一桌菜”叫 kernel;每一桌是一个 program,由 program_id 唯一标识。核心里用 tl.arange 生成块内编号,用 tl.load / tl.store 读写数据,用 tl.dot 做矩阵乘。编译器负责把块内的循环变成 GPU 上的向量化指令,并自动调度到所有计算单元上。

一桌处理多大的块,用网格公式描述:

有了“一桌菜”的思维,融合就变得自然:你只需要在这一桌的案板(SRAM)上连续完成多道工序,中间不换桌、不下楼。

G=⌈MBM⌉×⌈NBN⌉\displaystyle G = \left\lceil \frac{M}{B_M} \right\rceil \times \left\lceil \frac{N}{B_N} \right\rceil
  1. M / N — 输出规模输出矩阵的总行数和总列数,对应整个后厨要服务的桌数。
    M=输出矩阵行数,N=输出矩阵列数M = \text{输出矩阵行数}, \quad N = \text{输出矩阵列数}
  2. B_M / B_N — 块大小每个 program 一次处理的行数和列数,对应一桌坐几位客人。
    BM=每个块处理的行数,BN=每个块处理的列数B_M = \text{每个块处理的行数}, \quad B_N = \text{每个块处理的列数}
  3. 向上取整⌈⌉ 表示有余数就多开一桌,确保所有元素都有程序负责。
    ⌈x⌉=大于等于 x 的最小整数\left\lceil x \right\rceil = \text{大于等于 } x \text{ 的最小整数}
  4. 总 program 数 G两者的乘积就是整个 grid 的规模,GPU 会把它铺到所有 SM 上执行。
    G=总共需要启动的 program 数量G = \text{总共需要启动的 program 数量}

Linear算子剖析

Linear 层在神经网络里,相当于“调底料”:输入 X 是成筐的食材,权重 W 是配方表,偏置 b 是最后的点睛之盐。每一份输出,都是把某一行食材和一列配方逐项相乘再累加的结果。

Triton 内核里,我们把输出划分成 BLOCK_M x BLOCK_N 的块,每个 program 负责一块;外层沿输出网格铺开,内层按 BLOCK_K 分小段,用 tl.dot 计算并累加。累加器 acc 在寄存器里一直接力,直到 K 维全部算完。

底料调好了,传统做法是盛出来送回仓库,等下一个算子 GELU 再端上来。而我们要做的,是让它留在锅里直接调味——这就是下一节的 GELU 融合。

Ym,n=∑k=0K−1Xm,k Wn,k+bn\displaystyle Y_{m,n} = \sum_{k=0}^{K-1} X_{m,k} \, W_{n,k} + b_n
  1. X_{m,k} — 输入元素第 m 行输入的第 k 个特征,来自上一层激活后的结果。
    Xm,k=输入矩阵第 m 行第 k 列元素X_{m,k} = \text{输入矩阵第 } m \text{ 行第 } k \text{ 列元素}
  2. W_{n,k} — 权重元素第 n 个输出神经元对应的第 k 个权重。
    Wn,k=权重矩阵第 n 行第 k 列元素W_{n,k} = \text{权重矩阵第 } n \text{ 行第 } k \text{ 列元素}
  3. K 次累加把所有 k 从 0 到 K-1 的乘积加起来,这就是一次内积。
    ∑k=0K−1Xm,kWn,k=Xm,0Wn,0+Xm,1Wn,1+⋯\sum_{k=0}^{K-1} X_{m,k} W_{n,k} = X_{m,0}W_{n,0} + X_{m,1}W_{n,1} + \cdots
  4. b_n — 偏置给第 n 个输出加上一个可学习的常数项。
    +bn=每个输出神经元独立的偏置项+ b_n = \text{每个输出神经元独立的偏置项}

GELU激活融合

GELU 是一个“看人下菜”的调味师:输入为正,基本原样保留;输入为负,狠狠压小;输入接近 0,适当削弱。它不像 ReLU 那样一刀切(负的全变 0),而是平滑地筛出“这条食材值得保留”的概率。

由于 GELU 是逐元素的纯函数(输入一个数,输出一个数),它不需要跨行跨列的“社交”,是最适合融合的调味料。把它直接接在 Linear 的累加器后面,对 acc 的每个元素套公式即可。

在代码里这只是一行:acc = acc * 0.5 * (1.0 + tl.erf(acc * 0.70710678))。累加完成立刻调味,锅都不用挪。如果拆成两个 kernel,这一行前后就是两次 HBM 往返——你算算省了多少趟。

GELU(x)≈0.5x(1+tanh⁡(2π(x+0.044715x3)))\displaystyle \mathrm{GELU}(x) \approx 0.5x\left(1 + \tanh\left(\sqrt{\frac{2}{\pi}}\left(x + 0.044715x^3\right)\right)\right)
  1. x³ 微小修正括号里的 0.044715x³ 是对高斯累积函数的泰勒近似补充。
    x+0.044715x3=修正后的输入x + 0.044715x^3 = \text{修正后的输入}
  2. tanh 压缩把修正后的值压到 (-1, 1) 区间,模拟正态分布的尾部形状。
    tanh⁡(z)=ez−e−zez+e−z\tanh(z) = \frac{e^z - e^{-z}}{e^z + e^{-z}}
  3. 1 + tanh(……)把区间平移到 (0, 2),相当于一个“保留概率”的未缩放版本。
    1+tanh⁡(z)∈(0,2)1 + \tanh(z) \in (0, 2)
  4. 0.5x 整体缩放最后乘以 x 本身,保证输入越大、保留越多。
    0.5x(1+tanh⁡(z))≈x⋅Φ(x)0.5x\left(1 + \tanh(z)\right) \approx x \cdot \Phi(x)

QKV投影融合

Attention 需要三种酱料:查询 Q、键 K、值 V。对同一个输入 X,憨厚的旧做法是分三次投影:先调 Q,送回仓库;再调 K,送回仓库;最后调 V。三次都从同一个 X 出发,等于把一份底料端下楼三次。

聪明的主厨会把三张配方表横向拼成一张大表 W_qkv,一次矩阵乘同时得到 Q、K、V 三个酱料桶,再按列切分。由于 GPU 的矩阵乘单元天然擅长大批量计算,合并后三个小矩阵乘变成一个大矩阵乘,利用率往往更高。

切分好的 Q、K、V 同样不需要下楼——它们就在案板(SRAM)上,直接进入下一道注意力工序。

[Q∣K∣V]=XWqkv⊤+bqkv\displaystyle [Q \mid K \mid V] = X W_{qkv}^{\top} + b_{qkv}
  1. W_qkv — 拼接权重把 W_q、W_k、W_v 按列拼接得到的权重大矩阵。
    Wqkv=[Wq∣Wk∣Wv]∈R(3d)×KW_{qkv} = [W_q \mid W_k \mid W_v] \in \mathbb{R}^{(3d) \times K}
  2. 一次矩阵乘X 与 W_qkv 转置相乘,得到 (M, 3d) 的拼接结果,三份投影共用一次数据搬运。
    XWqkv⊤∈RM×3dX W_{qkv}^{\top} \in \mathbb{R}^{M \times 3d}
  3. b_qkv — 拼接偏置三个偏置向量拼接后,同样一次广播加完。
    bqkv=[bq∣bk∣bv]b_{qkv} = [b_q \mid b_k \mid b_v]
  4. 按列切片把结果切成 Q、K、V 三份,分别对应前 d、中 d、后 d 列。
    Q=H[:,:d],K=H[:,d:2d],V=H[:,2d:]Q = H[:, :d], \quad K = H[:, d:2d], \quad V = H[:, 2d:]

注意力内核融合

注意力就像给每位客人上菜前做一次“搭配确认”:Q 问“这道菜想要什么”,K 说“这道菜有什么”,两者一比得到相似度,再用 softmax 归一化成权重,最后按权重把 V(食材)组合起来。关键难点是:完整 softmax 需要看到整行所有分数,而块式计算一次只看一部分。

FlashAttention 的思路是“边尝边调整”:每算完一个分块,用当前最大值更新运行最大值 m,并用指数缩放修正旧的累加和,最后统一除。这样既不需要先算出完整的 QKᵀ 矩阵,也无需把整行分数写回 HBM,正是融合思想在注意力上的应用。

配合在线 softmax 的运行最大值技巧,即使完整的 QKᵀ 矩阵永远不在 HBM 里出现,数值结果也和标准实现完全一致。这就像厨师边尝边加盐——不需要等所有菜上齐才定咸淡。

Attention(Q,K,V)=softmax(QK⊤dk)V\displaystyle \mathrm{Attention}(Q,K,V) = \mathrm{softmax}\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V
  1. QKᵀ — 相似度矩阵查询与键两两内积,得到 (M, N) 的得分表。
    QK⊤=每对查询与键的点积得分QK^{\top} = \text{每对查询与键的点积得分}
  2. 除以 √d_k防止维度大时点积数值过大,把 softmax 推向梯度消失的饱和区。
    QK⊤dk=缩放后的得分\frac{QK^{\top}}{\sqrt{d_k}} = \text{缩放后的得分}
  3. softmax 归一化把每行得分变成和为 1 的权重。
    softmax(z)i=ezi∑jezj\mathrm{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}
  4. 乘以 V按权重把值向量加权累加,得到最终输出。
    softmax(⋅) V=加权的值向量组合\mathrm{softmax}(\cdot)\, V = \text{加权的值向量组合}

工程实践验收

新菜谱好不好,得上菜计时。工程上,融合 kernel 上线前要做三件事:测速、调参、验精度。缺一不可,否则可能“看起来快了,结果错了”。

测速用端到端延迟对比;调参主要调节 BLOCK_M / BLOCK_N / BLOCK_K 三个旋钮,让块大小与寄存器压力平衡;精度验证则把融合 kernel 与 PyTorch 参考实现逐元素比对,统计最大绝对误差。

回到后厨:同样的菜,出锅时间从三趟仓库变成一趟案板。指标不会说谎——把加速比、误差报告贴在墙上,整个团队都对“一锅出”心服口服。融合不是一个魔法,而是一套可测量、可验证的工程方法。

Speedup=TbeforeTafter\displaystyle \text{Speedup} = \frac{T_{\text{before}}}{T_{\text{after}}}
  1. T_before — 融合前耗时未融合的基线版本端到端耗时。
    Tbefore=多个 kernel 分别执行的端到端时间T_{\text{before}} = \text{多个 kernel 分别执行的端到端时间}
  2. T_after — 融合后耗时融合成单个 kernel 后的端到端耗时。
    Tafter=单个融合 kernel 的端到端时间T_{\text{after}} = \text{单个融合 kernel 的端到端时间}
  3. 加速比大于 1 说明融合生效;一般 Linear+GELU 融合可拿到 1.2-2x 的收益。
    Speedup>1⇒融合有效\text{Speedup} > 1 \Rightarrow \text{融合有效}

AI 应用场景

LLM 推理服务 Prefill 加速

阶段:模型推理部署

解决的问题:大模型生成每个 token 都要执行 QKV 投影和注意力计算。传统做法中,投影结果矩阵频繁写回 HBM,显存带宽被吃满,kernel 启动次数多达七八次。将 QKV 投影融合为一次大矩阵乘,再把 Attention 用在线 Softmax 分块执行后,中间矩阵不再落地,首 token 延迟显著下降。

具体案例:部署 7B 参数模型时,输入 2048 个 token 的 prefill 阶段,QKV 投影融合配合分块注意力,使首 token 延迟从约 1.2 秒降至 0.8 秒,吞吐提升约 30%。这就是每个 LLM 推理服务都值得做的“案板革命”。

高性能算子库内核开发

阶段:算子开发与优化

解决的问题:框架自带算子覆盖通用场景,但对融合 GELU 的多层感知机块、动态形状的注意力等特殊结构,存在明显冗余访存。手写 Triton 融合 kernel 能把“仓库往返”变成“案板流转”。

具体案例:为训练框架开发 fused Linear+GELU kernel,在 2000 次迭代的 MLP 基准上,单卡训练吞吐提升 18%,同时代码量只有 CUDA 方案的十分之一。该 kernel 被封装成自定义算子,经自动测试后合入内部算子库。

Python 演示

python61 行
import torchimport tritonimport triton.language as tl  @triton.jitdef linear_gelu_fused_kernel(    X_ptr, W_ptr, B_ptr, Y_ptr,    M, N, K,    stride_xm, stride_xk,    stride_wn, stride_wk,    stride_ym, stride_yn,    BLOCK_M: tl.constexpr,    BLOCK_N: tl.constexpr,    BLOCK_K: tl.constexpr,):    pid_m = tl.program_id(0)    pid_n = tl.program_id(1)    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)    offs_k = tl.arange(0, BLOCK_K)    x_ptrs = X_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk    w_ptrs = W_ptr + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)    for k in range(0, tl.cdiv(K, BLOCK_K)):        x = tl.load(x_ptrs)        w = tl.load(w_ptrs)        acc = tl.dot(x, w, acc)        x_ptrs += BLOCK_K * stride_xk        w_ptrs += BLOCK_K * stride_wk    b = tl.load(B_ptr + offs_n)    acc = acc + b[None, :]    acc = acc * 0.5 * (1.0 + tl.erf(acc * 0.70710678))    y_ptrs = Y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :] * stride_yn    tl.store(y_ptrs, acc)  def linear_gelu_fused(x, w, b):    M, K = x.shape    N = w.shape[0]    y = torch.empty((M, N), device=x.device, dtype=x.dtype)    grid = (triton.cdiv(M, 64), triton.cdiv(N, 64))    linear_gelu_fused_kernel[grid](        x, w, b, y, M, N, K,        x.stride(0), x.stride(1),        w.stride(0), w.stride(1),        y.stride(0), y.stride(1),        BLOCK_M=64, BLOCK_N=64, BLOCK_K=32,    )    return y  if __name__ == "__main__":    torch.manual_seed(0)    M, K, N = 256, 128, 256    x = torch.randn((M, K), device="cuda", dtype=torch.float32)    w = torch.randn((N, K), device="cuda", dtype=torch.float32)    b = torch.randn((N,), device="cuda", dtype=torch.float32)    y = linear_gelu_fused(x, w, b)    ref = torch.nn.functional.gelu(x @ w.T + b)    print("max abs err:", (y - ref).abs().max().item())
目的:用 Triton 手写一个 Linear+GELU 融合 kernel,并在 GPU 上与 PyTorch 参考实现做精度比对。
  1. @triton.jit把 Python 函数编译成 GPU kernel 的装饰器,相当于“一桌菜”的菜谱。
  2. tl.program_id(0)拿到当前 program 在网格中的坐标,即桌号。
  3. tl.arange(0, BLOCK_M)生成块内的行编号,对应这一桌的座位号。
  4. acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)初始化累加器,底料锅先烧热。
  5. acc = tl.dot(x, w, acc)分块矩阵乘并累加,内层循环沿 K 维接力。
  6. b = tl.load(B_ptr + offs_n)加载偏置向量,相当于最后的点睛之盐。
  7. acc = acc * 0.5 * (1.0 + tl.erf(acc * 0.70710678))GELU 累加后原地调味,中间结果不落地。
  8. tl.store(y_ptrs, acc)只写一次最终结果,不再为中间张量跑仓库。
  9. grid = (triton.cdiv(M, 64), triton.cdiv(N, 64))按 64 的块算出网格大小,多开一桌补余数。

常见误区

误区我以为算子融合就是在 Python 里先算 Linear 再算 GELU,把两段计算逻辑写在一起就算融合了,代码看着紧凑就行。

正解真正的融合是让中间结果留在 SRAM 或寄存器里,不写回 HBM。就算代码写在同一个 kernel 中,如果中间结果仍然落回全局内存,性能不会有本质提升。融合的关键是数据流,不是代码排版。

为何常见:名字里的“融合”容易让人联想到代码层面的合并,而忽略了它真正的目标是减少访存。

误区我总觉得 BLOCK_M 和 BLOCK_N 越大,一个块处理的活越多,线程利用率越高,性能就一定越好。

正解块太大会撑爆寄存器和共享内存,导致 occupancy 下降,甚至无法编译。要参考目标 GPU 的硬件规格,让块体积与资源容量匹配。案板就那么大,硬塞一吨食材只会放不下。

为何常见:“一次多干点”的直觉很自然,但忽略了 GPU 资源上限这个硬约束。

误区PyTorch 2.0 不是能自动融合算子吗?那我还花时间学 Triton 干嘛,交给编译器就好了。

正解torch.compile 对常见图模式有效,但对涉及切片重组的 QKV 投影、跨自定义算子的复杂融合,自动编译器的收益有限。手写 Triton 能精确控制块大小、数据布局和融合边界,是高级性能优化的必经之路。

为何常见:宣传材料常夸大自动编译器的适用范围,实际项目里“编译器没融好”的案例比比皆是。

误区GELU 有 erf 和 tanh 两种近似形式,我觉得算出来差不多,随便挑一个用,反正都是激活函数。

正解两种形式的数值差异大约在 1e-3 量级,训练时可接受;但推理时如果做精度对齐或量化校准,必须与参考实现保持一致。融合 kernel 里选哪种公式,直接影响端到端误差报告。

为何常见:误差在浮点上看起来很小,容易忽略“基准不一致”会引发的调试灾难。

面试题

Q1. 算子融合带来的最主要收益是什么?

单选题 · easy

  • A.减少中间结果的访存开销← 正确答案
  • B.减少模型的训练参数数目
  • C.让代码结构显得更加简洁
  • D.提升数据类型的表示精度

答案:A。融合最核心的收益是减少中间结果在 HBM 与计算单元之间的往返,对应后厨里“少跑仓库”。参数总量没变,代码风格只是副产物,精度也未改变,因此 B、C、D 都不对。

Q2. Linear 与 GELU 融合后,中间结果存储在哪里?

单选题 · mid

  • A.被写回 HBM 全局内存等待
  • B.留在 SRAM 供下一步使用← 正确答案
  • C.暂存在 CPU 主存队列里
  • D.直接丢弃再重新计算一次

答案:B。融合后 Linear 的累加结果留在 SRAM(案板)上,GELU 直接读取并原地变换。A 是未融合的旧行为,C 跨设备搬运更慢,D 丢弃数据会导致结果错误,故选 B。

Q3. 以下哪些场景适合做算子融合?

多选题 · mid

  • A.Linear 输出被 GELU 消费← 正确答案
  • B.QKV 三个投影共用输入← 正确答案
  • C.两个算子间数据量巨大← 正确答案
  • D.两个算子完全互不依赖

答案:ABC。融合适合“生产者—消费者”紧邻且中间数据量大的情况。A 是典型配对,B 共用输入可合并成一次投影,C 数据量越大省得越多;D 无依赖的算子合并没有收益,反而可能增加寄存器压力,不选。

Q4. Attention 中除以根号 d_k 的主要目的是什么?

单选题 · mid

  • A.防止 QK 点积过大进入饱和区← 正确答案
  • B.使点积结果保持精确为零
  • C.增加模型参数量提升表达力
  • D.减少键矩阵的显存占用大小

答案:A。点积随维度线性增大,除以根号 d_k 把它压回单位量级,避免 softmax 进入指数饱和区导致梯度消失。B、C、D 都与缩放的真实目的无关,故选 A。

Q5. Triton 中块大小参数的选择需要考虑哪些因素?

多选题 · easy

  • A.GPU 寄存器与共享内存容量← 正确答案
  • B.输入矩阵的具体形状尺寸← 正确答案
  • C.目标 GPU 的硬件架构版本← 正确答案
  • D.训练优化器的学习率大小

答案:ABC。块大小受硬件资源上限(寄存器/SRAM)、输入形状与架构特性共同约束,案板能放多少得看案板本身。学习率是训练超参数,与 kernel 的网格划分无关,D 不选。

Q6. FlashAttention 中在线 Softmax 解决的核心问题是什么?

单选题 · hard

  • A.无需保存完整分数矩阵即可归一化← 正确答案
  • B.让注意力分数结果更加精确
  • C.把 Softmax 替换成线性激活函数
  • D.消除 QKV 三个矩阵的参数冗余

答案:A。在线 Softmax 用运行最大值和运行累加和修正,使分块计算也能得到与全局 softmax 完全一致的结果,因此不需要把完整 (M, N) 分数矩阵留在 HBM。B、C、D 都不是它要解决的核心问题。

Q7. 融合后 GELU 计算发生在哪个时刻?

单选题 · mid

  • A.线性累加完成之后立刻执行← 正确答案
  • B.输入数据拷贝完成之前执行
  • C.权重矩阵初始化过程中执行
  • D.结果写回 HBM 之后才执行

答案:A。融合内核中 GELU 在 acc 完成所有 K 维累加并加上偏置后、写回 Y 之前执行,即“累加完立刻调味”。B、C 在计算前无从激活,D 写回后执行已失去融合意义,故选 A。

Q8. 验证融合算子正确性可以采用哪些方法?

多选题 · hard

  • A.与 PyTorch 参考实现逐元素比对← 正确答案
  • B.检查最大绝对误差是否可接受← 正确答案
  • C.用多组随机输入统计误差分布← 正确答案
  • D.直接查看代码缩进是否工整

答案:ABC。工程上融合 kernel 必须与参考实现逐元素比对、检查最大误差,并用多组随机输入验证误差分布稳定。D 只关乎代码风格,无法验证数值正确性,不选。