海边拾贝-FlashInfer

Attention Engine 可以被理解成“Attetion 算子库 + Attention 运行时”。有以下设计：

可拆分的 Attention 算子：提高了 GPU 内存带宽利用率

新的 KV Cache 管理抽象 CBSR：兼具 PagedAttention 和 RadixAttention 的优点，更 general

宏观上动态、微观上静态的调度运行时：动态的同时不损害静态抽象和收益

可定制的 Attention 算子框架、JIT：一些工程特性

FlashInfer 这篇工作在 2023 年就提出了，据作者所言，那个时候只有 FlashAttention1，还没有 FA2&3，FlashDecode 等工作，但是这些工作论文发得更早。

一、Background

FlashAttention

传统的 Attention 运算需要扫描 3 遍 GPU Memory 中的 Attention Logits 矩阵（len, len），计算密度低。

FlashAttention 提出了 1-pass 的 attention 算法，提高了计算密度，缓解了 GPU Memory 到 GPU Cache 的 IO 瓶颈：

FlashAttention1 不适应 Decode 场景

FlashAttention 是为了训练开发的，Queries 的长度往往很长（也就是 Queries 彩色方块的高度很高），这样可以充分利用 GPU 上的计算资源。

而在 decode 阶段，采用自回归生成，每次的 Queries 的长度就为 1（也就是 Queries 彩色方块高度只有 1），那么计算资源就不会得到充分利用（按照 FlashDecode 的说法，是不到 1%）。

更加形式化地去看，在 FlashAttention 中，设置 $l_q$ 是 Queries 的长度， $l_k$ 是 K 的长度，那么一次 Attention 计算的访存开销是 $O(l_q + l_k)$ ，计算开销是 $O(l_q l_k )$ ，则计算密度是：

$O(\frac{l_q l_k}{l_q + l_k}) = O(\frac{1}{\frac{1}{l_q} + \frac{1}{l_k}} )$

如果考虑 $l_{k}$ 是一个很大的值（长文本或者推理模型都会导致 $l_k$ 很大），那么计算密度约等于 $O(l_q)$ 。当 $l_q = 1$ 时就会导致计算资源利用不足。

Attention 的输入是动态变化的

而实际情况会更加复杂，query 的长度是会动态变化的，从应用场景区分，有 3 种：

放到 roofline 上来看

动态变化的 Queries 长度就对 Attention 的动态性提出了一定的要求。

此外优化长文本有一种经典的技术就是 KV Cache 稀疏，也就是 KV Cache 也会存在变化（不止是单调递增），如 NSA 就包含三种稀疏特性：

Attention 算子库既要利用特化的优势，又要有足够好的定制性。

二、Design

2.1 Split-K

原版的 FlashAttention 需要利用前缀和不断缩放校正（scale）局部结果，但是“前缀和”就意味着“顺序遍历”，而当 Queries 的长度较小时，就容易导致利用率不高。

通过调整算法，我们可以实现并行计算不同的 KV Cache Chunk：

所以为了达到这种 merge 的效果，我们需要记录每个 block 的一些运算结果，这些结果在文中被称为 Attention State。

每个 block $\mathcal{I}$ 需要记录两种 State，分别是 attention scale：

和 attention output：

有了这两个东西以后，我们就可以将 block $\mathcal{I}$ 和 block $\mathcal{J}$ 融合到一起了：

Split-K 算法改善了 FA1 在 Decode 场景下计算资源利用率不足的问题。

Split-K 策略也是 FlashDecode 的核心 Idea，对此 FI 的作者叶子豪解释道：

其实主要原因是我们跟 FA2 和 FlashDecoding 的开发几乎都是同期进行的，在 FlashAttention2 发布之前我们已经独立探索过了 FA2 中大部分的优化。而 FlashDecoding 我们在去年 8 月就已经有了较完整的实现和评测，不过我对 LLM 这个领域的内卷程度稍有低估没有及时推广，导致被抢发出来。抛开这些虚名而言，LLM Serving 还有很多工程上的问题需要解决，同志仍需努力。