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 瓶颈: