如果你认真研究 LLM 推理，那就去学 FlashAttention。它是现代 Transformer 性能背后最重要的优化之一。学习内容：IO 感知的注意力机制、分块与内核融合、SRAM 与 HBM 内存层级、在线 Softmax、KV 缓存优化、Warp 级并行、CUDA 内核、张量核心、内存带宽优化、长上下文推理、FP16、BF16 与 FP8 优化。GitHub: https://github.com/Dao-AILab/flash-attention… FlashAttention 论文: https://arxiv.org/abs/2205.14135 FlashAttention-2: https://tridao.me/publications/flash2/flash2.pdf… FlashAttention-3: https://tridao.me/blog/2024/flash3/… FlashAttention 通过避免大型中间注意力矩阵来减少 GPU 内存流量，实现 2–4 倍更快的注意力机制，并在不改变 Transformer 架构的情况下支持更长的上下文窗口。如果你想成为 AI 基础设施工程师、CUDA 工程师或 LLM 推理工程师，FlashAttention 是必读内容。

Dao-AILab/flash-attention

来源: https://github.com/Dao-AILab/flash-attention

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é 论文: https://arxiv.org/abs/2205.14135 IEEE Spectrum 文章 (https://spectrum.ieee.org/mlperf-rankings-2022) 关于我们使用 FlashAttention 提交至 MLPerf 2.0 基准测试。 FlashAttention FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning Tri Dao 论文: https://tridao.me/publications/flash2/flash2.pdf FlashAttention-2

使用方法

