FlashAttention

本仓库提供了 FlashAttention 和 FlashAttention-2 的官方实现，源自 以下论文。

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré
Paper: https://arxiv.org/abs/2205.14135
IEEE Spectrum article about our submission to the MLPerf 2.0 benchmark using FlashAttention.

FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Tri Dao

Paper: https://tridao.me/publications/flash2/flash2.pdf

用法

我们非常高兴看到 FlashAttention 在发布后如此短的时间内就被广泛采用。此页面 包含了一个 FlashAttention 被使用地点的部分列表。

FlashAttention 和 FlashAttention-2 可自由使用和修改（参见 LICENSE）。 如果您使用了 FlashAttention，请引用并致谢 FlashAttention。

FlashAttention-3 测试版发布

FlashAttention-3 针对 Hopper GPU（例如 H100）进行了优化。

博客文章：https://tridao.me/blog/2024/flash3/

论文：https://tridao.me/publications/flash3/flash3.pdf

这是一个用于在将其集成到仓库其余部分之前进行测试 / 基准测试的 beta 版本。

当前已发布：

• FP16 / BF16 前向和反向，FP8 前向

要求：H100 / H800 GPU，CUDA >= 12.3。

我们强烈建议使用 CUDA 12.8 以获得最佳性能。

安装方法：

cd hopper
python setup.py install

运行测试：

export PYTHONPATH=$PWD
pytest -q -s test_flash_attn.py

安装包安装完成后，您可以按以下方式导入：

from flash_attn_3 import flash_attn_interface
flash_attn_interface.flash_attn_func()

使用 uv 安装，在您的 pyproject.toml 中：

[project]
dependencies = [
    "flash-attn-3"
]

[tool.uv]
no-build-isolation = true

[tool.uv.sources]
flash-attn-3 = { git = "https://github.com/Dao-AILab/flash-attention", subdirectory = "hopper" }

FlashAttention-4 (CuTeDSL)

FlashAttention-4 使用 CuTeDSL 编写，并针对 Hopper 和 Blackwell GPU（例如 H100、B200）进行了优化。

安装方法如下：

pip install flash-attn-4

如果你使用的是 CUDA 13，我们建议安装 cu13 附加组件以获得最佳性能：

pip install "flash-attn-4[cu13]"

安装完成后，你可以按如下方式使用：

from flash_attn.cute import flash_attn_func

out = flash_attn_func(q, k, v, causal=True)

安装与功能

要求：

• CUDA 工具包或 ROCm 工具包
• PyTorch 2.2 及以上版本。
• packaging Python 包（pip install packaging）
• psutil Python 包（pip install psutil）
• ninja Python 包（pip install ninja） *
• Linux。从 v2.3.2 开始可能在 Windows 上也能工作（我们收到了一些积极的报告），但 Windows 编译仍需更多测试。如果你有关于如何为 Windows 设置预构建 CUDA 轮子的想法，请通过 Github issue 联系我们。

