PaddlePaddle / FlashMask

# ⚡ FlashMask
**Flexible and Efficient Attention with Column-wise Sparse Masks**
English | 简体中文
 Figure 1: FlashMask overview. (a) Supported mask types; (b) column-wise sparse representation; (c) efficient implementation 
 Figure 2: FlashMask-3 across the four layers of the execution stack 
## 📋 Contents
- News
- Overview
- Key Features
- Performance
- Installation
- Quick Start
- Citation
- Acknowledgement
- License
---
## 📢 News
- **2026-09-25 FlashMask-4**: Supports Hopper and Blackwell. Kernels are JIT-compiled at runtime with CuTe DSL, so compilation is fast. Flexible masks are available at every supported head dim, with full support for sink bias and deterministic computation, plus large head dims. Both PaddlePaddle and PyTorch interfaces are provided.
- **2026-09-25 FlashMask-3**: Supports Hopper and introduces mask-aware distributed attention with communication/computation overlap and load balancing.
- **2026-09-25 FlashMask**: Supports Ampere and introduces the column-wise sparse mask representation, which reduces memory complexity from $O(N^2)$ to $O(N)$ and uses mask sparsity to skip unnecessary computation.
---
## 🎯 Overview
FlashMask is an **efficient attention kernel library** for flexible sparse masks. It uses a column-wise sparse mask representation that encodes a 2D dense mask as 1D per-column row-index ranges, and it exploits mask sparsity to skip unnecessary computation.
This repository evolved from PaddlePaddle/flash-attention and is now maintained as the standalone FlashMask repository. It implements forward and backward kernels for **Hopper (SM90) and Blackwell (SM100)** with **CuTe DSL**, and provides both **PaddlePaddle and PyTorch** interfaces. The core code is JIT-compiled at runtime with CuTe DSL, so compilation is fast. Beyond flexible masks, it supports mask-aware distributed attention with communication/computation overlap and load balancing, and adds large head dim support on SM100.
### Formulation
$$
O = \mathrm{softmax}\left(\frac{QK^\mathsf{T}}{\sqrt{d}} + M\right)V
$$
where:
- $Q,\ K,\ V$ are the input tensors of the attention module;
- $d$ is the size of the last dimension of $Q$ (the head dim);
- $M$ is the column-wise sparse mask introduced by FlashMask.
---
## ✨ Key Features
FlashMask-4 provides three capabilities:
- memory-efficient, flexible, and fast masked attention kernels;
- mask-aware distributed attention with communication/computation overlap and load balancing;
- large head dim support.
All three ship in the same `flashmask` package. The core code is JIT-compiled at runtime with CuTe DSL, so compilation is fast, and **the PaddlePaddle and PyTorch frontends share the same kernel implementation**.
### Memory-efficient, flexible, and fast masked attention kernels
The table below compares **FlashMask-4** with the **`mask_mod` path of FlashAttention-4**, not with every FlashAttention-4 interface. `d` is the head dim of Q/K and `dv` is the head dim of V.
| d | dv | FlashMask SM90 | FlashMask SM100 | Flas