{"success":true,"results":[{"repo":"23silicon/flashattention","url":"https://github.com/23silicon/flashattention","readmeUrl":"https://raw.githubusercontent.com/23silicon/flashattention/HEAD/README.md","snippet":"# FlashAttention Pedagogical Implementation in CUDA C++\n\nThis repository contains a memory-optimized CUDA C++ implementation of the **FlashAttention** algorithm (Dao et al.). It demonstrates the progression from a naive memory-bound attention kernel to an optimized, tile-based implementation capable of linear context scaling on sequence lengths of 262k+ on limited hardware.\n\nThe project is structured into two main iterations:\n* `FlashAttentionFirstAttempt.cu`: Initial FP32 implementation. Functional but limited by bank conflicts and low occupancy. Contain the tiling and online softmax logic at the core of FlashAttention's increased processing capabilities.\n* `FlashAttention.cu`: Optimized FP16 implementation featuring strided memory access, increased tile sizes, and reduced instruction overhead. Allows it to outperform naive attention on large datasets, but still far behind Pytorch's Scaled Dot Product Attention (SDPA) due to its industry-grade optimizations.\n* `FlashAttentionBenchmark.ipynb`: Jupyter notebook for benchmarking kernel performance against naive CUDA and PyTorch SDPA.\n\n## Performance Benchmark\n*Hardware: NVIDIA T4 Tensor Core GPU | Hidden Dimension d=128 | Precision: FP16*\n\nThe optimized kernel (`FlashAttention.cu`) successfully overtakes the Naive baseline at sequence length $N=4096$ and scales linearly to $N=196,608+$ where Naive attention crashes due to Out-Of-Memory (OOM) errors.\n\n| N (Seq Len) | Naive (ms) | Flash (ms) | PyTorch (ms) | Speedup (vs Naive) | Notes |\n| :--- | :--- | :--- | :--- | :--- | :--- |\n| **1024** | 3.41 | 10.44 | 0.06 | 0.33x | Naive wins (L2 Cache hits) |\n| **4096** | 49.55 | 41.24 | 0.61 | **1.20x** | **Crossover Point** (Flash becomes faster) |\n| **16384** | 862.99 | 720.87 | 10.95 | **1.20x** | Naive hits HBM bandwidth wall |\n| **24576** | 2202.61 | 1499.46 | 26.54 | **1.47x** | Significant memory bandwidth savings |\n| **32768** | 4304.87 | 2786.46 | 44.46 | **1.54x** | **Peak Relative Speedup** & Last Naive Run |\n| **49152** | **OOM** | 6119.29 | 104.08 | **$\\infty$** | Naive self-attention can no longer run |\n| **131072** | **OOM** | 42834.59 | 817.84 | **$\\infty$** | Scaling to 128k context |\n| **262144** | **OOM** | 169369.66 | 3939.94 | **$\\infty$** | Scaling to **262k context** |\n\n## Optimization Journey\n\n### Iteration 1: `FlashAttentionFirstAttempt.cu`\nThe initial attempt implemented the core tiling logic and online softmax.\n* **Precision:** FP32 (Single Precision).\n* **Bottlenecks:**\n * **Bank Conflicts:** The row stride of 128 (512 bytes) caused 32-way shared memory bank conflicts, serializing memory reads.\n * **Low Occupancy:** Small tile sizes ($B_r=64$) resulted in low warp occupancy, preventing the GPU from hiding global memory latency.\n * **Memory Limit:** FP32 data size restricted the maximum tile size per block.\n\n### Iteration 2: `FlashAttention.cu` (Current Optimized Version)\nThis version addresses the architectural bottlenecks of the first attempt, resulting in a robust speedup.\n\n1. **Bank Conflict Resolution (Padding):**\n * **The Fix:** Introduced a \"dummy\" padding of 1 element per row to the Shared Memory allocation (`stride = d + 1`).\n * **The Result:** This shifts the memory addresses such that threads in a warp access distinct memory banks simultaneously. This restored full SRAM bandwidth, moving from 32-cycle serial reads to 1-cycle parallel reads.\n\n2. **Precision Shift (FP32 $\\to$ FP16):**\n * Switched from `float` to `half` precision, accelerating computations. This also effectively doubled the available Shared Memory capacity per block (2 bytes vs. 4 per element), allowing for larger tile sizes.\n\n3. **Increased Occupancy:**\n * Increased Query Tile Size ($B_r$) from 64 to 128.\n * This ensures 4 warps run per block, allowing the CUDA scheduler to execute math instructions on one warp while others stall on global memory loads.\n\n\"image\"\n\n## Comparison to State of the Art (PyTorch SDPA)\n\nWhile `FlashAttention.cu` beats the Naive implementation, it runs at approximately **2%** of the speed of PyTorch's native `scaled_dot_product_attention`.\n\n**Why the massive gap between my kernel and SDPA?**\n1. **Hardware Utilization (Most important):**\n * **My Kernel:** Uses **CUDA Cores** (Scalar `hadd`/`hmul`). It computes matrix multiplications by iterating through vectors one element at a time.\n * **PyTorch SDPA:** Uses **Tensor Cores** (Matrix `hmma`). These specialized hardware units perform $4 \\times 4$ matrix multiplications in a single clock cycle, providing vastly higher theoretical throughput.\n2. **Memory Pipelining:**\n * **My Kernel:** Uses synchronous \"Stop-and-Go\" loading. (Load Tile $\\to$ Wait $\\to$ Compute).\n * **PyTorch SDPA:** Uses asynchronous copying (`cp.async`) and multi-stage pipelining. It loads the *next* tile from HBM into registers while simultaneously computing the *current* tile in SRAM.\n\n## Future Optimization Roadmap\nTo bridge the gap between this implementation and production kernels:\n\n1. **Tensor Cores (`nvcuda::wmma`):** Rewrite the inner dot-product loops to use Warp Matrix Multiply Accumulate instructions instead of scalar math.\n2. **Software Pipelining:** Implement double-buffering to fetch data for iteration $i+1$ while computing iteration $i$.\n3. **Warp-Level Primitives:** Replace block-level reductions with `__shfl_down_sync` for faster softmax calculation.\n\n## Usage\nThe kernel is wrapped for Python usage via `torch.utils.cpp_extension`.\n\n```python\nimport torch\nimport flash_attn_cuda # The compiled extension\n\n# Shapes: (N, d) - Must be contiguous and FP16\nN, d = 131072, 128\nQ = torch.randn(N, d, device='cuda', dtype=torch.float16)\nK = torch.randn(N, d, device='cuda', dtype=torch.float16)\nV = torch.randn(N, d, device='cuda', dtype=torch.float16)\n\n# Output: (N, d)\noutput = flash_attn_cuda.run_flash(Q, K, V)\n```","license":null},{"repo":"pengzhangzhi/flash-attention-with-bias-triton","url":"https://github.com/pengzhangzhi/flash-attention-with-bias-triton","readmeUrl":"https://raw.githubusercontent.com/pengzhangzhi/flash-attention-with-bias-triton/HEAD/README.md","snippet":"# Flash Attention w. Bias Implementation in Triton\n\nA Triton implementation of Flash Attention with differential bias support. This implementation is inspired by the\noriginal [Flash Attention paper](https://arxiv.org/abs/2205.14135).\n\n\"image\"\n\n## Features\n\n- ⚑ Efficient Flash Attention implementation with Triton\n- 🎯 Support for matrix-form bias in both forward and backward passes\n- πŸ“Š Built-in benchmarking against PyTorch, PyTorch Compile, xFormers and SDPA implementations\n- πŸ” Comprehensive testing for correctness verification\n- πŸ“ˆ Performance optimizations for common use cases\n\n## Requirements\n\n- Python\n- PyTorch\n- Triton\n- xFormers\n- CUDA-capable GPU (tested on A100)\n\n## Installation\n\n```bash\npip install torch triton xformers\n```\n\n## Quick Start\n\n```python\nimport torch\nfrom benckmark.attention_func import attention_triton\n\n# Create input tensors\nbatch_size, seq_len, n_heads, head_dim = 2, 1024, 8, 64\nq = torch.randn(batch_size, seq_len, n_heads, head_dim, device='cuda', dtype=torch.float16)\nk = torch.randn_like(q)\nv = torch.randn_like(q)\nbias = torch.randn(batch_size, n_heads, seq_len, seq_len, device='cuda', dtype=torch.float16)\n\n# Run Flash Attention\noutput = attention_triton(q, k, v, bias, False)\n```\n\n## Performance\n\n- The implementation shows significant speedups compared to standard PyTorch attention, especially for longer sequence\n lengths:\n\n![Forward Pass Performance](benckmark/attention-comparison-batch2-head4-d32-fwd.png)\n![Backward Pass Performance](benckmark/attention-comparison-batch2-head4-d32-bwd.png)\n\n- Efficiency Benchmark of Attention Implementations (Forward, Batch=2, Heads=4, HeadDim=32)\n\n| Length | Triton (FLOPS) | PyTorch (FLOPS) | PyTorch-Compile (FLOPS) | xFormers (FLOPS) | PyTorch-SDPA (FLOPS) |\n|--------|----------------|-----------------|-------------------------|------------------|----------------------|\n| 256 | 5.542667 | 1.574715 | 1.943982 | 3.198991 | 3.192787 |\n| 512 | 11.284882 | 2.762875 | 5.355243 | 7.125481 | 7.176634 |\n| 1024 | 23.356817 | 2.547090 | 5.691676 | 13.810132 | 13.839005 |\n| 2048 | 30.678131 | 2.247875 | 4.738464 | 19.586347 | 19.654297 |\n| 4096 | 33.154529 | 2.586121 | 4.071432 | 18.912378 | 18.881439 |\n| 8192 | 32.935548 | 2.822556 | 4.118749 | 18.189223 | 18.278633 |\n\n- Efficiency Benchmark of Attention Implementations (Backward, Batch=2, Heads=4, HeadDim=32)\n\n| Length | Triton (FLOPS) | PyTorch (FLOPS) | PyTorch-Compile (FLOPS) | xFormers (FLOPS) | PyTorch-SDPA (FLOPS) |\n|--------|----------------|-----------------|-------------------------|------------------|----------------------|\n| 256 | 0.073699 | 0.062677 | 0.076821 | 0.077729 | 0.168817 |\n| 512 | 0.371210 | 0.242403 | 0.348554 | 0.306856 | 1.007847 |\n| 1024 | 1.257951 | 1.036895 | 1.181298 | 1.738618 | 6.938603 |\n| 2048 | 4.089672 | 2.213153 | 4.083681 | 4.527996 | 8.918591 |\n| 4096 | 10.732837 | 2.395701 | 6.392552 | 13.702077 | 13.685270 |\n| 8192 | 13.440012 | 2.485934 | 6.257983 | 13.159468 | 13.134775 |\n\n## Usage Notes\n\n### Optimal Performance\n\nFor best performance:\n\n- Make sequence lengths a multiple of 2\n- Use head dimensions ≀ 128\n- Test with your specific data shapes before production use\n- Consider disabling autotune if experiencing errors\n\n### Known Limitations\n\n- Currently tested primarily on A100 GPUs\n- Head dimensions must be ≀ 128\n- Supports only fp16 and bf16 dtypes\n- Autotune can occasionally introduce race conditions\n\n## Running Tests and Benchmarks\n\n```bash\n# Run correctness tests\npython flashattn_triton.py\n\n# This will run:\n# 1. Correctness tests comparing against PyTorch, PyTorch Compile, xFormers and SDPA\n# 2. Performance benchmarks for various sequence lengths\n```\n\n## Implementation Details\n\nThe implementation includes several key optimizations:\n\n- Efficient memory access patterns\n- Block-wise computation for better cache utilization\n- Optimized backward pass with gradient computation for bias\n- Configurable block sizes and warps for different hardware\n\n## Contributing\n\nWhile this is primarily an educational resource, contributions are welcome! Please feel free to:\n\n- Report issues\n- Suggest improvements\n- Add documentation\n- Share benchmark results on different hardware\n\n## Acknowledgments\n\n- @triDao for the original Flash Attention implementation","license":"Apache-2.0"},{"repo":"ssiu/flash-attention-turing","url":"https://github.com/ssiu/flash-attention-turing","readmeUrl":"https://raw.githubusercontent.com/ssiu/flash-attention-turing/HEAD/README.md","snippet":"# FlashAttention Turing\n\nThis repository provides an implementation of [FlashAttention](https://github.com/Dao-AILab/flash-attention) for the Turing architecture. \n\n## Features\n\nSupports:\n\n - fwd and bwd\n - head dim 64, 128\n - causal mask\n - gqa\n - varlen\n\nDoes not support:\n\n - dropout\n - local mask\n - kv cache\n\n## Performance\n\nWe currently have benchmarks for T4.\n\n### Forward pass \n\nUp to 2.19x and 1.95x faster than PyTorch's [Attention](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) for non-causal and causal workloads.\n\nOn Turing GPUs, PyTorch's Attention calls Memory-Efficient Attention from [xformers](https://github.com/facebookresearch/xformers) in the backend.\n\nFor long sequences, the forward kernel reaches up to 66% compute throughput.\n\n\"Forward\n\"Forward\n\n### Backward pass \n\nThe backward pass is split into two kernels: one for `dQ` and one for `dK` and `dV`.\n\nUp to 1.35x and 1.51x faster than PyTorch's Attention for non-causal and causal workloads.\n\nFor long sequences, the backward kernels reach up to 49% compute throughput for `dK` and `dV`, and 45% for `dQ`.\n\n\"Backward\n\"Backward\n\n\n## How to use FlashAttention\nThe main functions implement scaled dot product attention: `softmax(Q @ K^T * softmax_scale) @ V`.\n\n```\nfrom flash_attention_interface import (\n flash_attn_func,\n flash_attn_kvpacked_func,\n flash_attn_qkvpacked_func,\n flash_attn_varlen_func,\n flash_attn_varlen_kvpacked_func,\n flash_attn_varlen_qkvpacked_func,\n)\n```\n\n\nThe arguments for these functions differ from the standard FlashAttention Python API because this implementation does not support every feature yet. See `turing/flash_attention_interface.py` for the full function signatures and parameter descriptions.\n\n\n## Requirements\nWe tested this implementation with:\n\n- CUDA 12.4\n- PyTorch 2.8.0 and 2.5.1\n\n## Build notes\n\nInstall with:\n\n```bash\ngit clone https://github.com/ssiu/flash-attention-turing\ncd /path/to/flash-attention-turing\npip install torch setuptools ninja wheel\npip install -v .\n```\n\nTo run the test suite, install the test dependencies:\n\n```bash\npip install pytest numpy pandas\npytest -q\n```\n\nIf you enable Excel debug dumps in `test_flash_attn.py`, also install:\n\n```bash\npip install openpyxl\n```","license":null},{"repo":"tile-ai/TileOPs","url":"https://github.com/tile-ai/TileOPs/issues/403","pageType":"issue","number":403,"title":"[FEAT][FLASH_ATTENTION] implement flash attention operator ...","snippet":"Implementation Notes. Flash attention uses online softmax and tiled accumulation to avoid materializing the full attention matrix. Decode kernels are memory ...","license":null},{"repo":"gitctrlx/flash-attention-tutorial","url":"https://github.com/gitctrlx/flash-attention-tutorial","readmeUrl":"https://raw.githubusercontent.com/gitctrlx/flash-attention-tutorial/HEAD/README.md","snippet":"# Flash Attention Tutorial\r\n\r\nThis tutorial is designed for beginners to understand and implement Flash Attention in PyTorch. We follow the structure of the document [\"From Online Softmax to FlashAttention\"](https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf) by Zihao Ye, explaining each concept step by step, deriving the formulas, and providing PyTorch code snippets. We'll build from basic self-attention to the full tiled Flash Attention, including both forward and backward passes for a complete, differentiable PyTorch implementation.\r\n\r\nWe assume basic knowledge of PyTorch and transformers. All code is tested for correctness and uses CPU/GPU-agnostic devices (tested on CPU for simplicity, but works on CUDA). We'll use small dimensions for examples but scale to larger ones in the final implementation.\r\n\r\n## Introduction\r\n\r\nFlash Attention is a memory-efficient and fast algorithm for computing self-attention in transformers. It avoids materializing large intermediate matrices (e.g., the attention matrix) in GPU global memory (High Bandwidth Memory, HBM). Instead, it employs tiling and online softmax techniques to fuse operations into a single kernel, leveraging fast GPU shared memory (SRAM). This reduces I/O overhead between slow HBM and fast SRAM, enabling longer sequences without out-of-memory errors.\r\n\r\nA key limitation in standard attention implementations is the small size of SRAM (typically ~100KB per streaming multiprocessor on modern GPUs), which cannot store the full $O(L^2)$ attention matrix for long sequences $L$. As a result, computations require multiple HBM accesses: writing intermediates to HBM, then reading them back for subsequent operations. This I/O bottleneck slows down performance, especially for large $L$. Flash Attention addresses this by recomputing intermediates on-the-fly in SRAM using tiled blocks, minimizing HBM accesses to sub-quadratic levels.\r\n\r\nThe core challenge is that softmax is not associative (unlike addition in matrix multiplication), complicating tiling. We'll derive how online softmax enables associative updates and build to a full implementation, highlighting reductions in HBM accesses at each stage.\r\n\r\n## 1. Standard Self-Attention\r\n\r\nSelf-attention computes weighted sums of values based on query-key similarities. For simplicity, we initially ignore batches, heads, masks, and scaling (added later).\r\n\r\nThe formula is:\r\n\r\n$$\r\nO = \\mathrm{softmax}(Q K^T) V\r\n$$\r\n\r\nwhere $Q, K, V, O \\in \\mathbb{R}^{L \\times D}$, $L$ is the sequence length, and $D$ is the head dimension. Softmax is applied row-wise.\r\n\r\nThe standard implementation factorizes as:\r\n\r\n$$\r\nX = Q K^T, \\quad A = \\mathrm{softmax}(X), \\quad O = A V\r\n$$\r\n\r\nThis requires storing $X$ and $A$ ($O(L^2)$ memory), which is inefficient for large $L$. In terms of memory accesses, it involves at least three HBM round-trips: (1) compute and write $X$ to HBM, (2) read $X$ from HBM for softmax and write $A$ to HBM, (3) read $A$ from HBM to compute $O$.\r\n\r\nA basic (non-Flash) PyTorch implementation with batches, heads, mask, and scaling:\r\n\r\n```python\r\nimport torch\r\nimport torch.nn as nn\r\nimport math # Use math for sqrt to align with PyTorch style\r\n\r\ndef standard_attention(Q, K, V, mask=None):\r\n # Inputs: Q, K, V shape (batch_size, num_heads, seq_len, head_dim)\r\n head_dim = Q.shape[-1]\r\n scale = 1 / math.sqrt(head_dim)\r\n # Scale queries\r\n Q_scaled = Q * scale\r\n # Compute logits: (b, h, l, l)\r\n logits = torch.matmul(Q_scaled, K.transpose(-2, -1))\r\n \r\n if mask is not None:\r\n # Mask: (b, l) -> (b, 1, 1, l)\r\n key_mask = mask.unsqueeze(1).unsqueeze(1)\r\n logits = torch.where(key_mask > 0, logits, float('-inf'))\r\n \r\n # Softmax to get attention weights\r\n attn_weights = nn.functional.softmax(logits, dim=-1)\r\n # Weighted sum\r\n output = torch.matmul(attn_weights, V)\r\n return output\r\n```\r\n\r\nExample usage:\r\n\r\n```python\r\n# Small example tensors\r\nbatch_size, num_heads, seq_len, head_dim = 1, 1, 4, 8\r\nQ = torch.randn(batch_size, num_heads, seq_len, head_dim)\r\nK = torch.randn(batch_size, num_heads, seq_len, head_dim)\r\nV = torch.randn(batch_size, num_heads, seq_len, head_dim)\r\nmask = torch.ones(batch_size, seq_len)\r\nO = standard_attention(Q, K, V, mask)\r\nprint(O.shape) # torch.Size([1, 1, 4, 8])\r\n```\r\n\r\nWhile matrix multiplication can be tiled (due to associativity), softmax cannot be directly tiled without recomputation techniques.\r\n\r\n## 2. Safe Softmax\r\n\r\nFor a vector $x = [x_1, \\dots, x_N]$:\r\n\r\n$$\r\n\\mathrm{softmax}(x)_i = \\frac{e^{x_i}}{\\sum_{j=1}^N e^{x_j}}\r\n$$\r\n\r\nLarge $x_i$ can cause numerical overflow (e.g., in float16). Safe softmax subtracts the row max $m = \\max_j x_j$:\r\n\r\n$$\r\n\\mathrm{softmax}(x)_i = \\frac{e^{x_i - m}}{\\sum_{j=1}^N e^{x_j - m}}\r\n$$\r\n\r\nThis is a 3-pass algorithm, requiring three HBM reads if $x$ does not fit in SRAM: (1) find max, (2) compute exp and sum, (3) normalize. Each pass reloads $x$ from HBM.\r\n\r\nPseudo-code:\r\n\r\n```text\r\nm = -∞\r\nfor i = 1 to N:\r\n m = max(m, x_i) # Pass 1: global max\r\n\r\nd = 0\r\nfor i = 1 to N:\r\n d += exp(x_i - m) # Pass 2: denominator\r\n\r\nfor i = 1 to N:\r\n a_i = exp(x_i - m) / d # Pass 3: normalize\r\n```\r\n\r\nPyTorch implementation (row-wise for 2D tensor):\r\n\r\n```python\r\ndef safe_softmax(x):\r\n # x: (batch_size, seq_len)\r\n # Pass 1: Compute row maxes\r\n m = torch.max(x, dim=-1, keepdim=True)[0]\r\n # Pass 2: Exps and sum (denominator)\r\n exp_shifted = torch.exp(x - m)\r\n denom = torch.sum(exp_shifted, dim=-1, keepdim=True)\r\n # Pass 3: Normalize\r\n return exp_shifted / denom\r\n\r\n# Example\r\nx = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\r\nprint(safe_softmax(x))\r\n```\r\n\r\nThis is inefficient for large sequences, as multiple passes increase HBM I/O.\r\n\r\n## 3. Online Softmax\r\n\r\nTo reduce passes, introduce a surrogate $d_i' = \\sum_{j=1}^i e^{x_j - m_i}$ with the recurrence:\r\n\r\n$$\r\nd_i' = d_{i-1}' \\cdot e^{m_{i-1} - m_i} + e^{x_i - m_i}\r\n$$\r\n\r\nwhere $m_i = \\max(m_{i-1}, x_i)$. At $i=N$, $d_N' = \\sum e^{x_j - m_N}$. This enables a 2-pass algorithm (2 HBM reads if not in SRAM): one for statistics, one for normalization.\r\n\r\nPseudo-code:\r\n\r\n```text\r\nm = -∞\r\nd' = 0\r\nfor i = 1 to N: # Pass 1: Compute m_i, d_i'\r\n m_new = max(m, x_i)\r\n d' = d' * exp(m - m_new) + exp(x_i - m_new)\r\n m = m_new\r\n\r\nfor i = 1 to N: # Pass 2: Normalize\r\n a_i = exp(x_i - m) / d'\r\n```\r\n\r\nPyTorch implementation (row-wise):\r\n\r\n```python\r\ndef online_softmax(x):\r\n # x: (batch_size, seq_len)\r\n batch_size, seq_len = x.shape\r\n # Initialize\r\n m = torch.full((batch_size, 1), float('-inf'), device=x.device)\r\n d_prime = torch.zeros((batch_size, 1), device=x.device)\r\n # Pass 1: Online update (loop for clarity; vectorize in practice)\r\n for i in range(seq_len):\r\n x_i = x[:, i:i+1]\r\n m_new = torch.maximum(m, x_i)\r\n exp_diff = torch.exp(m - m_new)\r\n exp_term = torch.exp(x_i - m_new)\r\n d_prime = d_prime * exp_diff + exp_term\r\n m = m_new\r\n # Pass 2: Compute probabilities\r\n probs = torch.exp(x - m) / d_prime\r\n return probs\r\n\r\n# Example\r\nx = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\r\nprint(online_softmax(x)) # Matches safe_softmax\r\n```\r\n\r\nThis reduces I/O to two HBM accesses but still requires multiple passes.\r\n\r\n## 4. Flash Attention (Without Tiling)\r\n\r\nIn attention, we only need $O = A V$, not $A$ itself. For each row $k$, derive a one-pass recurrence using surrogates.\r\n\r\nNotation:\r\n\r\n- $x_j = Q[k, :] \\cdot K[j, :]$ (scalar logit for position $j$)\r\n- $o_i = \\sum_{j=1}^i a_j V[j, :]$ (partial output up to $i$)\r\n\r\nSurrogate $o_i' = \\sum_{j=1}^i e^{x_j - m_i} V[j, :] / d_i'$, with recurrence:\r\n\r\n$$\r\no_i' = o_{i-1}' \\cdot \\frac{d_{i-1}' e^{m_{i-1} - m_i}}{d_i'} + \\frac{e^{x_i - m_i}}{d_i'} V[i, :]\r\n$$\r\n\r\nThis enables a single loop (1 HBM access per block if tiled later), updating statistics on-the-fly.\r\n\r\nPseudo-code for one row:\r\n\r\n```text\r\nm = -∞\r\nd' = 0\r\no' = zeros(D)\r\nfor j = 1 to N:\r\n x_j = Q[k, :] @ K[j, :]\r\n m_new = max(m, x_j)\r\n exp_diff = exp(m - m_new)\r\n exp_term = exp(x_j - m_new)\r\n d'_new = d' * exp_diff + exp_term\r\n o' = o' * (d' * exp_diff / d'_new) + (exp_term / d'_new) * V[j, :]\r\n m = m_new\r\n d' = d'_new\r\nO[k, :] = o'\r\n```\r\n\r\nPyTorch implementation (batched, multi-head, no tiling; for illustration only):\r\n\r\n```python\r\ndef simple_flash_forward(Q, K, V, mask=None):\r\n # Inputs: (b, h, l, d); no tiling (inefficient for large l)\r\n b, h, l, d = Q.shape\r\n scale = 1 / math.sqrt(d)\r\n O = torch.zeros_like(Q)\r\n # Loop over batch*heads\r\n for batch_head_idx in range(b * h):\r\n # Flatten to (l, d)\r\n q_bh = Q.view(b * h, l, d)[batch_head_idx]\r\n k_bh = K.view(b * h, l, d)[batch_head_idx]\r\n v_bh = V.view(b * h, l, d)[batch_head_idx]\r\n m_bh = mask.view(b, l)[batch_head_idx // h] if mask is not None else torch.ones(l, device=Q.device)\r\n # Loop over query rows\r\n for row_idx in range(l):\r\n q_row = q_bh[row_idx] * scale # (d,)\r\n row_max = float('-inf')\r\n row_denom_prime = torch.tensor(0.0, device=Q.device)\r\n row_output_prime = torch.zeros(d, device=Q.device)\r\n # Loop over key columns (online update)\r\n for col_idx in range(l):\r\n if m_bh[col_idx] <= 0:\r\n continue\r\n logit = torch.dot(q_row, k_bh[col_idx]) # scalar\r\n new_max = max(row_max, logit)\r\n exp_diff = torch.exp(row_max - new_max)\r\n exp_term = torch.exp(logit - new_max)\r\n new_denom_prime = row_denom_prime * exp_diff + exp_term\r\n new_output_prime = row_output_prime * (row_denom_prime * exp_diff / new_denom_prime) + \\\r\n (exp_term / new_denom_prime) * v_bh[col_idx]\r\n row_max = new_max\r\n row_denom_prime = new_denom_prime\r\n row_output_prime = new_output_prime\r\n # Assign to output\r\n O.view(b * h, l, d)[batch_head_idx, row_idx] = row_output_prime\r\n return O\r\n```\r\n\r\n(Note: This scalar loop is for clarity and not efficient; tiling fuses operations in the next section.)\r\n\r\n## 5. Flash Attention (With Tiling)\r\n\r\nFor large $L$, tile $K$ and $V$ into blocks of size $B$ (fitting in SRAM). For each query row/block, load tiles, compute local max/exp/sum in SRAM, and update global statistics associatively. This ensures a single effective pass, with HBM accesses scaling as $O(L^2 D^2 / M)$ (where $M$ is SRAM size), far fewer than standard attention's $O(L^2 + L D)$.\r\n\r\nPseudo-code for one row (tiled):\r\n\r\n```text\r\nm = -∞\r\nd' = 0\r\no' = zeros(D)\r\nfor tile_idx = 1 to num_tiles:\r\n tile_start = (tile_idx - 1) * B\r\n tile_end = tile_start + B\r\n K_tile = K[tile_start:tile_end, :]\r\n V_tile = V[tile_start:tile_end, :]\r\n logits_tile = Q[k, :] @ K_tile.T # (B,)\r\n local_max = max(logits_tile)\r\n global_max = max(m, local_max)\r\n exp_global_diff = exp(m - global_max)\r\n exp_local_terms = exp(logits_tile - global_max)\r\n local_sum = sum(exp_local_terms)\r\n d'_new = d' * exp_global_diff + local_sum\r\n weighted_v_sum = (exp_local_terms / d'_new) @ V_tile # (D,)\r\n o' = o' * (d' * exp_global_diff / d'_new) + weighted_v_sum\r\n m = global_max\r\n d' = d'_new\r\nO[k, :] = o'\r\n```\r\n\r\nFull tiled forward in PyTorch (returns $O$, row sums $l$, row maxes $m$ for backward):\r\n\r\n```python\r\nBLOCK_SIZE = 1024 # Adjust based on SRAM\r\nNEG_INF = float('-inf')\r\nEPS = 1e-6 # For numerical stability\r\n\r\ndef flash_attention_forward(Q, K, V, mask=None):\r\n # Inputs: (b, h, l, d)\r\n b, h, l, d = Q.shape\r\n device = Q.device\r\n O = torch.zeros_like(Q)\r\n row_sums = torch.zeros(b, h, l, 1, device=device) # l (denominators)\r\n row_maxes = torch.full((b, h, l, 1), NEG_INF, device=device) # m\r\n \r\n # Block sizes (Q rows, KV columns)\r\n q_block_size = min(BLOCK_SIZE, l)\r\n kv_block_size = BLOCK_SIZE\r\n \r\n # Split into blocks\r\n Q_blocks = torch.split(Q, q_block_size, dim=2)\r\n K_blocks = torch.split(K, kv_block_size, dim=2)\r\n V_blocks = torch.split(V, kv_block_size, dim=2)\r\n mask = torch.ones(b, l, device=device) if mask is None else mask\r\n mask_blocks = torch.split(mask, kv_block_size, dim=1)\r\n \r\n num_q_blocks, num_kv_blocks = len(Q_blocks), len(K_blocks)\r\n \r\n # Split outputs for accumulation\r\n O_blocks = list(torch.split(O, q_block_size, dim=2))\r\n row_sums_blocks = list(torch.split(row_sums, q_block_size, dim=2))\r\n row_maxes_blocks = list(torch.split(row_maxes, q_block_size, dim=2))\r\n \r\n scale = 1 / math.sqrt(d)\r\n \r\n # Outer loop over KV tiles\r\n for kv_idx in range(num_kv_blocks):\r\n K_tile = K_blocks[kv_idx]\r\n V_tile = V_blocks[kv_idx]\r\n mask_tile = mask_blocks[kv_idx].unsqueeze(1).unsqueeze(1) # (b, 1, 1, block_size)\r\n \r\n # Inner loop over Q tiles\r\n for q_idx in range(num_q_blocks):\r\n Q_tile = Q_blocks[q_idx]\r\n curr_O = O_blocks[q_idx]\r\n curr_row_sums = row_sums_blocks[q_idx]\r\n curr_row_maxes = row_maxes_blocks[q_idx]\r\n \r\n # Compute block logits\r\n Q_tile_scaled = Q_tile * scale\r\n logits_block = torch.matmul(Q_tile_scaled, K_tile.transpose(-2, -1)) # (b, h, q_block, kv_block)\r\n logits_block = torch.where(mask_tile > 0, logits_block, NEG_INF)\r\n \r\n # Local max and exp\r\n local_max = torch.max(logits_block, dim=-1, keepdim=True)[0]\r\n exp_logits = torch.exp(logits_block - local_max)\r\n exp_logits = torch.where(mask_tile > 0, exp_logits, 0.0)\r\n \r\n # Local sum\r\n local_sum = torch.sum(exp_logits, dim=-1, keepdim=True) + EPS\r\n \r\n # Update global max and sum\r\n new_max = torch.maximum(curr_row_maxes, local_max)\r\n new_row_sums = torch.exp(curr_row_maxes - new_max) * curr_row_sums + \\\r\n torch.exp(local_max - new_max) * local_sum\r\n \r\n # Update output\r\n exp_max_diff = torch.exp(curr_row_maxes - new_max)\r\n exp_local_diff = torch.exp(local_max - new_max)\r\n weighted_v = torch.matmul(exp_logits, V_tile) # (b, h, q_block, d)\r\n O_blocks[q_idx] = (curr_row_sums * exp_max_diff / new_row_sums) * curr_O + \\\r\n (exp_local_diff / new_row_sums) * weighted_v\r\n \r\n # Store updated stats\r\n row_sums_blocks[q_idx] = new_row_sums\r\n row_maxes_blocks[q_idx] = new_max\r\n \r\n # Concatenate blocks\r\n O = torch.cat(O_blocks, dim=2)\r\n row_sums = torch.cat(row_sums_blocks, dim=2)\r\n row_maxes = torch.cat(row_maxes_blocks, dim=2)\r\n return O, row_sums, row_maxes\r\n```\r\n\r\n## 6. Flash Attention Backward Pass\r\n\r\nThe backward pass computes gradients $dQ$, $dK$, $dV$ without storing the $O(L^2)$ attention matrix. It recomputes block-wise in SRAM using saved row maxes $m$ and sums $l$ from forward, with similar tiling. This maintains memory efficiency, with HBM accesses $O(L^2 D^2 / M)$.\r\n\r\nKey derivations (from original paper):\r\n\r\n- $dV = A^T dO$\r\n- For softmax grad: $dS_{ij} = A_{ij} (dO_i^T V_j - D_i)$ where $D_i = dO_i^T O_i$\r\n- $dQ_i = dS_{i:} K^T$, $dK_j = dS_{:j}^T Q$\r\n\r\nPseudo-code (tiled, simplified from paper):\r\n\r\n```text\r\nfor kv_tile_idx = 1 to num_kv_tiles:\r\n Load K_tile, V_tile\r\n Init temp_dK_tile = zeros, temp_dV_tile = zeros\r\n for q_tile_idx = 1 to num_q_tiles:\r\n Load Q_tile, O_tile, dO_tile, row_sums_tile, row_maxes_tile\r\n Compute logits_block = scale * Q_tile @ K_tile.T\r\n Apply mask\r\n probs_block = exp(logits_block - row_maxes_tile) / row_sums_tile\r\n temp_dV_tile += probs_block.T @ dO_tile\r\n dP_block = dO_tile @ V_tile.T\r\n D_tile = row_sum(dO_tile * O_tile) # (q_block,)\r\n dS_block = probs_block * (dP_block - D_tile)\r\n dQ_tile += dS_block @ K_tile\r\n temp_dK_tile += dS_block.T @ Q_tile\r\n Write dK_tile, dV_tile\r\n```\r\n\r\nPyTorch implementation (tiled backward; assumes no dropout/mask for simplicity; extend as needed):\r\n\r\n```python\r\ndef flash_attention_backward(dO, Q, K, V, O, row_sums, row_maxes, mask=None):\r\n # Inputs: dO, Q, K, V, O (b, h, l, d); row_sums, row_maxes (b, h, l, 1)\r\n b, h, l, d = Q.shape\r\n device = Q.device\r\n scale = 1 / math.sqrt(d)\r\n \r\n dQ = torch.zeros_like(Q)\r\n dK = torch.zeros_like(K)\r\n dV = torch.zeros_like(V)\r\n \r\n q_block_size = min(BLOCK_SIZE, l)\r\n kv_block_size = BLOCK_SIZE\r\n \r\n # Splits (reuse forward logic)\r\n Q_blocks = torch.split(Q, q_block_size, dim=2)\r\n K_blocks = torch.split(K, kv_block_size, dim=2)\r\n V_blocks = torch.split(V, kv_block_size, dim=2)\r\n O_blocks = torch.split(O, q_block_size, dim=2)\r\n dO_blocks = torch.split(dO, q_block_size, dim=2)\r\n row_sums_blocks = torch.split(row_sums, q_block_size, dim=2)\r\n row_maxes_blocks = torch.split(row_maxes, q_block_size, dim=2)\r\n dQ_blocks = list(torch.split(dQ, q_block_size, dim=2))\r\n mask = torch.ones(b, l, device=device) if mask is None else mask\r\n mask_blocks = torch.split(mask, kv_block_size, dim=1)\r\n \r\n num_q_blocks, num_kv_blocks = len(Q_blocks), len(K_blocks)\r\n \r\n # Outer loop over KV tiles\r\n for kv_idx in range(num_kv_blocks):\r\n K_tile = K_blocks[kv_idx]\r\n V_tile = V_blocks[kv_idx]\r\n mask_tile = mask_blocks[kv_idx].unsqueeze(1).unsqueeze(1)\r\n \r\n temp_dK = torch.zeros_like(K_tile)\r\n temp_dV = torch.zeros_like(V_tile)\r\n \r\n # Inner loop over Q tiles\r\n for q_idx in range(num_q_blocks):\r\n Q_tile = Q_blocks[q_idx]\r\n O_tile = O_blocks[q_idx]\r\n dO_tile = dO_blocks[q_idx]\r\n row_sums_tile = row_sums_blocks[q_idx]\r\n row_maxes_tile = row_maxes_blocks[q_idx]\r\n \r\n # Recompute block probs\r\n Q_tile_scaled = Q_tile * scale\r\n logits_block = torch.matmul(Q_tile_scaled, K_tile.transpose(-2, -1))\r\n logits_block = torch.where(mask_tile > 0, logits_block, NEG_INF)\r\n probs_block = torch.exp(logits_block - row_maxes_tile) / (row_sums_tile + EPS)\r\n \r\n # dV accumulation\r\n temp_dV += torch.matmul(probs_block.transpose(-2, -1), dO_tile)\r\n \r\n # dP and D\r\n dP_block = torch.matmul(dO_tile, V_tile.transpose(-2, -1))\r\n D_tile = torch.sum(dO_tile * O_tile, dim=-1, keepdim=True) # (b, h, q_block, 1)\r\n \r\n # dS\r\n dS_block = probs_block * (dP_block - D_tile)\r\n \r\n # dQ accumulation\r\n dQ_blocks[q_idx] += torch.matmul(dS_block, K_tile)\r\n \r\n # dK accumulation\r\n temp_dK += torch.matmul(dS_block.transpose(-2, -1), Q_tile)\r\n \r\n # Write gradients\r\n dK_blocks = torch.split(dK, kv_block_size, dim=2)\r\n dV_blocks = torch.split(dV, kv_block_size, dim=2)\r\n dK_blocks[kv_idx].copy_(temp_dK)\r\n dV_blocks[kv_idx].copy_(temp_dV)\r\n \r\n dQ = torch.cat(dQ_blocks, dim=2)\r\n dK = torch.cat(torch.split(dK, kv_block_size, dim=2), dim=2) # Reassemble if needed\r\n dV = torch.cat(torch.split(dV, kv_block_size, dim=2), dim=2)\r\n return dQ, dK, dV\r\n```\r\n\r\n(Note: This is a simplified version without dropout; refer to the original paper for full details including dropout regeneration.)\r\n\r\nTo make it differentiable, wrap in a custom `torch.autograd.Function`.\r\n\r\n## Reference\r\n\r\n- **Tutorial Notes**: [From Online Softmax to FlashAttention](https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf) by Zihao Ye\r\n- **Original Paper**: [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135) by Tri Dao et al.\r\n- **Official Repository**: [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)\r\n- **PyTorch Example**: [shreyansh26/FlashAttention-PyTorch](https://github.com/shreyansh26/FlashAttention-PyTorch)\r\n\r\n## License\r\n\r\nMIT License","license":"MIT"},{"repo":"dao-ailab/flash-attention","url":"https://github.com/dao-ailab/flash-attention/issues/2439","pageType":"merged_pr","number":2439,"segmentCount":3,"title":"Conversation","snippet":"Add dropout support to CuTe DSL attention kernels\n` …\n\n`\n 0fae293\n`\n\n````\nImplements position-keyed Philox 4Γ—32Γ—10 dropout for the CuTe DSL forward\nand backward paths on SM80/SM90/SM100/SM120. Dropout mask is derived from\n(row, col) coordinates rather than thread/element index so that forward\nand backward independently regenerate the same mask via the same Philox\ncounter, which is what makes the gradient correct without storing the\nmask tensor.\n\n- `flash_attn/cute/dropout.py` (new) β€” Philox 4Γ—32Γ—10 in CuTe DSL, mask\n precompute/apply split for fused backward, per-element fallback for\n SM90 warp-group MMA whose threadβ†’element mapping differs from SM80's\n warp-level MMA. Uses FA2-style MMA-layout keying via `logical_divide`\n on the accumulator's N mode: each Philox call produces 8 Γ— 16-bit\n masks covering 8 elements, using all 128 random bits (addresses\n tridao's efficiency feedback).\n- `flash_attn/cute/{flash_fwd,flash_bwd}.py` + per-arch SM80/SM90/SM100\n variants β€” wire `p_dropout` through forward/backward, apply per-block\n inside the K/V loop; forward and backward tile sizes are matched when\n dropout is enabled so the Philox counter aligns.\n- `flash_attn/cute/interface.py` β€” `flash_attn_func(...)` and\n `flash_attn_varlen_func(...)` take `dropout_p` + optional\n `dropout_seed`; seed is auto-generated from `torch.randint` if not\n supplied; (lo, hi) halves of the 64-bit seed reach the kernel via\n compile_args; backward retrieves seed from autograd ctx.\n- `flash_attn/cute/utils.py` β€” `get_smem_store_atom()` arch guard\n extended to `arch >= 120` so SM120 doesn't pick up SM90's\n `StMatrix8x8x16bOp` warp-group instruction (it has warp-level MMA,\n not WGMMA β€” the output store was producing garbage before this fix).\n- Tests:\n - `tests/cute/test_dropout_grad_match.py` β€” tridao-style flash vs\n reference gradient match using the standard FA error metric\n (|flash βˆ’ ref_f32| ≀ 2 Β· |ref_bf16 βˆ’ ref_f32| + atol). 7 tests\n covering causal/non-causal, GQA, SM80/SM90/SM120 dispatch.\n - `tests/cute/test_dropout_correctness.py`, `test_dropout_e2e.py` β€”\n deterministic-seed, dropout-p-zero-equals-no-dropout, output-changes\n invariants, backward grad presence.\n - `test_dropout_standalone.py` (repo root) β€” minimal repro that runs\n without conftest pytest-ini dependencies.\n- `AI/DROPOUT_IMPLEMENTATION_NOTES.md` β€” design notes on position-keyed\n Philox, the SM90 per-element fallback, and the MMA-layout keying\n decision.\n\n1. d0d29cc β€” base dropout implementation (forward + backward, all arches)\n2. b4c98cb β€” forward/backward dropout mask consistency test\n3. 4d8d71c β€” SM120 StMatrix8x8x16bOp fix in utils.py + gradient match test\n4. b498ca9 β€” efficient Philox: 8 masks per call, 16-bit thresholds,\n group-cached (addresses tridao's \"4Γ—32 bits but only 16 used\")\n5. 1ff618d β€” FA2-style MMA-layout dropout keying + correct backward\n gradient\n6. b111322 β€” SM90 backward undefined `acc_shape_SdP` + ruff format\n7. e4ecab8 β€” SM90 per-element position-keyed Philox for WGMMA (warp-group\n thread layout differs from warp-level MMA)\n8. 3cab834 β€” per-element Python Philox reference for SM90/SM100 testing\n9. fa2ed75 β€” precompute/apply split, fused backward, tile tuning\n10. 1671230 β€” rename `apply_dropout_mask_sm100` β†’ `apply_dropout_mask_per_element`\n11. 70e7f81 β€” ruff formatting in dropout.py and flash_bwd_sm90.py\n\n`test_dropout_standalone.py` against the rebased branch on current\n`main` (`0bbb25a`):\n\n```\nPASS test_fwd_no_dropout baseline forward works\nPASS test_fwd_dropout_zero dropout_p=0.0 bytewise matches no-dropout\nPASS test_fwd_dropout dropout_p=0.5 changes output (max diff 0.9336)\nPASS test_fwd_deterministic same seed β†’ identical output\nPASS test_bwd_dropout q, k, v gradients computed correctly\n```\n\n`pre-commit run --files flash_attn/cute/*.py tests/cute/*.py\ntest_dropout_standalone.py` β†’ ruff check + ruff format both pass.\n\nEarlier on-hardware coverage (per the PR comment thread): johnnynunez\nvalidated SM90 on GH200 and SM110 on AGX Thor; SM120 validated on\nDGX Spark GB10. The SM90 per-element fallback (per Philox call) was\nadded after johnnynunez's GH200 run identified that the warp-group\nMMA register layout requires per-element rather than per-tile keying.\n\nCo-Authored-By: Claude Opus 4.7 (1M context) \n````\n\nthe\nfeat/cute-dsl-dropout\nbranch\nfrom\n[`e48c5d8`](https://github.com/Dao-AILab/flash-attention/commit/e48c5d86ce795b6ca8d9af4a887837e92c41761b) to\n[`0fae293`](https://github.com/Dao-AILab/flash-attention/commit/0fae2934ccf023f13a584ebde8453df5345640b9) [Compare](https://github.com/Dao-AILab/flash-attention/compare/e48c5d86ce795b6ca8d9af4a887837e92c41761b..0fae2934ccf023f13a584ebde8453df5345640b9) [last weekJune 2, 2026 19:56](https://github.com/Dao-AILab/flash-attention/pull/2439#event-26255307920)\n\nContributorAuthor\n\n| |\n| --- |\n| Rebased onto current `main` (resolves the merge conflict). All conflicts were API-compatibility with no change to dropout numerics:
- Adopted main's `cute.make_fragment` β†’ `cute.make_rmem_tensor` rename, including one call our backward path adds in `flash_bwd_sm100.py`.
- Kept both main's new MLA-specialization block and the dropout-seed init in `flash_attn_func` / `flash_attn_varlen_func`.
I also dropped two files that shouldn't have been in the PR β€” a root-level `test_dropout_standalone.py` scratch script and `AI/DROPOUT_IMPLEMENTATION_NOTES.md`. Coverage stays in `tests/cute/test_dropout_{correctness,e2e,grad_match}.py`.
The dropout code paths themselves are unchanged from the version the grad-match suite already passed, so the rebase carries no numerical delta. |\n\nThis file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.\n[Learn more about bidirectional Unicode characters](https://github.co/hiddenchars)\n\n[Show hidden characters](https://github.com/Dao-AILab/flash-attention/pull/2439)","license":"BSD-3-Clause"},{"repo":"alexzhang13/flashattention2-custom-mask","url":"https://github.com/alexzhang13/flashattention2-custom-mask","readmeUrl":"https://raw.githubusercontent.com/alexzhang13/flashattention2-custom-mask/HEAD/README.md","snippet":"# FlashAttention2 with Custom Masks 🎭\n**Note: This is an unofficial implementation of FlashAttention2.**\n\nFor efficiency purposes, the standard implementations of FlashAttention currently do not support **arbitrary custom masks**. \nTheir implementation of specific masks like causal masking for language modeling are implemented using branch logic to save memory. This repository is just a modified version of the tutorial Triton implementation of FlashAttention2 that allows the user\nto define a (batch of) custom mask. It modifies both the forward and backwards pass to handle custom masking (you can define a different mask per head and batch).\n \nOriginal Triton code: [https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html](https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html)\n\nSee the original thread: [https://github.com/Dao-AILab/flash-attention/issues/352](https://github.com/Dao-AILab/flash-attention/issues/352)\n\n## Quick Install\nCreate a Python environment (>=3.8) and install through pip:\n```\npip install flashattention2-custom-mask\n```\n\n## Example Setup\nThe relevant libraries needed to use the custom-mask FlashAttention2 kernel are below:\n```\npip install triton>=3.0.0\npip install torch\n```\n\n#### For Viewing Benchmarking Results\nOther libraries for evaluating the performance of the models is below. These are primarily for `test_benchmark.py`, which verifies the correctness of the implementation.\n```\npip install pytest\npip install matplotlib\npip install pandas\n```\nTo compare with the official FlashAttention and `xformers.ops.memory_efficient_attention` implementations, make sure to install both libraries separately (follow the instructions on these repositories).\n```\npip install flash-attn --no-build-isolation\npip3 install -U xformers --index-url https://download.pytorch.org/whl/cu121\n```\n\n## Testing Correctness\nThere are two `pytest` functions in `test_benchmark.py`, one that tests whether a reference implementation of multi-head attention with a causal mask matches the Triton version in both the forward pass and backwards pass gradients. The second tests whether the same implementation with **random** masks matches the Triton version. You can modify these tests to do more rigorous correctness tests and check with `pytest`.\n\n## Simple Example\nYou can insert this module into your standard attention pipeline.\n```python\nfrom fa2_custom_mask import flash_attention_custom_mask\n\nB, H, L, D = 4, 16, 4096, 64\nsm_scale = 1 / (D ** 0.5)\n\nfp32_q = torch.randn(B, H, L, D).float().cuda()\nfp32_k = torch.randn(B, H, L, D).float().cuda()\nfp32_v = torch.randn(B, H, L, D).float().cuda()\nmask = torch.randint(0, 2, (B, 1, L, L)).int().cuda()\nmask = torch.broadcast_to(mask, (B, H, L, L))\n\nout = flash_attention_custom_mask(fp32_q, fp32_k, fp32_v, mask=mask, sm_scale=sm_scale)\n...\nout.backward(loss)\n```\n\n## Benchmarking\nSimple benchmark against the base Triton implementation. In our custom mask version, we pass in the canonical causal mask as input (hence storing in global device memory). Running `test_benchmark.py`,\nwith batch size=4, # heads=16, hidden dim=64, and sequence length `N_CTX` ranging from 256 to 16384 in powers of 2. You can replicate the experiments by running\n```\npytest\npython test_benchmark.py\n```\n\n#### Causal Masks and No Masks Comparisons \nWe compare against the original experiments and original implementation, as well as the official FlashAttention and xformers implementation (note: there seems to be a versioning issue, so it's using a different implementation. I corrected the version in the later benchmarking experiments). \n![causal and no masking with flash attn](./data/results-causal-fa.png)\n \n#### Causal Masks and No Masks Comparisons (with Correct xfrormers version)\nWe compare against the original experiments and original implementation, as well as the xformers implementation. Notably, the original implementation does well for causal masking because of some pipelining tricks and ability to not have to store masks.\n![causal and no masking](./data/results-causal.png)\n#### Custom Masking Comparison\nWe compare directly to the [xformers memory efficient attention](https://facebookresearch.github.io/xformers/components/ops.html) which allows for custom masking. We generate random masks (fixed across the head dimension).\n![custom masking](./data/results-random.png)\n\n\n## Notes and Bugs\n1. This implementation only works on Ampere devices and up. I originally tried running it on a V100 (Volta) and it failed. \n2. You need to be on `triton>=3.0.0`, or it'll complain about permutation indices on the value vector pointer. The `torch` and `flash-attn` libraries may force you to install `triton=2.x.x`, but you can just re-install `triton>=3.0.0` and it should work. I may fix this manually in the future.\n * This is oddly specific, but I'm not able to have `flash-attn` and `xformers` at the same time. I had to run them separately and generate the plots.\n3. TODO: Add benchmarking for peak memory consumption and other efficiency metrics.\n\nIf time permits, I'm interested in making this implementation generalizable / changing the CUDA implementation for FA3 (if it's necessary of course). I also probably will run some more realistic workloads and see what happens.","license":"Apache-2.0"},{"repo":"dl-attention/flash-attention-1","url":"https://github.com/dl-attention/flash-attention-1","readmeUrl":"https://raw.githubusercontent.com/dl-attention/flash-attention-1/HEAD/README.md","snippet":"# FlashAttention\nThis repository provides the official implementation of FlashAttention and\nFlashAttention-2 from the\nfollowing papers.\n\n**FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness** \nTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré \nPaper: https://arxiv.org/abs/2205.14135 \nIEEE Spectrum [article](https://spectrum.ieee.org/mlperf-rankings-2022) about our submission to the MLPerf 2.0 benchmark using FlashAttention.\n![FlashAttention](assets/flashattn_banner.jpg)\n\n**FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning** \nTri Dao\n\nPaper: https://tridao.me/publications/flash2/flash2.pdf\n\n![FlashAttention-2](assets/flashattention_logo.png)\n\n\n## Usage\n\nWe've been very happy to see FlashAttention being widely adopted in such a short\ntime after its release. This [page](https://github.com/Dao-AILab/flash-attention/blob/main/usage.md)\ncontains a partial list of places where FlashAttention is being used.\n\nFlashAttention and FlashAttention-2 are free to use and modify (see LICENSE).\nPlease cite and credit FlashAttention if you use it.\n\n## Installation and features\n\nRequirements:\n- CUDA 11.4 and above.\n- PyTorch 1.12 and above.\n\nWe recommend the\n[Pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)\ncontainer from Nvidia, which has all the required tools to install FlashAttention.\n\nTo install:\n1. Make sure that PyTorch is installed.\n2. Make sure that `packaging` is installed (`pip install packaging`)\n3. Make sure that `ninja` is installed and that it works correctly (e.g. `ninja\n--version` then `echo $?` should return exit code 0). If not (sometimes `ninja\n--version` then `echo $?` returns a nonzero exit code), uninstall then reinstall\n`ninja` (`pip uninstall -y ninja && pip install ninja`). Without `ninja`,\ncompiling can take a very long time (2h) since it does not use multiple CPU\ncores. With `ninja` compiling takes 3-5 minutes on a 64-core machine.\n4. Then:\n```sh\npip install flash-attn --no-build-isolation\n```\nAlternatively you can compile from source:\n```sh\npython setup.py install\n```\n\nIf your machine has less than 96GB of RAM and lots of CPU cores, `ninja` might\nrun too many parallel compilation jobs that could exhaust the amount of RAM. To\nlimit the number of parallel compilation jobs, you can set the environment\nvariable `MAX_JOBS`:\n```sh\nMAX_JOBS=4 pip install flash-attn --no-build-isolation\n```\n\nInterface: `src/flash_attention_interface.py`\n\nFlashAttention-2 currently supports:\n1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). Support for Turing\n GPUs (T4, RTX 2080) is coming soon, please use FlashAttention 1.x for Turing\n GPUs for now.\n2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).\n3. All head dimensions up to 256. Head dim > 192 backward requires A100/A800 or H100/H800.\n\n\n## How to use FlashAttention\n\nThe main functions implement scaled dot product attention (softmax(Q @ K^T *\nsoftmax_scale) @ V):\n```python\nfrom flash_attn import flash_attn_qkvpacked_func, flash_attn_func\n```\n\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False):\n\"\"\"dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nArguments:\n qkv: (batch_size, seqlen, 3, nheads, headdim)\n dropout_p: float. Dropout probability.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\nReturn:\n out: (batch_size, seqlen, nheads, headdim).\n\"\"\"\n```\n\n```python\nflash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False):\n\"\"\"dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in KV must be divisible by the number of heads in Q.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\nArguments:\n q: (batch_size, seqlen, nheads, headdim)\n k: (batch_size, seqlen, nheads_k, headdim)\n v: (batch_size, seqlen, nheads_k, headdim)\n dropout_p: float. Dropout probability.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\nReturn:\n out: (batch_size, seqlen, nheads, headdim).\n\"\"\"\n```\n\nTo see how these functions are used in a multi-head attention layer (which\nincludes QKV projection, output projection), see the MHA [implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py).\n\n## Upgrading from FlashAttention (1.x) to FlashAttention-2\n\nThese functions have been renamed:\n- `flash_attn_unpadded_func` -> `flash_attn_varlen_func`\n- `flash_attn_unpadded_qkvpacked_func` -> `flash_attn_varlen_qkvpacked_func`\n- `flash_attn_unpadded_kvpacked_func` -> `flash_attn_varlen_kvpacked_func`\n\nIf the inputs have the same sequence lengths in the same batch, it is simpler\nand faster to use these functions:\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p, softmax_scale=None, causal=False)\n```\n```python\nflash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)\n```\n\n## Performance\n\nWe present expected speedup (combined forward + backward pass) and memory savings from using FlashAttention against PyTorch standard attention, depending on sequence length, on different GPUs (speedup depends on memory bandwidth - we see more speedup on slower GPU memory).\n\nWe currently have benchmarks for these GPUs:\n* [A100](#a100)\n* [H100](#h100)\n\n\n\n### A100\n\nWe display FlashAttention speedup using these parameters:\n* Head dimension 64 or 128, hidden dimension 2048 (i.e. either 32 or 16 heads).\n* Sequence length 512, 1k, 2k, 4k, 8k, 16k.\n* Batch size set to 16k / seqlen.\n\n#### Speedup\n\n![FlashAttention speedup on A100 80GB SXM5 with FP16/BF16](assets/flash2_a100_fwd_bwd_benchmark.png)\n\n#### Memory\n\n![FlashAttention memory](assets/flashattn_memory.jpg)\n\nWe show memory savings in this graph (note that memory footprint is the same no matter if you use dropout or masking).\nMemory savings are proportional to sequence length -- since standard attention has memory quadratic in sequence length, whereas FlashAttention has memory linear in sequence length.\nWe see 10X memory savings at sequence length 2K, and 20X at 4K.\nAs a result, FlashAttention can scale to much longer sequence lengths.\n\n### H100\n\n![FlashAttention speedup on H100 SXM5 with FP16/BF16](assets/flash2_h100_fwd_bwd_benchmark.png)\n\n## Full model code and training script\n\nWe have released the full GPT model\n[implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/models/gpt.py).\nWe also provide optimized implementations of other layers (e.g., MLP, LayerNorm,\ncross-entropy loss, rotary embedding). Overall this speeds up training by 3-5x\ncompared to the baseline implementation from Huggingface, reaching up to 225\nTFLOPs/sec per A100, equivalent to 72% model FLOPs utilization (we don't need\nany activation checkpointing).\n\nWe also include a training\n[script](https://github.com/Dao-AILab/flash-attention/tree/main/training) to\ntrain GPT2 on Openwebtext and GPT3 on The Pile.\n\n## Triton implementation of FlashAttention\n\nPhil Tillet (OpenAI) has an experimental implementation of FlashAttention in Triton:\nhttps://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py\n\nAs Triton is a higher-level language than CUDA, it might be easier to understand\nand experiment with. The notations in the Triton implementation are also closer\nto what's used in our paper.\n\nWe also have an experimental implementation in Triton that support attention\nbias (e.g. ALiBi):\nhttps://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py\n\n\n## Tests\nWe test that FlashAttention produces the same output and gradient as a reference\nimplementation, up to some numerical tolerance. In particular, we check that the\nmaximum numerical error of FlashAttention is at most twice the numerical error\nof a baseline implementation in Pytorch (for different head dimensions, input\ndtype, sequence length, causal / non-causal).\n\nTo run the tests:\n```sh\npytest -q -s tests/test_flash_attn.py\n```\n## When you encounter issues\n\nThis new release of FlashAttention-2 has been tested on several GPT-style\nmodels, mostly on A100 GPUs.\n\nIf you encounter bugs, please open a GitHub Issue!\n\n## Citation\nIf you use this codebase, or otherwise found our work valuable, please cite:\n```\n@inproceedings{dao2022flashattention,\n title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},\n author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\\'e}, Christopher},\n booktitle={Advances in Neural Information Processing Systems},\n year={2022}\n}\n@article{dao2023flashattention2,\n title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning,\n author={Dao, Tri},\n year={2023}\n}\n```","license":"BSD-3-Clause"},{"repo":"nebius/kvax","url":"https://github.com/nebius/kvax","readmeUrl":"https://raw.githubusercontent.com/nebius/kvax/HEAD/README.md","snippet":"# Kvax: fast and easy-to-use flash attention implementation for JAX\n\nKvax is an open-source library offering fast and efficient attention operations for the JAX framework. Built with [Flash Attention 2](https://arxiv.org/abs/2307.08691) algorithms implemented in the Triton language, it is optimised for high-performance attention computation with document masks and supports context parallelism. Kvax is designed to perform exceptionally well in distributed training scenarios on long sequences using FSDP/HSDP sharding.\n\nMore technical details in our blogpost: https://nebius.com/blog/posts/kvax-open-source-flash-attention-for-jax\n\n#### Table of Contents:\n- [Key Concepts of Kvax Implementation](#key-concepts-of-kvax-implementation)\n- [Kvax Features](#kvax-features)\n- [Kvax Results](#kvax-results)\n- [How to install](#how-to-install)\n- [How to use](#how-to-use)\n- [Package Description](#package-description)\n- [Benchmarks](#benchmarks)\n- [Limitations](#limitations)\n- [Contributing](#contributing)\n- [Citation](#citation)\n- [License](#license)\n\n## Key Concepts of Kvax Implementation\n\n### Document Mask Optimisation\n\nWhen training transformer models on long sequences, a significant amount of compute is spent on attention operations due to the quadratic complexity of the attention algorithm. [Flash Attention algorithm](https://github.com/Dao-AILab/flash-attention) offers hardware-specific optimisations to significantly reduce latency and memory requirements for these operations.\n\nDuring training on long sequences, dense packing is often used to maximise compute resource utilisation. In this approach, multiple data points are packed into a single sequence while avoiding cross-sequence attention contamination. The main idea is to calculate only the blocks of attention weights that include tokens which should attend to each other while skipping other blocks. Various methods can efficiently handle this, with [PyTorch's FlexAttention](https://pytorch.org/blog/flexattention/) being one example. Kvax takes a similar approach to achieve high performance in these scenarios.\n\n### Context Parallelism\n\nUsing long sequences during training can also lead to high GPU memory consumption for storing layer activations. Context parallelism helps solve this problem, speeding up the computations and reducing memory required for layer activations.\n\nThere are several approaches to implementing context parallelism for transformer architectures, such as [RingAttention](https://arxiv.org/abs/2310.01889) and all-gather based method. The all-gather based method, described in the [Llama 3 training paper](https://arxiv.org/abs/2407.21783), performs an all-gather on the key and value tensors, collecting tensors before attention computation due to their lower memory requirements enabled by [GQA](https://arxiv.org/abs/2305.13245). This method is particularly well-suited for document masks, and Kvax leverages it in its implementation.\n\n## Kvax Features\n\n- **Block-wise Attention Masks**: Like [FlexAttention](https://pytorch.org/blog/flexattention/), our implementation builds the attention mask once per forward-backward pass, reusing it across layers. Our high-performance Triton kernel builds this mask blockwise, and does not require `O(seq_len^2)` GPU memory.\n\n- **Optimised Memory Storage**: Kvax stores attention masks in block-wise format, requiring `3 * 4 * batch_size * seq_len // block_size * 4 bytes` (block_size is typically 64 or 128).\n\n- **Skipping Pad Tokens**: Kvax skips blocks consisting entirely of padding tokens. See the \"How to Use\" section for details on defining padding tokens.\n\n- **Context Parallelism**: Kvax balances tokens across GPUs to ensure equal attention operation loads, accounting for causal masks. This feature is described in [Llama 3 training paper](https://arxiv.org/abs/2407.21783) and fully integrates with document mask optimisations.\n\n## Kvax Results\n\n![Comparison of attention implementations with causal masks; forward pass only](assets/attn_doc.png)\n\n![Comparison of attention implementations with causal masks; forward + backward pass](assets/attn_doc_bwd.png)\n\nMore details on Kvax benchmarking and its results can be found in the [blogpost](https://nebius.com/blog/posts/kvax-open-source-flash-attention-for-jax#results).\n\n## How to install\n\nInstall the latest stable release from pip:\n\n```bash\npip install kvax\n```\n\n**Note: The automatically installed versions of Triton and JAX-Triton might not be compatible. If you encounter an error while running the provided benchmarks, please ensure that you install compatible versions manually. For benchmarking, we used `triton==3.1` and `jax-triton==0.2.0`.**\n\n\n## How to use\n\nFirst, ensure that the position of every padding token is marked with `PADDING_SEGMENT_ID` in the `query_segment_ids` and `kv_segment_ids` tensors:\n\n```python\nfrom kvax.utils import PADDING_SEGMENT_ID\n\n# In this example, the sequence length is 8, and there are 2 padding tokens.\npad_token_id = 128001\ninput_ids = [6151, 0, 52043, 710, 374, 1618, pad_token_id, pad_token_id]\nquery_segment_ids = [0, 0, 0, 0, 0, 0, PADDING_SEGMENT_ID, PADDING_SEGMENT_ID]\nkv_segment_ids = [0, 0, 0, 0, 0, 0, PADDING_SEGMENT_ID, PADDING_SEGMENT_ID]\n```\n\nThen, kvax functions can be used in the transformer code:\n\n```python\nimport flax.linen as nn\nfrom kvax.ops import (\n create_attention_mask,\n flash_attention,\n)\nfrom kvax.utils import (\n attention_specs,\n permute_tokens_context_parallelism,\n unpermute_tokens_context_parallelism,\n)\n\n\nclass AttentionLayer(nn.Module):\n def __call__(\n self,\n embedding,\n query_positions,\n query_segment_ids,\n kv_positions,\n kv_segment_ids,\n attn_mask,\n ):\n query, key, value = ...\n scale = ...\n\n # Call the Flash Attention op\n attn_out = flash_attention(\n query=query,\n key=key,\n value=value,\n query_positions=positions,\n query_segment_ids=segment_ids,\n kv_positions=kv_positions,\n kv_segment_ids=kv_segment_ids,\n mask=attn_mask,\n assume_sequential_positions=self.config.assume_sequential_positions,\n scale=scale,\n # Mesh is defined as a global context\n # mesh=mesh,\n )\n\n out = ...\n return out\n\n\nclass Transformer(nn.Module):\n ...\n def setup(self):\n self.attn_layers = [AttentionLayer(...) for _ in range(self.num_layers)]\n self.mlp_layers = ...\n\n def __call__(\n self,\n embedding,\n positions,\n segment_ids,\n ):\n # During inference, create kv_positions and kv_segment_ids from positions and segment_ids\n # For training they could be simply defined as:\n # kv_positions, kv_segment_ids = positions, segment_ids\n kv_positions, kv_segment_ids = self._maybe_cache(positions, segment_ids)\n\n # Permute input tokens to balance load between GPUs during context_parallelism\n if self._should_permute_input_tokens:\n embeddings, query_positions, query_segment_ids = permute_tokens_context_parallelism(\n (embeddings, positions, segment_ids),\n )\n\n # Call it once and then pass the mask into all attention blocks\n attention_mask = create_attention_mask(\n query_positions,\n query_segment_ids,\n kv_positions,\n kv_segment_ids,\n fwd_params=self.fa_config.fwd_params,\n bwd_params=self.fa_config.bwd_params,\n skip_pad_tokens=self.fa_config.skip_pad_tokens,\n calc_bwd_mask=True,\n # Mesh is defined as a global context\n # mesh=mesh,\n )\n\n # Call transformer's layers sequentially\n for attn_layer, mlp_layer in zip(self.attn_layers, self.mlp_layers):\n embedding = attn_layer(\n embedding,\n query_positions,\n query_segment_ids,\n kv_positions,\n kv_segment_ids,\n attention_mask,\n )\n embedding = mlp_layer(...)\n\n # Unpermute outputs\n if self._should_permute_input_tokens:\n embeddings = unpermute_tokens_context_parallelism(embeddings)\n\n logits = ...\n return logits\n\n\ndef training_loop(...):\n ...\n # Define mesh as a global context and axes sharding for query, key and value.\n # Can be called inside the Transformer class but before\n # the first call of create_attention_mask, flash_attention,\n # permute_tokens_context_parallelism or unpermute_tokens_context_parallelism.\n mesh = jax.sharding.Mesh(mesh_devices, mesh_names)\n with mesh, attention_specs(\n query_specs=(\"data\", \"context\", None, None),\n kv_specs=(\"data\", None, None, None),\n ):\n ...\n logits = Transformer(...)(\n embeddings,\n positions,\n segment_ids,\n )\n\n```\n\n## Package Description\n\n### Operations\n\n#### **`flash_attention`**\n\nThe function for attention operation, based on precomputed masks and input tensor sharding specifications. Should be used within the attention_specs context manager.\n\n**Arguments**:\n\n- `query`: Query tensor of shape `(batch_size, query_seq_length, num_heads, head_dim)`.\n- `key`: Key tensor of shape `(batch_size, kv_seq_length, num_kv_heads, head_dim)`.\n- `value`: Value tensor of shape `(batch_size, kv_seq_length, num_kv_heads, head_dim)`.\n- `query_positions`: A tensor of query positions with shape `(batch_size, query_seq_length)`. For sequential tokens, use `range(0, query_seq_length)`. This tensor is ignored if `assume_sequential_positions` is set to `True`.\n- `query_segment_ids`: A tensor with segment IDs for the query tokens, shaped `(batch_size, query_seq_length)`. Tokens from the same sequence must share the same segment ID. Segment IDs should be in the range `(0, max(int32))`. All padding tokens should be marked with `PADDING_SEGMENT_ID`.\n- `kv_positions`: A tensor of key/value positions with shape `(batch_size, kv_seq_length)`. For sequential tokens, use `range(0, kv_seq_length)`. This tensor is ignored if `assume_sequential_positions` is set to `True`.\n- `kv_segment_ids`: A tensor with segment IDs for the key/value tokens, shaped `(batch_size, kv_seq_length)`. Tokens from the same sequence must share the same segment ID. Segment IDs should be in the range `(0, max(int32))`. All padding tokens should be marked with `PADDING_SEGMENT_ID`.\n- `mask`: Precomputed block-wise mask from the `create_attention_mask` function.\n- `scale`: Scaling factor for attention scores. Default is `1.0`.\n- `fwd_params`: `FlashAttentionParamsConfig` for the forward pass. Defaults to predefined parameters for the GPU model.\n- `bwd_params`: `FlashAttentionParamsConfig` for the backward pass. Defaults to predefined parameters for the GPU model.\n- `assume_sequential_positions`: Assumes sequential token positions and skips loading `query_positions` and `kv_positions`. If set to `True`, the attention behaves the same as when `is_causal == True` in `jax.nn.dot_product_attention`. The default is `False`.\n- `memory_optimized_gqa_backward`: Enables memory-optimised gradient computation for grouped-query attention when set to `True`. This flag affects performance, making it slower, but can save GPU memory on activations during the backward pass if it becomes a bottleneck. It may be useful for small models with long contexts. The default is `False`.\n- `permute_tokens_for_load_balance`: Permutes tokens to achieve better load balancing across GPUs when set to True. Used only during context parallelism. For more details, refer to the [Llama 3 training paper](https://arxiv.org/abs/2407.21783). The default is True.\n- `debug`: Prints the low-level IR of the kernel when set to `True`. The default is `False`.\n- `mesh`: Device mesh configuration for distributed execution. If set to `None`, it uses the mesh from the global context. An exception is raised if `None` is provided and no mesh is available from the global context. The default is `None`.\n\n**Returns**:\nTensor with attention-weighted values.\n\n#### **`create_attention_mask`**\n\nThis function calculates attention masks for both forward and backward Flash Attention operations, using Triton kernels for block-wise computation.\n\n**Arguments**:\n\n- `query_positions`: A tensor of query positions with shape `(batch_size, query_seq_length)`. For sequential tokens, use `range(0, query_seq_length)`.\n- `query_segment_ids`: A tensor with segment IDs for the query tokens, shaped `(batch_size, query_seq_length)`. Tokens from the same sequence must share the same segment ID. Segment IDs should be in the range `(0, max(int32))`. All padding tokens should be marked with `PADDING_SEGMENT_ID`.\n- `kv_positions`: A tensor of key/value positions with shape `(batch_size, kv_seq_length)`. For sequential tokens, use `range(0, kv_seq_length)`.\n- `kv_segment_ids`: A tensor with segment IDs for the key/value tokens, shaped `(batch_size, kv_seq_length)`. Tokens from the same sequence must share the same segment ID. Segment IDs should be in the range `(0, max(int32))`. All padding tokens should be marked with `PADDING_SEGMENT_ID`.\n- `fwd_params`: `FlashAttentionParamsConfig` for the forward pass. Defaults to predefined parameters for the GPU model.\n- `bwd_params`: `FlashAttentionParamsConfig` for the backward pass. Defaults to predefined parameters for the GPU model.\n- `calc_bwd_mask`: Whether to calculate the attention masks for the backward pass. Default is `False`.\n- `skip_pad_tokens`: Whether to skip padding tokens in calculations of attention operation. If `True`, the blocks with padding tokens only will be skipped. Defaults is `True`.\n- `mesh`: Device mesh configuration for distributed execution. If set to `None`, it uses the mesh from the global context. An exception is raised if `None` is provided and no mesh is available from the global context. The default is `None`.\n\n**Returns**:\nThe forward attention mask and optionally attention masks for the backward pass if `calc_bwd_mask` is `True`.\n\n### Utilities\n\n#### **`FlashAttentionParamsConfig`**\n\nDataclass that contains parameters for the Flash Attention Triton kernel. Increasing `query_block_size` and `kv_block_size` can lead to better performance but requires more streaming multiprocessor register memory on the GPU.\n\n#### **`PADDING_SEGMENT_ID`**\n\nSegment ID for padding tokens. This value should correspond to the position of padding tokens in the `kv_segment_ids` and `query_segment_ids` tensors. See the 'How to use' section for an example.\n\n#### **`attention_specs`**\n\nA context manager for setting the attention specifications for `query` and `key`/`value` tensors. All other specifications in the Kvax are calculated based on these specifications.\n\n**Arguments**:\n\n- `query_specs`: Specifications for sharding the `query` tensor. Specs must have 4 dimensions and provide sharding dimensions for the following axes:
`(batch, query_sequence, heads, attention_head_dim)`\n- `kv_specs`: Specifications for sharding the `key`/`value` tensors. Specs must have 4 dimensions and provide sharding dimensions for the following axes:
`(batch, kv_sequence, kv_heads, attention_head_dim)`\n\n**Notes**:\n\n- Specs must have the same sharding dimensions for `batch` and `attention_head_dim`.\n- Typical values for tensor parallelism with Mesh with axes `\"data\"` and `\"model\"`:
\n`query_specs: (\"data\", None, \"model\", None)\nkv_specs: (\"data\", None, \"model\", None)`\n- Typical values for context parallelism with Mesh with axes `\"data\"` and `\"context\"`:
\n`query_specs: (\"data\", \"context\", None, None)\nkv_specs: (\"data\", None, None, None)`\n\n#### **`permute_tokens_context_parallelism`**\n\nA function to permute tokens across the sequence length `(axis==1)` to balance computation of the attention operation between GPUs for the causal mask case. For more details, please see the [Llama 3 training paper](https://arxiv.org/abs/2407.21783). For examples, please refer to the 'How to use' section.\n\n**Arguments**:\n\n- `inputs`: An input tensor or tuple of tensors to permute.\n- `mesh`: Device mesh configuration for distributed execution. If set to `None`, it uses the mesh from the global context. An exception is raised if `None` is provided and no mesh is available from the global context. The default is `None`.\n\n**Returns**:\nPermuted tensor or tuple of tensors.\n\n#### **`unpermute_tokens_context_parallelism`**\n\nA function to unpermute tokens across the sequence length `(axis==1)` after the `permute_tokens_context_parallelism` function to return them to their original order. For examples, please refer to the 'How to use' section.\n\n**Arguments**:\n\n- `inputs`: An input tensor or tuple of tensors to unpermute.\n- `mesh`: Device mesh configuration for distributed execution. If set to `None`, it uses the mesh from the global context. An exception is raised if `None` is provided and no mesh is available from the global context. The default is `None`.\n\n**Returns**:\nA tensor or tuple of tensors with tokens in their original order.\n\n## Benchmarks\n\n**Note**: Before benchmarking, you need to install the required dependencies. First, install the [GPU version of JAX](https://jax.readthedocs.io/en/latest/installation.html). After that, you can install required dependencies:\n\n```bash\npip install -e .[dev]\n```\n\n**Note: The automatically installed versions of Triton and JAX-Triton might not be compatible. If you encounter an error while running the provided benchmarks, please ensure that you install compatible versions manually. For benchmarking, we used `triton==3.1` and `jax-triton==0.2.0`.**\n\nBenchmarking CuDNN implementation vs our implementation:\n\n```bash\n# Forward with only 1 segment\npython3 benchmarks.py mha\n\n# Forward with 3 segments\npython3 benchmarks.py mha --num-segments 3\n\n# Forward+backward with 3 segments and 1000 pad tokens\npython3 benchmarks.py mha_bwd --num-segments 3 --num-pad-tokens 1000\n\n# Forward with 3 segments and 1000 pad tokens with printing attention mask\npython3 benchmarks.py mha --num-segments 3 --num-pad-tokens 1000 --show-attention-mask\n```\n\nBenchmarking context vs tensor parallelism on our implementation:\n\n```bash\n# Forward with only 1 segment with token permutation enabled\npython3 benchmarks.py mha_cp\n\n# Forward+backward with 12 segments with token permutation disabled\npython3 benchmarks.py mha_cp_bwd --num-segments 3 --permute-tokens-for-load-balance false\n\n```\n\n## Limitations\n\n- Bias is not supported.\n- Sliding window, [ALiBi](https://arxiv.org/abs/2108.12409), and custom masks are not implemented.\n- Context parallelism does not support sharding across kv_sequence as in [RingAttention](https://arxiv.org/abs/2310.01889).\n\n## Contributing\n\nCommunity contributions are welcome. For more detailed information, please refer to the [contributing guidelines](CONTRIBUTING.md).\n\n## Citation\n\nPlease cite as:\n\n```\nSkvortsov et al., \"Kvax: Fast and easy-to-use Flash Attention implementation for JAX\", Nebius blog, 2025.\n```\n\nBibTeX citation:\n```\n@article{skvortsov2025kvax,\n title={Kvax: Fast and easy-to-use Flash Attention implementation for JAX},\n author={Skvortsov, Sergei and Fisin, Filipp and Trofimova, Maria and Yangel, Boris},\n year={2025},\n journal={Nebius blog},\n note={}\n}\n```\n## License\n\nThis project is licensed under the Apache License, Version 2.0. See the [LICENSE](LICENSE) file for details.\n\n---\nΒ© Nebius BV, 2025","license":"Apache-2.0"},{"repo":"vllm-project/flash-attention","url":"https://github.com/vllm-project/flash-attention","readmeUrl":"https://raw.githubusercontent.com/vllm-project/flash-attention/HEAD/README.md","snippet":"# FlashAttention\nThis repository provides the official implementation of FlashAttention and\nFlashAttention-2 from the\nfollowing papers.\n\n**FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness** \nTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher RΓ© \nPaper: https://arxiv.org/abs/2205.14135 \nIEEE Spectrum [article](https://spectrum.ieee.org/mlperf-rankings-2022) about our submission to the MLPerf 2.0 benchmark using FlashAttention.\n![FlashAttention](assets/flashattn_banner.jpg)\n\n**FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning** \nTri Dao\n\nPaper: https://tridao.me/publications/flash2/flash2.pdf\n\n![FlashAttention-2](assets/flashattention_logo.png)\n\n\n## Usage\n\nWe've been very happy to see FlashAttention being widely adopted in such a short\ntime after its release. This [page](https://github.com/Dao-AILab/flash-attention/blob/main/usage.md)\ncontains a partial list of places where FlashAttention is being used.\n\nFlashAttention and FlashAttention-2 are free to use and modify (see LICENSE).\nPlease cite and credit FlashAttention if you use it.\n\n\n## FlashAttention-3 beta release\nFlashAttention-3 is optimized for Hopper GPUs (e.g. H100). \n\nBlogpost: https://tridao.me/blog/2024/flash3/\n\nPaper: https://tridao.me/publications/flash3/flash3.pdf\n\n![FlashAttention-3 speedup on H100 80GB SXM5 with FP16](assets/flash3_fp16_fwd.png)\n\nThis is a beta release for testing / benchmarking before we integrate that with\nthe rest of the repo.\n\nCurrently released:\n- FP16 / BF16 forward and backward, FP8 forward\n\nRequirements: H100 / H800 GPU, CUDA >= 12.3.\n\nWe highly recommend CUDA 12.8 for best performance.\n\nTo install:\n```sh\ncd hopper\npython setup.py install\n```\nTo run the test:\n```sh\nexport PYTHONPATH=$PWD\npytest -q -s test_flash_attn.py\n```\nOnce the package is installed, you can import it as follows:\n```python\nfrom flash_attn_3 import flash_attn_interface\nflash_attn_interface.flash_attn_func()\n```\n\nTo install using `uv`, in your `pyproject.toml`:\n\n```toml\n[project]\ndependencies = [\n \"flash-attn-3\"\n]\n\n[tool.uv]\nno-build-isolation = true\n\n[tool.uv.sources]\nflash-attn-3 = { git = \"https://github.com/Dao-AILab/flash-attention\", subdirectory = \"hopper\" }\n```\n\n## FlashAttention-4 (CuTeDSL)\n\nFlashAttention-4 is written in CuTeDSL and optimized for Hopper and Blackwell GPUs (e.g. H100, B200).\n\nTo install:\n```sh\npip install flash-attn-4\n```\n\nIf you're on CUDA 13, we recommend installing with the `cu13` extra for best performance:\n```sh\npip install \"flash-attn-4[cu13]\"\n```\n\nOnce installed, you can use it as follows:\n```python\nfrom flash_attn.cute import flash_attn_func\n\nout = flash_attn_func(q, k, v, causal=True)\n```\n\n## Installation and features\n**Requirements:**\n- CUDA toolkit or ROCm toolkit\n- PyTorch 2.2 and above.\n- `packaging` Python package (`pip install packaging`)\n- `psutil` Python package (`pip install psutil`)\n- `ninja` Python package (`pip install ninja`) *\n- Linux. Might work for Windows starting v2.3.2 (we've seen a few positive [reports](https://github.com/Dao-AILab/flash-attention/issues/595)) but Windows compilation still requires more testing. If you have ideas on how to set up prebuilt CUDA wheels for Windows, please reach out via Github issue.\n\n\\* Make sure that `ninja` is installed and that it works correctly (e.g. `ninja\n--version` then `echo $?` should return exit code 0). If not (sometimes `ninja\n--version` then `echo $?` returns a nonzero exit code), uninstall then reinstall\n`ninja` (`pip uninstall -y ninja && pip install ninja`). Without `ninja`,\ncompiling can take a very long time (2h) since it does not use multiple CPU\ncores. With `ninja` compiling takes 3-5 minutes on a 64-core machine using CUDA toolkit.\n\n**To install:**\n```sh\npip install flash-attn --no-build-isolation\n```\nAlternatively you can compile from source:\n```sh\npython setup.py install\n```\n\nIf your machine has less than 96GB of RAM and lots of CPU cores, `ninja` might\nrun too many parallel compilation jobs that could exhaust the amount of RAM. To\nlimit the number of parallel compilation jobs, you can set the environment\nvariable `MAX_JOBS`:\n```sh\nMAX_JOBS=4 pip install flash-attn --no-build-isolation\n```\n\n**Interface:** `src/flash_attention_interface.py`\n\n### NVIDIA CUDA Support\n**Requirements:**\n- CUDA 12.0 and above.\n\nWe recommend the\n[Pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)\ncontainer from Nvidia, which has all the required tools to install FlashAttention.\n\nFlashAttention-2 with CUDA currently supports:\n1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). For Turing GPUs (T4, RTX 2080), see the separate [flash-attention-turing](https://github.com/ssiu/flash-attention-turing) repo, which supports a core subset of FlashAttention features on Turing.\n2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).\n3. All head dimensions up to 256. ~~Head dim > 192 backward requires A100/A800 or H100/H800~~. Head dim 256 backward now works on consumer GPUs (if there's no dropout) as of flash-attn 2.5.5.\n\n### AMD ROCm Support\nROCm version has two backends. There is [composable_kernel](https://github.com/ROCm/composable_kernel) (ck) which is the default backend and a [Triton](https://github.com/triton-lang/triton) backend. They provide an implementation of FlashAttention-2.\n\n**Requirements:**\n- ROCm 6.0 and above.\n\nWe recommend the\n[Pytorch](https://hub.docker.com/r/rocm/pytorch)\ncontainer from ROCm, which has all the required tools to install FlashAttention.\n\n#### Composable Kernel Backend\nFlashAttention-2 ROCm CK backend currently supports:\n1. MI200x, MI250x, MI300x, MI355x, and RDNA 3/4 GPUs.\n2. Datatype fp16 and bf16\n3. Both forward's and backward's head dimensions up to 256.\n\n#### Triton Backend\nThe Triton implementation of [Flash Attention](https://tridao.me/publications/flash2/flash2.pdf) supports AMD's CDNA (MI200, MI300) and RDNA GPUs using fp16, bf16, and fp32 datatypes. It provides forward and backward passes with causal masking, variable sequence lengths, arbitrary Q/KV sequence lengths and head sizes, MQA/GQA, dropout, rotary embeddings, ALiBi, paged attention, and FP8 (via the Flash Attention v3 interface). Sliding window attention is currently a work in progress.\n\nThe Triton backend kernels are provided by the [aiter](https://github.com/ROCm/aiter) package, included as a git submodule at `third_party/aiter` and automatically installed during setup.\n\nTo install, first get PyTorch for ROCm from https://pytorch.org/get-started/locally/, then install Flash Attention:\n```sh\ncd flash-attention\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pip install --no-build-isolation .\n```\n\nTo use a specific aiter commit (e.g., for testing or development):\n```sh\ncd flash-attention\ncd third_party/aiter && git fetch origin && git checkout && cd ../..\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pip install --no-build-isolation .\n```\n\nTo run the tests (note: full suite takes hours):\n```sh\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pytest tests/test_flash_attn_triton_amd.py\n```\n\nThe Triton backend uses a default kernel configuration optimized for determinism and reasonable performance across workloads. For peak throughput, enable `FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=\"TRUE\"` to search for optimal settings, which incurs a one-time warmup cost.\n\nAlternativly, if _not_ autotuning, `FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON` may be used to set a single triton config overriding the hardcoded defaults for `attn_fwd`. E.g.\n```sh\nFLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON='{\"BLOCK_M\":128,\"BLOCK_N\":64,\"waves_per_eu\":1,\"PRE_LOAD_V\":false,\"num_stages\":1,\"num_warps\":8}'\n```\n\nFor a quick start with Docker:\n```dockerfile\nFROM rocm/pytorch:latest\n\nWORKDIR /workspace\n\n# build flash attention with triton backend\nRUN git clone https://github.com/Dao-AILab/flash-attention &&\\ \n cd flash-attention &&\\\n FLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pip install --no-build-isolation .\n\n# set working dir\nWORKDIR /workspace/flash-attention\n\n# set env variable to use triton backend\nENV FLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\"\n```\n\nBuild and run:\n```sh\ndocker build -t flash-attn-triton .\ndocker run -it --network=host --user root --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --shm-size 16G --device=/dev/kfd --device=/dev/dri flash-attn-triton\n```\n\n## How to use FlashAttention\n\nThe main functions implement scaled dot product attention (softmax(Q @ K^T *\nsoftmax_scale) @ V):\n```python\nfrom flash_attn import flash_attn_qkvpacked_func, flash_attn_func\n```\n\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False,\n window_size=(-1, -1), alibi_slopes=None, deterministic=False):\n\"\"\"dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\nArguments:\n qkv: (batch_size, seqlen, 3, nheads, headdim)\n dropout_p: float. Dropout probability.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to\n the attention score of query i and key j.\n deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n which is slightly slower and uses more memory. The forward pass is always deterministic.\nReturn:\n out: (batch_size, seqlen, nheads, headdim).\n\"\"\"\n```\n\n```python\nflash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False,\n window_size=(-1, -1), alibi_slopes=None, deterministic=False):\n\"\"\"dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n q: (batch_size, seqlen, nheads, headdim)\n k: (batch_size, seqlen, nheads_k, headdim)\n v: (batch_size, seqlen, nheads_k, headdim)\n dropout_p: float. Dropout probability.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n is added to the attention score of query i and key j.\n deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n which is slightly slower and uses more memory. The forward pass is always deterministic.\nReturn:\n out: (batch_size, seqlen, nheads, headdim).\n\"\"\"\n```\n\n```python\ndef flash_attn_with_kvcache(\n q,\n k_cache,\n v_cache,\n k=None,\n v=None,\n rotary_cos=None,\n rotary_sin=None,\n cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,\n cache_batch_idx: Optional[torch.Tensor] = None,\n block_table: Optional[torch.Tensor] = None,\n softmax_scale=None,\n causal=False,\n window_size=(-1, -1), # -1 means infinite context window\n rotary_interleaved=True,\n alibi_slopes=None,\n):\n \"\"\"\n If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from\n k and v. This is useful for incremental decoding: you can pass in the cached keys/values from\n the previous step, and update them with the new keys/values from the current step, and do\n attention with the updated cache, all in 1 kernel.\n\n If you pass in k / v, you must make sure that the cache is large enough to hold the new values.\n For example, the KV cache could be pre-allocated with the max sequence length, and you can use\n cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.\n\n Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be\n rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\n If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos\n and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\n If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at\n indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).\n\n See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.\n\n Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\n than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\n For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\n If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\n For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n 1 1 1 1 0\n 1 1 1 1 1\n If seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n 0 0\n 0 0\n 0 0\n 1 0\n 1 1\n If the row of the mask is all zero, the output will be zero.\n\n If window_size != (-1, -1), implements sliding window local attention. Query at position i\n will only attend to keys between\n [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\n Note: Does not support backward pass.\n\n Arguments:\n q: (batch_size, seqlen, nheads, headdim)\n k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n page_block_size must be a multiple of 256.\n v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate\n k with k_cache, starting at the indices specified by cache_seqlens.\n v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.\n rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding\n to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.\n rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.\n cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the\n KV cache.\n block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.\n cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.\n If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].\n If the indices are not distinct, and k and v are provided, the values updated in the cache\n might come from any of the duplicate indices.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.\n If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,\n rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1\n (i.e. GPT-NeoX style).\n alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n is added to the attention score of query i and key j.\n\n Return:\n out: (batch_size, seqlen, nheads, headdim).\n \"\"\"\n```\n\nTo see how these functions are used in a multi-head attention layer (which\nincludes QKV projection, output projection), see the MHA [implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py).\n\n### Using with πŸ€— Kernels\n\nIf your hardware environment belongs to any of the above-mentioned, you can also use the [`kernels` library](https://github.com/huggingface/kernels)\nto use Flash Attention 2 and 3 right away.\n\n```py\n# pip install kernels\n\nfrom kernels import get_kernel\n\n# FA2\nfa_module = get_kernel(\"kernels-community/flash-attn2\", version=1)\nflash_attn_func = fa_module.flash_attn_func\n\n# FA3\nfa3_module = get_kernel(\"kernels-community/flash-attn3\", version=1)\nflash_attn_func = fa3_module.flash_attn_func\n```\n\n## Changelog\n\n### 2.0: Complete rewrite, 2x faster\nUpgrading from FlashAttention (1.x) to FlashAttention-2\n\nThese functions have been renamed:\n- `flash_attn_unpadded_func` -> `flash_attn_varlen_func`\n- `flash_attn_unpadded_qkvpacked_func` -> `flash_attn_varlen_qkvpacked_func`\n- `flash_attn_unpadded_kvpacked_func` -> `flash_attn_varlen_kvpacked_func`\n\nIf the inputs have the same sequence lengths in the same batch, it is simpler\nand faster to use these functions:\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)\n```\n```python\nflash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)\n```\n### 2.1: Change behavior of causal flag\n\nIf seqlen_q != seqlen_k and causal=True, the causal mask is aligned to the\nbottom right corner of the attention matrix, instead of the top-left corner.\n\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 =\nmasked out) is: \nv2.0: \n 1 0 0 0 0 \n 1 1 0 0 0 \nv2.1: \n 1 1 1 1 0 \n 1 1 1 1 1 \n\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is: \nv2.0: \n 1 0 \n 1 1 \n 1 1 \n 1 1 \n 1 1 \nv2.1: \n 0 0 \n 0 0 \n 0 0 \n 1 0 \n 1 1 \nIf the row of the mask is all zero, the output will be zero.\n\n### 2.2: Optimize for inference\n\nOptimize for inference (iterative decoding) when query has very small sequence\nlength (e.g., query sequence length = 1). The bottleneck here is to load KV\ncache as fast as possible, and we split the loading across different thread\nblocks, with a separate kernel to combine results.\n\nSee the function `flash_attn_with_kvcache` with more features for inference\n(perform rotary embedding, updating KV cache inplace).\n\nThanks to the xformers team, and in particular Daniel Haziza, for this\ncollaboration.\n\n### 2.3: Local (i.e., sliding window) attention\n\nImplement sliding window attention (i.e., local attention). Thanks to [Mistral\nAI](https://mistral.ai/) and in particular TimothΓ©e Lacroix for this\ncontribution. Sliding window was used in the [Mistral 7B](https://mistral.ai/news/announcing-mistral-7b/) model.\n\n### 2.4: ALiBi (attention with linear bias), deterministic backward pass.\n\nImplement ALiBi (Press et al., 2021). Thanks to Sanghun Cho from Kakao Brain for this contribution.\n\nImplement deterministic backward pass. Thanks to engineers from [Meituan](www.meituan.com) for this contribution.\n\n### 2.5: Paged KV cache.\n\nSupport paged KV cache (i.e., [PagedAttention](https://arxiv.org/abs/2309.06180)).\nThanks to @beginlner for this contribution.\n\n### 2.6: Softcapping.\n\nSupport attention with softcapping, as used in Gemma-2 and Grok models.\nThanks to @Narsil and @lucidrains for this contribution.\n\n### 2.7: Compatibility with torch compile\n\nThanks to @ani300 for this contribution.\n\n## Performance\n\nWe present expected speedup (combined forward + backward pass) and memory savings from using FlashAttention against PyTorch standard attention, depending on sequence length, on different GPUs (speedup depends on memory bandwidth - we see more speedup on slower GPU memory).\n\nWe currently have benchmarks for these GPUs:\n* [A100](#a100)\n* [H100](#h100)\n\n\n\n### A100\n\nWe display FlashAttention speedup using these parameters:\n* Head dimension 64 or 128, hidden dimension 2048 (i.e. either 32 or 16 heads).\n* Sequence length 512, 1k, 2k, 4k, 8k, 16k.\n* Batch size set to 16k / seqlen.\n\n#### Speedup\n\n![FlashAttention speedup on A100 80GB SXM5 with FP16/BF16](assets/flash2_a100_fwd_bwd_benchmark.png)\n\n#### Memory\n\n![FlashAttention memory](assets/flashattn_memory.jpg)\n\nWe show memory savings in this graph (note that memory footprint is the same no matter if you use dropout or masking).\nMemory savings are proportional to sequence length -- since standard attention has memory quadratic in sequence length, whereas FlashAttention has memory linear in sequence length.\nWe see 10X memory savings at sequence length 2K, and 20X at 4K.\nAs a result, FlashAttention can scale to much longer sequence lengths.\n\n### H100\n\n![FlashAttention speedup on H100 SXM5 with FP16/BF16](assets/flash2_h100_fwd_bwd_benchmark.png)\n\n## Full model code and training script\n\nWe have released the full GPT model\n[implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/models/gpt.py).\nWe also provide optimized implementations of other layers (e.g., MLP, LayerNorm,\ncross-entropy loss, rotary embedding). Overall this speeds up training by 3-5x\ncompared to the baseline implementation from Huggingface, reaching up to 225\nTFLOPs/sec per A100, equivalent to 72% model FLOPs utilization (we don't need\nany activation checkpointing).\n\nWe also include a training\n[script](https://github.com/Dao-AILab/flash-attention/tree/main/training) to\ntrain GPT2 on Openwebtext and GPT3 on The Pile.\n\n## Triton implementation of FlashAttention\n\nPhil Tillet (OpenAI) has an experimental implementation of FlashAttention in Triton:\nhttps://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py\n\nAs Triton is a higher-level language than CUDA, it might be easier to understand\nand experiment with. The notations in the Triton implementation are also closer\nto what's used in our paper.\n\nWe also have an experimental implementation in Triton that support attention\nbias (e.g. ALiBi):\nhttps://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py\n\n\n## Tests\nWe test that FlashAttention produces the same output and gradient as a reference\nimplementation, up to some numerical tolerance. In particular, we check that the\nmaximum numerical error of FlashAttention is at most twice the numerical error\nof a baseline implementation in Pytorch (for different head dimensions, input\ndtype, sequence length, causal / non-causal).\n\nTo run the tests:\n```sh\npytest -q -s tests/test_flash_attn.py\n```\n## When you encounter issues\n\nThis new release of FlashAttention-2 has been tested on several GPT-style\nmodels, mostly on A100 GPUs.\n\nIf you encounter bugs, please open a GitHub Issue!\n\n## Tests\nTo run the tests:\n```sh\npytest tests/test_flash_attn_ck.py\n```\n\n## Citation\nIf you use this codebase, or otherwise found our work valuable, please cite:\n```\n@inproceedings{dao2022flashattention,\n title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},\n author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\\'e}, Christopher},\n booktitle={Advances in Neural Information Processing Systems (NeurIPS)},\n year={2022}\n}\n@inproceedings{dao2023flashattention2,\n title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},\n author={Dao, Tri},\n booktitle={International Conference on Learning Representations (ICLR)},\n year={2024}\n}\n```","license":"BSD-3-Clause"}],"warnings":["The research index GitHub search (GET /v2/search/research/github, legacy GET /v2/research/github) is deprecated and stops responding after 2026-11-03. Use GET or POST /v2/search/developer instead: it searches GitHub issues, pull requests and READMEs plus curated documentation sources, returns matched passages, and adds filters for repo, language, license and stars. Response changes: 'snippet' becomes 'passages', results gain an 'id', and there is no score breakdown and no web fallback result type."],"replacement":"/v2/search/developer"}