* 请确保已安装 ninja 并且其工作正常（例如 ninja --version then echo $? should return exit code 0). If not (sometimes ninja --version then echo $? 返回非零退出码），然后卸载并重新安装 ninja（pip uninstall -y ninja && pip install ninja）。如果没有 ninja， 编译可能需要很长时间（2 小时），因为它不使用多个 CPU 核心。使用 ninja 时，在 64 核机器上使用 CUDA 工具包编译只需 3-5 分钟。

安装方法：

pip install flash-attn --no-build-isolation

或者你可以从源代码编译：

python setup.py install

如果你的机器内存少于 96GB 且拥有大量 CPU 核心，ninja 可能会 运行过多的并行编译任务，从而耗尽内存。要 限制并行编译任务的数量，你可以设置环境变量 MAX_JOBS：

MAX_JOBS=4 pip install flash-attn --no-build-isolation

接口： src/flash_attention_interface.py

NVIDIA CUDA 支持

要求：

• CUDA 12.0 及以上版本。

我们推荐 Nvidia 的 Pytorch 容器，其中包含安装 FlashAttention 所需的所有工具。

支持 CUDA 的 FlashAttention-2 目前支持：

1. Ampere、Ada 或 Hopper GPU（例如 A100、RTX 3090、RTX 4090、H100）。对于 Turing GPU（T4、RTX 2080），请参阅单独的 flash-attention-turing 仓库，该仓库在 Turing 上支持 FlashAttention 功能的核心子集。
2. 数据类型 fp16 和 bf16（bf16 需要 Ampere、Ada 或 Hopper GPU）。
3. 所有头维度最高可达 256。头维度 > 192 的反向传播需要 A100/A800 或 H100/H800。自 flash-attn 2.5.5 起，头维度 256 的反向传播现在可以在消费级 GPU 上运行（如果没有 dropout）。

AMD ROCm 支持

ROCm 版本有两个后端。有 composable_kernel (ck)，它是默认后端，以及一个 Triton 后端。它们提供了 FlashAttention-2 的实现。

要求：

• ROCm 6.0 及以上版本。

我们推荐 ROCm 的 Pytorch 容器，其中包含安装 FlashAttention 所需的所有工具。

Composable Kernel 后端

FlashAttention-2 ROCm CK 后端目前支持：

1. MI200x、MI250x、MI300x、MI355x 以及 RDNA 3/4 GPU。
2. 数据类型 fp16 和 bf16
3. 前向和反向的头维度最高可达 256。

Triton 后端

Triton 实现的 Flash Attention 支持 AMD 的 CDNA (MI200, MI300) 和 RDNA GPU，使用 fp16、bf16 和 fp32 数据类型。它提供带有因果掩码、可变序列长度、任意 Q/KV 序列长度和头大小、MQA/GQA、dropout、旋转位置编码、ALiBi、分页注意力以及 FP8（通过 Flash Attention v3 接口）的前向和反向传播。滑动窗口注意力目前仍在开发中。

Triton 后端内核由 aiter 包提供，作为 git 子模块包含在 third_party/aiter 中，并在安装过程中自动安装。

要安装，请首先从 https://pytorch.org/get-started/locally/ 获取适用于 ROCm 的 PyTorch，然后安装 Flash Attention：

cd flash-attention
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .

要使用特定的 aiter 提交（例如用于测试或开发）：

cd flash-attention
cd third_party/aiter && git fetch origin && git checkout <commit-sha> && cd ../..
FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .

要运行测试（注意：完整测试套件需要数小时）：

FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pytest tests/test_flash_attn_triton_amd.py

Triton 后端使用针对确定性和跨工作负载合理性能进行优化的默认内核配置。为了获得峰值吞吐量，请启用 FLASH_ATTENTION_TRITON_AMD_AUTOTUNE="TRUE" 以搜索最佳设置，这将产生一次性的预热成本。

或者，如果 不 进行自动调优，可以使用 FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON 来设置单个 triton 配置，以覆盖 attn_fwd 的硬编码默认值。例如：

FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON='{"BLOCK_M":128,"BLOCK_N":64,"waves_per_eu":1,"PRE_LOAD_V":false,"num_stages":1,"num_warps":8}'

使用 Docker 快速入门：

FROM rocm/pytorch:latest

WORKDIR /workspace

# build flash attention with triton backend
RUN git clone https://github.com/Dao-AILab/flash-attention &&\ 
    cd flash-attention &&\
    FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .

# set working dir
WORKDIR /workspace/flash-attention

# set env variable to use triton backend
ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"

构建并运行：

docker build -t flash-attn-triton .
docker run -it --network=host --user root --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --shm-size 16G --device=/dev/kfd --device=/dev/dri flash-attn-triton

如何使用 FlashAttention

主要函数实现缩放点积注意力（softmax(Q @ K^T * softmax_scale) @ V）：

from flash_attn import flash_attn_qkvpacked_func, flash_attn_func

flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False,
                          window_size=(-1, -1), alibi_slopes=None, deterministic=False):
"""dropout_p should be set to 0.0 during evaluation
If Q, K, V are already stacked into 1 tensor, this function will be faster than
calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
of the gradients of Q, K, V.
If window_size != (-1, -1), implements sliding window local attention. Query at position i
will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
Arguments:
    qkv: (batch_size, seqlen, 3, nheads, headdim)
    dropout_p: float. Dropout probability.
    softmax_scale: float. The scaling of QK^T before applying softmax.
        Default to 1 / sqrt(headdim).
    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
    window_size: (left, right). If not (-1, -1), implements sliding window local attention.
    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
        the attention score of query i and key j.
    deterministic: bool. Whether to use the deterministic implementation of the backward pass,
        which is slightly slower and uses more memory. The forward pass is always deterministic.
Return:
    out: (batch_size, seqlen, nheads, headdim).
"""

flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False,
                window_size=(-1, -1), alibi_slopes=None, deterministic=False):
