cuflash

📚 Portfolio map: https://github.com/holtwood-systems/portfolio

cuflash — 从零手写的 CUDA FlashAttention：标量 → WMMA Tensor Core 前向、FlashDecoding/Split-KV、Roofline 性能分析；FP16/BF16/FP32 前反向，sm_70–sm_90。

文档 · API 参考 · 更新日志

🎯 项目简介

cuflash 是一个从零手写的 CUDA FlashAttention 深度实现：从正确起步，经标量内核 → WMMA Tensor Core 前向的优化迭代，形成有数据支撑的性能叙事（教学可读性为底色，内核深度为作品定位，见 ROADMAP）。

在 GPU 学习主线（五主仓）中，本仓库负责 CUDA C++ FlashAttention 前向/反向的专项深挖；CUDA 基础、Triton 实现、完整推理运行时和 Serving 控制面分别由其他主仓承担。整体顺序见 LEARNING_PATH.md（meta 仓）。

项目状态

状态：stable（当前源码与最新 Release 为 0.6.0；只修正确性 bug 与文档，处于维护收敛阶段）
权威入口：公开头文件、实现、测试和用户文档
定位：从零手写的 kernel 深度作品——教学可读性为底色，优化迭代叙事为定位（见 ROADMAP）
当前重点：删除过时流程层、收紧文档并修复长尾缺陷，而不是扩展新功能

⚡ 核心亮点（30 秒速览）

维度	内容
内核矩阵	标量 kernel（FP32 / 回退路径）+ WMMA Tensor Core 前向（FP16 `sm_70+` / BF16 `sm_80+`，`m16n16k16`，FP32 accumulate）+ FlashDecoding/Split-KV + 完整反向
算法	tiling + online softmax + FA2 延迟归一化：不物化 O(N²) 注意力矩阵，辅助内存 O(N)，反向复用 L=logsumexp
代码规模	~2.8K 行核心 CUDA 实现（`src/` + `include/`），770 行 Google Benchmark
正确性验证	11 个单元测试文件（数值稳定性/causal mask/dtype/边界压力）+ 集成测试 + PyTorch SDPA 对比测试 + compute-sanitizer
架构覆盖	sm_70（V100）→ sm_90（H100），实测硬件 sm_86
实测快照	RTX 3060 Laptop：FP16 Tensor Core 前向 seq=4096/d=64 23.5 ms，同形状 FP32 标量 177 ms（约 7.5×，同仓库同硬件对比，详见性能）

数字口径：所有性能数字均标注硬件与环境，完整协议与约束见 docs/performance/benchmarks.md。本仓库遵循「每个数字必须可复现、可归档」的证据标准，不引用无原始产物支持的跨 GPU 历史数据。

为什么选择 cuflash？

挑战	解决方案
📚 学习 FlashAttention	清晰、文档完善的 CUDA 内核，逐步的算法实现
🔬 研究与实验	修改注意力机制，无需复杂的框架依赖
🔧 便于集成	C++ API 配合 C ABI 绑定，可通过 ctypes 接入 Python
⚡ GPU 优化	多架构支持，从 V100 (sm_70) 到 H100 (sm_90)

✨ 主要特性

特性	说明
⚡ O(N) 辅助内存	不物化完整 O(N²) 注意力矩阵
🔢 多精度支持	FP32、FP16、BF16 的前向和反向路径
🔁 完整训练	完整的前向/反向传播，包含梯度计算
🎭 因果掩码	内置自回归模型支持（GPT 风格）
🔧 易于集成	简洁的 C++ API + C ABI，便于 Python ctypes 集成
🏎️ 多架构	优化的 CUDA 内核，支持 sm_70 → sm_90（V100 → H100）
🧪 全面测试	单元测试、集成测试、压力测试、PyTorch 对比测试
📊 性能基准	Google Benchmark 集成，用于性能追踪
📚 完整文档	VitePress 文档站点，覆盖算法、API、性能与故障排除

与同类库对比

特性	cuflash	PyTorch SDPA	xFormers	FlashAttention-2
教学代码	✅ 清晰简洁	❌ 复杂难读	❌ 复杂难读	⚠️ 中等
自定义修改	✅ 容易	⚠️ 困难	⚠️ 困难	⚠️ 困难
无需框架依赖	✅ 是	❌ PyTorch	❌ PyTorch	❌ PyTorch/Cutlass
Python 绑定	✅ ctypes	✅ 原生	✅ 原生	✅ PyTorch
训练支持	✅ 完整	✅ 完整	✅ 完整	✅ 完整
BF16 支持	✅ 是	✅ 是	✅ 是	✅ 是

