想象你是一家超级中央厨房的行政总厨。厨房里有上万名厨师,每个厨师只负责切一根葱、翻一次锅,他们记不住全局菜谱;仓库(显存)里堆满食材,但取食材很慢;每个厨师小组的备料台(SRAM)又小又快。你不可能站在每个厨师耳边下达逐条指令,于是请来一位叫 Triton 的传令官:你把“一盘菜”的批处理做法写在菜谱里,他自动把任务分给所有厨师。Softmax 是给每道菜计算“受欢迎热度”的归一化步骤;FlashAttention 则是为几千位宾客准备酒宴时,不让成堆的中间匹配单堆满仓库,而是分盘在备料台上现做现并。这篇文章带你跟着 Triton 的块级菜谱,从两趟取料的朴素 Softmax,升级到一趟完成的在线 Softmax,再看懂分块流水的 FlashAttention。

岗位关联

对 AI 从业者而言,这相当于从“只会调模型”走向“能优化底层算子”的翻身课,也是通往大模型推理加速、高性能 AI 编译方向的核心基本功。

目录
  1. 可视化演示
  2. 块级编程模型:一托盘一托盘地写菜谱
  3. 内存层次解析:仓库、备料台与合并领取
  4. 朴素 Softmax:来回仓库的两趟取料
  5. 在线 Softmax:边洗边切的一趟流水
  6. FlashAttention:分盘流水的大宴会
  7. 性能调优要点:调整备料台的人手和托盘
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

通过在线Softmax、FlashAttention分块和GPU内存合并访问三个模式,展示Triton块级编程模型的核心优化技术。

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

在线Softmax算法 流式处理:逐元素更新max和sum,无需全量加载 数值稳定:减去当前最大值后再exp,防止溢出 m_new = max(m_old, x_i); s_new = s_old*e^(m_old-m_new) + e^(x_i-m_new) 在线算法将Softmax从两遍降为一遍,减少内存访问

块级编程模型:一托盘一托盘地写菜谱

中央厨房里最怕的,是总厨站在厨师耳边说“你,切第一根葱;你,切第二根葱”。上万名厨师需要的是批量指令:把一托盘葱花一起切好。Triton 的块级编程模型就是这个思路——你不再手写每个线程做什么,而是写“一块数据”整体做什么。一个 kernel 就像一张托盘菜谱,GPU 上的每个 program instance 负责处理托盘中的一块。

在 Triton 中,这个“托盘”就是 block,也叫 tile。你可以用 tl.arange 生成下标,用 tl.load 一次加载一整块数据,用 tl.store 一次写回一整块结果。编译器会帮你把块操作切成线程束能并行执行的小任务,你不用像 CUDA 那样手动管理 blockIdx.x 和 threadIdx.x 的每个细节。这就是块级编程模型最核心的抽象:把“多点并行”交给编译器,把“块级流水”留给你。

为什么这个抽象重要?因为 GPU 的算力来自成千上万个线程,但让人类逐个管理线程就像指挥十万只蚂蚁排队,既慢又容易错。Triton 把“蚂蚁队列”统一成“托盘”,你只需要决定托盘多大、往哪送。回到厨房里,总厨不再喊“一号厨师切一根葱”,而是喊“三号备料台切一百根葱”——所有厨师同时动手,效率立刻上来。

B=X[i:i+BLOCK,:],Y=op(B),Xout[i:i+BLOCK,:]=Y\displaystyle B = X[i:i+\text{BLOCK}, :], \quad Y = \text{op}(B), \quad X_{\text{out}}[i:i+\text{BLOCK}, :] = Y
  1. 切出块从大张量 X 中按行切片,得到一个小张量 B,它就是当前这道菜的一托盘原料。
    B=X[i:i+BLOCK,:]B = X[i:i+\text{BLOCK}, :]
  2. 块上计算对整块 B 施加某个算子,比如指数、求和、矩阵乘法,结果仍是块级张量 Y。
    Y=op(B)Y = \text{op}(B)
  3. 写回结果把处理好的 Y 放回目标张量的对应位置,完成一次托盘流转。
    Xout[i:i+BLOCK,:]=YX_{\text{out}}[i:i+\text{BLOCK}, :] = Y

内存层次解析:仓库、备料台与合并领取