"""dropout_p should be set to 0.0 during evaluation
Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
If window_size != (-1, -1), implements sliding window local attention. Query at position i
will only attend to keys between
[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.

Arguments:
    q: (batch_size, seqlen, nheads, headdim)
    k: (batch_size, seqlen, nheads_k, headdim)
    v: (batch_size, seqlen, nheads_k, headdim)
    dropout_p: float. Dropout probability.
    softmax_scale: float. The scaling of QK^T before applying softmax.
        Default to 1 / sqrt(headdim).
    causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
    window_size: (left, right). If not (-1, -1), implements sliding window local attention.
    alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
        (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
        is added to the attention score of query i and key j.
    deterministic: bool. Whether to use the deterministic implementation of the backward pass,
        which is slightly slower and uses more memory. The forward pass is always deterministic.
Return:
    out: (batch_size, seqlen, nheads, headdim).
"""

def flash_attn_with_kvcache(
    q,
    k_cache,
    v_cache,
    k=None,
    v=None,
    rotary_cos=None,
    rotary_sin=None,
    cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,
    cache_batch_idx: Optional[torch.Tensor] = None,
    block_table: Optional[torch.Tensor] = None,
    softmax_scale=None,
    causal=False,
    window_size=(-1, -1),  # -1 means infinite context window
    rotary_interleaved=True,
    alibi_slopes=None,
):
    """
    If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
    k and v. This is useful for incremental decoding: you can pass in the cached keys/values from
    the previous step, and update them with the new keys/values from the current step, and do
    attention with the updated cache, all in 1 kernel.

    If you pass in k / v, you must make sure that the cache is large enough to hold the new values.
    For example, the KV cache could be pre-allocated with the max sequence length, and you can use
    cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.

    Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be
    rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
    If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos
    and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.
    If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at
    indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).

    See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.

    Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
    than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
    For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
    0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.

    If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
    For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
        1 1 1 1 0
        1 1 1 1 1
    If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
        0 0
        0 0
        0 0
        1 0
        1 1
    If the row of the mask is all zero, the output will be zero.

    If window_size != (-1, -1), implements sliding window local attention. Query at position i
    will only attend to keys between
    [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.

    Note: Does not support backward pass.

    Arguments:
        q: (batch_size, seqlen, nheads, headdim)
        k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
            page_block_size must be a multiple of 256.
        v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,
            or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)
        k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate
            k with k_cache, starting at the indices specified by cache_seqlens.
        v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.
        rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding
            to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.
        rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.
        cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the
            KV cache.
        block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.
        cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.
            If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].
            If the indices are not distinct, and k and v are provided, the values updated in the cache
                 might come from any of the duplicate indices.
        softmax_scale: float. The scaling of QK^T before applying softmax.
            Default to 1 / sqrt(headdim).
        causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
        window_size: (left, right). If not (-1, -1), implements sliding window local attention.
        rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.
            If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,
            rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1
            (i.e. GPT-NeoX style).
        alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
            (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
            is added to the attention score of query i and key j.

    Return:
        out: (batch_size, seqlen, nheads, headdim).
    """

要了解这些函数如何在多头注意力层（包括 QKV 投影、输出投影）中使用，请参阅 MHA 实现。

使用 🤗 Kernels

如果你的硬件环境属于上述任何一种，你也可以使用 kernels 库 立即使用 Flash Attention 2 和 3。

# pip install kernels

from kernels import get_kernel

# FA2
fa_module = get_kernel("kernels-community/flash-attn2", version=1)
flash_attn_func = fa_module.flash_attn_func

# FA3
fa3_module = get_kernel("kernels-community/flash-attn3", version=1)
flash_attn_func = fa3_module.flash_attn_func

更新日志

2.0：完全重写，速度提升 2 倍

从 FlashAttention (1.x) 升级到 FlashAttention-2

以下函数已重命名：

• flash_attn_unpadded_func -> flash_attn_varlen_func
• flash_attn_unpadded_qkvpacked_func -> flash_attn_varlen_qkvpacked_func
• flash_attn_unpadded_kvpacked_func -> flash_attn_varlen_kvpacked_func

如果同一批次中的输入具有相同的序列长度，使用以下函数更简单 且更快：

flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)

flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)

2.1：更改 causal 标志的行为

如果 seqlen_q != seqlen_k 且 causal=True，因果掩码将对齐到 注意力矩阵的右下角，而不是左上角。

例如，如果 seqlen_q = 2 且 seqlen_k = 5，因果掩码（1 = 保留，0 = 被掩码）为：
v2.0：
1 0 0 0 0
1 1 0 0 0
v2.1：
1 1 1 1 0
1 1 1 1 1

如果 seqlen_q = 5 且 seqlen_k = 2，因果掩码为：
v2.0：
1 0
1 1
1 1
1 1
1 1
v2.1：
0 0
0 0
0 0
1 0
1 1
如果掩码的某一行全为零，输出将为零。