选择 cuflash 的场景：希望理解、修改或嵌入 FlashAttention，同时避免繁重的依赖。

🧭 项目边界（IN / OUT）

IN（本仓库负责）：

FlashAttention 前向 + 反向
FP16 / BF16 / FP32 多精度，causal mask
FlashDecoding（decode 阶段 KV 分块并行）
优化叙事与 benchmark（Roofline、kernel 深挖）

OUT（明确不做，见对应仓库）：

GEMM 基础与 CUDA 编程学习 → cuda-foundations
Triton 版算子（参考实现）→ trifuse
完整推理运行时（模型加载/采样/生成）→ tiny-llm

🏗️ 架构与实现

内核分层

                cuflash::flash_attention_forward   （API 层：src/api/）
                        │  参数校验 / dtype 分派 / launch 配置
        ┌───────────────┴────────────────┐
        ▼                                ▼
  WMMA Tensor Core 路径              标量 kernel 路径
  (FP16 sm_70+ / BF16 sm_80+)        (FP32，及其他架构回退)
  S = Q·Kᵀ 与 O += P·V 走            纯 CUDA core，
  m16n16k16 fragment，               与 WMMA 共享同一套
  FP32 accumulate                    online softmax 逻辑
        └───────────────┬────────────────┘
                        ▼
      online softmax + FA2 延迟归一化（src/kernels/online_softmax.cu）
      逐 K/V 块更新 m、l，O 累加后延迟除以 l；L = m + log(l) 落盘，
      反向据此重计算 P，避免保存 O(N²) 中间矩阵
                        ▼
      decode 场景（query_len = 1）：FlashDecoding / Split-KV
      Phase 1 各 KV chunk 并行计算局部 (m, l, O)
      → Phase 2 按 flash rescaling 规则归并：m = max m_c，
        l = Σ l_c·e^{m_c−m}，O = Σ O_c·e^{m_c−m} / l

关键实现事实

Q/K/V tile 常驻 shared memory：BLOCK_M × HEAD_DIM / BLOCK_N × HEAD_DIM 布局，extern __shared__ 一次性划分（见 flash_attention_forward_wmma.cu）。
launch 越界修复：grid 展平到一维（num_q_blocks × batch_heads），避免 B*H > 65535 时超出 gridDim.y 上限——附带回归测试（ROADMAP E2a）。
多精度一条管线：FP32/FP16/BF16 通过 TypeAdapter 走同一 kernel 模板，BF16 fragment 需 sm_80+，不满足时 launcher 自动回退标量路径。
完整反向：src/backward/flash_attention_backward_typed.cu（600 行），dQ/dK/dV 全路径，数值稳定性由专项测试覆盖。

深度解读见文档站：算法详解 · 内核深潜 · FlashDecoding 设计 · Roofline 分析

🚀 快速开始

环境要求

GPU: NVIDIA GPU，计算能力 7.0+（V100、RTX 20/30/40、A100、H100）
CUDA Toolkit: 11.8 或更高版本（推荐 12.x；sm_90/Hopper 需要 11.8+；注意 CUDA 13.x 已移除 sm_70 支持）
CMake: 3.18 或更高版本（使用 CMake preset 需 3.20+）
编译器: GCC 7+、Clang 5+ 或 MSVC 2017+（需要 C++17 支持）

方案 1：源码编译

# 克隆仓库
git clone https://github.com/holtwood-systems/cuflash.git
cd cuflash

# 使用预设构建（Release 模式）
cmake --preset release
cmake --build --preset release

# 运行测试
ctest --preset release --output-on-failure

方案 2：Docker（推荐用于快速测试）

# 构建 Docker 镜像
docker build -t cuflash .

# 运行（需要 GPU 支持）
docker run --gpus all -it cuflash

# 容器内运行基准测试
./cuflash_bench

第一个 C++ 程序

#include <cuda_runtime.h>
#include "cuflash/flash_attention.h"
#include <iostream>
#include <vector>
#include <cmath>
#include <random>

// 辅助函数：用随机值初始化设备内存
void init_random(float* d_ptr, size_t n, unsigned seed = 42) {
    std::mt19937 gen(seed);
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    std::vector<float> h_data(n);
    for (auto& v : h_data) v = dist(gen);
    cudaMemcpy(d_ptr, h_data.data(), n * sizeof(float), cudaMemcpyHostToDevice);
}