中央厨房的仓库(显存 HBM/DRAM)可以放几百公斤食材,但离灶台远,取一趟要很久;每个厨师小组的备料台(SRAM)只有巴掌大,但就在手边,拿取极快。GPU 程序最怕的不是算得慢,而是反复去仓库搬食材。理解内存层次,就是知道哪些食材该放备料台,哪些必须送回仓库。

Triton 虽然隐藏了线程,但没有隐藏内存带宽的现实。tl.load 把数据从全局显存搬到 SRAM 或寄存器,tl.store 再搬回去。如果一块数据需要反复使用,就应该尽量留在备料台上,而不是每次用都去仓库取。另一个关键概念是合并访问:当同一排厨师同时伸手,如果他们的手在仓库传送带上相邻的位置,管理员可以一次搬一整排食材;如果大家东一个西一个,管理员就要跑好几趟。

合并访问的实战启示是:让相邻线程访问相邻地址。比如按行优先存储的矩阵,块内下标连续的一段就会触发合并访问;如果随意跳着取,硬件会拆成多次小事务。厨房里的规则很简单:食材按货架顺序摆,厨师也按货架顺序站,搬运次数最少。

I=FLOPsBytes\displaystyle I = \frac{\text{FLOPs}}{\text{Bytes}}
  1. 分子是计算量FLOPs 表示浮点运算次数,比如做多少次乘加,对应厨房里切了多少刀。
    FLOPs\text{FLOPs}
  2. 分母是搬移量Bytes 表示从全局内存读写了多少字节,对应从仓库搬了多少斤食材。
    Bytes\text{Bytes}
  3. 比值是算术强度I 越大,说明每次搬食材都伴随大量计算,更可能是计算瓶颈;I 越小,说明大部分时间都在搬运,是内存瓶颈。
    I=FLOPsBytesI = \frac{\text{FLOPs}}{\text{Bytes}}

朴素 Softmax:来回仓库的两趟取料

每道菜要打出“受欢迎程度”的权重。Softmax 要做的事是:给一堆数值,把它们变成非负、加起来等于 1 的权重。总厨的朴素做法是:先派一位厨师跑到仓库,把所有热度值看一遍,找出最高热度;再回来告诉所有人“以最高热度为基准,算指数”;最后再统一归一化。这个流程至少需要两趟仓库往返。

为什么非要找最大值?因为指数函数增长太快。如果某个热度是 1000,e^1000 在计算机里直接变成无穷大,后面的归一化全崩了。减去最大值后,最大的指数变成 e^0 = 1,其他项都小于 1,数值就安全了。朴素 Softmax 虽然数学正确,但在 GPU 上有一个致命问题:如果向量很长,中间结果装不进备料台,就必须写回仓库,带宽被白白浪费。

如果你只需要算一次 Softmax,两趟仓库还能忍;但在大模型里,每个 token 的注意力都要算 Softmax,形状是 N×N,N 可能是几千甚至几十万。这时候还写完整中间矩阵,就像为了做一道菜把整个厨房的食材全搬出来再搬回去,仓库门口立刻堵死。

softmax(xi)=exi−m∑j=1Nexj−m,m=max⁡1≤j≤Nxj\displaystyle \text{softmax}(x_i) = \frac{e^{x_i-m}}{\sum_{j=1}^N e^{x_j-m}}, \quad m=\max_{1\le j \le N} x_j
  1. 找最大值先扫描整个向量,得到最大值 m,它是所有指数计算的“温度基准”。
    m=max⁡1≤j≤Nxjm = \max_{1\le j \le N} x_j
  2. 减最大值把每个元素都减去 m,让最大项变成 0,避免指数上溢。
    xi−mx_i - m
  3. 算指数对平移后的值取指数,得到每个元素的非负权重。
    exi−me^{x_i - m}
  4. 求分母把所有指数结果加起来,得到归一化常数,也就是这桌菜的总热度。
    ∑j=1Nexj−m\sum_{j=1}^N e^{x_j - m}
  5. 归一化每个指数除以总热度,得到最终权重,所有权重加起来等于 1。
    exi−m∑j=1Nexj−m\frac{e^{x_i - m}}{\sum_{j=1}^N e^{x_j - m}}

在线 Softmax:边洗边切的一趟流水

总厨想出一个聪明办法:与其先跑一趟仓库找最大值,不如让厨师一边看食材一边更新手里的“当前最高温度”和“累计热度总和”。每见到一个新热度,就重新比较;如果最高温度变了,旧累计热度必须做一个等比缩放,否则新旧比例会错位。这就是在线 Softmax,也叫 streaming softmax。