2.2：针对推理进行优化

针对推理（迭代解码）场景进行优化，此时查询具有非常小的序列 长度（例如，查询序列长度 = 1）。这里的瓶颈是以尽可能快的速度 加载 KV 缓存，我们将加载过程拆分到不同的线程块中，并使用单独的 内核来组合结果。

参见函数 flash_attn_with_kvcache，它具有更多针对推理的功能 （执行旋转嵌入，就地更新 KV 缓存）。

感谢 xformers 团队，特别是 Daniel Haziza，为此合作做出的贡献。

2.3：局部（即滑动窗口）注意力

实现滑动窗口注意力（即局部注意力）。感谢 Mistral AI 以及 Timothée Lacroix 对此贡献。滑动窗口在 Mistral 7B 模型中被使用。

2.4: ALiBi（带线性偏置的注意力），确定性反向传播。

实现 ALiBi（Press et al., 2021）。感谢来自 Kakao Brain 的 Sanghun Cho 对此贡献。

实现确定性反向传播。感谢来自 Meituan 的工程师对此贡献。

2.5: 分页 KV 缓存。

支持分页 KV 缓存（即 PagedAttention）。 感谢 @beginlner 对此贡献。

2.6: 软截断（Softcapping）。

支持带软截断的注意力，如 Gemma-2 和 Grok 模型中所使用的。 感谢 @Narsil 和 @lucidrains 对此贡献。

2.7: 与 torch compile 的兼容性

感谢 @ani300 对此贡献。

性能

我们展示了使用 FlashAttention 相对于 PyTorch 标准注意力所获得的预期加速（前向 + 反向传播总和）和内存节省，这取决于序列长度，并在不同的 GPU 上进行了测试（加速效果取决于内存带宽 - 在较慢的 GPU 内存上我们看到更多的加速）。

我们目前拥有以下 GPU 的基准测试：

• A100
• H100

A100

我们使用以下参数展示 FlashAttention 的加速效果：

• 头维度 64 或 128，隐藏维度 2048（即 32 或 16 个头）。
• 序列长度 512、1k、2k、4k、8k、16k。
• 批大小设置为 16k / seqlen。

加速

内存

我们在该图中展示了内存节省情况（请注意，无论使用 dropout 还是 masking，内存占用都是相同的）。 内存节省与序列长度成正比——因为标准注意力的内存与序列长度呈二次方关系，而 FlashAttention 的内存与序列长度呈线性关系。 在序列长度为 2K 时，我们观察到 10 倍的内存节省，在 4K 时为 20 倍。 因此，FlashAttention 可以扩展到更长的序列长度。

H100

完整模型代码与训练脚本

我们已发布完整的 GPT 模型 实现。 我们还提供了其他层的优化实现（例如，MLP、LayerNorm、 交叉熵损失、旋转位置编码）。总体而言，与 Huggingface 的基线实现相比， 这将训练速度提升了 3-5 倍， 在每块 A100 上达到 225 TFLOPs/sec，相当于 72% 的模型 FLOPs 利用率（我们不需要 任何激活检查点）。

我们还包含一个训练 脚本，用于 在 Openwebtext 上训练 GPT2，在 The Pile 上训练 GPT3。

FlashAttention 的 Triton 实现

Phil Tillet (OpenAI) 有一个 FlashAttention 的 Triton 实验性实现： https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py

由于 Triton 是比 CUDA 更高级的语言，它可能更容易 理解和实验。Triton 实现中的符号也 更接近我们论文中使用的符号。

我们还有一个支持注意力 偏置（例如 ALiBi）的 Triton 实验性实现： https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py

测试

我们测试了 FlashAttention 产生的输出和梯度与参考 实现相同，在一定的数值容差范围内。具体而言，我们检查 FlashAttention 的最大数值误差至多为 Pytorch 中基线实现 数值误差的两倍（针对不同的头维度、输入 dtype、序列长度、因果 / 非因果）。

要运行测试：

pytest -q -s tests/test_flash_attn.py

当你遇到问题时

FlashAttention-2 的这个新版本已在多个 GPT 风格的 模型上进行了测试，主要在 A100 GPU 上。

如果你遇到 bug，请打开一个 GitHub Issue！

测试

要运行测试：

pytest tests/test_flash_attn_ck.py

引用

如果你使用了此代码库，或以其他方式认为我们的工作有价值，请引用：

@inproceedings{dao2022flashattention,
  title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
  author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
  booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
  year={2022}
}
@inproceedings{dao2023flashattention2,
  title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
  author={Dao, Tri},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2024}
}