我们非常高兴地看到 FlashAttention 在发布后短时间内被广泛采用。此页面 (https://github.com/Dao-AILab/flash-attention/blob/main/usage.md) 列出了 FlashAttention 被使用的部分场所。 FlashAttention 和 FlashAttention-2 可自由使用和修改（见 LICENSE）。如果使用 FlashAttention，请引用并致谢。

FlashAttention-3 beta 版本

FlashAttention-3 针对 Hopper GPU（如 H100）进行了优化。 博客文章: https://tridao.me/blog/2024/flash3/ 论文: https://tridao.me/publications/flash3/flash3.pdf FlashAttention-3 在 H100 80GB SXM5 上的 FP16 加速 这是一个 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（我们看到了少量正面反馈 (https://github.com/Dao-AILab/flash-attention/issues/595)），但 Windows 编译仍需更多测试。如果你有关于为 Windows 设置预构建 CUDA wheels 的想法，请通过 GitHub issue 联系我们。 * 请确保 ninja 已安装且正常工作（例如 ninja --version 然后 echo $? 应返回退出码 0）。如果不正常（有时 ninja --version 然后 echo $? 返回非零退出码），卸载并重新安装 ninja (pip uninstall -y ninja && pip install ninja)。没有 ninja，编译可能需要很长时间（2 小时），因为它不会使用多个 CPU 核心。使用 ninja 和 CUDA 工具包，在 64 核机器上编译只需 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 (https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) 容器，它包含了安装 FlashAttention 所需的所有工具。 目前支持 CUDA 的 FlashAttention-2：

1. Ampere、Ada 或 Hopper GPU（如 A100、RTX 3090、RTX 4090、H100）。对于 Turing GPU（T4、RTX 2080），请参考单独的 flash-attention-turing (https://github.com/ssiu/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 (https://github.com/ROCm/composable_kernel) (ck)，另一个是 Triton (https://github.com/triton-lang/triton) 后端。它们提供了 FlashAttention-2 的实现。 要求：

• ROCm 6.0 及以上版本。我们推荐使用 ROCm 的 Pytorch (https://hub.docker.com/r/rocm/pytorch) 容器，它包含了安装 FlashAttention 所需的所有工具。

Composable Kernel 后端

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

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

Triton 后端

Flash Attention (https://tridao.me/publications/flash2/flash2.pdf) 的 Triton 实现支持 AMD 的 CDNA（MI200、MI300）和 RDNA GPU，使用 fp16、bf16 和 fp32 数据类型。它提供前向和反向传播，支持因果掩码、可变序列长度、任意 Q/KV 序列长度和头大小、MQA/GQA、dropout、旋转位置编码、ALiBi、分页注意力和 FP8（通过 Flash Attention v3 接口）。滑动窗口注意力目前正在开发中。 Triton 后端内核由 aiter (https://github.com/ROCm/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> && 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
# 使用 triton 后端构建 flash attention
RUN git clone https://github.com/Dao-AILab/flash-attention && \
    cd flash-attention && \
    FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" pip install --no-build-isolation .
# 设置工作目录
WORKDIR /workspace/flash-attention
# 设置环境变量以使用 triton 后端
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 应设为 0.0
    如果 Q、K、V 已经堆叠成 1 个张量，此函数将比在 Q、K、V 上调用 flash_attn_func 更快，因为反向传播避免了显式拼接 Q、K、V 的梯度。
    如果 window_size != (-1, -1)，实现滑动窗口局部注意力。位置 i 的查询只会关注 [i - window_size[0], i + window_size[1]] 范围内的键（包括边界）。
    参数：
        qkv: (batch_size, seqlen, 3, nheads, headdim)
        dropout_p: float。Dropout 概率。
        softmax_scale: float。应用 softmax 前 QK^T 的缩放。默认为 1 / sqrt(headdim)。
        causal: bool。是否应用因果注意力掩码（例如，自回归建模）。
        window_size: (left, right)。如果不是 (-1, -1)，实现滑动窗口局部注意力。
        alibi_slopes: (nheads,) 或 (batch_size, nheads), fp32。向查询 i 和键 j 的注意力分数添加 (-alibi_slope * |i - j|) 偏置。
        deterministic: bool。是否使用确定性的反向传播实现，这会稍微慢一些且使用更多内存。前向传播始终是确定性的。
    返回：
        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 应设为 0.0
    通过传入比 Q 更少头的 KV 来支持多查询和分组查询注意力 (MQA/GQA)。注意，Q 中的头数必须能被 KV 中的头数整除。
    例如，如果 Q 有 6 个头，K、V 有 2 个头，则 Q 的头 0、1、2 将关注 K、V 的头 0，Q 的头 3、4、5 将关注 K、V 的头 1。
    如果 window_size != (-1, -1)，实现滑动窗口局部注意力。位置 i 的查询只会关注 [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] 范围内的键（包括边界）。
    参数：
        q: (batch_size, seqlen, nheads, headdim)
        k: (batch_size, seqlen, nheads_k, headdim)
        v: (batch_size, seqlen, nheads_k, headdim)
        dropout_p: float。Dropout 概率。
        softmax_scale: float。应用 softmax 前 QK^T 的缩放。默认为 1 / sqrt(headdim)。
        causal: bool。是否应用因果注意力掩码（例如，自回归建模）。
        window_size: (left, right)。如果不是 (-1, -1)，实现滑动窗口局部注意力。
        alibi_slopes: (nheads,) 或 (batch_size, nheads), fp32。向查询 i 和键 j 的注意力分数添加 (-alibi_slope * |i + seqlen_k - seqlen_q - j|) 偏置。
        deterministic: bool。是否使用确定性的反向传播实现，这会稍微慢一些且使用更多内存。前向传播始终是确定性的。
    返回：
        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 表示无限上下文窗口
    rotary_interleaved=True,
    alibi_slopes=None,
):
    """如果 k 和 v 不为 None，k_cache 和 v_cache 将会被原地更新，使用 k 和 v 中的新值。
    这对于增量解码很有用：你可以传入上一步缓存的键/值，并用当前步骤的新键/值更新它们，然后对更新后的缓存执行注意力，所有这些都在 1 个内核中完成。
    如果你传入 k / v，必须确保缓存足够大以容纳新值。例如，KV 缓存可以用最大序列长度预分配，你可以使用 cache_seqlens 跟踪批次中每个序列的当前序列长度。
    如果传入 rotary_cos 和 rotary_sin，还会应用旋转位置编码。
    键 @k 将在索引 cache_seqlens, cache_seqlens + 1 等处被 rotary_cos 和 rotary_sin 旋转。
    如果 causal 或 local（即 window_size != (-1, -1)），查询 @q 将在索引 cache_seqlens, cache_seqlens + 1 等处被旋转。
    如果不是 causal 且不是 local，查询 @q 将仅在索引 cache_seqlens 处被旋转（即我们认为 @q 中的所有 token 都位于 cache_seqlens 位置）。
    参见 tests/test_flash_attn.py::test_flash_attn_kvcache 了解如何使用该函数的示例。
    通过传入比 Q 更少头的 KV 来支持多查询和分组查询注意力 (MQA/GQA)。注意，Q 中的头数必须能被 KV 中的头数整除。
    例如，如果 Q 有 6 个头，K、V 有 2 个头，则 Q 的头 0、1、2 将关注 K、V 的头 0，Q 的头 3、4、5 将关注 K、V 的头 1。
    如果 causal=True，因果掩码与注意力矩阵的右下角对齐。例如，如果 seqlen_q = 2 且 seqlen_k = 5，因果掩码（1 = 保留，0 = 掩盖）为：
        1 1 1 1 0
        1 1 1 1 1
    如果 seqlen_q = 5 且 seqlen_k = 2，因果掩码为：
        0 0
        0 0
        0 0
        1 0
        1 1
    如果掩码的行全为零，输出将为零。
    如果 window_size != (-1, -1)，实现滑动窗口局部注意力。位置 i 的查询只会关注 [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] 范围内的键（包括边界）。
    注意：不支持反向传播。
    参数：
        q: (batch_size, seqlen, nheads, headdim)
        k_cache: 如果没有 block_table，则为 (batch_size_cache, seqlen_cache, nheads_k, headdim)；如果有 block_table（即分页 KV 缓存），则为 (num_blocks, page_block_size, nheads_k, headdim)
                 page_block_size 必须是 256 的倍数。
        v_cache: 如果没有 block_table，则为 (batch_size_cache, seqlen_cache, nheads_k, headdim)；如果有 block_table（即分页 KV 缓存），则为 (num_blocks, page_block_size, nheads_k, headdim)
        k [可选]: (batch_size, seqlen_new, nheads_k, headdim)。如果提供，我们将 k 与 k_cache 拼接，起始索引由 cache_seqlens 指定。
        v [可选]: (batch_size, seqlen_new, nheads_k, headdim)。类似于 k。
        rotary_cos [可选]: (seqlen_ro, rotary_dim / 2)。如果提供，我们对 k 和 q 应用旋转位置编码。仅当传入 k 和 v 时适用。rotary_dim 必须能被 16 整除。
        rotary_sin [可选]: (seqlen_ro, rotary_dim / 2)。类似于 rotary_cos。
        cache_seqlens: int，或 (batch_size,)，dtype torch.int32。KV 缓存的序列长度。