它维护两个滚动变量:当前最大值 m_k 和当前分母 d_k。处理新元素 x_k 时,先更新最大值;然后旧分母乘以 exp(m_{k-1} - m_k) 完成缩放,再加上新元素的指数贡献。看起来多了一次乘法,但换来的是只遍历一遍数据,不需要保存完整中间向量。对于 FlashAttention,这意味着可以一边分块读取数据,一边增量更新输出。

在线 Softmax 的妙处在于,它和朴素 Softmax 在数学上是完全等价的,只是把“先全部看完再算”改成了“边看边算,温度变了就整体校准”。厨师再也不用为了找最高温跑一趟仓库,只需要在备料台上维护一张小纸条,所有中间状态都能留在手边。

\displaystylem_k = \max(m_{k-1}, x_k), \quadd_k = d_{k-1} e^{m_{k-1} - m_k} + e^{x_k - m_k}, \quadO_k = \frac{O_{k-1} d_{k-1} e^{m_{k-1} - m_k} + e^{x_k - m_k}}{d_k}
  1. 更新最大值新元素与旧最大值取较大者,得到当前块或当前元素的标准温度。
    mk=max⁡(mk−1,xk)m_k = \max(m_{k-1}, x_k)
  2. 缩放旧分母如果最大值变大了,旧分母必须以 e^{m_{k-1} - m_k} 缩回到新标准下,保持比例一致。
    dk−1emk−1−mkd_{k-1} e^{m_{k-1} - m_k}
  3. 加入新指数把新元素的指数项加到分母中,形成新的累计分母。
    dk=dk−1emk−1−mk+exk−mkd_k = d_{k-1} e^{m_{k-1} - m_k} + e^{x_k - m_k}
  4. 更新输出旧输出同样缩放,再加上新元素的贡献,最后除以新分母,得到当前归一化的加权结果。
    Ok=Ok−1dk−1emk−1−mk+exk−mkdkO_k = \frac{O_{k-1} d_{k-1} e^{m_{k-1} - m_k} + e^{x_k - m_k}}{d_k}

FlashAttention:分盘流水的大宴会

模型中的 Attention 要计算“每个查询和每个键的匹配程度”。如果序列长度是 N,朴素实现会生成一个 N×N 的打分矩阵,这个矩阵在长文本下大得惊人,写回显存再读出来会消耗大量带宽。FlashAttention 的解决思路,就像大宴会上不让所有匹配单据堆满仓库:把 Q 切成若干行块,把 K、V 切成若干列块,每次只在备料台上计算一小块“查询 × 键”,立刻做 Softmax,再乘对应的 V,最后把小块结果累加到输出上。

每一步都依赖在线 Softmax:先取当前 Q 块和 K 块,算出局部打分矩阵 S_ij;然后对这个局部矩阵做缩放 Softmax;再把得到的概率乘 V 块,累加到输出 O_i。整个过程里,S_ij 和 P_ij 只存在于片上 SRAM,不需要写回全局显存。反向传播时还可以用重计算技巧重新算一遍 S_ij,从而省下保存完整注意力矩阵的显存。

FlashAttention 把“全局 Softmax”拆成了“局部 Softmax + 在线合并”,所以最终结果和标准 Attention 几乎完全一致,只是省去了 N×N 中间矩阵的搬运。厨房里的画面是:每一桌客人的匹配评分都在备料台上现算、现乘、现合并,最后端出的是已经分好的成品菜盘,而不是成箱的中间单据。

\displaystyleS_{ij}=\frac{Q_iK_j^\top}{\sqrt{d}},\quadP_{ij}=\mathrm{softmax}(S_{ij}),\quadO_i=\sum_{j=1}^{T}P_{ij}V_j
  1. 分块打分取出第 i 个 Q 块和第 j 个 K 块,做矩阵乘法,得到一个小块的相似度矩阵。
    Sij=QiKj⊤S_{ij}=Q_iK_j^\top
  2. 缩放稳定除以 sqrt(d),防止维度变大后点积数值过大,Softmax 进入饱和区。
    Sijd\frac{S_{ij}}{\sqrt{d}}
  3. 块内 Softmax对该小块的行做在线 Softmax,得到概率块 P_ij,注意这里不跨整个序列。
    Pij=softmax(Sij)P_{ij}=\mathrm{softmax}(S_{ij})
  4. 乘 V 累加把 P_ij 与对应的 V 块相乘,累加到 O_i;因为使用在线 Softmax,之前的累加结果会被正确缩放。
    Oi=∑j=1TPijVjO_i=\sum_{j=1}^{T}P_{ij}V_j

