MMDec, 2023

CUDA 内核融合研究案例：在 NVIDIA Hopper 架构上使用 CUTLASS 库实现 FlashAttention-2

PDF
TL;DR我们提供了 FlashAttention-2 前向传递的优化实现，使用了自定义融合的 CUDA 内核，以适应 NVIDIA Hopper 架构，并使用开源的 CUTLASS 库编写。在此过程中，我们解释了将在线 softmax 与连续的 GEMM 内核融合在一起的挑战和技术，利用 Hopper 特定的 Tensor Memory Accelerator（TMA）和 Warpgroup Matrix-Multiply-Accumulate（WGMMA）指令，定义和转换 CUTLASS 布局和张量，重叠复制和 GEMM 操作，并选择 Q、K 和 V 注意力矩阵的最优瓦片大小，同时平衡寄存器压力和共享内存利用率。在单个 H100 PCIe GPU 上进行的对比性测试中，针对某些常见的超参数选择，我们观察到与针对上一代 NVIDIA Ampere 架构进行优化的 FlashAttention-2 版本相比，FLOPs/s 高出 20-50%。
发现论文，激发创造
  • FlashAttention-2：更快的注意力计算，更好的并行和工作分割

    使用 FlashAttention 和 FlashAttention-2 算法来提高模型的效率

    Jul, 2023

  • NVIDIA Tensor Core 可编程性、性能和精度

    本文介绍了 NVIDIA Volta GPU 微架构中一种特殊的处理器 Tensor Core，它能够在每个时钟周期内对 4x4 矩阵执行一个矩阵乘积累加操作，经过测试，我们发现 NVIDIA Tensor Core 在 Tesla V100 GPU 上能够以混合精度提供 83 Tflops/s 的处理速度，使得 HPC 应用的矩阵乘法能够获得显著的性能提升。

    Mar, 2018

  • SWAT：基于 FPGA 的可扩展和高效窗口注意力 Transformer 加速

    提出了一种基于 FPGA 的加速器设计，SWAT，通过最大程度地利用稀疏性实现可扩展的性能，相比基准 FPGA 加速器，在延迟和能效方面改进了 22 倍和 5.7 倍，并且比基于 GPU 的解决方案提高了 15 倍的能效。

    May, 2024

  • 精简注意力：面向 Transformer 解码阶段的硬件感知可扩展注意力机制

    LeanAttention 是一种可扩展的自注意力计算技术，通过重新设计解码阶段的执行流程，将自注意力机制的实现扩展到具有挑战性的长上下文长度情况，以并行计算的方式提供 2.6 倍的平均注意力执行加速和最多 8.33 倍的速度提升。

    May, 2024

  • FlashAttention: 带 IO 感知的快速、节省内存的精确注意力机制

    提出了 FlashAttention，一种 IO 感知的确切关注算法。FlashAttention 使用平铺减少 GPU 高带宽内存（HBM）和 GPU 片上 SRAM 之间的内存读取 / 写入次数，并可扩展为块状稀疏关注。FlashAttention 使 Transformers 速度提高，使其具有更长的上下文并获得更高质量的模型，以及实现了 Path-X 挑战的首个 Transformers。

    May, 2022

  • 大规模序列中通过稀疏闪存注意力加速因果注意力

    我们扩展了 FlashAttention 以容纳一大类稀疏性注意力模式，其中包括关键 / 查询删除和基于哈希的注意力。即使具有相对较低的稀疏度，我们的方法在序列长度增加时可见地改进 FlashAttention 的性能。不牺牲困惑度的情况下，我们将 transformer 语言模型的训练速度分别提高了 2.0 倍和 3.3 倍，对于长度为 8k 和 16k 的序列。

    Jun, 2023

  • 基于元素的注意力层：一种优化选项

    本文提出了一种新的注意力机制，并通过将 Dot-Product Attention 中的矩阵乘法转换成数组乘法来实现元素级别的注意力，从而在保持实验精度的同时减少了网络的参数量。在 Fashion MNIST 和 CIFAR10 数据集上的分类实验中，该方法的结果表明，相对于 VGG 类似结构，神经网络降低了 97% 的参数量，但在 Fashion MNIST 数据集上的分类准确率达到了 92%，而在 CIFAR10 数据集上的分类准确率仍相当于 VGG-like 架构的 60%。

    Feb, 2023

  • SpAtten: 级联记号与头部修剪的高效稀疏注意力架构

    提出了 SpAtten，这是一种高效的算法 - 架构协同设计，利用标记稀疏性、头部稀疏性和量化机会，通过级联剪枝和渐进式量化来减少注意力计算和内存访问，实验结果证明其在 30 个基准测试中平均减少 DRAM 访问 10 倍，并且相对于其他加速器和处理器能达到巨大的加速和能耗降低。

    Dec, 2020

  • 随机尖峰注意力：在尖峰网络中利用随机计算加速注意力

    通过利用随机计算（SC）执行基于脉冲神经网络（SNN）的 Transformer 模型的点积注意力机制，实现高准确率和能源效率。

    Feb, 2024

  • 硬件高效训练的门控线性注意力变换器

    通过引入数据依赖的门控机制，我们开发了一种硬件高效的并行形式，使得门控线性注意力（GLA）Transformers 在适度规模的语言建模中表现竞争力，并在训练速度上与基于 CUDA 优化的 FlashAttention-2 相媲美。

    Dec, 2023