int main() {
    const int B = 2, H = 8, N = 1024, D = 64;
    const float scale = 1.0f / std::sqrt(static_cast<float>(D));
    const size_t qkv_size = B * H * N * D;
    const size_t l_size = B * H * N;
    
    // 分配设备内存
    float *d_Q, *d_K, *d_V, *d_O, *d_L;
    cudaMalloc(&d_Q, qkv_size * sizeof(float));
    cudaMalloc(&d_K, qkv_size * sizeof(float));
    cudaMalloc(&d_V, qkv_size * sizeof(float));
    cudaMalloc(&d_O, qkv_size * sizeof(float));
    cudaMalloc(&d_L, l_size * sizeof(float));
    
    // 用随机数据初始化
    init_random(d_Q, qkv_size, 1);
    init_random(d_K, qkv_size, 2);
    init_random(d_V, qkv_size, 3);
    
    // 使用因果掩码计算 FlashAttention
    auto err = cuflash::flash_attention_forward(
        d_Q, d_K, d_V, d_O, d_L,
        B, H, N, D, scale,
        true  // 因果掩码
    );
    
    if (err != cuflash::FlashAttentionError::SUCCESS) {
        std::cerr << "错误: " << cuflash::get_error_string(err) << std::endl;
        return 1;
    }
    
    std::cout << "FlashAttention 前向传播成功完成！" << std::endl;
    
    // 释放资源
    cudaFree(d_Q); cudaFree(d_K); cudaFree(d_V);
    cudaFree(d_O); cudaFree(d_L);
    return 0;
}

📖 更多示例: 请参见 examples/ 目录，包含反向传播和 Python 集成的完整程序。

Python 集成

import ctypes
import numpy as np
import torch

# 加载动态库
lib = ctypes.CDLL("./build/release/libcuflash.so")

# 定义 API
lib.cuflash_attention_forward_f32.argtypes = [
    ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p,
    ctypes.c_void_p, ctypes.c_void_p,
    ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int,
    ctypes.c_float, ctypes.c_bool, ctypes.c_void_p
]
lib.cuflash_attention_forward_f32.restype = ctypes.c_int

# 使用 PyTorch 准备数据
B, H, N, D = 2, 8, 1024, 64
Q = torch.randn(B, H, N, D, dtype=torch.float32, device='cuda')
K = torch.randn(B, H, N, D, dtype=torch.float32, device='cuda')
V = torch.randn(B, H, N, D, dtype=torch.float32, device='cuda')
O = torch.empty_like(Q)
L = torch.empty(B, H, N, dtype=torch.float32, device='cuda')

# 调用 cuflash
scale = 1.0 / np.sqrt(D)
result = lib.cuflash_attention_forward_f32(
    ctypes.c_void_p(Q.data_ptr()),
    ctypes.c_void_p(K.data_ptr()),
    ctypes.c_void_p(V.data_ptr()),
    ctypes.c_void_p(O.data_ptr()),
    ctypes.c_void_p(L.data_ptr()),
    B, H, N, D, scale, True, None
)

assert result == 0, f"FlashAttention 失败，错误码 {result}"
print(f"输出形状: {O.shape}, 平均值: {O.mean().item():.4f}")

📖 完整 Python 示例: 参见 examples/python_binding.py

📊 性能

FlashAttention 通过分块与在线 softmax 避免物化完整注意力矩阵。实际延迟、吞吐和显存收益取决于 GPU、CUDA、dtype、形状和 causal 模式；仓库不维护无法持续复测的固定性能数字，请在目标硬件上运行自带 benchmark。

指标口径：FlashAttention kernel 用 CUDA Event 计时（排除 host 开销）；带宽指标为模型化 logical HBM（前向：Q/O 各一次 + K/V 按 Q block 重载，见 algorithm），并非 ncu 实测的物理带宽。请勿把 LogicalHBM GB/s 当作硬件实测 HBM 带宽。

运行基准测试

cmake --preset release
cmake --build --preset release
./build/release/cuflash_bench

本机实测快照（RTX 3060 Laptop，sm_86，2026-08-18，commit `6860cbc`）

非正式快照：未归档成套原始 JSON，仅作 sanity check 与量级参考；环境、复现命令与完整表格见 benchmarks.md §1.5。配置：batch=1, heads=8，CUDA Event kernel-only 计时，10 预热 + 30 采样中位数。

形状	FP16 (Tensor Core)	FP32 (标量)	FP32 causal
seq=1024, d=64	1.76 ms	12.0 ms	4.55 ms
seq=4096, d=64	23.5 ms	177 ms	56.7 ms
seq=4096, d=128	84.1 ms	379 ms	186 ms

