$ show Dao-AILab/flash-attention
FlashAttention is a high-performance library implementing IO-aware exact attention for transformer models, with CUDA kernels targeting NVIDIA Ampere (SM80), Hopper (SM90), and Blackwell (SM100) architectures, plus a Python CuTe DSL implementation (FA4), a Composable Kernel backend for AMD, and Triton fallbacks. The library exposes PyTorch-integrated Python APIs (flash_attn_func, flash_attn_varlen_func, flash_attn_with_kvcache) backed by custom CUDA extensions with autograd support. It also ships fused ops (dense layers, layer norm, cross-entropy), model definitions with HuggingFace state-dict remapping (GPT, BERT, LLaMA, Falcon, Baichuan, OPT, GPT-J, GPT-NeoX, ViT), and a PyTorch Lightning + Hydra training pipeline for GPT-style models. The dominant character is a systems library, though it has a hybrid research flavor: experimental kernels for unreleased hardware, benchmark suites, and a training sub-package all coexist in one monorepo. The architecture is a layered monolith: a Python API layer dispatches to multiple kernel backends (CUDA C++ for FA2, CUTLASS/CuTe for FA3, Python CuTe DSL for FA4, CK for AMD, Triton for portability). The core algorithmic contributions are the online softmax with rescaling, tiled attention with shared-memory staging, warp-specialized producer/consumer pipelines using TMA and GMMA, persistent tile schedulers, and deterministic backward passes with semaphore-based reductions. The training sub-package is a secondary component demonstrating the kernels in full GPT training runs with tensor/sequence parallelism. The codebase is actively maintained with extensive CI/CD (including Apptainer-based FA4 CI), a large test suite covering correctness against reference implementations and HuggingFace models, and performance benchmarks. However, it carries significant complexity: multiple overlapping backends with duplicated logic and inconsistent feature support, a complex build system with many environment variables, a monolithic gpt.py model file, and deprecated components (csrc/layer_norm is explicitly deprecated in favor of Triton). The kernel code is world-class and worth studying; the surrounding infrastructure is functional but organically grown.
The LLM identified 6 files worth highlighting for reuse:
Parameterized online softmax classes (Softmax, SoftmaxSm100) implementing the rescaling algorithm that makes FlashAttention memory-efficient and numerically stable. Self-contained, architecture-parameterized, and directly reusable in any tiled attention or streaming-softmax kernel. Demonstrates the non-obvious rescaling math (row-wise max tracking, log-sum-exp accumulation, finalization) that is the algorithmic heart of FlashAttention.
A family of tile schedulers (SingleTileScheduler, DynamicPersistentTileScheduler, VarlenDynamicPersistentTileScheduler) that abstract block-to-SM work distribution for persistent GPU kernels, including L2 cache swizzling and variable-length sequence handling. Cleanly separated from the attention math, making it reusable for any persistent-kernel workload that needs dynamic work stealing or static tile assignment.
A JIT persistent cache with file locking and source fingerprinting that avoids recompilation of identical CuTe DSL kernel configurations across processes. Solves a recurring problem (concurrent JIT compilation in distributed training) with a portable, dependency-light design. Extractable into any project that JIT-compiles GPU kernels or other expensive artifacts.
Distributed communication utilities including all_gather/reduce_scatter with custom autograd functions and an FP16 gradient compression hook for all-reduce. Self-contained and reusable across any PyTorch distributed training codebase. Demonstrates non-obvious patterns for avoiding communication overhead in sequence-parallel and tensor-parallel settings.
Block sparsity configuration normalization and deterministic backward dQ write-order computation for sparse attention. Solves the subtle problem of ensuring bit-identical gradient writes when multiple sparse blocks alias the same output positions. Includes broadcast-pattern detection, subtile factor inference, and context-specific error messages. Reusable for any sparse-tensor kernel that needs deterministic reduction ordering.
Autotuned Triton fused linear+activation kernels (forward and backward) with IO-bound configuration selection. Demonstrates a clean pattern for writing portable, autotuned GPU kernels that fuse elementwise activations into GEMM epilogues. Extractable as a template for any fused GEMM+activation workload on NVIDIA or AMD GPUs.