AI Infra:FlashAttention,具有I/O感知的快速且内存高效的精确注意力机制
一、FlashAttention 简介FlashAttention 是由 Stanford Hazy Research 团队提出的高效注意力机制实现,核心目标是在不牺牲准确率的前提下最大化注意力计算效率,尤其适用于 GPU 上的训练与推理。https://github.com/Dao-AILab/flash-attentionPaper: https://tridao.me/publications/flash3/flash3.pdf1.1 核心原理将传统的多次读取/写入的注意力计算过程,融合为一次 GPU kernel 调用使用 tile-based 的块级处理方式,在寄存器中缓存中间...