在 NVIDIA Ampere 架构A100 /sm_80时代Tensor Core 的运算指令以mma.sync为代表。然而到了 Hopper 架构H100 /sm_90FlashAttention-3FA3能够实现 FP16 算力翻倍、FP8 算力突破 PFLOPS 级别的核心秘诀就在于放弃了mma.sync全面采用了全新的WGMMAWarp Group Matrix Multiply-Accumulate异步矩阵乘指令。WGMMA 绝不仅仅是“一条新的汇编指令”它代表了 GPU执行粒度、数据流路径与异步流水线的一次范式革命。一、 核心对比Amperemma.syncvs HopperWGMMA为了直观展现变革我们先来看两者的底层差异对比维度Amperemma.sync(sm_80)HopperWGMMA(sm_90)驱动粒度 (Granularity)Single Warp32 个线程Warp Group4 个连续 Warp共 128 个线程同步范式 (Execution)阻塞同步Synchronous完全异步Asynchronous Execution数据源 (Operand A Source)必须从寄存器RF读取支持直接从 Shared Memory (SRAM) 读取寄存器压力 (Register File)极大必须先 SRAM→\to→Reg再计算极小SRAM 驱动省去矩阵AAA的寄存器指令发射开销高32 线程频繁发射mma.sync极低128 线程作为一个硬件整体发射一次二、 深度拆解一执行粒度的跃迁Single Warp→\to→Warp Group1. Ampere 的 Warp 级粒度 (mma.sync)在 Ampere 架构中Tensor Core 的最小驱动单位是1 个 Warp32 线程。每 32 个线程协同分配一个微小矩阵切片如16×8×1616 \times 8 \times 1616×8×16。为了计算一个较大的 Tile如64×6464 \times 6464×64SM 需要频繁为各个 Warp 派发大量的mma.sync指令。硬件指令译码器Instruction Fetch/Decode和发射槽Issue Slots面临巨大的负载。2. Hopper 的 Warp Group 级粒度 (wgmma.mma_async)Hopper 硬件原生引入了Warp Group硬件抽象——由4 个连续且对齐的 Warp共 128 个线程组成一个统一的调度单位。硬件级大矩阵切片128 个线程作为一个整体一条指令即可驱动 Tensor Core 执行高达64×256×3264 \times 256 \times 3264×256×32FP8或64×128×1664 \times 128 \times 1664×128×16FP16的单步大 GEMM。极致的指令发射效率指令开销降低为原来的14\frac{1}{4}41指令发射管线Instruction Pipeline被大幅释放彻底告别了前端译码瓶颈。Ampere (32 Threads): [ Warp 0 ] ──► mma.sync ──► Tensor Core \ 频繁发射小粒度指令 [ Warp 1 ] ──► mma.sync ──► Tensor Core ├─ 译码开销大 [ Warp 2 ] ──► mma.sync ──► Tensor Core / Hopper (128 Threads): ┌─────────────────────────────────────────┐ │ Warp Group (Warp 0 1 2 3) │ ──► wgmma.mma_async ──► Hopper Tensor Cores └─────────────────────────────────────────┘ (一条指令驱动 128 线程大切片)三、 深度拆解二数据路径的革命SRAM-Driven GEMM这是 WGMMA 给 CUDA 编程带来的最大物理红利——直接省掉了一半的寄存器占用1. 传统的mma.sync数据路径必须经过寄存器在 A100 上即便数据已经通过cp.async到了 SRAMTensor Core 依然无法直接读取 SRAM。必须经历以下步骤LDG/cp.async:Global Memory (HBM)→\to→Shared Memory (SRAM)LDS(Load Shared):Shared Memory (SRAM)→\to→通用寄存器 (Register File)-- 致命瓶颈mma.sync:通用寄存器→\to→Tensor Core 执行 GEMM代价矩阵AAA和矩阵BBB的数据必须双双保存在通用寄存器中。在 Attention 中为了维持高占有率Occupancy寄存器迅速被打爆Register Spilling导致 Tile Size 无法做大。2. WGMMA 的 SRAM 直读数据路径SRAM-DrivenHopper 架构在硬件层面将 Shared Memory (SRAM) 与 Tensor Core 的输入管线直接相连TMA:Global Memory (HBM)→\to→Shared Memory (SRAM)WGMMA:Tensor Core直接从 SRAM 读取矩阵AAA甚至矩阵BBB结果直接累加到 Output 寄存器[ Ampere mma.sync Data Path ] HBM ──► SRAM ──► [ Register File ] ──► Tensor Core ──► Accumulator Reg [ Hopper WGMMA Data Path ] HBM ──► SRAM ──┬──────────────┐ │ (Direct Read)│ └──────────► Tensor Core ──► Accumulator RegFA3 收益在 FlashAttention-3 中矩阵QQQ或KKK可以完全停留在 Shared Memory 中无需为其分配任何通用寄存器节省出来的极大量寄存器空间可以全部用来存放矩阵乘法的累加结果Accumulator Registers或者增大矩阵 Tile 的尺寸如将 TileNNN从 64 扩大到 128/256从而指数级提升计算密度。四、 深度拆解三完全异步与非阻塞执行Async Pipelines1.mma.sync的“假异步”与同步停顿虽然名字叫mma.sync但它本质上是阻塞式同步指令当一个 Warp 发射mma.sync时该 Warp 的流水线必须等待 Tensor Core 完成计算或至少完成数据准备后才能继续向下发射后续无关指令例如 Softmax 计算。计算与指令流难以解耦导致依赖冲突Dependency Stalls。2. WGMMA 的硬件异步队列wgmma.mma_async是纯粹的非阻塞异步指令发射即返回128 个线程组成的 Warp Group 发射一条wgmma.mma_async后指令被推入硬件级的 WGMMA 异步队列Warp Group无需等待计算完成指令立刻返回。后台并行计算Tensor Core 在后台静默读取 SRAM 并进行 GEMM 矩阵乘法。协同计算Overlap在 Tensor Core 异步计算矩阵乘法的同时CUDA 线程Vector Core / ALU可以立刻去计算上一轮 Tile 的 Softmax如求exp\expexp、Sum 或 Scaling实现了GEMM 与非 GEMM 算子的完全重叠统一屏障等待当必须依赖 GEMM 结果时只需要调用wgmma.wait_group指令进行组级等待即可。// PTX 级别的 WGMMA 异步调用伪代码逻辑wgmma.mma_async.sync.aligned.m64n128k16...// 发射异步 GEMM (Tile 0)不阻塞wgmma.mma_async.sync.aligned.m64n128k16...// 发射异步 GEMM (Tile 1)不阻塞// 【关键重叠区】CUDA Vector Core 同时在计算 Softmax完全掩盖 GEMM 计算开销compute_softmax_on_vector_core(...);wgmma.wait_group0;// 等待所有后台 WGMMA 组计算完毕再使用累加器结果五、 总结FA3 如何以 WGMMA 为基石建立榨干 H100 的物理管线FlashAttention-3FA3正是将TMA、Warp Specialization、mbarrier 与 WGMMA融为一体Producer Warp发射TMA指令把 HBM 的Q,K,VQ, K, VQ,K,V零开销拉到 SRAMmbarrier硬件自动通知数据到齐Consumer Warp Group发射wgmma.mma_async指令驱动 Tensor Core直接读取 SRAM执行QKTQ K^TQKT和PVP VPV矩阵乘法Vector Core利用 WGMMA 的非阻塞特性在后台 GEMM 计算的同时在同一空间内交错执行 Softmax 归一化。通过放弃传统的 Warp 级同步mma.sync拥抱 Warp Group 级异步 SRAM 直读的WGMMAFlashAttention-3 彻底移除了寄存器瓶颈与指令发射瓶颈将 H100 GPU 的物理计算效率推向了理论极致。