性能调优要点:调整备料台的人手和托盘

有了正确内核还不够,总厨还要调流水线参数。第一是托盘大小:块太大,备料台放不下,编译器会把数据临时溢出到仓库,反而更慢;块太小,每个厨师组只做一点点事,调度开销占比上升。第二是厨师组配置,Triton 里的 num_warps 决定一个 block 用多少个线程束,像决定一个备料台配几个人。人太少切得慢,人太多抢备料台。

性能调优不能靠猜,要用 Profiler 看瓶颈。如果 kernel 是 memory-bound,重点优化合并访问、减少 HBM 读写;如果是 compute-bound,重点优化矩阵分块和计算指令。Triton 还会通过 tl.constexpr 让块大小在编译期固定,方便编译器做循环展开和寄存器分配,这些都属于“编译器提示”。一个常见经验是:先保证结果正确,再用 Profiler 数据调 BLOCK 和 num_warps。

厨房运营同样看“流动效率”:不是厨师越多越好,而是让每个备料台都恰好放得下当前托盘、每个厨师组都不闲置。Triton 把很多繁琐的底层调度交给编译器,但块大小、线程束配置、数据布局仍然需要你像总厨一样拍板。掌握了这些参数,你就真正从“会写 kernel”进阶到“会调 kernel”。

Occupancy=active_warpsmax_warps_per_SM×100%\displaystyle \text{Occupancy} = \frac{\text{active\_warps}}{\text{max\_warps\_per\_SM}}\times 100\%
  1. 活跃线程束某个时刻每个 SM 上真正在跑的 warp 数量,它代表备料台上同时忙碌的厨师组。
    active_warps\text{active\_warps}
  2. 硬件上限每个 SM 最多能容纳的 warp 数,由寄存器数量和 SRAM 容量共同决定。
    max_warps_per_SM\text{max\_warps\_per\_SM}
  3. 占用率两者相除得到占用率。占用率太低,隐藏不了访存延迟;占用率过高,又可能挤爆寄存器。
    Occupancy=active_warpsmax_warps_per_SM\text{Occupancy} = \frac{\text{active\_warps}}{\text{max\_warps\_per\_SM}}

AI 应用场景

大模型推理引擎的 Attention 算子

阶段:模型部署与推理优化

解决的问题:传统 Attention 把 N×N 打分矩阵写回显存,长上下文场景下 GPU 带宽成为瓶颈,首 token 延迟和吞吐都受影响。FlashAttention 通过分块 + 在线 Softmax 减少 HBM 往返,让长序列推理成为可能。

具体案例:vLLM、TensorRT-LLM 等推理框架中,FlashAttention 风格算子被用于处理 32K 甚至 128K 上下文。开启后 KV cache 访存压力明显下降,长文本生成时的每秒 token 数显著提升。

长序列训练中的显存急救

阶段:模型训练与显存优化

解决的问题:训练 Bert、GPT 等模型时,注意力矩阵的显存占用随序列长度平方增长。AI 工程师使用 memory-efficient attention,在前向传播中不保存完整 Softmax 结果,反向传播时重计算,从而把显存占用降到线性。

具体案例:在序列长度 8192 的 fine-tune 任务中,使用 flash-attn 库替换原始多头注意力后,单卡 batch size 翻倍,且训练精度与原始实现基本一致。

自研算子的性能 Profiling

阶段:性能分析与算子调优

解决的问题:即使有了 Triton,初学者仍可能写出带宽极低的 kernel。通过 Nsight Compute 等 Profiler 查看 memory-bound 还是 compute-bound,再结合在线 Softmax 和分块技巧迭代优化,是算子落地的常态。

具体案例:某团队用 Triton 实现自定义 attention 时,最初带宽利用率不到 20%;调整块大小和 num_warps 后利用率提升到 60% 以上,端到端推理提速 1.8 倍。

Python 演示

