Dao-AILab/flash-attention

一、项目定位

FlashAttention 是 Dao-AILab 维护的精确注意力（exact attention）加速库，核心卖点是 IO-aware 的 tiling 实现：在 HBM 与 SRAM 之间按 tile 搬运与计算，把标准 Attention 的 HBM 读写压到理论下界，从而在不牺牲数值精度的前提下显著提速并省显存。仓库目前同时承载多代产品：FA2 在顶层 csrc/flash_attn（CUDA，Ampere/Ada/Hopper），FA3 在 hopper/（Hopper-only，FP16/BF16/FP8），而活跃开发已切到 FA4——用 CuTeDSL（NVIDIA CUTLASS DSL）写成、在运行时编译为 PTX/CUBIN，面向 Hopper (SM90) 与 Blackwell (SM100/SM110)。AGENTS.md 与 CLAUDE.md 明确写明「active development is on FA4 in flash_attn/cute/」，commit 热点也集中在 flash_attn/cute（近 43 次提交）、tests/cute（15 次）与 hopper/。FA4 以 flash-attn-4 包发布，最新 release 为 fa4-v4.0.0.beta33（2026-09-30），仍处于 beta 阶段。README 对 FA3 标注为「beta release for testing / benchmarking」，FA4 则仍在快速迭代（SM100 支持、MLA、FP8、varlen 等特性在最近提交中频繁出现）。整体定位：面向大模型训练/推理的底层 attention 原语层，被 PyTorch生态（HuggingFace、xFormers、SDXL/DiT 训练等）广泛依赖。

解决什么问题、给谁用

标准 scaled dot-product attention 在长序列下的瓶颈是 QK^T 与 softmax 中间矩阵必须完整写入 HBM，导致 (seq_len² · d) 的显存占用与带宽受限的读写。FlashAttention 通过 tiling + 在线 softmax（recompute P 而非保存完整注意力矩阵）把显存降到 O(N·d) 并减少 HBM 访问。FA4 进一步把实现迁移到 CuTeDSL，以便直接利用 Hopper 的 TMA、WGMMA、named barrier 等硬件原语，并在同一套 Python DSL 上同时覆盖 SM90 与 SM100/SM110。目标用户：(1) 大模型训练框架开发者（需要更快、更省显存的 attention 后端，例如 GPT、LLaMA、DiT 训练）；(2) 推理/服务侧工程师（decode、varlen、paged-KV、MLA 等场景）；(3) 硬件/编译器方向的研究者（CuTeDSL 内核模板、kernel 生成、性能调优）。典型场景：替换 torch.nn.functional.scaled_dot_product_attention 后端；为扩散模型视频生成提供长上下文注意力；在 H100/B200 上做 FP8/MLA decode；在 AMD MI300 上通过 CK 或 Triton 后端跑 FA2。

1) 大模型训练框架与研究人员——用 flash_attn_func / flash_attn_qkvpacked_func 替换原生 attention，降低显存、缩短 step time；2) 推理/服务工程师——使用 FA3/FA4 的 varlen、paged-KV、MLA、FP8 接口做低延迟 decode；3) 硬件性能工程师——在 flash_attn/cute/ 与 hopper/ 中研究 CuTeDSL 内核、TMA/WGMMA 利用、SM90/SM100 性能调优；4) AMD ROCm 用户——通过 csrc/flash_attn_ck 或 third_party/aiter Triton 后端在 MI200/MI300/RDNA 上获得 FA2 等价实现。

同类项目与差别

• xFormers (Meta)：提供 memory-efficient attention 与可组合算子块，含 CUTLASS / Triton 实现，偏上层算子库与 Torch 集成；FlashAttention 更聚焦于单一 attention 原语的极致性能与多代硬件覆盖，FA4 直接下沉到 CuTeDSL 生成 PTX/CUBIN。
• SDPA (PyTorch scaled_dot_product_attention)：PyTorch 原生后端，调用 FlashAttention / Memory-Efficient Attention / Math 实现；FlashAttention 是 SDPA 所调用的底层后端之一，本仓库提供 SDPA 尚未原生覆盖的 FP8、MLA、paged-KV、varlen 等特性。
• CUTLASS / CuTeDSL (NVIDIA)：是 FA4 的底层 DSL 与模板库；本仓库在其之上构建 attention 专用内核，并暴露 Python 接口 flash_attn.cute.flash_attn_func。
• aiter (ROCm)：为 AMD GPU 提供 Triton 版 FlashAttention 后端，已作为 third_party/aiter 子模块被本仓库集成进 ROCm Triton 路径。
• FlashInfer：同样聚焦 decode / ragged / paged-KV / MLA，但更偏推理运行时与 kernel 库；FlashAttention 覆盖更广的训练+推理场景，且是 FA4 在 CuTeDSL 上同时服务 SM90/SM100。

核心能力

阶段：成熟（FA2）+ beta（FA3）+ 快速迭代 beta（FA4）。FA2 是业界事实标准；FA3 README 自标 beta；FA4 以 flash-attn-4 包发布，最新 tag 为 fa4-v4.0.0.beta33，commit 集中在 SM100/MLA/varlen 特性，仍属 beta。

