AI技术总监级拆解大模型第17讲Tensor Core、HBM、SRAM 与 FlashAttention——大模型为什么会被“搬数据”拖慢本文承接第16讲《GPU 与 CUDA》继续往 GPU 内部深入Tensor Core 为什么能加速矩阵计算HBM、Cache、Shared Memory 和 Register 为什么要分层为什么算法减少内存读写有时比单纯增加计算单元更有效FlashAttention 又是怎样重写 Attention 的计算流程的本文会覆盖Tensor Core 与混合精度、GPU 内存层次、算术强度、Roofline 模型、矩阵分块、Kernel Fusion、FlashAttention 与在线 Softmax、正确的性能测量方法、训练和推理中的性能差异。一、先给结论GPU 不是只要“算得快”就够了在理解大模型的训练和推理时大家经常关注显卡的 TFLOPS、显存容量和显存带宽。这些参数很重要但它们不能单独决定真实性能。一个 GPU Kernel 想跑得快至少需要三件事配合计算单元够高效例如用 Tensor Core 执行适合的矩阵乘加。数据供给跟得上不要让计算单元长时间等待 HBM 中的数据。算法与实现合理让数据尽可能被复用减少不必要的中间张量、显存往返和 Kernel 启动。用一句话概括本讲大模型性能既是计算问题也是数据移动问题。以矩阵乘法为例[CAB,\qquad A\in\mathbb{R}^{M\times K},\ B\in\mathbb{R}^{K\times N}]数学表达式只有一个 (AB)但高性能实现要回答A 和 B 从哪里加载每个元素会复用多少次中间结果保存在哪里由哪类硬件单元执行怎样避免不同线程反复搬运同一份数据这些才是 GPU Kernel 的核心工程问题。二、Tensor Core为什么矩阵运算有专用硬件2.1 从普通乘加到矩阵乘加一个矩阵乘法的输出元素为[C_{ij}\sum_{k0}^{K-1}A_{ik}B_{kj}]每个输出都包含一串乘法和加法。若矩阵较大整体会出现大量规则、重复、可并行的乘加工作。这类负载非常适合专门为矩阵运算设计的执行单元。NVIDIA 的 Tensor Core 是部分 GPU 架构提供的矩阵计算硬件可加速矩阵乘加等操作。典型的矩阵乘加可以抽象成[DA\times BC]这里的 MMAMatrix Multiply-Accumulate通常以硬件支持的 tile 形状处理矩阵片段而不是要求硬件一次吞下整个大矩阵。具体指令、tile 尺寸、数据类型和布局限制因架构而异。2.2 Tensor Core 与 CUDA Core 的关系项目CUDA CoreTensor Core定位NVIDIA GPU 中执行常见算术指令的计算资源面向矩阵乘加等操作的专用计算单元擅长工作多类普通算术与指令执行高吞吐矩阵运算尤其常见于深度学习负载是否可互相简单替代不能用数量直接换算成另一类硬件的性能需要适合的算子、精度、形状与软件路径才能发挥优势工程关键点指令效率、调度、寄存器与访存数据格式、矩阵布局、tile、累加方式、数据供给不要把 Tensor Core 理解成“所有计算都更快的 CUDA Core”。它是一类针对特定计算形态优化的硬件。某个程序若主要由小规模标量操作、复杂控制流或频繁访存组成即使 GPU 有大量 Tensor Core也不代表这些单元一定能被高效利用。2.3 为什么峰值 TFLOPS 不等于真实性能厂商公布的峰值通常对应特定架构、数据类型、指令条件甚至特定稀疏性或精度设置。程序能否接近峰值取决于算子是不是 Tensor Core 擅长的矩阵乘加输入 dtype、布局、尺寸和对齐是否适配是否有足够多的独立工作让硬件持续运行数据是否能及时从内存层次中供给Kernel 实现、编译器和高层库是否选中了高效路径是否受同步、launch overhead、寄存器或共享内存资源限制。因此峰值算力是理论天花板不是应用性能承诺。做工程评估时必须测真实任务而不是只把 TFLOPS 当成答案。三、混合精度FP32、FP16、BF16 与 TF32数值格式会同时影响内存占用、计算吞吐和误差。常见格式如下格式特点使用时的关注点FP32单精度浮点范围和尾数精度较均衡存储与计算成本可能更高部分硬件的 FP32 吞吐低于低精度矩阵路径FP1616 位浮点内存占用较低很多 GPU 支持高吞吐矩阵计算指数范围较窄训练时要关注溢出、下溢与梯度缩放策略BF1616 位格式指数位数与 FP32 相同尾数精度更低动态范围通常更适合深度学习训练但相对粗糙的尾数仍会带来舍入误差TF32NVIDIA 部分架构用于 Tensor Core 路径的计算格式/模式不应简单视作独立的、通用的存储 dtype实际行为依硬件、框架与开关设置而定INT8、FP8 等某些场景中可提升吞吐、减少内存占用需要相应量化/缩放方法验证任务精度与数值稳定性支持情况依架构而异3.1 三种“精度”不能混为一谈输入精度A、B 元素以什么格式提供给运算。累加精度多个乘积相加时累加器采用什么精度。输出精度结果最终以什么格式保存。矩阵乘法的输入、累加和输出格式可以不同。混合精度的目标是合理利用高速计算路径同时控制误差而不是不加判断地把一切都变成最低精度。3.2 FP16 与 BF16 怎么理解二者都是 16 位浮点但位的分配不同。FP16 的尾数相对更多、指数范围较小BF16 的指数位与 FP32 相同尾数较少。因此BF16 通常具有更大的动态范围而 FP16 在一些数值范围内能表达更细的数值间隔。这不是说 BF16 永远优于 FP16也不是说 FP16 一定更准。要结合模型中的数值分布、硬件支持、训练稳定性和最终评估判断。数据格式是性能设计的一部分也是一项需要测试的数值决策。四、GPU 内存层次HBM 为什么很重要片上存储为什么更重要这里先区分几个名字HBM 是一种高带宽内存技术GPU 的设备全局内存常被泛称为显存但并非每块 GPU 都使用 HBM。GDDR 是另一类常见显存技术。GPU 计算单元 │ Register寄存器线程私有 │ Shared Memory / L1 Cache片上 │ L2 Cache片上 │ HBM / Device Global Memory 容量大承担主要数据存储这张图用于建立直觉不是所有 GPU 都严格按相同方式组织的物理结构。具体缓存大小、共享关系、带宽、延迟以及 SRAM 的划分随架构而变。4.1 HBM容量和带宽是两种能力HBM 的目标之一是提供很高的内存带宽帮助 GPU 把数据持续送到计算单元。工程上需要区分容量GB模型参数、激活、中间张量、优化器状态、KV Cache 等能否放得下。带宽GB/s 或 TB/s单位时间能搬运多少数据。延迟一次访问等待多久。带宽高并不代表单次访问没有等待成本。有的工作负载首先受显存容量约束有的则是容量够但数据吞吐跟不上。优化前必须先认清是哪种约束。4.2 Register、Shared Memory、L1/L2 Cache 的区别Register寄存器通常是线程私有的高速存储资源保存当前线程运算中使用的值。寄存器很宝贵一个 Kernel 每线程使用寄存器过多可能会减少一个 SM 能同时驻留的 Warp/Block 数若使用压力超过可用寄存器资源也可能产生 spill把部分值放到较慢的本地内存路径。Shared Memory共享内存是程序可显式管理的片上存储通常由同一个 Thread Block 中的线程共同访问。它常用于把全局内存中的 tile 载入片上再让多个线程重复使用。它并不是“所有线程都能任意访问的一块全 GPU 共享内存”可见范围和同步规则受到编程模型限制。L1/L2 Cache是缓存结构主要由硬件根据访问行为管理。L1 通常与 SM 关系密切L2 常由多个 SM 共享但具体设计依架构而变。Cache 可以帮助复用数据却不等于程序能保证所有访问都命中缓存。4.3 SRAM 不等于 Shared MemorySRAM 是一种存储器实现技术Register、Cache、Shared Memory 是硬件结构或编程模型抽象概念层次不同。GPU 上很多片上存储结构采用 SRAM 实现但不能把“SRAM、Shared Memory、Cache、寄存器”当作四个同义词。4.4 存储层次的根本矛盾更靠近计算单元的存储一般更快但容量和资源受限HBM 容量大、带宽高但数据经过更长的访问路径。高性能 Kernel 的基本策略是从 HBM 把当前要用的数据取出来暂存在合适的片上层次尽量让同一份数据服务多个运算减少不必要的读写和中间结果落地在寄存器、共享内存和并行度之间取得平衡。五、矩阵分块与数据复用把数据“拿来多用几次”考虑矩阵乘法[C_{ij}\sum_{k0}^{K-1}A_{ik}B_{kj}]朴素实现可能对每一个输出元素逐次计算。这样A 和 B 中许多元素会被多个输出使用。如果每次使用都重新从 HBM 获取内存流量就会增大。5.1 Tiling分块做了什么把大矩阵分成多个小块A tile [BM × BK] B tile [BK × BN] │ │ └──── 加载到片上 ───┘ │ 多次乘加并累加 │ C tile [BM × BN] │ 写回显存典型流程如下从 HBM 载入 A 的一个 tile 和 B 的一个 tile在 Shared Memory 或其他合适的片上路径中暂存一份 tile 被多个线程或多次乘加复用将部分结果累加到寄存器或其他合适位置遍历 K 维度后把 C 的结果块写回。这是一种概念化流程。真实高性能 GEMM 可能使用多级分块、寄存器分块、流水加载、异步拷贝、专门的 MMA 指令、复杂布局及多阶段流水线。5.2 为什么分块通常有帮助假设 A 的一个元素会参与多个输出B 的一个元素也会参与多个输出。若能在它们仍留在片上时完成多次乘加就能以更少的数据搬运完成相同工作提升数据复用率。但分块不是免费午餐tile 太小复用不充分还可能导致更多循环与调度开销tile 太大寄存器和共享内存占用增高可能压低并行度数据加载与计算安排不当会出现等待或同步成本不合适的访存布局可能造成访问不合并或共享内存 Bank Conflict。所以最佳 tile 取决于 GPU 架构、矩阵形状、精度、Kernel 实现和资源限制不存在一个适用于所有任务的万能尺寸。5.3 一个矩阵乘法工作量的估算对 (M\times K) 的 A 与 (K\times N) 的 B常用的浮点操作量近似为[\text{FLOPs}\approx 2MKN]这里把一次乘法和一次加法按两个浮点操作计数。实际 FLOPs 统计口径需保持一致。若 A、B、C 都以 16 位元素保存理想化地假定 A 和 B 各读取一次、C 写回一次则最低数据量粗略为[2(MKKNMN)\ \text{Bytes}]因为每个 16 位元素占 2 Bytes。这个估算忽略缓存、额外读写、中间累加格式、对齐及真实 Kernel 的流量它只是帮助理解“计算量与数据量之比”不能直接当作实测带宽流量。六、算术强度与 Roofline判断 GPU 到底卡在哪里6.1 算术强度算术强度Arithmetic Intensity定义为[I\frac{\text{FLOPs}}{\text{内存传输的 Bytes}}]单位是 FLOPs/Byte。它衡量每搬运一个字节数据平均完成多少浮点运算。算术强度较低数据搬运相对多可能是带宽瓶颈算术强度较高数据被较充分复用更有机会受计算吞吐限制。注意算术强度需要说明“算了多少”和“统计了哪一级存储之间的数据传输”。以 HBM 为分母和以 Cache/片上存储之间的传输为分母得到的含义不一样。6.2 Roofline 模型简化的性能上界是[P\leq\min(P_{\text{peak}},BW\times I)]其中(P)实际性能常用 FLOPs/s 表示(P_{\text{peak}})目标数据类型和硬件路径下的理论计算峰值(BW)对应内存层次的可用带宽(I)对应层次口径下的算术强度。这个模型帮助区分两种主要上界计算单元的峰值或数据传输能力乘以算术强度。实际 Kernel 还会受到启动与同步开销、并行度、缓存命中、指令吞吐等限制。6.3 手算一个例子假设某 Kernel 执行 (2\times10^{12}) FLOPs传输 (2\times10^{12}) Bytes则[I1\ \text{FLOP/Byte}]假设对应内存层次能提供的带宽是 (1\ \text{TB/s}10^{12}\ \text{Bytes/s})带宽约束给出的性能上界是[BW\times I10{12}\times110{12}\ \text{FLOPs/s}1\ \text{TFLOP/s}]即使这张 GPU 的理论算力远高于 1 TFLOP/s程序也可能受带宽上界限制无法接近计算峰值。这里的数字仅用于演示方法不对应某个实际 GPU 的指标。6.4 用 Roofline 反推优化策略观察到的主要瓶颈首先考虑什么不该直接下的结论Compute-bound合适的 Tensor Core 路径、精度、矩阵形状、指令与 tile“再增加显存带宽一定有效”Memory-bound减少 HBM 往返、合并访存、分块复用、融合中间操作“GPU 核心少所以慢”Launch-bound合并细小算子、减小框架调度开销、批处理、图捕获等“换更强显卡就能线性加速”资源限制/占用率不足检查寄存器、Shared Memory、Block 配置、spill 与 Warp 分歧“占用率越高就一定越快”七、Kernel Fusion为什么少写一次中间张量可能比多算几次更值考虑Y GELU(X W b)一个简单实现可能先算矩阵乘法再加偏置再执行 GELU。若各步骤分成不同 Kernel中间结果有可能写到 HBM再被后续 Kernel 读回来。执行步骤越碎启动次数和中间数据流量可能越多。Kernel FusionKernel 融合将相邻操作合并到一个 Kernel 或更紧密的执行路径中目标是减少Kernel 启动与调度成本中间张量写入 HBM 的次数后续重新读取中间张量的流量部分重复的地址计算与执行开销。但融合有资源代价更复杂的 Kernel 可能占用更多寄存器和共享内存指令路径可能变长甚至降低并行度。“融合”是优化手段不是无条件收益保证。必须在目标形状与目标 GPU 上测量。八、FlashAttention用算法改写内存访问方式FlashAttention 是本讲的重点因为它很好地展示了一个原则有时更快的实现不是改变数学目标而是改变完成数学计算时的数据移动方式。8.1 标准 Attention 的计算式缩放点积 Attention[\operatorname{Attention}(Q,K,V)\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V]各符号含义(Q)Query(K)Key(V)Value(d_k)Key 向量的维度(QK^T)每个 Query 与各个 Key 的匹配分数除以 (\sqrt{d_k})控制分数尺度Softmax把分数归一化成权重与 V 相乘按权重混合 Value 信息。如果序列长度为 (N)单个头的分数矩阵通常有 (N\times N) 个元素概率矩阵也有同阶大小。仅保存一个 (N\times N) 中间矩阵在元素数上就随 (N^2) 增长。长上下文的中间张量读写会成为明显成本。8.2 FlashAttention 的洞察FlashAttention 的核心是IO-aware exact attention考虑 HBM 与 GPU 片上 SRAM 之间的读写使用分块与在线 Softmax减少完整分数矩阵、概率矩阵的物化materialization以及往返 HBM 的数据量。概念流程如下传统流程 Q、K → 计算完整 S QKᵀ / √d → 保存/读取 S ↓ Softmax 得到 P ↓ 保存/读取 P ↓ P × V → 输出 FlashAttention 思路 取 Q tile K/V tile ↓ 在片上算当前分数块 ↓ 在线更新 Softmax 状态和输出累加器 ↓ 读取下一块 K/V继续合并 ↓ 输出最终结果不必把完整 S 和 P 写回 HBM这里的关键不只是“分块”而是分块执行时仍能正确完成 Softmax 的归一化并把不同分块的贡献逐步合并。8.3 在线 Softmax如何不保存完整分数矩阵Softmax 对一行分数 (s_j) 的定义为[\operatorname{softmax}(s_j)\frac{e^{s_j}}{\sum_t e^{s_t}}]直接计算时需要知道整行的归一化分母。FlashAttention 按 tile 分批处理分数但用一个可更新状态来保持计算正确性。为避免指数溢出常用行最大值进行平移。对某一行假设已有状态(m)已处理部分的最大分数(l)平移后指数项的和(o)指数权重作用于 V 后的累加向量。当前新 tile 的分数向量为 (s)更新最大值[m’\max(m,\max(s))]旧累加量的缩放系数为[\alphae^{m-m’}]当前 tile 的指数权重为[pe^{s-m’}]于是归一化分母状态和输出累加器更新为[l’\alpha l\sum_j p_j][o’\alpha opV_{\text{tile}}]当所有 tile 处理完后输出为[O\frac{o}{l}]这是便于理解的逐行表达。实际 FlashAttention Kernel 还必须处理多维张量布局、Mask、因果注意力、并行归约、数值稳定性以及 GPU 特定指令等细节。8.4 FlashAttention 解决什么不解决什么解决或改善减少完整注意力分数与概率矩阵的物化降低 HBM 与片上存储之间的读写量增强分块内的数据复用在合适的上下文长度、批量、硬件与实现下降低 Attention 的运行时间和显存压力。没有直接解决没有让稠密 Attention 所需的全部算术工作自动变成线性复杂度不保证任何形状、任何硬件、任何 batch 都以相同倍率加速不保证浮点计算结果逐 bit 完全相同不取代量化、KV Cache、并行训练或调度等其他优化。对固定 head dimension 而言稠密 Attention 的主要矩阵计算仍随序列长度呈二次增长。FlashAttention 主要改变的是IO 复杂度与中间结果存储方式而不是消除全部的成对 Token 计算。8.5 它为什么能成为一堂重要的系统课传统优化视角可能只问“乘法能不能更快”FlashAttention 还会问“这次乘法的输入是否已经在片上中间结果是否非得写回 HBM能不能只保存足够的统计量最后直接得到输出”这就是算法、硬件存储层次和数值计算协同设计的典型例子。九、训练和推理为什么有不同的性能画像9.1 训练训练包含前向计算、反向传播和优化器更新。大型矩阵乘法往往提供大量可并行计算因此 Tensor Core、合适精度、GEMM 实现和数据并行效率很重要。同时激活、梯度和优化器状态也可能带来显存容量与带宽压力。9.2 Prefill 与 Decode大模型推理通常可粗分为两段Prefill处理输入提示词的阶段一次处理一段已有序列往往能形成较大矩阵计算Attention 的分块实现和矩阵效率通常很重要。Decode逐步生成后续 Token 的阶段。每一步的工作量、batch、模型参数读取、KV Cache 访问和调度成本都会影响表现。低 batch 的自回归 Decode 常更容易受内存流量制约但实际瓶颈还取决于系统配置与并发规模。不要绝对化地说“Prefill 一定算力受限Decode 一定带宽受限”。这是常见倾向不是放之四海皆准的定律。9.3 FlashAttention 与 KV Cache 不是一回事FlashAttention改善 Attention 计算过程中的 IO 和中间结果处理方式。KV Cache在自回归生成时保存历史 Token 的 Key/Value避免每次新生成 Token 都重新计算历史 Token 的 K/V 投影。两者优化对象不同可以在同一系统里配合使用。十、实践怎样正确测量 GPU Kernel10.1 为什么不能随便用 CPU 计时GPU 操作经常是异步提交的。Python 发起一次矩阵乘法后CPU 可能很快返回但 GPU 还没执行完。只用普通 CPU 墙钟统计 Python 调用耗时可能主要测到的是任务提交时间。正确的基准至少应做到检查 GPU 是否可用进行预热避免首次初始化开销干扰使用 CUDA Event 或适当同步测量 GPU 时间固定数据类型、形状、batch 和测试条件多轮重复并报告稳定统计量验证输出误差避免把错误结果当作“更快”需要深入分析时使用 Profiler而不是凭感觉猜瓶颈。10.2 PyTorch 的 GPU Event 示例importtorchifnottorch.cuda.is_available():raiseRuntimeError(当前环境没有可用的 CUDA GPU)atorch.randn(2048,2048,devicecuda,dtypetorch.float16)btorch.randn(2048,2048,devicecuda,dtypetorch.float16)# 预热for_inrange(10):ca b torch.cuda.synchronize()starttorch.cuda.Event(enable_timingTrue)endtorch.cuda.Event(enable_timingTrue)start.record()for_inrange(100):ca b end.record()torch.cuda.synchronize()avg_msstart.elapsed_time(end)/100print(f平均单次耗时{avg_ms:.3f}ms)这个例子展示测量方式不是严谨的跨 GPU 基准。实际实验需考虑 warm-up 数量、重复次数、功耗与频率状态、输入形状、同步点和误差验证若要比较不同方法必须保持测试条件一致。10.3 用什么工具Nsight Systems看整体时间线检查 CPU、GPU、Kernel 启动、同步、传输之间的关系。Nsight Compute深入观察单个 Kernel 的硬件性能指标辅助判断内存访问、计算吞吐、占用率和资源限制。框架 Profiler观察 PyTorch 等框架里的算子时间、调用次数、内存使用与调用栈。Profiler 的数据需要结合硬件架构解释。一个指标偏低并不自动说明某个单独原因要结合 Kernel 特性和多个指标交叉判断。十一、常见性能问题与排查方向症状/现象可能原因建议检查理论算力很高实测远低于峰值算子没走高效矩阵路径、矩阵规模过小、数据供给不足核对 dtype、形状、库实现和 Kernel 选择查看 Compute 指标HBM 流量很大重复读写中间结果、访存不合并、数据复用差检查 tile、融合、访问布局与缓存行为Kernel 很多且都很短细碎算子多、启动开销占比高分析时间线考虑融合或批处理加大 tile 后反而更慢寄存器/共享内存占用升高驻留 Block 减少数据路径变差测试不同 tile关注资源占用和 spillShared Memory 加了却没有加速复用不足、同步太多、Bank Conflict、原本缓存已有效比较实际传输、同步和耗时不要只看代码“用上了共享内存”低精度后速度提升很小原瓶颈在带宽、启动或其他算子未走目标硬件路径Profiling检查算子覆盖、形状及误差GPU 使用率看似不高小 batch、CPU 供给不足、等待同步、任务短或采样统计误差看完整时间线与请求规模不要只凭单一利用率百分比判断十二、把整讲串起来高性能 GPU 计算的设计链数学工作量与数据依赖 ↓ 选择适合的数据格式和计算路径 ↓ Tensor Core / 其他计算单元 ↓ 分块与数据复用尽量让数据在片上多用几次 ↓ 控制 HBM 读写、Kernel 启动、同步及资源占用 ↓ 用 Profiler 确认瓶颈 ↓ 根据测量结果修改并复测这条链的重点是不要孤立地优化一个部件。Tensor Core 再快如果得不到数据也会空等Shared Memory 再快如果同步和搬运成本超过收益也不一定有用FlashAttention 的价值就在于它把数学计算、片上存储和 IO 成本放在一起重新设计。十三、本讲总结需要牢牢记住的 10 个结论Tensor Core 是面向矩阵乘加等运算的专用单元必须有合适的数据类型、矩阵形状和软件路径。CUDA Core 数量、Tensor Core 峰值或 TFLOPS 都不能单独决定真实应用性能。FP16、BF16、TF32、FP32 的范围、精度和硬件路径不同混合精度要同时考虑速度与数值稳定性。HBM、Cache、Shared Memory、Register 在容量、速度、访问范围和控制方式上各有分工。SRAM 是存储技术概念Shared Memory 是编程模型中的存储抽象不可混为一谈。分块通过数据复用减少重复显存流量但也增加资源与同步方面的权衡。算术强度和 Roofline 模型帮助初步判断是计算瓶颈还是带宽瓶颈。Kernel Fusion 可能减少启动和中间张量读写但也可能带来寄存器压力等副作用。FlashAttention 通过分块与在线 Softmax 降低 Attention 的 IO 压力它没有自动消除稠密 Attention 的二次级别算术工作。GPU 优化必须基于正确的基准测试和 Profiler不要靠猜测或宣传峰值下结论。十四、自测题用自己的话说明 CUDA Core 与 Tensor Core 的定位差异。为什么 HBM 容量、HBM 带宽与访问延迟不是同一个指标Register、Shared Memory、L1/L2 Cache 分别适合处理什么为什么数据复用可以减少 HBM 流量请画出一个矩阵 tile 的数据流。写出算术强度定义并解释 Roofline 模型的两个上界。为什么更大的 tile 可能降低性能至少说出两个原因。标准 Attention 为什么会出现 (N\times N) 的中间矩阵FlashAttention 如何通过在线 Softmax 得到最终输出FlashAttention 优化的是哪一部分成本它没有消除什么怎样正确测量 PyTorch GPU 运算时间参考资料NVIDIA,CUDA C Programming Guidehttps://docs.nvidia.com/cuda/cuda-c-programming-guide/NVIDIA,CUDA Refresher: The CUDA Programming Modelhttps://developer.nvidia.com/blog/cuda-refresher-cuda-programming-model/Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré,FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesshttps://arxiv.org/abs/2205.14135NVIDIA,Nsight Compute Documentationhttps://docs.nvidia.com/nsight-compute/最后一句话Tensor Core 决定一部分矩阵计算能有多快内存层次决定数据能否及时到达分块与融合决定数据是否被充分复用而 FlashAttention 证明了——重视 IO 成本有时比单纯追求更多计算更有价值。下一讲将继续讨论多 GPU 训练当一张 GPU 装不下、算不完时如何把模型、数据和计算分布到多张 GPU 上。
AI技术总监级拆解大模型第17讲Tensor Core、HBM、SRAM 与 FlashAttention——大模型为什么会被“搬数据”拖慢本文承接第16讲《GPU 与 CUDA》继续往 GPU 内部深入Tensor Core 为什么能加速矩阵计算HBM、Cache、Shared Memory 和 Register 为什么要分层为什么算法减少内存读写有时比单纯增加计算单元更有效FlashAttention 又是怎样重写 Attention 的计算流程的本文会覆盖Tensor Core 与混合精度、GPU 内存层次、算术强度、Roofline 模型、矩阵分块、Kernel Fusion、FlashAttention 与在线 Softmax、正确的性能测量方法、训练和推理中的性能差异。一、先给结论GPU 不是只要“算得快”就够了在理解大模型的训练和推理时大家经常关注显卡的 TFLOPS、显存容量和显存带宽。这些参数很重要但它们不能单独决定真实性能。一个 GPU Kernel 想跑得快至少需要三件事配合计算单元够高效例如用 Tensor Core 执行适合的矩阵乘加。数据供给跟得上不要让计算单元长时间等待 HBM 中的数据。算法与实现合理让数据尽可能被复用减少不必要的中间张量、显存往返和 Kernel 启动。用一句话概括本讲大模型性能既是计算问题也是数据移动问题。以矩阵乘法为例[CAB,\qquad A\in\mathbb{R}^{M\times K},\ B\in\mathbb{R}^{K\times N}]数学表达式只有一个 (AB)但高性能实现要回答A 和 B 从哪里加载每个元素会复用多少次中间结果保存在哪里由哪类硬件单元执行怎样避免不同线程反复搬运同一份数据这些才是 GPU Kernel 的核心工程问题。二、Tensor Core为什么矩阵运算有专用硬件2.1 从普通乘加到矩阵乘加一个矩阵乘法的输出元素为[C_{ij}\sum_{k0}^{K-1}A_{ik}B_{kj}]每个输出都包含一串乘法和加法。若矩阵较大整体会出现大量规则、重复、可并行的乘加工作。这类负载非常适合专门为矩阵运算设计的执行单元。NVIDIA 的 Tensor Core 是部分 GPU 架构提供的矩阵计算硬件可加速矩阵乘加等操作。典型的矩阵乘加可以抽象成[DA\times BC]这里的 MMAMatrix Multiply-Accumulate通常以硬件支持的 tile 形状处理矩阵片段而不是要求硬件一次吞下整个大矩阵。具体指令、tile 尺寸、数据类型和布局限制因架构而异。2.2 Tensor Core 与 CUDA Core 的关系项目CUDA CoreTensor Core定位NVIDIA GPU 中执行常见算术指令的计算资源面向矩阵乘加等操作的专用计算单元擅长工作多类普通算术与指令执行高吞吐矩阵运算尤其常见于深度学习负载是否可互相简单替代不能用数量直接换算成另一类硬件的性能需要适合的算子、精度、形状与软件路径才能发挥优势工程关键点指令效率、调度、寄存器与访存数据格式、矩阵布局、tile、累加方式、数据供给不要把 Tensor Core 理解成“所有计算都更快的 CUDA Core”。它是一类针对特定计算形态优化的硬件。某个程序若主要由小规模标量操作、复杂控制流或频繁访存组成即使 GPU 有大量 Tensor Core也不代表这些单元一定能被高效利用。2.3 为什么峰值 TFLOPS 不等于真实性能厂商公布的峰值通常对应特定架构、数据类型、指令条件甚至特定稀疏性或精度设置。程序能否接近峰值取决于算子是不是 Tensor Core 擅长的矩阵乘加输入 dtype、布局、尺寸和对齐是否适配是否有足够多的独立工作让硬件持续运行数据是否能及时从内存层次中供给Kernel 实现、编译器和高层库是否选中了高效路径是否受同步、launch overhead、寄存器或共享内存资源限制。因此峰值算力是理论天花板不是应用性能承诺。做工程评估时必须测真实任务而不是只把 TFLOPS 当成答案。三、混合精度FP32、FP16、BF16 与 TF32数值格式会同时影响内存占用、计算吞吐和误差。常见格式如下格式特点使用时的关注点FP32单精度浮点范围和尾数精度较均衡存储与计算成本可能更高部分硬件的 FP32 吞吐低于低精度矩阵路径FP1616 位浮点内存占用较低很多 GPU 支持高吞吐矩阵计算指数范围较窄训练时要关注溢出、下溢与梯度缩放策略BF1616 位格式指数位数与 FP32 相同尾数精度更低动态范围通常更适合深度学习训练但相对粗糙的尾数仍会带来舍入误差TF32NVIDIA 部分架构用于 Tensor Core 路径的计算格式/模式不应简单视作独立的、通用的存储 dtype实际行为依硬件、框架与开关设置而定INT8、FP8 等某些场景中可提升吞吐、减少内存占用需要相应量化/缩放方法验证任务精度与数值稳定性支持情况依架构而异3.1 三种“精度”不能混为一谈输入精度A、B 元素以什么格式提供给运算。累加精度多个乘积相加时累加器采用什么精度。输出精度结果最终以什么格式保存。矩阵乘法的输入、累加和输出格式可以不同。混合精度的目标是合理利用高速计算路径同时控制误差而不是不加判断地把一切都变成最低精度。3.2 FP16 与 BF16 怎么理解二者都是 16 位浮点但位的分配不同。FP16 的尾数相对更多、指数范围较小BF16 的指数位与 FP32 相同尾数较少。因此BF16 通常具有更大的动态范围而 FP16 在一些数值范围内能表达更细的数值间隔。这不是说 BF16 永远优于 FP16也不是说 FP16 一定更准。要结合模型中的数值分布、硬件支持、训练稳定性和最终评估判断。数据格式是性能设计的一部分也是一项需要测试的数值决策。四、GPU 内存层次HBM 为什么很重要片上存储为什么更重要这里先区分几个名字HBM 是一种高带宽内存技术GPU 的设备全局内存常被泛称为显存但并非每块 GPU 都使用 HBM。GDDR 是另一类常见显存技术。GPU 计算单元 │ Register寄存器线程私有 │ Shared Memory / L1 Cache片上 │ L2 Cache片上 │ HBM / Device Global Memory 容量大承担主要数据存储这张图用于建立直觉不是所有 GPU 都严格按相同方式组织的物理结构。具体缓存大小、共享关系、带宽、延迟以及 SRAM 的划分随架构而变。4.1 HBM容量和带宽是两种能力HBM 的目标之一是提供很高的内存带宽帮助 GPU 把数据持续送到计算单元。工程上需要区分容量GB模型参数、激活、中间张量、优化器状态、KV Cache 等能否放得下。带宽GB/s 或 TB/s单位时间能搬运多少数据。延迟一次访问等待多久。带宽高并不代表单次访问没有等待成本。有的工作负载首先受显存容量约束有的则是容量够但数据吞吐跟不上。优化前必须先认清是哪种约束。4.2 Register、Shared Memory、L1/L2 Cache 的区别Register寄存器通常是线程私有的高速存储资源保存当前线程运算中使用的值。寄存器很宝贵一个 Kernel 每线程使用寄存器过多可能会减少一个 SM 能同时驻留的 Warp/Block 数若使用压力超过可用寄存器资源也可能产生 spill把部分值放到较慢的本地内存路径。Shared Memory共享内存是程序可显式管理的片上存储通常由同一个 Thread Block 中的线程共同访问。它常用于把全局内存中的 tile 载入片上再让多个线程重复使用。它并不是“所有线程都能任意访问的一块全 GPU 共享内存”可见范围和同步规则受到编程模型限制。L1/L2 Cache是缓存结构主要由硬件根据访问行为管理。L1 通常与 SM 关系密切L2 常由多个 SM 共享但具体设计依架构而变。Cache 可以帮助复用数据却不等于程序能保证所有访问都命中缓存。4.3 SRAM 不等于 Shared MemorySRAM 是一种存储器实现技术Register、Cache、Shared Memory 是硬件结构或编程模型抽象概念层次不同。GPU 上很多片上存储结构采用 SRAM 实现但不能把“SRAM、Shared Memory、Cache、寄存器”当作四个同义词。4.4 存储层次的根本矛盾更靠近计算单元的存储一般更快但容量和资源受限HBM 容量大、带宽高但数据经过更长的访问路径。高性能 Kernel 的基本策略是从 HBM 把当前要用的数据取出来暂存在合适的片上层次尽量让同一份数据服务多个运算减少不必要的读写和中间结果落地在寄存器、共享内存和并行度之间取得平衡。五、矩阵分块与数据复用把数据“拿来多用几次”考虑矩阵乘法[C_{ij}\sum_{k0}^{K-1}A_{ik}B_{kj}]朴素实现可能对每一个输出元素逐次计算。这样A 和 B 中许多元素会被多个输出使用。如果每次使用都重新从 HBM 获取内存流量就会增大。5.1 Tiling分块做了什么把大矩阵分成多个小块A tile [BM × BK] B tile [BK × BN] │ │ └──── 加载到片上 ───┘ │ 多次乘加并累加 │ C tile [BM × BN] │ 写回显存典型流程如下从 HBM 载入 A 的一个 tile 和 B 的一个 tile在 Shared Memory 或其他合适的片上路径中暂存一份 tile 被多个线程或多次乘加复用将部分结果累加到寄存器或其他合适位置遍历 K 维度后把 C 的结果块写回。这是一种概念化流程。真实高性能 GEMM 可能使用多级分块、寄存器分块、流水加载、异步拷贝、专门的 MMA 指令、复杂布局及多阶段流水线。5.2 为什么分块通常有帮助假设 A 的一个元素会参与多个输出B 的一个元素也会参与多个输出。若能在它们仍留在片上时完成多次乘加就能以更少的数据搬运完成相同工作提升数据复用率。但分块不是免费午餐tile 太小复用不充分还可能导致更多循环与调度开销tile 太大寄存器和共享内存占用增高可能压低并行度数据加载与计算安排不当会出现等待或同步成本不合适的访存布局可能造成访问不合并或共享内存 Bank Conflict。所以最佳 tile 取决于 GPU 架构、矩阵形状、精度、Kernel 实现和资源限制不存在一个适用于所有任务的万能尺寸。5.3 一个矩阵乘法工作量的估算对 (M\times K) 的 A 与 (K\times N) 的 B常用的浮点操作量近似为[\text{FLOPs}\approx 2MKN]这里把一次乘法和一次加法按两个浮点操作计数。实际 FLOPs 统计口径需保持一致。若 A、B、C 都以 16 位元素保存理想化地假定 A 和 B 各读取一次、C 写回一次则最低数据量粗略为[2(MKKNMN)\ \text{Bytes}]因为每个 16 位元素占 2 Bytes。这个估算忽略缓存、额外读写、中间累加格式、对齐及真实 Kernel 的流量它只是帮助理解“计算量与数据量之比”不能直接当作实测带宽流量。六、算术强度与 Roofline判断 GPU 到底卡在哪里6.1 算术强度算术强度Arithmetic Intensity定义为[I\frac{\text{FLOPs}}{\text{内存传输的 Bytes}}]单位是 FLOPs/Byte。它衡量每搬运一个字节数据平均完成多少浮点运算。算术强度较低数据搬运相对多可能是带宽瓶颈算术强度较高数据被较充分复用更有机会受计算吞吐限制。注意算术强度需要说明“算了多少”和“统计了哪一级存储之间的数据传输”。以 HBM 为分母和以 Cache/片上存储之间的传输为分母得到的含义不一样。6.2 Roofline 模型简化的性能上界是[P\leq\min(P_{\text{peak}},BW\times I)]其中(P)实际性能常用 FLOPs/s 表示(P_{\text{peak}})目标数据类型和硬件路径下的理论计算峰值(BW)对应内存层次的可用带宽(I)对应层次口径下的算术强度。这个模型帮助区分两种主要上界计算单元的峰值或数据传输能力乘以算术强度。实际 Kernel 还会受到启动与同步开销、并行度、缓存命中、指令吞吐等限制。6.3 手算一个例子假设某 Kernel 执行 (2\times10^{12}) FLOPs传输 (2\times10^{12}) Bytes则[I1\ \text{FLOP/Byte}]假设对应内存层次能提供的带宽是 (1\ \text{TB/s}10^{12}\ \text{Bytes/s})带宽约束给出的性能上界是[BW\times I10{12}\times110{12}\ \text{FLOPs/s}1\ \text{TFLOP/s}]即使这张 GPU 的理论算力远高于 1 TFLOP/s程序也可能受带宽上界限制无法接近计算峰值。这里的数字仅用于演示方法不对应某个实际 GPU 的指标。6.4 用 Roofline 反推优化策略观察到的主要瓶颈首先考虑什么不该直接下的结论Compute-bound合适的 Tensor Core 路径、精度、矩阵形状、指令与 tile“再增加显存带宽一定有效”Memory-bound减少 HBM 往返、合并访存、分块复用、融合中间操作“GPU 核心少所以慢”Launch-bound合并细小算子、减小框架调度开销、批处理、图捕获等“换更强显卡就能线性加速”资源限制/占用率不足检查寄存器、Shared Memory、Block 配置、spill 与 Warp 分歧“占用率越高就一定越快”七、Kernel Fusion为什么少写一次中间张量可能比多算几次更值考虑Y GELU(X W b)一个简单实现可能先算矩阵乘法再加偏置再执行 GELU。若各步骤分成不同 Kernel中间结果有可能写到 HBM再被后续 Kernel 读回来。执行步骤越碎启动次数和中间数据流量可能越多。Kernel FusionKernel 融合将相邻操作合并到一个 Kernel 或更紧密的执行路径中目标是减少Kernel 启动与调度成本中间张量写入 HBM 的次数后续重新读取中间张量的流量部分重复的地址计算与执行开销。但融合有资源代价更复杂的 Kernel 可能占用更多寄存器和共享内存指令路径可能变长甚至降低并行度。“融合”是优化手段不是无条件收益保证。必须在目标形状与目标 GPU 上测量。八、FlashAttention用算法改写内存访问方式FlashAttention 是本讲的重点因为它很好地展示了一个原则有时更快的实现不是改变数学目标而是改变完成数学计算时的数据移动方式。8.1 标准 Attention 的计算式缩放点积 Attention[\operatorname{Attention}(Q,K,V)\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V]各符号含义(Q)Query(K)Key(V)Value(d_k)Key 向量的维度(QK^T)每个 Query 与各个 Key 的匹配分数除以 (\sqrt{d_k})控制分数尺度Softmax把分数归一化成权重与 V 相乘按权重混合 Value 信息。如果序列长度为 (N)单个头的分数矩阵通常有 (N\times N) 个元素概率矩阵也有同阶大小。仅保存一个 (N\times N) 中间矩阵在元素数上就随 (N^2) 增长。长上下文的中间张量读写会成为明显成本。8.2 FlashAttention 的洞察FlashAttention 的核心是IO-aware exact attention考虑 HBM 与 GPU 片上 SRAM 之间的读写使用分块与在线 Softmax减少完整分数矩阵、概率矩阵的物化materialization以及往返 HBM 的数据量。概念流程如下传统流程 Q、K → 计算完整 S QKᵀ / √d → 保存/读取 S ↓ Softmax 得到 P ↓ 保存/读取 P ↓ P × V → 输出 FlashAttention 思路 取 Q tile K/V tile ↓ 在片上算当前分数块 ↓ 在线更新 Softmax 状态和输出累加器 ↓ 读取下一块 K/V继续合并 ↓ 输出最终结果不必把完整 S 和 P 写回 HBM这里的关键不只是“分块”而是分块执行时仍能正确完成 Softmax 的归一化并把不同分块的贡献逐步合并。8.3 在线 Softmax如何不保存完整分数矩阵Softmax 对一行分数 (s_j) 的定义为[\operatorname{softmax}(s_j)\frac{e^{s_j}}{\sum_t e^{s_t}}]直接计算时需要知道整行的归一化分母。FlashAttention 按 tile 分批处理分数但用一个可更新状态来保持计算正确性。为避免指数溢出常用行最大值进行平移。对某一行假设已有状态(m)已处理部分的最大分数(l)平移后指数项的和(o)指数权重作用于 V 后的累加向量。当前新 tile 的分数向量为 (s)更新最大值[m’\max(m,\max(s))]旧累加量的缩放系数为[\alphae^{m-m’}]当前 tile 的指数权重为[pe^{s-m’}]于是归一化分母状态和输出累加器更新为[l’\alpha l\sum_j p_j][o’\alpha opV_{\text{tile}}]当所有 tile 处理完后输出为[O\frac{o}{l}]这是便于理解的逐行表达。实际 FlashAttention Kernel 还必须处理多维张量布局、Mask、因果注意力、并行归约、数值稳定性以及 GPU 特定指令等细节。8.4 FlashAttention 解决什么不解决什么解决或改善减少完整注意力分数与概率矩阵的物化降低 HBM 与片上存储之间的读写量增强分块内的数据复用在合适的上下文长度、批量、硬件与实现下降低 Attention 的运行时间和显存压力。没有直接解决没有让稠密 Attention 所需的全部算术工作自动变成线性复杂度不保证任何形状、任何硬件、任何 batch 都以相同倍率加速不保证浮点计算结果逐 bit 完全相同不取代量化、KV Cache、并行训练或调度等其他优化。对固定 head dimension 而言稠密 Attention 的主要矩阵计算仍随序列长度呈二次增长。FlashAttention 主要改变的是IO 复杂度与中间结果存储方式而不是消除全部的成对 Token 计算。8.5 它为什么能成为一堂重要的系统课传统优化视角可能只问“乘法能不能更快”FlashAttention 还会问“这次乘法的输入是否已经在片上中间结果是否非得写回 HBM能不能只保存足够的统计量最后直接得到输出”这就是算法、硬件存储层次和数值计算协同设计的典型例子。九、训练和推理为什么有不同的性能画像9.1 训练训练包含前向计算、反向传播和优化器更新。大型矩阵乘法往往提供大量可并行计算因此 Tensor Core、合适精度、GEMM 实现和数据并行效率很重要。同时激活、梯度和优化器状态也可能带来显存容量与带宽压力。9.2 Prefill 与 Decode大模型推理通常可粗分为两段Prefill处理输入提示词的阶段一次处理一段已有序列往往能形成较大矩阵计算Attention 的分块实现和矩阵效率通常很重要。Decode逐步生成后续 Token 的阶段。每一步的工作量、batch、模型参数读取、KV Cache 访问和调度成本都会影响表现。低 batch 的自回归 Decode 常更容易受内存流量制约但实际瓶颈还取决于系统配置与并发规模。不要绝对化地说“Prefill 一定算力受限Decode 一定带宽受限”。这是常见倾向不是放之四海皆准的定律。9.3 FlashAttention 与 KV Cache 不是一回事FlashAttention改善 Attention 计算过程中的 IO 和中间结果处理方式。KV Cache在自回归生成时保存历史 Token 的 Key/Value避免每次新生成 Token 都重新计算历史 Token 的 K/V 投影。两者优化对象不同可以在同一系统里配合使用。十、实践怎样正确测量 GPU Kernel10.1 为什么不能随便用 CPU 计时GPU 操作经常是异步提交的。Python 发起一次矩阵乘法后CPU 可能很快返回但 GPU 还没执行完。只用普通 CPU 墙钟统计 Python 调用耗时可能主要测到的是任务提交时间。正确的基准至少应做到检查 GPU 是否可用进行预热避免首次初始化开销干扰使用 CUDA Event 或适当同步测量 GPU 时间固定数据类型、形状、batch 和测试条件多轮重复并报告稳定统计量验证输出误差避免把错误结果当作“更快”需要深入分析时使用 Profiler而不是凭感觉猜瓶颈。10.2 PyTorch 的 GPU Event 示例importtorchifnottorch.cuda.is_available():raiseRuntimeError(当前环境没有可用的 CUDA GPU)atorch.randn(2048,2048,devicecuda,dtypetorch.float16)btorch.randn(2048,2048,devicecuda,dtypetorch.float16)# 预热for_inrange(10):ca b torch.cuda.synchronize()starttorch.cuda.Event(enable_timingTrue)endtorch.cuda.Event(enable_timingTrue)start.record()for_inrange(100):ca b end.record()torch.cuda.synchronize()avg_msstart.elapsed_time(end)/100print(f平均单次耗时{avg_ms:.3f}ms)这个例子展示测量方式不是严谨的跨 GPU 基准。实际实验需考虑 warm-up 数量、重复次数、功耗与频率状态、输入形状、同步点和误差验证若要比较不同方法必须保持测试条件一致。10.3 用什么工具Nsight Systems看整体时间线检查 CPU、GPU、Kernel 启动、同步、传输之间的关系。Nsight Compute深入观察单个 Kernel 的硬件性能指标辅助判断内存访问、计算吞吐、占用率和资源限制。框架 Profiler观察 PyTorch 等框架里的算子时间、调用次数、内存使用与调用栈。Profiler 的数据需要结合硬件架构解释。一个指标偏低并不自动说明某个单独原因要结合 Kernel 特性和多个指标交叉判断。十一、常见性能问题与排查方向症状/现象可能原因建议检查理论算力很高实测远低于峰值算子没走高效矩阵路径、矩阵规模过小、数据供给不足核对 dtype、形状、库实现和 Kernel 选择查看 Compute 指标HBM 流量很大重复读写中间结果、访存不合并、数据复用差检查 tile、融合、访问布局与缓存行为Kernel 很多且都很短细碎算子多、启动开销占比高分析时间线考虑融合或批处理加大 tile 后反而更慢寄存器/共享内存占用升高驻留 Block 减少数据路径变差测试不同 tile关注资源占用和 spillShared Memory 加了却没有加速复用不足、同步太多、Bank Conflict、原本缓存已有效比较实际传输、同步和耗时不要只看代码“用上了共享内存”低精度后速度提升很小原瓶颈在带宽、启动或其他算子未走目标硬件路径Profiling检查算子覆盖、形状及误差GPU 使用率看似不高小 batch、CPU 供给不足、等待同步、任务短或采样统计误差看完整时间线与请求规模不要只凭单一利用率百分比判断十二、把整讲串起来高性能 GPU 计算的设计链数学工作量与数据依赖 ↓ 选择适合的数据格式和计算路径 ↓ Tensor Core / 其他计算单元 ↓ 分块与数据复用尽量让数据在片上多用几次 ↓ 控制 HBM 读写、Kernel 启动、同步及资源占用 ↓ 用 Profiler 确认瓶颈 ↓ 根据测量结果修改并复测这条链的重点是不要孤立地优化一个部件。Tensor Core 再快如果得不到数据也会空等Shared Memory 再快如果同步和搬运成本超过收益也不一定有用FlashAttention 的价值就在于它把数学计算、片上存储和 IO 成本放在一起重新设计。十三、本讲总结需要牢牢记住的 10 个结论Tensor Core 是面向矩阵乘加等运算的专用单元必须有合适的数据类型、矩阵形状和软件路径。CUDA Core 数量、Tensor Core 峰值或 TFLOPS 都不能单独决定真实应用性能。FP16、BF16、TF32、FP32 的范围、精度和硬件路径不同混合精度要同时考虑速度与数值稳定性。HBM、Cache、Shared Memory、Register 在容量、速度、访问范围和控制方式上各有分工。SRAM 是存储技术概念Shared Memory 是编程模型中的存储抽象不可混为一谈。分块通过数据复用减少重复显存流量但也增加资源与同步方面的权衡。算术强度和 Roofline 模型帮助初步判断是计算瓶颈还是带宽瓶颈。Kernel Fusion 可能减少启动和中间张量读写但也可能带来寄存器压力等副作用。FlashAttention 通过分块与在线 Softmax 降低 Attention 的 IO 压力它没有自动消除稠密 Attention 的二次级别算术工作。GPU 优化必须基于正确的基准测试和 Profiler不要靠猜测或宣传峰值下结论。十四、自测题用自己的话说明 CUDA Core 与 Tensor Core 的定位差异。为什么 HBM 容量、HBM 带宽与访问延迟不是同一个指标Register、Shared Memory、L1/L2 Cache 分别适合处理什么为什么数据复用可以减少 HBM 流量请画出一个矩阵 tile 的数据流。写出算术强度定义并解释 Roofline 模型的两个上界。为什么更大的 tile 可能降低性能至少说出两个原因。标准 Attention 为什么会出现 (N\times N) 的中间矩阵FlashAttention 如何通过在线 Softmax 得到最终输出FlashAttention 优化的是哪一部分成本它没有消除什么怎样正确测量 PyTorch GPU 运算时间参考资料NVIDIA,CUDA C Programming Guidehttps://docs.nvidia.com/cuda/cuda-c-programming-guide/NVIDIA,CUDA Refresher: The CUDA Programming Modelhttps://developer.nvidia.com/blog/cuda-refresher-cuda-programming-model/Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré,FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesshttps://arxiv.org/abs/2205.14135NVIDIA,Nsight Compute Documentationhttps://docs.nvidia.com/nsight-compute/最后一句话Tensor Core 决定一部分矩阵计算能有多快内存层次决定数据能否及时到达分块与融合决定数据是否被充分复用而 FlashAttention 证明了——重视 IO 成本有时比单纯追求更多计算更有价值。下一讲将继续讨论多 GPU 训练当一张 GPU 装不下、算不完时如何把模型、数据和计算分布到多张 GPU 上。
