flash_attn_interface
flash_attn/flash_attn_interface.py 1 文件 + csrc/flash_attn 系列 .cu/.h | FA2 时代的 Python → C++ 包装层,暴露 flash_attn_func/flash_attn_varlen_func/flash_attn_qkvpacked_func 等,内部通过 torch.ops 加载 csrc/flash_attn 编译出的 CUDA op。 入口:flash_attn_func, flash_attn_varlen_func, flash_attn_qkvpacked_func, flash_attn_with_kvcache 依赖:csrc/flash_attn (FA2 C++ 扩展) README 与 usage.md 主要示例入口;PyTorch 训练代码(HuggingFace/Megatron)大多走这里。 |
cute_interface
flash_attn/cute/interface.py 1 文件(模块总 110 文件) | FA4 的 Python 入口,把 Q/K/V、causal、window、alibi、paged-kv、MLA 等参数转成 CuTeDSL kernel 调用。 入口:flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache (FA4 路径) 依赖:flash_attn/cute/flash_fwd.py, flash_bwd.py, mla_*, tile_scheduler.py, prepare_scheduler.py CLAUDE.md 与 AGENTS.md 指定为当前活跃入口;与 flash_attn_interface 并存,分别对应 FA4 / FA2。 |
flash_fwd
flash_attn/cute/flash_fwd.py ~1 主文件 + 多个变体 flash_fwd_sm90/sm100/sm120.py、flash_fwd_mla_sm100.py 等 | FA4 前向主 kernel(SM90/SM100/SM120、pack-gqa、paged-kv、varlen),组装 TMA/UMMA/MMA、mask、softmax、scheduler。 入口:flash_fwd, FlashAttentionForwardSm90, FlashAttentionForwardSm100, FlashAttentionForwardSm120 等(需验证具体类名) 依赖:tile_scheduler.py, prepare_scheduler.py, mask.py, softmax.py, pack_gqa.py, paged_kv.py, pipeline.py, named_barrier.py, ampere_helpers.py, blackwell_helpers.py 提交热点区(hd256 fwd、S/P ping-pong、decode 修复、MLA fwd)集中在此。 |
flash_bwd
flash_attn/cute/flash_bwd.py ~1 主文件 + flash_bwd_sm90/sm100/sm120.py、flash_bwd_mla_sm100.py、sm100_hd256_2cta_fmha_backward_*.py | FA4 反向主 kernel(标准 + MLA + hd256 2CTA),把 dQ/dK/dV 拆成 DK/DV 与 DQ 两阶段或融合。 入口:flash_bwd, FlashAttentionBackwardSm90/SM100/SM120, BlackwellFusedMultiHeadAttentionBackwardDKDVKernel, BlackwellFusedMultiHeadAttentionBackwardDQKernel 依赖:tile_scheduler.py, mask.py, softmax.py, pipeline.py, copy_utils.py, fast_math.py 最近提交集中在 hd256 backward 统一、deterministic=True、2CTA dQaccum load、S/P ping-pong。 |
tile_scheduler
flash_attn/cute/tile_scheduler.py 1 大文件 | FA4 的 tile 调度核心:static/dynamic/CLC/persistent、varlen、2CTA 命名化调度、与 CLC 硬件调度器(ClcDynamicPersistentTileScheduler)交互。 入口:SchedulerState, ClcSchedulerState, DynamicPersistentSchedulerState, SingleTileVarlenScheduler, Sm100FmhaClcDynamicTileScheduler, Sm100FmhaStaticTileScheduler, compute_sm100_fmha_varlen_grid, SchedulingMode 依赖:cutlass.utils.ClcDynamicPersistentTileScheduler, quack.cute_dsl_utils.ParamsBase, flash_attn.cute.utils 提交热点(named 2CTA schedules、SM100_TMEM_CAPACITY_COLUMNS、varlen grid)直接改这里。 |
prepare_scheduler
flash_attn/cute/prepare_scheduler.py 1 文件 | FA4 的 host 侧调度预处理:把 (batch, seqlen, nheads, tile) 变成 cu_total_m_blocks / cu_total_splits_m_blocks / blocks_to_batch_idx 等元数据 tensor,供 kernel 启动。 入口:FlashPrepareScheduler, SchedulerMetadataTensorsTorch 依赖:flash_attn.cute.utils, cutlass.cute.FastDivmodDivisorV2 注释明确说明是 hopper/flash_prepare_scheduler.cu 的 CuTeDSL 重写版;varlen/paged-kv 路径必过。 |
flash_attn_ck
csrc/flash_attn_ck/ ~7 文件 | FA2 的 ROCm Composable Kernel 后端(AMD GPU),提供 mha_fwd/mha_bwd/mha_varlen_fwd/mha_varlen_bwd/mha_fwd_kvcache。 入口:flash_api.cpp 中注册的 mha_fwd/mha_bwd/mha_varlen_fwd/mha_bwd/kvcache 依赖:composable_kernel 子树 README 指定 ROCm CK 为默认后端;与 csrc/flash_attn(CUDA)并列。 |
hopper_fa3
hopper/ 507 文件 | FA3(Hopper beta)完整实现:C++ 版的 SM90 TMA/GMMA 前向/反向、softcap、pack-gqa、paged-kv、split-kv、MLA decode、benchmark。 入口:flash_api.cpp, flash_fwd_kernel_sm90.h, flash_bwd_kernel_sm90.h, mainloop_fwd_sm90_tma_gmma_ws.hpp, mainloop_bwd_sm90_tma_gmma_ws.hpp, generate_kernels.py 依赖:csrc/cutlass(老 C++ 子树), 部分 csrc/flash_attn 共用工具 README 标记为 beta;仍被 H100/H800 用户和 benchmark 使用,但活跃开发已切到 FA4。 |
csrc_fa2
csrc/flash_attn/ ~197 文件(含子树) | FA2(Ampere/Ada)CUDA 实现,仍是 PyTorch 默认 pip install flash-attn 的编译目标。 入口:flash_api.cpp, flash_fwd_kernel.h, flash_bwd_kernel.h, generate_kernels.py, src/flash_fwd_*.cu, src/flash_bwd_*.cu 依赖:csrc/cutlass, csrc/composable_kernel(子树) 与 flash_attn_interface 配套;支持 SM80/SM90(非 TMA 路径)。 |
model_layers
flash_attn/modules/, flash_attn/layers/, flash_attn/models/ ~30 文件 | 把 flash attention 包装成 nn.Module(MHA、embedding、rotary、block),并提供 GPT/LLaMA/Baichuan/Falcon/BERT/ViT 等模型级替换。 入口:modules/mha.py:MultiheadAttention, layers/rotary.py, models/gpt.py, models/llama.py, models/bert.py, models/vit.py 依赖:flash_attn_interface 或 cute_interface(取决于 import 路径) 训练/推理集成的主要扩展点;新模型通常在这里加一个 models/xxx.py。 |
benchmarks
benchmarks/ 12 文件 | 性能对标:FA2/FA3/FA4、SM90、MLA paged-kv、sparse MLA fwd、varlen scheduling、ALiBi、GEMM、CLC bench。 入口:benchmark_flash_attention.py, benchmark_mla_paged_kv.py, benchmark_sparse_mla_fwd.py, benchmark_varlen_sched.py, bench_sm90.py, clc_bench.py 依赖:flash_attn_interface, cute_interface, hopper_fa3 README 与发布说明的性能数字来源;新 kernel 提 PR 时通常要补对应 benchmark。 |
tests
tests/ 55 文件 | 回归测试:FA2/FA3/FA4、Triton-AMD、MLA、varlen、mask/score mod、block sparsity。 入口:tests/cute/test_flash_attn.py, test_flash_attn_mla.py, test_flash_attn_varlen.py, test_mask_mod.py, test_score_mod.py, test_block_sparsity.py, test_flash_attn.py, test_flash_attn_ck.py 依赖:各代接口 AGENTS.md 给出 FakeTensor 两-pass 编译/执行工作流;新扩展必须补 test。 |
training
training/ 145 文件 | 分布式训练集成:GPT/GPT3 训练脚本、FSDP/TP/PP/EP 配置、Dockerfile。 入口:run.py, configs/, src/ 依赖:flash_attn_interface 或 cute_interface 与 Megatron/DeepSpeed 生态衔接;新模型/新分布式策略可在此加配置验证。 |
utils_and_helpers
flash_attn/utils/, flash_attn/cute/utils.py, flash_attn/cute/copy_utils.py, flash_attn/cute/fast_math.py, flash_attn/cute/barrier.py, flash_attn/cute/named_barrier.py, flash_attn/cute/pipeline.py, flash_attn/cute/mma_sm100_desc.py ~15 文件 | 公共工具:分布式 wrapper、generation、pretrained、benchmark、CuTeDSL 的 barrier/pipeline/divmod/TMEM/PTX 辅助。 入口:utils/distributed.py, utils/generation.py, cute/utils.py:constexpr 工具, cute/pipeline.py:PipelineStateSimple, cute/barrier.py:ld_acquire/wait_eq/arrive_inc 依赖:cutlass.cute, quack-kernels, apache-tvm-ffi 跨代共享;FA4 的 producer/consumer pipeline 与 named barrier 都在此实现。 |
aiter_rocm
third_party/aiter 子模块 | ROCm Triton 后端子模块,提供 FA2 级别的 Triton kernel(CDNA + RDNA、FP8、sliding-window WIP)。 入口:aiter 包内 Triton kernel(具体符号需验证) 依赖:Triton, PyTorch ROCm README 与 FLASH_ATTENTION_TRITON_AMD_ENABLE 流程说明其用途;默认不参与 CUDA 构建。 |