NLP高频面试题（二十）——flash attention原理

FlashAttention是一种针对Transformer模型中自注意力机制的优化算法，旨在提高计算效率并降低内存占用，特别适用于处理长序列任务。
在Transformer架构中，自注意力机制的计算复杂度和内存需求随着序列长度的平方增长。这意味着当处理较长序列时，计算和内存负担会显著增加，导致模型训练和推理的效率降低。

FlashAttention的核心思想

FlashAttention通过以下关键技术来优化自注意力机制：

1. 分块计算（Tiling）：将输入序列划分为较小的块（tiles），并在每个块上独立执行注意力计算。这种方法减少了对高带宽内存（HBM）的读写操作，因为计算可以在更接近处理单元的片上高速缓存（SRAM）中进行，从而提高了数据访问效率。

2. 重计算策略（Recomputation）：在反向传播阶段，选择性地重新计算前向传播中未存储的中间结果，而不是将所有中间结果都保存在内存中。这种策略减少了内存占用，同时通过权衡计算和内存使用来优化整体性能。

FlashAttention的实现细节

在具体实现中，FlashAttention采用以下步骤：

• 前向传播：对于每个输入块，依次加载查询（Q）、键（K）和值（V）矩阵的相关部分到片上高速缓存中，执行注意力计算，生成输出。计算完成后，丢弃不再需要的中间结果，以释放内存。