1 FlashAttention

FlashAttention是一种优化Transformer模型计算效率和内存使用的技术。它通过减少存储访问开销（Memory Access Cost，MAC），而非降低FLOPS（浮点运算次数），来提升性能。

2 前述知识点

涉及到内存访问，肯定与计算的硬件架构有关系。

从GPU架构进行解析，参考如下博客：
大模型推理加速技术的学习路线是什么
首先，我们将探讨GPU架构，特别是其内存层次结构。我们确定了两个重要模式：计算限制（compute bound）和内存限制（memory bound），并讨论了大规模Transformer推理受内存限制的原因。大部分优化都基于Transformer推理受内存限制这一基本事实，例如只要我们提高FLOP利用率，就能提高效率。

2.1 GPU架构

GPU架构总体如下图所示：

基础部分：DRAM（动态随机存取存储器）、L2缓存和SM（流处理器单元）

• 与CPU对比
  • SM类似于CPU核心，但具有更高级的并行性；
  • L2缓存和DRAM类似于CPU的L2缓存和DRAM
  • 在Flash Attention论文中，L2缓存被称为SRAM（静态随机存取存储器）
• A100 80G SXM
  • 08个SM，DRAM容量为80GB，有40M L2缓存

SM内部包含什么？

• L1缓存：指令和数据
• 张量核心：进行矩阵乘法运算的地方。回想一下，神经网络计算基本上就是巨大批量的矩阵乘法。

GPU编程基础

在执行model.generate(prompt)时，我们进行以下操作：

• 内存访问:
  • 从高带宽内存（HBM）加载模型权重 -> L2缓存 -> 传输到SM（流处理器单元）
• 计算:
  • 在SM中执行矩阵乘法，SM请求张量核心执行计算
• A100:
  • 108个SM，DRAM容量为80G，40M L2缓存
  • bf16张量核心：每秒312万亿浮点运算（TFLOPS）
  • DRAM内存带宽为2039GB/秒 = 2.039T/秒
• 如果模型很大，我们将其分割到多个GPU上，比如两个由NVLink连接的GPU
  • NVLink 300GB/秒 = 0.3T/秒
  • 我们大致观察了速度层次结构。尽管不能直接比较，但它们的数量级差异是我们需要优化的主要方面：
  • 312T（SM计算） > 2.03T（DRAM内存访问） > 0.3T=300G（NVLink跨设备通信） > 60G（PCIe跨设备通信）
• 这意味着，如果我们希望速度更快，我们应该尽力：
  • 充分利用SM
  • 减少单个GPU的内存访问（因为它比计算慢得多），减少GPU之间的通信（因为它甚至比内存访问还要慢）。

计算限制与内存限制

如何确定我们是否充分利用了SM呢？我们通过以下方式检查是否计算或内存限制：

定义每字节GPU操作 = flop / 内存带宽

• A100 = 312 / 2.039

• 定义计算强度 = 计算 / 内存访问

• 如果计算强度大，说明程序更会受到计算限制；如果计算强度较小，则更受内存限制。
  

• 增加批次大小会将行为从内存限制变为计算限制。

• 内核融合：减少了内存访问操作，因为我们将多个操作合并为一个操作。

2.2 Transformer推理

内存布局

正如我们所看到的，为了在bf16格式下运行一个13B模型，我们大约只有10GB的内存来存储kv缓存。这意味着：

• 不能使用太大型的批次（尽管我们希望使用更大的批次大小以提高效率）
• 也不能处理太长的序列，尽管我们确实希望能够处理长度为100k的序列。

3 FlashAttention的策略

FlashAttention的核心策略包括：

• Tiling（平铺/切分）：将注意力矩阵分解成更小的子矩阵，分别计算，确保每个子矩阵的大小适合SRAM（静态随机存取存储器）的存储能力，从而减少对HBM（高带宽内存）的访问。
• Recomputation（重算）：在反向传播时，不存储所有中间状态，而是在需要时重新计算，节省内存。
• 分块SoftMax：解决标准SoftMax在分块计算中的问题，确保整个Flash Attention的正确性。
• 优化显存交换：减少SRAM与HBM之间的数据交换，加速计算。
  这些策略共同作用，使FlashAttention在保持计算精度的同时，显著提高计算速度和内存效率

4 Ascend 上的FlashAttention

昇腾异构计算架构CANN针对昇腾AI处理器的片上内存和缓存大小，以及数据搬运通路，基于Ascend C算子编程语言优化实现FlashAttention融合算子，充分利用片上缓存，提升Attention处理性能。根据实测，在一些典型场景中CANN的FlashAttention算子相比小算子取得了5倍以上的性能提升，开发者可直接调用相关算子API接口使能大模型极致性能优化。

可参考：
基于Ascend C的FlashAttention算子性能优化最佳实践-技术干货-昇腾社区
​www.hiascend.com/developer/techArticles/20240607-1?envFlag=1