一种通过IO感知的分块计算方式降低Transformer注意力机制显存占用和计算复杂度的技术,是实现超大上下文窗口的关键技术之一
FlashAttention等优化技术已降低实际内存开销,但计算量的基本规律未改变
FlashAttention由斯坦福大学的Tri Dao等人提出,通过将注意力计算分块在GPU片上SRAM中完成,将内存访问次数降低一个数量级,是数学上精确等价的实现
许多第三方库如自定义CUDA kernel、FlashAttention等高性能注意力实现不支持MPS
FlashAttention通过重新设计注意力计算的内存访问模式,将显存占用从与序列长度平方成正比降低为线性关系