python43 行
import numpy as np def flash_attention(Q, K, V, block_size=32):    n = Q.shape[0]    d_k = Q.shape[1]    O = np.zeros_like(Q, dtype=np.float32)    m = np.full(n, -np.inf, dtype=np.float32)    l = np.zeros(n, dtype=np.float32)     for j in range(0, n, block_size):        K_b = K[j:j+block_size]      # Key block        V_b = V[j:j+block_size]      # Value block         S = Q @ K_b.T / np.sqrt(d_k) # Block scores         m_new = np.maximum(m, S.max(axis=1))        rescale = np.exp(m - m_new)        P = np.exp(S - m_new[:, None])         l = l * rescale + P.sum(axis=1)        O = O * rescale[:, None] + P @ V_b        m = m_new     O = O / l[:, None]    return O def naive_attention(Q, K, V):    d_k = Q.shape[1]    S = Q @ K.T / np.sqrt(d_k)    P = np.exp(S - S.max(axis=1, keepdims=True))    P = P / P.sum(axis=1, keepdims=True)    return P @ V if __name__ == "__main__":    rng = np.random.default_rng(0)    n, d = 128, 16    Q = rng.normal(size=(n, d)).astype(np.float32)    K = rng.normal(size=(n, d)).astype(np.float32)    V = rng.normal(size=(n, d)).astype(np.float32)     out = flash_attention(Q, K, V, block_size=16)    ref = naive_attention(Q, K, V)    print("max diff:", np.max(np.abs(out - ref)))
目的:用纯 NumPy 验证 FlashAttention 的分块在线 Softmax 逻辑,跑通后可在 Triton 中把相同结构翻译为块级 kernel。
  1. S = Q @ K_b.T / np.sqrt(d_k)只计算当前 Key 块的点积,不生成完整 N×N 打分矩阵。
  2. m_new = np.maximum(m, S.max(axis=1))在旧最大值和新块最大值之间取更大者,这就是在线更新。
  3. rescale = np.exp(m - m_new)当最大值变大时,旧状态必须缩放到新标准下。
  4. l = l * rescale + P.sum(axis=1)滚动累加每行的分母。
  5. O = O * rescale[:, None] + P @ V_b旧输出缩放后叠加当前 Value 块贡献。
  6. O = O / l[:, None]所有块处理完后统一归一化。

常见误区

误区以为写出块级代码后编译器会把所有性能问题解决,于是随意选择块大小和访问顺序,结果 kernel 仍然很慢。

正解Triton 只是隐藏了线程调度,没有消除内存带宽和 SRAM 容量的物理约束。块太大可能溢出到显存,访问不连续会破坏合并访问。编译器负责把块操作翻译好,但“拿多大的托盘、按什么顺序进货”仍需你决定。

为何常见:因为块级 API 看起来像写普通 NumPy,容易让人忽略底层还有显存搬运。

误区认为减去最大值只是为了推导方便,实际计算中不剪也无所谓,反正归一化后结果一样。

正解如果输入是 1000,np.exp(1000) 会直接上溢成 inf,后面的归一化全部变成 inf / inf = nan。减去最大值后最大指数是 e^0 = 1,计算才稳定。这不是可选的数学优化,而是数值稳定的必需步骤。

为何常见:很多教材只写 Softmax 公式,不强调浮点溢出,初学者便以为最大值项可有可无。

误区听到“分块”“在线”“增量”这些词,就以为 FlashAttention 牺牲精度换速度,结果和标准 Attention 不一样。

正解在线 Softmax 与标准 Softmax 在数学上等价,只是把全局最大值和全局分母拆成滚动更新。FlashAttention 的精度差异仅来自浮点累加顺序,通常只有 1e-6 量级,不会改变模型效果。

为何常见:性能优化常伴随精度取舍,所以人们默认“快的东西是近似的”。

误区认为 BLOCK 设得越大,单个块处理的数据越多,效率自然越高。

正解块大小受 SRAM 容量限制。块太大会导致寄存器溢出或 shared memory 放不下,反而触发额外显存访问;块太小又会让调度开销变大。最佳块大小需要结合矩阵形状、数据类型和 num_warps 实测。

为何常见:直觉上“批量大 = 效率高”,但 GPU 的片上资源有限,批量和资源之间必须平衡。

面试题

Q1. 在 Triton 中,一个 kernel 的基本操作单元是什么?

单选题 · easy

  • A.单个线程操作
  • B.块级张量操作← 正确答案
  • C.单个寄存器操作
  • D.一条机器指令

答案:B。Triton 的核心抽象是块,也就是 block/tile。开发者写的是对整个块张量的操作,Triton 编译器再自动拆分到线程。选项 A 是 CUDA 的编程方式,C 和 D 离得更远。

Q2. GPU 全局内存的合并访问是指什么?

单选题 · easy

  • A.相邻线程访问相邻地址← 正确答案
  • B.单线程访问连续大段
  • C.线程间共享私有缓存
  • D.把数据复制到多处备份

答案:A。合并访问指一个线程束中的相邻线程在同一个指令周期访问相邻的全局内存地址,硬件可以一次完成整段搬运。这能显著减少访存事务次数,是 GPU 性能的关键。

Q3. 朴素 Softmax 在 GPU 上通常需要两遍扫描,第二遍主要做什么?

单选题 · easy

  • A.找出整个向量的最大值
  • B.累加指数并完成归一化← 正确答案
  • C.对向量做一个随机打乱
  • D.把向量从显存搬回 CPU

答案:B。第一遍找最大值,第二遍计算指数和分母,再做归一化。如果数据量太大不能留在 SRAM,中间结果必须写回显存,导致额外带宽消耗。B 准确描述第二遍的核心任务。

Q4. Softmax 计算中减去最大值的主要目的是什么?

单选题 · mid

  • A.让所有输出都变成正数
  • B.防止指数函数发生上溢← 正确答案
  • C.减少计算量的整体规模
  • D.让总和自动等于一

答案:B。指数函数增长极快,输入较大时 e^x 会溢出成 inf。减去最大值后最大指数为 e^0=1,数值稳定。输出非正和总和为一都是 Softmax 的性质,不是减最大值的目的。

Q5. 在线 Softmax 相比朴素 Softmax 的核心优势是什么?

单选题 · mid

  • A.不需要保存完整中间向量← 正确答案
  • B.可以完全避免指数运算
  • C.能得到更高数学精度
  • D.只适用于纯量化模型

答案:A。在线 Softmax 用滚动最大值和滚动分母替代“先全部扫描再统一计算”,因此可以流式处理数据,不需要保存完整中间结果。它不是不用指数,也不是精度更高。

Q6. FlashAttention 分块计算中包含哪些步骤?

多选题 · mid

  • A.对 Q 和 K 分块做乘法← 正确答案
  • B.块内执行在线 Softmax← 正确答案
  • C.对整行向量随机采样
  • D.用旧分母更新输出结果← 正确答案

答案:ABD。FlashAttention 每次取 Q 块和 K 块做乘法得到局部分数,然后做在线 Softmax,再用旧分母和旧最大值更新输出。随机采样与注意力计算无关。

Q7. 两个分块合并时,旧最大值为 m_old,新最大值为 m_new,旧分母 l_old 的缩放因子是?

单选题 · hard

  • A.exp(m_old - m_new)← 正确答案
  • B.exp(m_new - m_old)
  • C.(m_old + m_new) / 2
  • D.d_new / d_old

答案:A。旧分母是在旧最大值标准下计算的,现在要换到更大的新最大值标准下,每个指数项都要乘以 exp(m_old - m_new)。因为 m_old <= m_new,所以该因子小于等于 1。

Q8. 为什么 FlashAttention 能显著减少 HBM 访问?

单选题 · hard

  • A.不写回完整的 N×N 打分矩阵← 正确答案
  • B.完全不需要全局显存
  • C.把 Q K V 全部放进寄存器
  • D.减少数学计算的总次数

答案:A。FlashAttention 的分块计算把打分矩阵留在 SRAM 中,只用最终输出 O 与全局显存交互,因此避免写回 N×N 中间矩阵。它仍需读 Q、K、V,也不是完全不用全局显存。

Q9. Triton Kernel 性能调优时,通常需要关注哪些因素?

多选题 · hard

  • A.块大小与 SRAM 容量匹配← 正确答案
  • B.num_warps 和线程束配置← 正确答案
  • C.CPU 上的 Python 版本
  • D.数据访问是否连续合并← 正确答案

答案:ABD。块大小影响 SRAM 占用,num_warps 影响并行度和资源竞争,数据连续性影响合并访问。Python 版本与 GPU kernel 性能基本无关,除非 CPU 端已成为瓶颈。

Q10. FlashAttention 反向传播中的重计算技巧有哪些作用?

多选题 · hard

  • A.节省正向计算时间
  • B.减少保存中间张量显存← 正确答案
  • C.反向时重新计算得分← 正确答案
  • D.让结果变成近似算法

答案:BC。重计算指的是反向传播时不保存完整注意力分数,而是重新算一遍 S 和 P,从而减少显存占用。它并不节省正向计算时间,也不会让结果变成近似值。