CuFlash-Attn从零实现的 CUDA FlashAttention
技术白皮书 · O(N) 内存 · FP32/FP16 · 前向与反向
Skip to content
v0.5.1稳定版
FP32/16/BF16数值路径
WMMA前向 Tensor Core
C++ / CUDA最小依赖
验证状态
公开数字只在原始命令、硬件/软件、commit 与结果产物齐全时引用。历史跨 GPU 表格已隔离为不可审计快照； 当前优先级是以同一输入契约复测 PyTorch SDPA、官方 FlashAttention 与本实现，并归档 JSON 和 profiler 产物。
查看结果发布门槛 →快速开始
5 分钟内构建并运行：
bash
git clone https://github.com/open-infra-ai/cuflash-attn.git
cd cuflash-attn
cmake --preset release
cmake --build --preset release
ctest --preset release --output-on-failurecpp
#include "cuflash/flash_attention.h"
auto err = cuflash::flash_attention_forward(
d_Q, d_K, d_V, d_O, d_L,
batch_size, num_heads, seq_len, head_dim,
scale, true, stream
);python
import ctypes
lib = ctypes.CDLL("./build/release/libcuflash_attn.so")
lib.cuflash_attention_forward_f32(
q_ptr, k_ptr, v_ptr, o_ptr, l_ptr,
B, H, N, D, scale, True, None
)快速开始Preset 构建与第一步算法详解分块、online softmax、重计算Kernel 逐行解读共享内存、warp 调度API 参考完整 C++ 与 C ABI 文档基准测试本机快照与正式结果发布门槛相关工作论文与实现对比
核心参考文献
FlashAttention — Dao et al., NeurIPS 2022.
arXiv:2205.14135
arXiv:2205.14135
FlashAttention-2 — Dao, ICLR 2024.
arXiv:2307.08691
arXiv:2307.08691
Online Softmax — Milakov & Gimelshein.
arXiv:1805.02867
arXiv:1805.02867