读法：同一仓库、同一硬件下，Tensor Core 路径相对标量路径约 7.5×（23.5 vs 177 ms）—— 这是「优化迭代叙事」的直接证据；绝对值与 PyTorch SDPA 的对比口径见完整文档，本仓库不把教学级实现包装成生产级性能。

姊妹实现对照：Triton 版（RTX 4090 实测）

同一 FlashAttention 学习主线上，trifuse 提供 Triton 参考实现（分块 + online softmax + exp2 折叠 + causal 对角短路）， 2026-11 在 2×RTX 4090 48GB 上实测：

配置（seq=16K, BF16）	Triton FA-2	参照
causal vs naive PyTorch 实现	24.7× 加速	峰值显存 17.35 GB → 0.19 GB（约 91×）
causal vs PyTorch SDPA (flash 后端)	3.90 ms vs 3.70 ms（~95% 持平）	vs SDPA 最大误差 4.9e-04
非 causal vs SDPA	6.52 ms vs 6.70 ms（反超 ~3%）	同上

两个实现互为印证：Triton 用于算法快速验证（online softmax、exp2 折叠等技巧的迭代）， CUDA C++ 版下沉到 Tensor Core 管线与完整反向。面试中可讲清两条路径的取舍—— 这也正是上文「项目边界（IN / OUT）」划分的由来。

🎯 面试考点映射

高频面试题	本仓库的对应物
FlashAttention 为什么快？显存为什么是 O(N)？	不减 FLOPs、减 HBM 读写（O(N²)→O(N)）；算法详解有读写量推导
推导 online softmax：两块拼接送进来时 m/l 怎么更新？	`src/kernels/online_softmax.cu` + 单元测试
FA2 的延迟归一化（deferred normalization）是什么？	前向只落盘 O 与 L=logsumexp，反向据此重算 P，不保存注意力矩阵
WMMA / Tensor Core 怎么编程？	`flash_attention_forward_wmma.cu`：m16n16k16 fragment、FP32 accumulate、warp 级 Q 行分配
decode 阶段为什么慢？Split-KV 怎么做？	Q=1 时每 token 要扫全部 KV；`flash_decoding.cu` 两阶段归并 + 设计文档
数值稳定性怎么验证？	`test_numerical_stability.cu`、dtype/causal mask/边界压力测试
CUDA 工程坑（launch 上限、shared memory 布局）？	`B*H > 65535` 的 `gridDim.y` 越界修复 + 回归测试（ROADMAP E2a）
优化无效怎么办？（负结果文化）	causal-boundary-skip.md：边界块特化实测 ±2%，诚实记录为噪声内负结果，不改写成加速
为什么比 PyTorch SDPA 慢？差距来自哪？	教学实现 vs cuDNN 生产管线的差距来源（warp specialization、TMA），以及追赶路径（kernel-deep-dive）

📖 文档

快速链接

资源	链接
📘 完整文档	https://holtwood-systems.github.io/cuflash/
🔌 API 参考	API 文档
🧠 算法详解	深入理解 FlashAttention
🔧 构建指南	从源码构建
❓ 故障排除	常见问题与解决方案

⚙️ 配置

支持的参数

参数	值	说明
`head_dim`	32, 64, 128	内核优化必需
数据类型	FP32 (`float`), FP16 (`half`), BF16 (`__nv_bfloat16`)	前向和反向都支持
因果掩码	可选	运行时启用/禁用
批大小	≥ 1	任意正整数
序列长度	≥ 1	优化用于 1K-16K+
头数	≥ 1	任意正整数

GPU 架构支持

架构	计算能力	示例 GPU
Volta	sm_70	V100
Turing	sm_75	RTX 2080 Ti
Ampere	sm_80, sm_86	A100, RTX 3090
Ada Lovelace	sm_89	RTX 4090
Hopper	sm_90	H100

默认构建目标: sm_80, sm_86（A100 + RTX 30xx）。可配置编译 sm_70–sm_90；仓库当前实测硬件为 sm_86（RTX 3060）。

自定义使用: cmake --preset release -DCMAKE_CUDA_ARCHITECTURES="90"

🏗️ 项目结构