技术栈：Python；C++；CUDA；CuTeDSL (nvidia-cutlass-dsl>=4.6.2)；Composable Kernel (ROCm)；Triton (via aiter)；PyTorch (>=2.2)；ninja；apache-tvm-ffi；quack-kernels>=0.5.3；CMake / setuptools (setup.py)；GitHub Actions CI (.github/workflows)；pre-commit

规模：仓库目录含 10 个顶层目录；csrc/ 197 项（含 composable_kernel、cutlass、flash_attn、flash_attn_ck、fused_dense_lib、layer_norm）；hopper/ 507 项；flash_attn/ 110 项；tests/ 53 项；training/ 145 项；benchmarks/ 12 项。最近提交集中在 flash_attn/cute（43 次热点）、tests/cute（15 次）、csrc/flash_attn_ck（4 次）。FA4 release 已发到 fa4-v4.0.0.beta33（2026-09-30），保持约每周一次 beta 节奏（beta31→beta32→beta33 各隔一周）。star 数未提供，需验证。

二、架构与代码地图

仓库按「代际 × 硬件目标」自然分成五层。 1) 接入层（Python API）：flash_attn/__init__.py、flash_attn/cute/__init__.py、flash_attn/cute/interface.py、hopper/flash_attn_interface.py、hopper/flash_attn_3/。职责是把 flash_attn_func / flash_attn_varlen_func 这些用户可见 API 翻译成内部 kernel 调用，处理 dtype、causal、softcap、varlen、paged KV、MLA 等开关。 2) 调度层（tile / split / batch 调度）：flash_attn/cute/prepare_scheduler.py、flash_attn/cute/tile_scheduler.py、flash_attn/cute/seqlen_info.py、flash_attn/cute/block_info.py、flash_attn/cute/cu_blocks_kernel.py，以及 hopper/flash_prepare_scheduler.cu、hopper/heuristics.h。职责是把 (batch, heads, seqlen) 切成 CTA tile、决定 split-KV 数、生成 cu_seqlens / cu_blocks / SchedulerMetadataTensorsTorch，并为 SM100 的 CLC / dynamic-persistent 调度器喂状态。 3) 执行层（CuTeDSL kernel 主体）：flash_attn/cute/flash_fwd*.py、flash_attn/cute/flash_bwd*.py、flash_attn/cute/flash_fwd_mla*.py、flash_attn/cute/flash_bwd_mla*.py、flash_attn/cute/sm100_hd256_2cta_fmha_backward_*.py。这是 FA4 的核心——用 CuTeDSL 在 Python 里写 TMA + WGMMA + named-barrier 的 forward/backward/MLA kernel。 4) 硬件抽象层（pipeline / barrier / copy / mma 描述）：flash_attn/cute/pipeline.py、flash_attn/cute/named_barrier.py、flash_attn/cute/copy_utils.py、flash_attn/cute/mma_sm100_desc.py、flash_attn/cute/barrier.py、flash_attn/cute/ampere_helpers.py、flash_attn/cute/blackwell_helpers.py。封装 Hopper/Blackwell 的 async pipeline、TMA、bulk-copy、UMMA/TMEM 描述符，让 kernel 层只描述「怎么算」。 5) 工具层（build / test / bench / debug）：setup.py、flash_attn/cute/cute_dsl_ptxas.py、flash_attn/cute/cute_dsl_utils.py、tests/cute/、benchmarks/、AI/*.md、tools/。负责编译到 PTX/CUBIN、持久化 kernel cache、two-pass 测试、SASS 级调试。 FA2（csrc/flash_attn/）与 FA3（hopper/）是上一代 C++/CUDA 实现，仍被 ROCm/CK 与老卡维护，但提交热点已切到 FA4 的 flash_attn/cute。

一次调用怎么流过这些模块

一次典型的 FA4 推理/训练前向调用链如下： 1) 用户调用 flash_attn.cute.interface.flash_attn_func(q, k, v, causal=..., softmax_scale=...)。interface 先做 dtype/shape 对齐（BF16/FP16，headdim 对齐到 8），判断走 standard、varlen、paged 还是 MLA 路径，并读取 env（FLASH_ATTENTION_MLA_1CTA、FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED 等）。 2) interface 调用 flash_attn.cute.prepare_scheduler.FlashPrepareScheduler，在 host 侧把 (batch, seqlen_q, seqlen_k, nheads, nheads_kv, headdim) 切成 M/N tile，计算 num_m_blocks、num_n_blocks、split-KV 数（num_splits_dynamic），产出 SchedulerMetadataTensorsTorch（含 cu_total_m_blocks、cu_total_splits_m_blocks、blocks_to_batch_idx、tile_count_semaphore）。对 varlen 还会先启动 CuSeqlensToBlocksKernel（cu_blocks_kernel.py）把 cu_seqlens 转成 block 偏移。 3) 根据目标架构（SM90/SM100/SM120）与 headdim、causal、MLA 标志，interface 选择具体 kernel：standard 走 flash_fwd.py 的 flash_fwd_sm{90,100,120}；MLA 走 flash_fwd_mla_sm100.py；block-sparse 走 block_sparsity.py + topk_gather_kv.py。 4) kernel 启动前先由 flash_attn.cute.cute_dsl_ptxas 把 @cute.jit 函数编译成 PTX/CUBIN（若命中磁盘缓存 /tmp/${USER}/flash_attention_cute_dsl_cache/ 则跳过编译）。编译使用 nvidia-cutlass-dsl + apache-tvm-ffi，支持 FakeTensor 模式（不分配 GPU 显存）。 5) kernel 执行时，tile_scheduler.py 构造 SchedulerState（create_clc 或 create_dynamic_persistent），把硬件 CLC 或软件 dynamic-persistent 调度器初始化；pipeline.py 的 PipelineStateSimple / PipelineTmaAsync / PipelineUmmaAsync 管理 producer-consumer 环形缓冲。kernel 主体用 copy_utils.py 的 TMA bulk-copy 把 Q/K/V tile 从 HBM 搬到 SMEM，用 mma_sm100_desc.py 构造 WGMMA/UMMA 描述符做 QK^T 与 PV 矩阵乘，softmax.py 做在线 softmax（recompute P，不保存完整注意力矩阵），mask.py 注入 causal/local/ALiBi mask，named_barrier.py 做 warp-group 同步。 6) 输出 out（与 Q 同形状）与 softmax_lse（log-sum-exp，形状 b×h×seqlen_q，给 backward 用）被写回 HBM。 7) backward 时，flash_bwd.py / flash_bwd_mla_sm100.py 复用 forward 的 scheduler 元数据，用 flash_bwd_preprocess 把 dout 与 out 结合，再启动 dK/dV kernel（BlackwellFusedMultiHeadAttentionBackwardDKDVKernel）与 dQ kernel（BlackwellFusedMultiHeadAttentionBackwardDQKernel），最后 flash_bwd_postprocess 把部分和归约成最终梯度。 关键调度点：interface 的路径选择、prepare_scheduler 的 tile/split 决策、tile_scheduler 的 CLC vs dynamic-persistent 选择、pipeline 的 producer-consumer 同步。关键数据结构：SchedulerMetadataTensorsTorch、SchedulerState / ClcSchedulerState、PipelineStateSimple、cute.Tensor（Q/K/V/out/lse）。

关键类型与函数

扩展点

• 新模型（如新架构的 attention 变体）：在 flash_attn/cute/ 下新增 flash_fwd_<variant>_sm100.py + flash_bwd_<variant>_sm100.py，并在 interface.py 里加派发分支；MLA 变体可参照 flash_fwd_mla_sm100.py 的 1CTA/2CTA/dense/sparse 分支
• 新 kernel（新 tiling 策略 / 新 split-KV 启发式）：修改 flash_attn/cute/prepare_scheduler.py 的 FlashPrepareScheduler 与 tile_scheduler.py 的 Sm100FmhaClcDynamicTileSchedulerParams；新调度模式可在 SchedulingMode 枚举中加项
• 新调度策略（如新的 dynamic-persistent 或 CLC 参数）：在 flash_attn/cute/tile_scheduler.py 扩展 SchedulerState.create_dynamic_persistent / create_clc，并在 kernel 的 __init__ 中消费新 SchedulerState
• 新后端（新 GPU 架构，如未来 SM110+）：在 flash_attn/cute/ 下新增 flash_fwd_sm<cc>.py、flash_bwd_sm<cc>.py，复用 pipeline.py / copy_utils.py / mma_sm100_desc.py；在 interface.py 的架构检测处加分支；setup.py 的 cuda_archs() 默认列表需同步更新（当前 80;90;100;110;120）
• 新硬件原语（新 TMA/UMMA/TMEM 用法）：在 flash_attn/cute/pipeline.py、copy_utils.py、mma_sm100_desc.py、named_barrier.py、barrier.py 中加封装，kernel 层只调用这些封装
• 新稀疏 / 新 mask 策略：在 flash_attn/cute/block_sparsity.py、compute_block_sparsity.py、mask.py 中加新实现，并在 flash_fwd.py / flash_fwd_mla_sm100.py 的 dispatch 处接新分支
• ROCm / CK 后端新特性：在 csrc/flash_attn_ck/ 修改 mha_fwd.cpp / mha_bwd.cpp / flash_common.cpp，README 已给出 FLASH_ATTENTION_TRITON_AMD_ENABLE 等 env 开关
• 调试 / SASS 分析：在 AI/ 加 .md 文档，用 tools/sass_diff.py、AI/parse_clc_log.py、AI/racecheck_repro_*.py 复现问题

最近在动的地方

• flash_attn/cute/：近 43 次提交，FA4 活跃开发中心：SM100 hd256 fwd/bwd、MLA 1CTA/2CTA、ping-pong、deterministic、varlen 支持
• tests/cute/：15 次提交，two-pass 测试、FakeTensor 编译缓存、新 kernel 覆盖
• flash_attn/cute/flash_fwd_sm100.py：最近提交：hd256 fwd unify、S/P ping-pong 重叠 QK 与 softmax、length 从 BlockInfo 推导
• flash_attn/cute/sm100_hd256_2cta_fmha_backward_*.py：最近提交：deterministic=True、stack 回收、seqused_q/k 支持
• flash_attn/cute/flash_fwd_mla_sm100.py：最近提交：1CTA MLA forward（dense + sparse top-k）、MLA dispatch heuristic、sparse MLA backward at any head count
• flash_attn/cute/tile_scheduler.py：CLC / dynamic-persistent 调度核心，SM100 性能关键
• csrc/flash_attn_ck/：4 次提交：修复 num_splits heuristic（原本永远返回 1）
• AI/SPARSE_MLA_*.md：3 份文档，记录 sparse MLA 的 1CTA/64H/recompute-P 设计决策

建议阅读顺序

1. README.md + AGENTS.md + CLAUDE.md（定位仓库：FA4 在 flash_attn/cute/，活跃开发；FA2 在 csrc/，FA3 在 hopper/）
2. flash_attn/cute/__init__.py + flash_attn/cute/interface.py（用户入口：flash_attn_func / flash_attn_varlen_func）
3. flash_attn/cute/prepare_scheduler.py + flash_attn/cute/seqlen_info.py + flash_attn/cute/block_info.py（host 侧 tile/split 决策）
4. flash_attn/cute/tile_scheduler.py（CLC / dynamic-persistent 调度器）
5. flash_attn/cute/pipeline.py + flash_attn/cute/named_barrier.py + flash_attn/cute/copy_utils.py + flash_attn/cute/mma_sm100_desc.py（硬件抽象层）
6. flash_attn/cute/flash_fwd.py + flash_attn/cute/flash_fwd_sm100.py（SM100 standard forward kernel）
7. flash_attn/cute/flash_bwd.py + flash_attn/cute/flash_bwd_sm100.py（SM100 standard backward kernel）
8. flash_attn/cute/flash_fwd_mla_sm100.py + flash_attn/cute/flash_bwd_mla_sm100.py（MLA 变体）
9. flash_attn/cute/sm100_hd256_2cta_fmha_backward_*.py（SM100 hd256 2CTA backward，最近热点）
10. tests/cute/test_flash_attn.py + flash_attn/cute/testing.py + tools/two_pass_tests.sh（测试基础设施与 two-pass 工作流）

三、本地跑起来（没有 GPU 的 Mac）

安装

1. # 1. 创建独立 Python 环境（推荐 Python 3.10–3.12，与 PyTorch 2.2+ 兼容）
2. conda create -n fa4 python=3.11 -y && conda activate fa4
4. # 2. 安装 PyTorch（Apple Silicon 用 MPS 后端，CPU 用默认）
5. pip install torch torchvision torchaudio
7. # 3. 克隆仓库（含子模块 third_party/aiter）
8. git clone --recurse-submodules https://github.com/Dao-AILab/flash-attention.git
9. cd flash-attention
11. # 4. 安装 FA4 的 Python 依赖（仅 CPU/MPS，不编译 CUDA kernel）
12. # 注意：nvidia-cutlass-dsl 依赖 CUDA 运行时，在 Mac 上无法真正编译 kernel，
13. # 但可以安装包以导入 Python 模块、阅读 CuTeDSL 代码。
14. pip install -e "flash_attn/cute[dev]" 2>&1 | tee install.log
16. # 5. 若上一步因 CUDA 依赖失败，则仅安装可运行部分：
17. pip install einops apache-tvm-ffi quack-kernels # 跳过 nvidia-cutlass-dsl
18. pip install -e flash_attn/cute # 无 extras，可能仍会尝试编译，需验证
20. # 6. 验证 Python 导入（不触发 kernel 编译）
21. python -c "from flash_attn.cute import flash_attn_func; print('import ok')"
23. # 7. 安装测试/开发工具
24. pip install pytest pytest-xdist pre-commit ruff
26. # 关键环境变量（Mac 上无 GPU 时设置）
27. export FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE # 跳过 CUDA 编译，仅保留 Python 层
28. export FLASH_ATTENTION_FORCE_BUILD=FALSE # 不强制本地编译

哪些路径能真跑

['【可在 CPU/MPS 上执行的代码路径】', '- flash_attn/cute/__init__.py、flash_attn/cute/interface.py：Python 调度层，可 import 并调用高层 API（但底层 kernel 会失败）。', '- flash_attn/cute/prepare_scheduler.py、flash_attn/cute/tile_scheduler.py、flash_attn/cute/seqlen_info.py、flash_attn/cute/block_info.py：调度逻辑，可在 CPU 上跑单元测试（部分测试用 FakeTensorMode）。', '- flash_attn/cute/pipeline.py、flash_attn/cute/named_barrier.py、flash_attn/cute/mask.py、flash_attn/cute/softmax.py：抽象描述，可在无 GPU 环境下 import 和调试。', '- flash_attn/flash_attn_triton.py、flash_attn/flash_attn_triton_og.py：Triton 实现，理论上可在 AMD GPU 或 CPU 上跑，但 Mac 上无 Triton 后端支持，仅能读代码。', '- tests/cute/test_*.py 中标记 @maybe_fake_tensor_mode(USE_FAKE_TENSOR) 的测试：用 FakeTensorMode 编译 kernel，不分配 GPU 内存（见 AGENTS.md 两阶段测试说明）。', '- tests/test_util.py、tests/test_rotary.py：纯 CPU 测试，无 GPU 依赖。', '', '【只能在 GPU 上执行的代码路径】', '- flash_attn/cute/flash_fwd*.py、flash_attn/cute/flash_bwd*.py、flash_attn/cute/flash_fwd_mla*.py、flash_attn/cute/flash_bwd_mla*.py：CuTeDSL kernel 主体，编译到 PTX/CUBIN，必须 NVIDIA GPU。', '- flash_attn/cute/sm100_hd256_2cta_fmha_backward_*.py、flash_attn/cute/sm100_*、flash_attn/cute/sm120_*：SM100/SM120 专用 kernel，需 Blackwell GPU。', '- csrc/flash_attn/src/*.cu：FA2 CUDA kernel，需 Ampere/Ada/Hopper GPU。', '- csrc/flash_attn_ck/*.cpp：ROCm Composable Kernel 实现，需 AMD GPU。', '- hopper/ 下所有 .cu、.hpp、.h：FA3 Hopper-only 实现，需 H100/H800。', '- benchmarks/*.py、hopper/benchmark_*.py：基准测试，需真实 GPU。', '', '【Mac 上的实际策略】', '- 本地用 FakeTensorMode 跑编译阶段测试（验证 kernel 能否编译通过），用 Colab 免费 T4 做最终执行验证。', '- 阅读代码时重点关注 flash_attn/cute/ 下的调度与 kernel 逻辑，用 print 或 logging 调试 Python 层。']

最小可运行

1. # 冒烟测试 1：验证 Python 导入与高层 API（CPU/MPS，不触发 kernel 编译）
2. python -c "from flash_attn.cute import flash_attn_func; from flash_attn.cute.interface import flash_attn_varlen_func; print('API import OK')"
4. # 冒烟测试 2：用 FakeTensorMode 编译 FA4 forward kernel（无 GPU 内存分配，仅验证编译）
5. FLASH_ATTENTION_FAKE_TENSOR=1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest tests/cute/test_flash_attn.py -k "test_flash_attn_output" -x --no-header 2>&1 | head -50
7. # 冒烟测试 3：运行纯 CPU 测试（test_util、test_rotary）
8. pytest tests/test_util.py tests/test_rotary.py -x -q
10. # 冒烟测试 4：在 Colab T4 上执行真实 forward/backward（需 GPU）
11. # 先安装：pip install flash-attn-4
12. # 然后运行：
13. python -c "
14. import torch
15. from flash_attn.cute import flash_attn_func
16. q = torch.randn(2, 1024, 8, 16, device='cuda', dtype=torch.float16)
17. k = torch.randn(2, 1024, 8, 16, device='cuda', dtype=torch.float16)
18. v = torch.randn(2, 1024, 8, 16, device='cuda', dtype=torch.float16)
19. out = flash_attn_func(q, k, v, causal=True)
20. print('FA4 forward OK, shape:', out.shape)
21. "
23. # 冒烟测试 5：运行完整 FA4 测试套件（需 GPU，Colab T4 约 30–60 分钟）
24. FLASH_ATTENTION_FAKE_TENSOR=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest tests/cute/test_flash_attn.py -x -q
26. # 冒烟测试 6：验证调度器与 block_info 逻辑（CPU）
27. python -c "
28. from flash_attn.cute.prepare_scheduler import FlashPrepareScheduler
29. from flash_attn.cute.block_info import BlockInfo
30. scheduler = FlashPrepareScheduler(num_warps=4, tile_m=128, tile_n=128, nheads=8, nheads_kv=8, headdim=128)
31. print('Scheduler init OK')
32. "

测试

['测试框架：pytest + pytest-xdist（并行执行），部分测试用 @maybe_fake_tensor_mode 装饰器支持 FakeTensorMode。', '', '测试目录：', '- tests/cute/test_flash_attn.py：FA4 forward/backward 主测试（smoke_test 2/5）。', '- tests/cute/test_flash_attn_mla.py：MLA（Multi-head Latent Attention）测试，覆盖 dense/sparse/paged KV。', '- tests/cute/test_flash_attn_varlen.py：变长序列（varlen）测试。', '- tests/cute/test_mask_mod.py、tests/cute/test_score_mod.py：mask 与 score 修饰符测试。', '- tests/cute/test_block_sparsity.py：块稀疏注意力测试。', '- tests/test_flash_attn.py：FA2 测试（需 GPU）。', '- tests/test_flash_attn_ck.py：ROCm CK 后端测试（需 AMD GPU）。', '- tests/test_flash_attn_triton_amd.py：Triton AMD 后端测试（需 AMD GPU）。', '- tests/test_util.py、tests/test_rotary.py：纯 CPU 工具测试。', '', '只跑 CPU 子集（Mac 无 GPU）：', 'pytest tests/test_util.py tests/test_rotary.py -x -q # 纯 CPU，秒级完成', 'pytest tests/cute/test_flash_attn.py -k "test_block_info or test_seqlen_info or test_scheduler" -x -q # 调度逻辑测试', '', 'FakeTensorMode 编译阶段测试（无 GPU 内存分配，但会打开 CUDA 上下文）：', 'FLASH_ATTENTION_FAKE_TENSOR=1 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest -n 4 tests/cute/test_flash_attn.py -x # 编译所有 kernel，约 5–15 分钟', '', '完整测试套件（需 GPU，Colab T4）：', 'FLASH_ATTENTION_FAKE_TENSOR=0 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 pytest tests/cute/ -x -q # 30–60 分钟', '', '耗时估算：', '- CPU 子集（test_util + test_rotary）：< 10 秒。', '- FakeTensorMode 编译 pass（test_flash_attn.py）：5–15 分钟（取决于 CPU 核心数）。', '- GPU 执行 pass（test_flash_attn.py）：20–40 分钟（T4）。', '- 完整 tests/cute/ 套件：30–60 分钟（T4）。']

调试

• 【断点与日志】
• - Python 层断点：在 flash_attn/cute/interface.py、flash_attn/cute/prepare_scheduler.py、flash_attn/cute/tile_scheduler.py 中用 import pdb; pdb.set_trace() 或 IDE 调试。
• - CuTeDSL kernel 内调试：用 cute.printf 或 print 在 kernel 内输出（需 GPU 执行）。
• - 日志开关：flash_attn/cute/fa_logging.py 提供日志工具，设置 logging.basicConfig(level=logging.DEBUG) 查看调度信息。
• 【环境变量】
• - FLASH_ATTENTION_FAKE_TENSOR=1：启用 FakeTensorMode，编译 kernel 但不执行（AGENTS.md 明确说明）。
• - FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1：启用持久化磁盘缓存，路径 /tmp/${USER}/flash_attention_cute_dsl_cache/。
• - FLASH_ATTENTION_TEST_EXPECT_CACHED=1：执行阶段测试时，若 kernel 未命中缓存则失败（用于 CI 验证缓存完整性）。
• - FLASH_ATTENTION_TEST_COUNT_COMPILES=1：仅统计编译次数，不执行测试。
• - FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE：跳过 CUDA 编译，仅保留 Python 层（setup.py 中定义）。
• - FLASH_ATTENTION_FORCE_BUILD=TRUE：强制本地编译，忽略预编译 wheel。
• - CUDA_VISIBLE_DEVICES=<gpu_id>：限制可见 GPU，用于多 GPU 环境隔离。
• 【Profiling 入口】
• - flash_attn/cute/benchmark.py、flash_attn/cute/benchmark_flash_attention_fp8.py：FA4 基准测试脚本。
• - benchmarks/bench_sm90.py、benchmarks/benchmark_flash_attention.py：FA2/FA3 基准测试。
• - hopper/benchmark_attn.py、hopper/benchmark_flash_attention_fp8.py：FA3 基准测试。
• - 用 nsys profile 或 nvprof 对 Python 脚本做 GPU profiling（需 NVIDIA GPU）。
• 【常见调试场景】
• - 若 import flash_attn.cute 失败，检查 nvidia-cutlass-dsl 是否安装，或降级到仅 Python 层。
• - 若 FakeTensorMode 编译失败，检查 FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED 是否启用，或删除缓存目录重试。
• - 若 kernel 执行结果不对，用 FLASH_ATTENTION_TEST_EXPECT_CACHED=1 确认是否使用了过期缓存。

CI

['CI 框架：GitHub Actions（.github/workflows/），含 bump beta tag 脚本（.github/scripts/bump_beta_tag.py）。', '', 'CI 主要跑：', '- 代码格式检查：ruff format（最近 commit fb97d25 为 ruff format 提交）。', '- 静态检查：pre-commit（.pre-commit-config.yaml 配置）。', '- 编译测试：FLASH_ATTENTION_FAKE_TENSOR=1 编译 pass + FLASH_ATTENTION_FAKE_TENSOR=0 执行 pass（见 tools/two_pass_tests.sh）。', '- 完整测试套件：pytest tests/cute/ 在多个 GPU 上并行（-n 256 workers）。', '', 'PR 会被以下检查卡住：', '- ruff format 格式不符（最近 commit fb97d25 显示仓库已强制 ruff 格式化）。', '- pre-commit 钩子失败（isort、ruff、mypy 等）。', '- FakeTensorMode 编译失败（kernel 无法编译）。', '- GPU 执行阶段测试失败（数值精度、shape 不匹配等）。', '- FLASH_ATTENTION_TEST_EXPECT_CACHED=1 时 kernel 未命中缓存（CI 要求所有 kernel 在编译 pass 中已缓存）。']

坑

• 【坑 1：Mac 上无法真正编译或执行 CUDA kernel】
• Apple Silicon 无 NVIDIA GPU，nvidia-cutlass-dsl 依赖 CUDA 运行时，只能安装 Python 包但无法编译 PTX/CUBIN。本地开发只能读代码、跑 FakeTensorMode 编译验证，最终执行必须用 Colab T4 或远程 GPU。
• 【坑 2：setup.py 默认尝试 CUDA 编译
• 直接 pip install -e . 会触发 CUDA 编译，在 Mac 上会失败。必须设置 FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE 或仅安装 flash_attn/cute 子包。
• 【坑 3：CuTeDSL 缓存导致过期 binary
• FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 启用磁盘缓存后，修改 kernel 代码可能仍命中旧缓存。调试时删除 /tmp/${USER}/flash_attention_cute_dsl_cache/ 或禁用缓存。
• 【坑 4：FakeTensorMode 与真实执行的差异
• FakeTensorMode 只验证编译，不验证数值精度。编译通过的 kernel 在 GPU 上仍可能数值错误（见 AI/ 目录下的 DEBUG_*.md 文档）。
• 【坑 5：环境变量拼写与大小写
• FLASH_ATTENTION_FAKE_TENSOR、FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED 等变量名拼写错误会静默失效，导致测试行为不符合预期。
• 【坑 6：子模块 third_party/aiter 未初始化
• ROCm Triton 后端依赖 third_party/aiter，git clone 时需 --recurse-submodules，否则 FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE 安装会失败。
• 【坑 7：测试装饰器 @maybe_fake_tensor_mode 缺失
• 新增测试函数必须加 @maybe_fake_tensor_mode(USE_FAKE_TENSOR) 装饰器，并在数据依赖检查前用 if is_fake_mode(): return 退出，否则 CI 两阶段测试会失败（AGENTS.md 明确说明）。
• 【坑 8：FA2/FA3/FA4 三代代码共存，易混淆
• 仓库同时包含 FA2（csrc/flash_attn/）、FA3（hopper/）、FA4（flash_attn/cute/），活跃开发在 FA4。修改代码前确认目标代际，避免在 FA2 上修复 FA4 的 bug。
• 【坑 9：Windows 支持仍 experimental
• README 提到 Windows 编译从 v2.3.2 开始可能工作，但需更多测试。Mac 上无官方 Windows 支持，但 Linux 容器或远程开发是可行替代。
• 【坑 10：ninja 编译 RAM 不足
• README 明确警告：若机器 RAM < 96GB 且 CPU 核心多，ninja 可能耗尽内存。Mac 16GB 应设置 MAX_JOBS=2 或 MAX_JOBS=4 限制并行编译任务。

四、维护者与社区

高频迭代。FA4 当前处于 beta 阶段，最近三个 release（fa4-v4.0.0.beta31/32/33）分别在 2026-09-16、2026-09-23、2026-09-30 发布，约每周一个 beta。commit 集中在 flash_attn/cute（近 43 次提交热点）、tests/cute（15 次）与 hopper/，最近提交（2026-10-04）仍在推进 SM100 hd256 backward 与 varlen Q/decode 的修复与统一。FA2/FA3 仍有维护提交（如 split-KV combine 归一化、Windows 构建修复），但开发重心已明确切到 FA4。

流程与 Review 风格

仓库根目录没有独立的 CONTRIBUTING.md，贡献流程主要靠 AGENTS.md / CLAUDE.md 与 setup.py 环境变量驱动。CLAUDE.md 明确：活跃开发在 flash_attn/cute/，安装用 pip install -e "flash_attn/cute[dev]"，测试用 pytest tests/cute/test_flash_attn.py 等；并给出「two-pass testing」流程（先用 FLASH_ATTENTION_FAKE_TENSOR=1 并行编译，再用缓存执行）。pre-commit 配置存在（.pre-commit-config.yaml），且 PR #2940 专门做 ruff format，表明代码风格用 ruff 强制格式化。没有看到 CLA/DCO 文件。Issue 模板与 PR 模板未在根目录显式出现，但 .github/workflows 与 .github/actions 存在，CI 会跑构建与测试。贡献者应先在 GitHub Issue 复现/讨论（如 #2957、#2969 带具体 GPU 与 head_dim 条件），再提 PR。

Review 风格偏「硬核内核工程师」：紧贴具体 GPU 架构、head_dim、varlen、GQA、deterministic 等细节。Issue 标题高度具体（如「[CuTe, SM100] dsink loses precision: dpsum uses the bf16-rounded O」「[CuTe, SM90] Paged KV with page_size != tile_n: wrong output at head_dim 80/96 (non-causal), illegal memory access at 160/224」），要求复现条件精确到 dtype、head_dim、page_size、causal 开关。PR 拆分细致（如 #2964–#2968 把「Unify hd256 backward」拆成 5 个子 PR：tests first → 2CTA schedules → opt-in → GQA/varlen → delete dedicated kernels），说明维护者偏好小步、可审查的增量提交。响应速度：高频提交与 weekly beta release 表明核心维护者几乎每日在线；Issue 指派（Johnsonms、tzadouri、micmelesse）显示会按模块分派。对数值精度、illegal memory access、deterministic 类 Issue 反应最快。

渠道

• GitHub Issues: https://github.com/Dao-AILab/flash-attention/issues
• GitHub Pull Requests: https://github.com/Dao-AILab/flash-attention/pulls
• GitHub repository: https://github.com/Dao-AILab/flash-attention
• README 中 usage 页: https://github.com/Dao-AILab/flash-attention/blob/main/usage.md
• README 中 flash-attention-turing 仓库: https://github.com/ssiu/flash-attention-turing
• README 中 PyTorch NGC 容器: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch
• README 中 ROCm PyTorch 容器: https://hub.docker.com/r/rocm/pytorch
• README 中 PyTorch 本地安装: https://pytorch.org/get-started/locally/
• README 中 aiter 仓库: https://github.com/ROCm/aiter
• README 中 FlashAttention 论文: https://arxiv.org/abs/2205.14135
• README 中 FlashAttention-3 博客: https://tridao.me/blog/2024/flash3/
• README 中 FlashAttention-3 论文: https://tridao.me/publications/flash3/flash3.pdf
• README 中 IEEE Spectrum 文章: https://spectrum.ieee.org/mlperf-rankings-2022

这里的规矩

• 得体：提 Issue 时明确标注 GPU 型号（SM90/SM100/SM120/ROCm）、head_dim、dtype、causal/varlen/paged_kv 开关、是否 deterministic；复现脚本最小化。
• 得体：提 PR 时遵循现有拆分习惯（tests first → 小步增量 → 最终 delete dedicated kernels），并在 PR 标题标注代际与硬件前缀，如 [CuTe, SM100]、[FA2]、[ROCm/CK]。
• 得体：本地用 two-pass testing（FLASH_ATTENTION_FAKE_TENSOR=1 编译 + 缓存执行）验证，确保 kernel 编译通过且测试命中缓存（FLASH_ATTENTION_TEST_EXPECT_CACHED=1）。
• 得体：代码风格遵循 ruff（PR #2940 已强制），提交前跑 pre-commit。
• 得体：在 RFC 类 Issue（如 #2456 head_dim=256/512、#2915 sparse MLA 64-head）先讨论方案再动笔。
• 不得体：在 Mac 上无 GPU 时直接跑 pytest tests/cute/test_flash_attn.py 期望全部通过——必须用 FakeTensorMode 或 two-pass 流程。
• 不得体：跨代际大改（如同时动 FA2/FA3/FA4）而不拆 PR；维护者明显偏好按代际与硬件分界线独立提交。
• 不得体：在 Issue 里只写「性能慢/报错」却不给 GPU 架构、head_dtype、seqlen、batch 等关键信息——这类 Issue 会被要求补充复现条件。
• 不得体：未经讨论直接提「重构整个 scheduler」类 PR；调度器（prepare_scheduler / tile_scheduler / seqlen_info）是核心路径，改动需先 RFC。

维护者现在最想要的帮助

• SM100/SM120 (Blackwell) 支持完善：近期提交与 Issue 密集指向 SM100（B200/SM103）的 hd256 backward、varlen Q、decode、sink 精度、illegal memory access 等问题（#2950、#2951、#2952、#2969、#2860、#2852），是当前最活跃的开发线。
• MLA (Multi-head Latent Attention) 扩展：#2938 刚合入 1CTA MLA forward (dense + sparse top-k) 与 dispatch heuristic；#2915 提出 native 64-head sparse MLA 设计；MLA backward 与 paged KV 仍在推进。
• head_dim=256/512 统一：#2456 RFC 与 #2964–#2968 系列 PR 正在把 hd256 backward 统一为默认并删除专用 kernel；hd512 仍在讨论。
• 数值精度与 deterministic：#2969 dsink precision、#2946 learnable-sink logits overflow、#2244 FA3 deterministic backward 非确定性、#2957 Paged KV 数值错误——精度类 bug 优先级高。
• ROCm / CK / aiter 后端：#2933 gfx905 asm backward、#2963 gfx12 BF16 D256 split-head backward、#2901 num_splits heuristic 修复——micmelesse 在推动，欢迎 ROCm 后端贡献。
• Python dispatch 开销优化：#2426 记录 FA4 ~20–25us Python dispatch overhead，需要熟悉 flash_attn/cute/interface.py 与 torch.library 的人做 dispatch 层瘦身。
• 跨平台构建（Windows / CUDA 13 / PyTorch 2.15）：#2872、#2799、#2861 显示 Windows 与最新 PyTorch/CUDA 适配仍在打磨。
• FA2 维护与 split-KV：#2961「Normalize split-KV combine scales」等表明 FA2 仍在被训练框架依赖，需要人维护 split-KV combine 的正确性。
• 测试基础设施：two-pass testing 工具（tools/two_pass_tests.sh）、CI matrix（#2976 publish.yml tag filter）需要人维护与扩展。

五、切入方案

它现在缺什么（你无 GPU 也能补）

第一批 PR

怎么知道自己站住了

• 首次 PR 在 2 周内被合并，且 review 来自核心维护者（Johnsonms / tzadouri / Tri Dao）。
• 在 2–3 个 Issue 下收到维护者回复「go ahead」或「assigned to you」。
• 被邀请加入项目的 Discord/Slack 贡献者频道，或收到 triage 权限邀请。
• 自己的测试/文档 PR 成为后续相关 PR 的参照模板（被引用或复用）。
• 在 flash_attn/cute 方向的 Issue 讨论中能提供代码级别的诊断意见（如指出某 bug 在 prepare_scheduler.py 的哪一行）。
• Colab T4 上的 benchmark 基线被社区引用（如在 Issue 讨论或 PR 中作为性能参照）。
• 被维护者提名为 flash_attn/cute/tests 或 flash_attn/cute/docs 方向的 reviewer。
• 连续 3 个月每月至少 1 个 PR 被合并，且范围从文档/测试逐步扩展到调度器逻辑或 kernel 修复。