cuflash/
├── benchmarks/                 # 性能基准测试（Google Benchmark）
├── cmake/                      # CMake 模块和打包配置
├── docs/                       # VitePress 文档站点（中文）
│   ├── .vitepress/             # 站点配置与主题
│   ├── public/                 # 静态资源（logo、favicon）
│   ├── guide/                  # 快速开始指南
│   ├── design/                 # 设计决策与内核深度解读
│   ├── performance/            # 基准测试与 Roofline 分析
│   └── research/               # 相关工作与参考文献
├── examples/                   # 完整使用示例
│   ├── basic_usage.cu          # C++ 基础示例
│   └── python_binding.py       # Python ctypes 示例
├── include/cuflash/            # 公共 API 头文件
│   ├── flash_attention.h       # 主 API，包含 C++ 和 C ABI
│   ├── export.h                # 可见性宏
│   └── version.h.in            # 版本头文件模板
├── src/                        # 实现代码
│   ├── api/                    # API 调度层
│   ├── forward/                # 前向传播内核实现
│   ├── backward/               # 反向传播内核实现
│   └── kernels/                # 内部内核工具（.cuh）
├── tests/                      # 测试套件
│   ├── unit/                   # 单元测试（10 个文件）
│   ├── integration/            # 集成测试 + PyTorch 对比
│   └── package_smoke/          # 包冒烟测试
├── CMakeLists.txt              # 主构建配置
├── CMakePresets.json           # 构建预设（release、debug、asan）
├── Dockerfile                  # 容器构建
└── .github/workflows/          # CI/CD 工作流
    ├── ci.yml                  # 多配置构建矩阵与测试
    ├── codeql.yml              # 安全扫描
    ├── pages.yml               # 文档部署
    └── release.yml             # 发布自动化

🧪 测试与质量

测试分类

# 运行所有测试
ctest --preset release --output-on-failure

# 运行特定测试类别
ctest --preset release -R ForwardTest    # 前向传播测试
ctest --preset release -R BackwardTest   # 反向传播测试
ctest --preset release -R StressTest     # 压力与边界测试
ctest --preset release -R PyTorch        # PyTorch 对比测试（需要 GPU + PyTorch）

代码质量工具

✅ clang-format: 自动化代码格式化（CI 强制执行）
✅ clang-tidy: 静态分析，50+ 检查（host 源文件在 CI 中强制执行）
✅ CodeQL: 每周安全扫描
✅ Sanitizers: AddressSanitizer & UBSan（CI 构建矩阵）+ compute-sanitizer（GPU workflow）

# 使用 AddressSanitizer 构建
cmake --preset debug-asan
cmake --build --preset debug-asan
ctest --preset debug-asan

🤝 贡献

欢迎贡献。请保持改动聚焦、明确，并围绕当前 CUDA 库边界推进。

开始贡献

克隆仓库 并确认当前范围
使用 preset 构建，先验证环境可用
编写或更新测试 覆盖你的改动
同步更新文档，确保行为、API 用法和工作流描述一致
提交拉取请求，附上简明摘要和验证说明

开发工作流

# 提交前格式化代码
find . -name "*.cu" -o -name "*.cuh" -o -name "*.cpp" -o -name "*.h" | xargs clang-format -i

# 本地运行测试
cmake --preset release && cmake --build --preset release
ctest --preset release --output-on-failure

# 可选：运行 clang-tidy（host 源文件）
./scripts/run_clang_tidy.sh build/release

📋 详细指南: 请参见 CONTRIBUTING.md

📄 许可证

本项目采用 MIT 许可证。

📚 参考文献

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao 等，NeurIPS 2022
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning — Dao，ICLR 2024

📈 版本历史

详细的版本历史和更新请参见 CHANGELOG.md。

_{用 ❤️ 打造的 CUDA FlashAttention 深度作品}
_{从零手写 · CUDA C++ · 开源}

Name		Name	Last commit message	Last commit date
.agents		.agents
.github		.github
.vscode		.vscode
benchmarks		benchmarks
cmake		cmake
docs		docs
examples		examples
include/cuflash		include/cuflash
scripts		scripts
src		src
tests		tests
.clang-format		.clang-format
.clang-tidy		.clang-tidy
.clangd		.clangd
.dockerignore		.dockerignore
.gitignore		.gitignore
.rgignore		.rgignore
AGENTS.md		AGENTS.md
CHANGELOG.md		CHANGELOG.md
CMakeLists.txt		CMakeLists.txt
CMakePresets.json		CMakePresets.json
CONTEXT.md		CONTEXT.md
CONTRIBUTING.md		CONTRIBUTING.md
Dockerfile		Dockerfile
LICENSE		LICENSE
README.md		README.md
ROADMAP.md		ROADMAP.md
package.json		package.json

Latest commit

History

Folders and files

Repository files navigation