{"success":true,"results":[{"resultType":"repo_readme","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```"},{"resultType":"repo_readme","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```"},{"resultType":"repo_readme","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"},{"resultType":"github_history","repo":"huggingface/transformers","url":"https://github.com/huggingface/transformers/issues/33522","pageType":"merged_pr","number":33522,"segmentCount":1,"title":"Conversation","snippet":"## Conversation\n\n# What does this PR do?\n\nThis PR adds preliminary support for Flash Attention 3.\n\n- `is_flash_attn_3_available` required a workaround in `_is_package_available` as `package_version = importlib.metadata.version(pkg_name)` fails with `importlib.metadata.PackageNotFoundError: No package metadata was found for flash_attn_interface`.\n- `_supports_flash_attn_3` and `_check_and_enable_flash_attn_3` added to `modeling_utils.py`, near duplicate of `_check_and_enable_flash_attn_2`.\n- ~~`_flash_attention_3_forward` implemented in `modeling_flash_attention_3_utils.py`~~`_flash_attention_forward` is now a unified interface for FAv2 and FAv3 controlled by `use_flash_attn_3` which is passed from `FlashAttention` classes based on `config._attn_implementation == \"flash_attention_3\"`.\n\n - Currently FAv3 does not support dropout, ~~sliding window~~ (edit: sliding window is now supported) or softcap, and in FAv3 `flash_attn_func`/`flash_attn_varlen_func` return a tuple.\n - `attention_mask is not None` and `position_ids is not None` paths depend on `_upad_input` and `prepare_fa2_from_position_ids` respectively, ~~these are duplicated from `modeling_flash_attention_utils.py`~~ and are not included in FAv3 package therefore FAv3 depends on `flash_attn`, this is reflected in `is_flash_attn_3_available` which checks for `is_flash_attn_2_available`.\n - In the remaining path FAv3 supports FP8, this PR currently uses environment variable `FLASH_ATTENTION_3_FP8` for this purpose, we can probably add something like `attention_kwargs` to model forwards to control this, or maybe another `_attn_implementation` type `flash_attention_3_fp8`, best to get reviews first and consensus on the best way to do it\\[1\\]\n- ~~`flash_attention_3` is added to Llama with `LlamaFlashAttention3`, similar to `LlamaFlashAttention2` with unsupported options like dropout and sliding window removed.~~ Edit: added to other models, see comment below.\n- ~~`_update_causal_mask` is updated in various models due to `utils/check_copies.py`, and `_supports_flash_attn_3` is added in to some other models already for the same reason.~~ See comment below.\n\nFixes [#33373](https://github.com/huggingface/transformers/issues/33373)\n\n## Todo\n\n- ~~Test `attention_mask is not None` and `position_ids is not None` paths~~\n- ~~Implement FlashAttention3 classes for other models~~ Done.\n- FP8 usage\\[1\\]\n- ~~Documentation~~ Partly done.\n- Benchmarks would be nice\n\n## Notes\n\nLlama tested on H100 SXM with:\n\n```\nimport torch\nfrom transformers import AutoTokenizer, LlamaForCausalLM\n\ntokenizer = AutoTokenizer.from_pretrained('NousResearch/Hermes-3-Llama-3.1-8B', trust_remote_code=True)\nmodel = LlamaForCausalLM.from_pretrained(\n \"NousResearch/Hermes-3-Llama-3.1-8B\",\n torch_dtype=torch.float16,\n device_map=\"auto\",\n attn_implementation=\"flash_attention_3\"\n)\n\nprompts = [\\\n \"\"\"<|im_start|>system\\\nYou are a sentient, superintelligent artificial general intelligence, here to teach and assist me.<|im_end|>\\\n<|im_start|>user\\\nWrite a short story about Goku discovering kirby has teamed up with Majin Buu to destroy the world.<|im_end|>\\\n<|im_start|>assistant\"\"\",\\\n ]\n\nfor chat in prompts:\n print(chat)\n input_ids = tokenizer(chat, return_tensors=\"pt\").input_ids.to(\"cuda\")\n generated_ids = model.generate(input_ids, max_new_tokens=750, temperature=0.8, repetition_penalty=1.1, do_sample=True, eos_token_id=tokenizer.eos_token_id)\n response = tokenizer.decode(generated_ids[0][input_ids.shape[-1]:], skip_special_tokens=True, clean_up_tokenization_space=True)\n print(f\"Response: {response}\")\n```\n\n(shortened) responses\n\nFP16:\n\n```\nIn the vast expanse of the universe, there existed a celestial realm where beings of extraordinary powers roamed freely. One such being was Goku, the legendary warrior known for his indomitable spirit and unconquerable will.\n\nOne fateful day, as Goku trained under the golden sun, he sensed an unusual disturbance in the cosmic energy. Puzzled by this anomaly, Goku rushed to the source of the disturbance and arrived at the edge of a hidden dimension.\n\nTo his horror, Goku witnessed Kirby and Majin Buu working together, devising devious plans to obliterate entire planets. Their combined strength was formidable, their intentions sinister, and their alliance unprecedented.\n```\n\nFP8:\n\n```\nIn the vast expanse of the universe, there existed a planet called Planet Vegeta, where the powerful Saiyan warrior, Goku, lived among his friends in the city of Earth.\n\nOne fateful day, Goku was training on the top of a mountain when he sensed an unusual energy signature. His spidey senses immediately tingled with suspicion.\n\n\"Who could that be?\" he wondered aloud as he leapt into the sky, soaring towards the source of the disturbance.\n\nAs Goku arrived at the scene, he discovered something utterly shocking - Kirby, the infamous villain from another galaxy, had formed an alliance with Majin Buu, the mischievous yet formidable entity who had caused Goku so much trouble in the past.\n```\n\n~~All other models will be tested after I've finished adding FlashAttention3 classes.~~ Other models have been tested, see comment below.\n\n## Who can review?\n\ncc [@ArthurZucker](https://github.com/ArthurZucker)\n\nArthurZucker and vasqu reacted with hooray emoji\n\n- ![tada](https://github.githubassets.com/assets/1f389-36899a2cb781.png)2 reactions\n\nthe\nflash-attention-3\nbranch\n3 times, most recently\nfrom\n[`b6afd63`](https://github.com/huggingface/transformers/commit/b6afd6351b838fc33398a444a05f6e487b692c1b) to\n[`0976545`](https://github.com/huggingface/transformers/commit/09765453c6ccca9360e3c62015f5424d142a1681) [Compare](https://github.com/huggingface/transformers/compare/b6afd6351b838fc33398a444a05f6e487b692c1b..09765453c6ccca9360e3c62015f5424d142a1681) [2 years agoSeptember 17, 2024 14:28](https://github.com/huggingface/transformers/pull/33522#event-14292407087)\n\nContributorAuthor\n\n| |\n| --- |\n| `FlashAttention3` classes added to the models that had to be modified due to `utils/check_copies.py`. The following models do not currently support Flash Attention and were only modified due to `utils/check_copies.py`: `bloom`, `codegen`, `gpt_neox_japanese`, `idefics`, `persimmon`.
~~There are more models that support FAv2, these will be done next.~~
~~Note that Sliding Window should be supported soon, after [Dao-AILab/flash-attention#1233](https://github.com/Dao-AILab/flash-attention/pull/1233)~~ |\n\nthe\nflash-attention-3\nbranch\n2 times, most recently\nfrom\n[`5aa58ab`](https://github.com/huggingface/transformers/commit/5aa58ab5fa127c4d551f6da773e74f56b7e3929b) to\n[`7ae105e`](https://github.com/huggingface/transformers/commit/7ae105e0f83f54d2ed55780969b567f8ddbc20f7) [Compare](https://github.com/huggingface/transformers/compare/5aa58ab5fa127c4d551f6da773e74f56b7e3929b..7ae105e0f83f54d2ed55780969b567f8ddbc20f7) [2 years agoSeptember 17, 2024 17:49](https://github.com/huggingface/transformers/pull/33522#event-14295064933)\n\nContributorAuthor\n\n| |\n| --- |\n| All models supporting FAv2 should now have FAv3 classes.
~~The following models will need sliding window adding back in when available:
`chameleon/modeling_chameleon.py`
`gemma/modeling_gemma.py`
`gemma2/modeling_gemma2.py`
`granite/modeling_granite.py`
`idefics2/modeling_idefics2.py`
`jamba/modeling_jamba.py`
`llama/modeling_llama.py`
`mistral/modeling_mistral.py`
`mixtral/modeling_mixtral.py`
`nemotron/modeling_nemotron.py`
`phi3/modeling_phi3.py`
`qwen2/modeling_qwen2.py`
`qwen2_moe/modeling_qwen2_moe.py`
`qwen2_vl/modeling_qwen2_vl.py`
`starcoder2/modeling_starcoder2.py`~~
~~There are some areas that use `config._attn_implementation == \"flash_attention_2\"` that I'll update next.~~ |\n\nthe\nflash-attention-3\nbranch\nfrom\n[`7ae105e`](https://github.com/huggingface/transformers/commit/7ae105e0f83f54d2ed55780969b567f8ddbc20f7) to\n[`bd6e9e7`](https://github.com/huggingface/transformers/commit/bd6e9e72b5ec434f829cddf9b4f699c19c75e3b7) [Compare](https://github.com/huggingface/transformers/compare/7ae105e0f83f54d2ed55780969b567f8ddbc20f7..bd6e9e72b5ec434f829cddf9b4f699c19c75e3b7) [2 years agoSeptember 17, 2024 18:33](https://github.com/huggingface/transformers/pull/33522#event-14295615854)\n\nContributorAuthor\n\n| |\n| --- |\n| All occurrences of `config._attn_implementation == ...`/`self._use_flash_attention_` and other mentions of `\"flash_attention_2\"` should now be updated with FAv3 versions. This should be about it on the modeling side with the exception of sliding window.
Documentation and tests will be done next. |\n\nthe\nflash-attention-3\nbranch\nfrom\n[`bd6e9e7`](https://github.com/huggingface/transformers/commit/bd6e9e72b5ec434f829cddf9b4f699c19c75e3b7) to\n[`fbf9bec`](https://github.com/huggingface/transformers/commit/fbf9beccc08e24abbb7468d0d48a5b2c1fbf232e) [Compare](https://github.com/huggingface/transformers/compare/bd6e9e72b5ec434f829cddf9b4f699c19c75e3b7..fbf9beccc08e24abbb7468d0d48a5b2c1fbf232e) [2 years agoSeptember 17, 2024 20:14](https://github.com/huggingface/transformers/pull/33522#event-14296932956)\n\nContributorAuthor\n\n| |\n| --- |\n| Some documentation and all tests are updated for FAv3. I'll run the tests on a H100 instance then mark this as ready for (initial) review. |\n\nthe\nflash-attention-3\nbranch\nfrom\n[`fbf9bec`](https://github.com/huggingface/transformers/commit/fbf9beccc08e24abbb7468d0d48a5b2c1fbf232e) to\n[`27edb62`](https://github.com/huggingface/transformers/commit/27edb6246b10d5b0d6adb6dbf2a8278541baaaa0) [Compare](https://github.com/huggingface/transformers/compare/fbf9beccc08e24abbb7468d0d48a5b2c1fbf232e..27edb6246b10d5b0d6adb6dbf2a8278541baaaa0) [2 years agoSeptember 17, 2024 23:16](https://github.com/huggingface/transformers/pull/33522#event-14299220509)\n\nContributorAuthor\n\n| |\n| --- |\n| Generally FAv3 tests are failing due to the small configurations used: `RuntimeError: Only support head size 64, 128, and 256 for now`.
Instead I've tested the majority of models from their examples, with a few exceptions like Gemma and Mistral that I need to request access to, and particularly large models such as Jamba that my instance doesn't have space to download.
All of the tested models with examples are ok, with the exception of `HuggingFaceM4/idefics2-8b-base`:
```
File \"/workspace/transformers/src/transformers/modeling_flash_attention_3_utils.py\", line 118, in _upad_input
query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: too many values to unpack (expected 4)
```
However this error also occurs with `flash_attention_2`.
StableLM models are currently not supported due to num\\_attention\\_heads/`Only support head size 64, 128, and 256 for now`.
I've attached test reports, the numerical accuracy failures may need special care as per `flash-attention/hopper/test_flash_attn.py` [here](https://github.com/Dao-AILab/flash-attention/blob/af314d400663fe895199b0586a9f1f718b1d7b79/hopper/test_flash_attn.py#L182-L191) and [here](https://github.com/Dao-AILab/flash-attention/blob/af314d400663fe895199b0586a9f1f718b1d7b79/hopper/test_flash_attn.py#L388-L393)
[test\\_report.zip](https://github.com/user-attachments/files/17036304/test_report.zip) |\n\nArthurZucker reacted with hooray emoji\n\n- ![tada](https://github.githubassets.com/assets/1f389-36899a2cb781.png)1 reaction\n\n**[ArthurZucker](https://github.com/ArthurZucker)**\n\nreviewed\n\n[on Sep 17, 2024Sep 17, 2024](https://github.com/huggingface/transformers/pull/33522#pullrequestreview-2311312396)\n\n[View reviewed changes](https://github.com/huggingface/transformers/pull/33522/files)\n\n### ![@ArthurZucker](https://avatars.githubusercontent.com/u/48595927?s=48&v=4)**[ArthurZucker](https://github.com/ArthurZucker)** left a comment\n\nThe reason will be displayed to describe this comment to others. [Learn more](https://docs.github.com/articles/managing-disruptive-comments/#hiding-a-comment).\n\nWowowo super nice initiative thanks! 🔥\n\nIMO since we already abstracted the flash attention API, let's try to keep it in `flashAttentionLlama` but maybe support `flash_attention_3` in the `attn_implementation` for example! WDYT?\n\nComment thread[src/transformers/models/llama/modeling\\_llama.py](https://github.com/huggingface/transformers/pull/33522/files#diff-06392bad3b9e97be9ade60d4ac46f73b6809388f4d507c2ba1384ab872711c51)\nOutdated\n\n|\n|\n\n### ![@ArthurZucker](https://avatars.githubusercontent.com/u/48595927?s=48&v=4)**[ArthurZucker](https://github.com/ArthurZucker)** [on Sep 17, 2024Sep 17, 2024](https://github.com/huggingface/transformers/pull/33522\\#discussion_r1764183358)\n\nThe reason will be displayed to describe this comment to others. [Learn more](https://docs.github.com/articles/managing-disruptive-comments/#hiding-a-comment).\n\nHey! As far as I can tell, the only diff is the forward function right?\n\n### ![@hlky](https://avatars.githubusercontent.com/u/106811348?s=48&v=4)**[hlky](https://github.com/hlky)** [on Sep 17, 2024Sep 17, 2024](https://github.com/huggingface/transformers/pull/33522\\#discussion_r1764219393)\n\nContributorAuthor\n\nThe reason will be displayed to describe this comment to others. [Learn more](https://docs.github.com/articles/managing-disruptive-comments/#hiding-a-comment).\n\nYeah the difference between FlashAttention2 classes and FlashAttention3 is just the forward function, and lack of dropout/sliding window/softcap for FAv3. As you suggest we could support v3 in the existing classes instead using `config.attn_implementation` to select the appropriate function, happy to make this change if you think that's better.\n\nComment thread[src/transformers/modeling\\_flash\\_attention\\_3\\_utils.py](https://github.com/huggingface/transformers/pull/33522/files#diff-d37407494bc6ab849e0fb7d6cdf04811a078a9f142b0bb4e074fe4c3df68f2ae)\nOutdated\n\n|\n|\n\n### ![@ArthurZucker](https://avatars.githubusercontent.com/u/48595927?s=48&v=4)**[ArthurZucker](https://github.com/ArthurZucker)** [on Sep 17, 2024Sep 17, 2024](https://github.com/huggingface/transformers/pull/33522\\#discussion_r1764185201)\n\nThe reason will be displayed to describe this comment to others. [Learn more](https://docs.github.com/articles/managing-disruptive-comments/#hiding-a-comment).\n\nlet's maybe replace flash\\_attention\\_forward by this one when flash attention3 is available WDYT?\n\n### ![@hlky](https://avatars.githubusercontent.com/u/106811348?s=48&v=4)**[hlky](https://github.com/hlky)** [on Sep 17, 2024Sep 17, 2024](https://github.com/huggingface/transformers/pull/33522\\#discussion_r1764214411)\n\nContributorAuthor\n\nThe reason will be displayed to describe this comment to others. [Learn more](https://docs.github.com/articles/managing-disruptive-comments/#hiding-a-comment).\n\nAFAIK FAv3 will be for Hopper GPUs only\n\nthe\nflash-attention-3\nbranch\nfrom\n[`27edb62`](https://github.com/huggingface/transformers/commit/27edb6246b10d5b0d6adb6dbf2a8278541baaaa0) to\n[`ba268ef`](https://github.com/huggingface/transformers/commit/ba268efc2fedcd2ad0be5c0f812d3ea895a7458f) [Compare](https://github.com/huggingface/transformers/compare/27edb6246b10d5b0d6adb6dbf2a8278541baaaa0..ba268efc2fedcd2ad0be5c0f812d3ea895a7458f) [2 years agoSeptember 18, 2024 16:31](https://github.com/huggingface/transformers/pull/33522#event-14310624234)\n\nContributorAuthor\n\n| |\n| --- |\n| I've replaced the `FlashAttention3` classes and integrated it into existing `FlashAttention2` classes. `self._flash_attn_3 = self.config._attn_implementation == \"flash_attention_3\"` is added to control which version to use with `if self._flash_attn_3:`.
I've renamed the `FlashAttention2` classes to just `FlashAttention` as `ATTENTION_CLASSES` would look strange doing e.g. `\"flash_attention_3\": Qwen2VLFlashAttention2`
Note that while I was checking all `Attention` classes contain config I added some missing type annotations to the config parameters, I then had to add a few more type due to `Copied from`.
In `src/transformers/models/qwen2_vl/modeling_qwen2_vl.py` I've changed `VisionFlashAttention` and `VisionSdpaAttention` to subclass `VisionAttention` and added `config` as a parameter, this was needed for `self.config._attn_implementation == \"flash_attention_3\"`.
We could simplify the changes to `FlashAttention` classes further by creating a wrapper for both `_flash_attention_forward` and `_flash_attention_3_forward` with the above `_flash_attn_3` as a parameter.
Checks like `config._attn_implementation != \"flash_attention_2\"`/`config._attn_implementation == \"flash_attention_2\"` could also be changed to something like `\"flash_attention\" not in config._attn_implementation`/`\"flash_attention\" in config._attn_implementation`. |\n\nthe\nflash-attention-3\nbranch\n3 times, most recently\nfrom\n[`4473129`](https://github.com/huggingface/transformers/commit/447312990634f31c594e332a3050f88e690ba405) to\n[`4a34da8`](https://github.com/huggingface/transformers/commit/4a34da8eee0d6927011827c5e2490483a93a3c6c) [Compare](https://github.com/huggingface/transformers/compare/447312990634f31c594e332a3050f88e690ba405..4a34da8eee0d6927011827c5e2490483a93a3c6c) [2 years agoSeptember 20, 2024 10:51](https://github.com/huggingface/transformers/pull/33522#event-14338162418)\n\nContributorAuthor\n\n| |\n| --- |\n| Sliding window is now supported. |\n\n| |\n| --- |\n| Very great work 🚀 just a passerby who looked into the code :)
> We could simplify the changes to FlashAttention classes further by creating a wrapper for both \\_flash\\_attention\\_forward and \\_flash\\_attention\\_3\\_forward with the above \\_flash\\_attn\\_3 as a parameter.
I'd be very pro this. It kinda looks misleading now with `_flash_attention_forward` which is meant to model the FA2 forward. If we keep it as a unified interface and an FA3 flag (to give control over the implementation used), it makes more sense imo.
Personal preference: I'd go a step further and move the `modeling_flash_attention_3_utils.py` file's content into the general `modeling_flash_attention_utils.py` file together with the combined interface.
> Checks like config.\\_attn\\_implementation != \"flash\\_attention\\_2\"/config.\\_attn\\_implementation == \"flash\\_attention\\_2\" could also be changed to something like \"flash\\_attention\" not in config.\\_attn\\_implementation/\"flash\\_attention\" in config.\\_attn\\_implementation.
Seems reasonable to me. Makes the code less verbose too.
Lastly, `_check_and_enable_flash_attn_3` in `modeling_utils.py` could benefit from a compute capability check since FA3 has a hard dependency on hopper gpus (9.0 iirc).
Edit: Maybe raising a value error / warning if dropout or similar values are passed, would also be nice since now it's just silently ignoring them. |\n\nthe\nflash-attention-3\nbranch\nfrom\n[`4a34da8`](https://github.com/huggingface/transformers/commit/4a34da8eee0d6927011827c5e2490483a93a3c6c) to\n[`9bcbe3f`](https://github.com/huggingface/transformers/commit/9bcbe3f38c3f27cc81f393115dbe06bf24ee8aa3) [Compare](https://github.com/huggingface/transformers/compare/4a34da8eee0d6927011827c5e2490483a93a3c6c..9bcbe3f38c3f27cc81f393115dbe06bf24ee8aa3) [2 years agoSeptember 20, 2024 18:06](https://github.com/huggingface/transformers/pull/33522#event-14343788543)\n\nContributorAuthor\n\n| |\n| --- |\n| I've removed `modeling_flash_attention_3_utils.py` and unified the interface for FAv2 and FAv3 in `_flash_attention_forward`.
I'll wait for input from a maintainer on changing the checks (`config._attn_implementation != \"flash_attention_2\"/config._attn_implementation == \"flash_attention_2\"`) in case there's some preference.
~~Compute capability check would indeed be useful, I'll add this in the next commit.~~ Done.
I assume this won't be merged until FAv3 is out of beta, at which point dropout and softcap should hopefully be supported, if not then I agree we should add an error/warning if they're used with FAv3. |\n\nvasqu and ArthurZucker reacted with thumbs up emoji\n\n- ![+1](https://github.githubassets.com/assets/1f44d-41cb66fe1e22.png)2 reactions\n\nthe\nflash-attention-3\nbranch\n4 times, most recently\nfrom\n[`113afe4`](https://github.com/huggingface/transformers/commit/113afe42e0ef4b0c0726677b35e4cdfd312327d9) to\n[`a7f521c`](https://github.com/huggingface/transformers/commit/a7f521ca9880fcf4d5e1b8ec06f994687105e7a0) [Compare](https://github.com/huggingface/transformers/compare/113afe42e0ef4b0c0726677b35e4cdfd312327d9..a7f521ca9880fcf4d5e1b8ec06f994687105e7a0) [2 years agoSeptember 26, 2024 11:36](https://github.com/huggingface/transformers/pull/33522#event-14414856066)\n\nthe\nflash-attention-3\nbranch\nfrom\n[`a7f521c`](https://github.com/huggingface/transformers/commit/a7f521ca9880fcf4d5e1b8ec06f994687105e7a0) to\n[`d50593a`](https://github.com/huggingface/transformers/commit/d50593aae7145330637d06995bc6f9804d8a39ca) [Compare](https://github.com/huggingface/transformers/compare/a7f521ca9880fcf4d5e1b8ec06f994687105e7a0..d50593aae7145330637d06995bc6f9804d8a39ca) [2 years agoOctober 2, 2024 20:04](https://github.com/huggingface/transformers/pull/33522#event-14494639909)\n\n`\n flash-attention-3\n`\n\n`\n 66eefe5\n`\n\nthe\nflash-attention-3\nbranch\nfrom\n[`d50593a`](https://github.com/huggingface/transformers/commit/d50593aae7145330637d06995bc6f9804d8a39ca) to\n[`66eefe5`](https://github.com/huggingface/transformers/commit/66eefe571611798b2f2b183e994f1e2b13744134) [Compare](https://github.com/huggingface/transformers/compare/d50593aae7145330637d06995bc6f9804d8a39ca..66eefe571611798b2f2b183e994f1e2b13744134) [2 years agoOctober 2, 2024 20:08](https://github.com/huggingface/transformers/pull/33522#event-14494684079)\n\n| |\n| --- |\n| Why not pull this request? I've been using it for months without a problem and have to hold back my transformers version to do so. |\n\n| |\n| --- |\n| [@bn999](https://github.com/bn999) Just my thoughts but I think it's mainly due to fa3 still being in beta and stable support so far isn't really given - there is especially a lot of development in the last couple of weeks to support more architectures, head dims etc.
Also, on another note, there also has been a refactor on the attn implementation side of transformers which most likely would force adjustments here. |\n\n| |\n| --- |\n| [@vasqu](https://github.com/vasqu) To your first point, I keep up with the fa3 changes, and so far none have broken my use case. Also, if something does break, the user can just fall back to using fa2 as they do now.
However, I can't comment on your second point. |\n\n**[ArthurZucker](https://github.com/ArthurZucker)**\n\nreviewed\n\n[on Feb 13, 2025Feb 13, 2025](https://github.com/huggingface/transformers/pull/33522#pullrequestreview-2614323436)\n\n[View reviewed changes](https://github.com/huggingface/transformers/pull/33522/files/66eefe571611798b2f2b183e994f1e2b13744134)\n\n### ![@ArthurZucker](https://avatars.githubusercontent.com/u/48595927?s=48&v=4)**[ArthurZucker](https://github.com/ArthurZucker)** left a comment\n\nThe reason will be displayed to describe this comment to others. [Learn more](https://docs.github.com/articles/managing-disruptive-comments/#hiding-a-comment).\n\nReviving this! [@hlky](https://github.com/hlky) , now that we have a simpler way to integrate it for all models thanks to [#35235](https://github.com/huggingface/transformers/pull/35235) do you want to help get this merged! 🤗 🚀\n\nContributorAuthor\n\n| |\n| --- |\n| [@ArthurZucker](https://github.com/ArthurZucker) Yes I can pick this up again 🚀 |\n\nvasqu, ArthurZucker, and zye1996 reacted with hooray emoji\n\n- ![tada](https://github.githubassets.com/assets/1f389-36899a2cb781.png)3 reactions\n\nmentioned this pull request\n[on Feb 14, 2025Feb 14, 2025](https://github.com/huggingface/transformers/pull/33522#ref-pullrequest-2852946465)\n\n[Flash Attention v3\\\\\n#36190](https://github.com/huggingface/transformers/pull/36190)\n\nClosed\n\n[on Apr 15, 2025Apr 15, 2025](https://github.com/huggingface/transformers/pull/33522#event-17268417185)\n\ndeleted the\n\nflash-attention-3\n\nbranch\n\n[last yearApril 15, 2025 12:30](https://github.com/huggingface/transformers/pull/33522#event-17268417807)\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/huggingface/transformers/pull/33522)"},{"resultType":"repo_readme","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"},{"resultType":"repo_readme","repo":"damienjose/cuda-flashattention","url":"https://github.com/damienjose/cuda-flashattention","readmeUrl":"https://raw.githubusercontent.com/damienjose/cuda-flashattention/HEAD/README.md","snippet":"# Flash Attention: Fast and Memory-Efficient Attention Mechanism\n\n## Overview\nFlash Attention is a fast, memory-efficient, and IO-aware implementation of attention mechanisms used in deep learning models. \nThis project showcases a simplified version of Flash Attention on GPUs using CUDA with randomly initialized matrix inputs, demonstrating how significant performance improvements can be achieved with optimized GPU compute, over traditional attention implementations.\n\n## Key Features\n1. **Fast**: Speeds up training for large models like BERT and GPT by up to 3x.\n2. **Memory-Efficient**: Reduces memory complexity from \\(O(N^2)\\) to \\(O(N)\\) by leveraging hardware memory hierarchies.\n3. **Exact**: Ensures accuracy without approximation.\n4. **IO-Aware**: Optimizes memory access and communication for modern GPUs.\n\n## Implementation Highlights\nFlash Attention is implemented in three versions:\n1. **CPU Implementation**: Provides a baseline implementation using sequential processing.\n2. **Naive GPU Implementation**: Uses CUDA global memory for parallel computations.\n3. **Optimized GPU Implementation**: Leverages shared memory, memory tiling, thread coarsening, and memory coalescing for maximum efficiency.\n\n## Key Optimizations\n1. **Memory Tiling**: Divides computation into manageable tiles to enhance memory locality.\n2. **Thread Coarsening**: Increases the work assigned per thread to reduce overhead.\n3. **Shared Memory**: Uses fast on-chip shared memory to minimize global memory access.\n4. **Memory Coalescing**: Ensures efficient memory access patterns for threads.\n\n## Experimental Results\nPerformance benchmarks were conducted on an NVIDIA GeForce RTX 4050 Laptop GPU with varying configurations. Below are the latency results for 1024 x 1024 matrices:\n\n| Implementation | Latency |\n|---------------------------|------------|\n| CPU | 9,427 ms |\n| GPU (Naive) | 38.78 ms |\n| GPU (Optimized Shared) | 14.12 ms |\n\nThe optimized implementation achieves a **3x speedup** compared to the naive GPU version and a **600x improvement** over the CPU baseline.\n\n## Source Code Structure\n- **CPU Implementation**: Implements scaled dot-product attention using sequential processing.\n- **Naive GPU Implementation**: Implements CUDA kernels using global memory for QK^T, softmax, and output computation.\n- **Optimized GPU Implementation**: Implements CUDA kernels using shared memory, tiling, and other optimizations for improved performance.\n\n## How to Build and Run\n1. Clone this repository:\n ```bash\n git clone https://github.com/damienjose/cuda-flashattention\n cd cuda-flashattention\n ```\n\n2. Set up your CUDA environment:\n - Install the latest CUDA toolkit.\n - Update the `sm` and `compute` settings in the Visual Studio project properties, specific for your GPU. See https://www.truehost.com/what-is-compute-capability-of-a-gpu/\n\n3. Build the project:\n ```bash\n nvcc -o cuda-flashattention kernel.cu -arch=sm_80\n ```\n\n4. Run the program:\n ```bash\n ./cuda-flashattention\n ```\n\n## Profiling and Performance Analysis\n- The naive GPU implementation is memory-bound due to frequent global memory accesses.\n- The optimized shared memory implementation transforms memory-bound operations into compute-bound ones.\n- Profiling tools like NVIDIA Nsight can be used for further performance analysis.\n\n## References\n1. [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135)\n2. [Understanding Flash-Attention and Flash-Attention-2](https://towardsai.net/p/artificial-intelligence/understanding-flash-attention-and-flash-attention-2-the-path-to-scale-the-context-lenght-of-language-models)\n3. [GitHub - Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)\n\n## Authors\n- Bryan Zhou ([bryanzhou2000@gmail.com](mailto:bryanzhou2000@gmail.com))\n- Damien Jose ([damien.jose@gmail.com](mailto:damien.jose@gmail.com))\n\n## Acknowledgments\nThis project was done as part of the AUT24 GPU Compute course at the University of Washington. We thank Professor Colin N. Reinhardt and TAs Arnab Karmakar, Xiaxi Shen for guiding us throughout this course. We also appreciate the authors of Flash Attention for their contributions and open sourcing their code."},{"resultType":"repo_readme","repo":"shlokvfx/mini-attention","url":"https://github.com/shlokvfx/mini-attention","readmeUrl":"https://raw.githubusercontent.com/shlokvfx/mini-attention/HEAD/Readme.md","snippet":"# Mini-Attention\n\n**FP16 Flash Attention 2 reimplemented from scratch in CUDA C++ — K7 beats cuDNN by 35% at N=512 on RTX 5090 (SM120). 139.5 TFLOPS, 272.5 GB/s DRAM bandwidth. SM100 (B200) WIP.**\n\n---\n\n## Why CUDA C++, not Triton?\n\nTriton's codegen exposes no control over `ldmatrix.x4.trans` scheduling, per-warp swizzle patterns, or async-copy double buffering — the three levers needed to close the gap with cuDNN at short sequences on SM86/SM120. The final kernel is hand-tuned PTX wrapped in CUDA C++. See [`notes/why_mma_sync.md`](notes/why_mma_sync.md).\n\n---\n\n## Kernels\n\n| Folder | Description |\n|--------|-------------|\n| [`kernels/flash_attn`](kernels/flash_attn/) | Flash Attention 2 from scratch: K1–K17 CUDA optimization progression on SM86, ported to SM120 Blackwell with Bc=128 / L2::256B / fast-reciprocal tuning |\n| [`kernels/standard`](kernels/standard/) | MHA, GQA, MQA variants — CUDA, Triton, and PyTorch reference implementations |\n\n---\n\n## Architectural notes\n\nDesign decisions are documented in [`notes/`](notes/):\n\n| File | Question answered |\n|------|-------------------|\n| [`notes/sm120_vs_sm100.md`](notes/sm120_vs_sm100.md) | Why SM120 ≠ SM100 — TMEM, `tcgen05.mma`, and what's actually on the RTX 5090 |\n| [`notes/why_mma_sync.md`](notes/why_mma_sync.md) | Why `mma.sync` instead of WGMMA — ISA availability on consumer Blackwell |\n| [`notes/pipeline_design.md`](notes/pipeline_design.md) | Tile sizes, double buffering, swizzle patterns, L2::256B rationale |\n\n---\n\n## Correctness\n\nAll K1–K7 pass numerical validation against PyTorch SDPA reference. Zero local memory spilling on every kernel.\n\n```bash\ncd /mnt/d/GITHUB/Mini-Attention\nLD_LIBRARY_PATH=/root/fa_env/lib/python3.12/site-packages/torch/lib \\\n /root/fa_env/bin/python tests/test_correctness.py\n```\n\n```\nConfig: B=2 H=8 N=512 D=128 fp16 tolerance=0.001\nReference: torch.nn.functional.scaled_dot_product_attention\n\n K1 max|out−ref| = 1.22e-04 [PASS]\n K2 max|out−ref| = 1.22e-04 [PASS]\n K3 max|out−ref| = 1.22e-04 [PASS]\n K4 max|out−ref| = 1.22e-04 [PASS]\n K5 max|out−ref| = 1.22e-04 [PASS]\n K6 max|out−ref| = 1.22e-04 [PASS]\n K7 max|out−ref| = 2.44e-04 [PASS]\n\nAll K1–K7 passed.\n```\n\n---\n\n## Benchmarks\n\nBaseline = `SDPBackend.CUDNN_ATTENTION` (fastest available attention on SM120). All runs on RTX 5090 (SM120).\n\n### Latency — K1–K7 vs cuDNN\n\n```\n N cuDNN (ms) K1 (ms) K3 (ms) K7 (ms) K7/cuDNN\n------------------------------------------------------------------\n 512 0.247 0.341 0.185 0.185 133.2%\n 1024 0.739 1.260 0.662 0.658 112.2%\n 2048 2.515 4.887 2.521 2.588 97.2%\n 4096 4.945 9.797 5.084 5.201 95.1%\n 8192 9.904 19.696 10.261 10.310 96.1%\n```\n\n**K7 beats cuDNN by 12–35% at N≤1024. Reaches 97–100% of cuDNN at N≥2048.**\n\nFull progression (% of cuDNN — lower = faster, >100% = faster than cuDNN):\n\n```\nConfig N=512 N=1024 N=2048 N=4096 N=8192\nK1 74.1% 59.6% 51.8% 50.4% 50.3%\nK2 121.6% 106.2% 93.5% 93.5% 92.4%\nK3 134.7% 114.4% 99.0% 97.3% 96.4%\nK4 132.3% 112.4% 97.5% 96.6% 94.6%\nK5 131.3% 110.6% 96.2% 94.4% 94.0%\nK6 132.6% 112.0% 97.7% 96.1% 95.0%\nK7 134.7% 114.4% 99.8% 99.2% 97.3%\n```\n\n![Latency vs N](figures/flash_attn_latency_vs_N.png)\n![Speedup vs SDPA](figures/flash_attn_speedup.png)\n\n### torch.profiler — B=1 H=16 N=1024 D=128, RTX 5090\n\n```\nK1 (base async-only): avg CUDA time: 0.005 ms\nK7 (L2::256B + frcp + Bc=128): avg CUDA time: 0.003 ms\nSpeedup K7 vs K1: 1.67×\n```\n\n### NCU-equivalent metrics — B=2 H=8 N=1024 D=128, RTX 5090 (419 TFLOPS FP16 peak, 1792 GB/s)\n\n```\nK1 (base)\n Duration: 112.55 µs | Compute: 18.2% ( 76.3 TFLOPS) | DRAM BW: 149.1 GB/s\n Local memory spill: 0\n\nK7 (SM120: L2::256B + fast reciprocal + Bc=128 tiles)\n Duration: 61.57 µs | Compute: 33.3% (139.5 TFLOPS) | DRAM BW: 272.5 GB/s\n Local memory spill: 0\n\nK7 vs K1: 1.83× speedup | K7 vs cuDNN at N=1024: beats by ~35%\n```\n\n![SM120 benchmark results](figures/bench_sm120_results.png)\n\n---\n\n## Profiling\n\n```bash\n# torch.profiler — CPU-side kernel timeline\nLD_LIBRARY_PATH=... python kernels/flash_attn/cuda/sm86/torch_profile.py\n\n# NCU-equivalent analytical metrics\nLD_LIBRARY_PATH=... python kernels/flash_attn/cuda/sm86/ncu_metrics.py\n\n# Per-kernel NCU report\nLD_LIBRARY_PATH=... python kernels/flash_attn/cuda/sm86/ncu_profile.py\n```\n\n---\n\n## WIP\n\n- **SM100 / B200 Blackwell** — `tcgen05.mma` TMEM-backed kernel in progress, pending hardware access. Consumer SM120 does not expose `tcgen05.mma` (datacenter Blackwell only). See [`notes/sm120_vs_sm100.md`](notes/sm120_vs_sm100.md).\n\n---\n\n## PR History\n\n- [PR #10](https://github.com/ShlokVFX/Mini-Attention/pull/10) — FlashAttention SM120 Blackwell: L2::256B, fast reciprocal, Bc=128 tiles + cuDNN baseline (RTX 5090)\n- [PR #9](https://github.com/ShlokVFX/Mini-Attention/pull/9) — SageAttention SM86: INT8 QK CUDA kernel + video quality validation\n- [PR #8](https://github.com/ShlokVFX/Mini-Attention/pull/8) — Simplified headers (K1–7) + K17s benchmark\n- [PR #7](https://github.com/ShlokVFX/Mini-Attention/pull/7) — Remove build artifacts, add gitignore\n- [PR #6](https://github.com/ShlokVFX/Mini-Attention/pull/6) — README overhaul: NCU metrics, FP16 from-scratch section, figures\n- [PR #5](https://github.com/ShlokVFX/Mini-Attention/pull/5) — Flash Attention FP16 from scratch: 16 CUDA iterations + inline pipeline (SM86)\n- [PR #3](https://github.com/ShlokVFX/Mini-Attention/pull/3) — Flash Attention v2: float4 LDS.128 + FP16 shmem (SM86)\n- [PR #1](https://github.com/ShlokVFX/Mini-Attention/pull/1) — SageAttention, StreamingLLM, PagedAttention, MoE kernels (SM86)\n\n---\n\n## Citations\n\n- Flash Attention 2 — Dao et al. 2023 · https://arxiv.org/abs/2307.08691\n- GQA — Ainslie et al. 2023 · https://arxiv.org/abs/2305.13245"},{"resultType":"repo_readme","repo":"tspeterkim/flash-attention-minimal","url":"https://github.com/tspeterkim/flash-attention-minimal","readmeUrl":"https://raw.githubusercontent.com/tspeterkim/flash-attention-minimal/HEAD/README.md","snippet":"# flash-attention-minimal\nA minimal re-implementation of Flash Attention with CUDA and PyTorch. \nThe official [implementation](https://github.com/Dao-AILab/flash-attention) can be quite daunting for a CUDA beginner\n(like myself), so this repo tries to be small and educational.\n\n* The entire forward pass is written in ~100 lines in `flash.cu`.\n* The variable names follow the notations from the original [paper](https://arxiv.org/abs/2205.14135).\n\n## Usage\n### Prerequisite\n* PyTorch (with CUDA)\n* `Ninja` for loading in C++\n\n### Benchmark\nCompare the wall-clock time between manual attention and minimal flash attention:\n```\npython bench.py\n```\n\nSample output on a [T4](https://aws.amazon.com/ec2/instance-types/g4/):\n```\n=== profiling manual attention ===\n...\nSelf CPU time total: 52.389ms\nSelf CUDA time total: 52.545ms\n\n=== profiling minimal flash attention === \n... \nSelf CPU time total: 11.452ms\nSelf CUDA time total: 3.908ms\n```\nSpeed-up achieved! \n\n### I don't have a GPU\nTry out this [online colab demo](https://colab.research.google.com/gist/tspeterkim/143bc7be7a845656817cf94c5228598e/demo-flash-attention-minimal.ipynb).\n\n## Caveats\n* No backward pass! To be honest, I found it a lot more complex than the forward pass, which was enough to show the\nuse of shared memory to avoid large N^2 read/writes.\n* In the inner loop, I assign each thread to a row of the output matrix. This differs from the original implementation.\n* This thread-per-row simplification makes the matrix multiplications very slow. This is probably why for longer \nsequences and larger block sizes, this gets slower than the manual implementation.\n* Q,K,Vs are in float32, unlike the original implementation which uses float16.\n* The block size is [fixed](https://github.com/tspeterkim/flash-attention-minimal/blob/9b7ca8ef4e6afdbfeb149a9cd488c8dea9af9ad6/flash.cu#L85) at compile time to 32.\n\n## Todos\n- [ ] Add backward pass\n- [x] Speed up matmults -> Check out the `tensorcore` branch\n- [ ] Dynamically set block size"},{"resultType":"github_history","repo":"oobabooga/textgen","url":"https://github.com/oobabooga/textgen/issues/2427","pageType":"merged_pr","number":2427,"segmentCount":1,"title":"Conversation","snippet":"## Conversation\n\nThis PR fixes the flash attention implementation that is currently in the repo and adds the `--flash-attention` flag to activate it.\n\nTo use it, install flash attention:\n\n```\npip install ninja\npip install flash-attn\n```\n\nIf that doesn't work (it should only take < 20 minutes), you need to compile from source with ninja:\n\n```\npip install ninja\ngit clone https://github.com/HazyResearch/flash-attention.git\ncd flash-attention\npython setup.py install\n```\n\nIf downstream users confirm it is working for them, I will add the pip install to the requirements.txt\n\nWith the fixed implementation, flash attention works as expected:\n\nRTX 3090 (24 GB VRAM)\n\ntsumeone SuperCOT 30B 128 groupsize CUDA:\n\n```\nMax new tokens: 1\n3001 tokens\nw/ flash attention: OOM\n\n2614 tokens:\nw/ flash attention: 23.7 GB\n```\n\nllama 13B with an 8 Rank LoRA applied:\n\n```\nMax new tokens: 1\n3001 tokens\nw flash attention: 12.7 GB\nw/o flash attention: 15.3 GB\n\n4002 tokens\nw flash attention: ~14 GB\nw/o flash attention: 17.7 GB\n```\n\nIn fact, I can get 6300 tokens on 13B.\n\nThe current implementation in the repo uses PyTorch's scaled\\_dot\\_product to perform the attention call. This method falls back to standard attention when flash attention or memory efficient attention is not available -- so what's wrong? PyTorch's implementation is somewhat old, and at the time, flash attention did not support attention masking in the forward pass, so trying to run it with the attention mask fails and falls back to normal attention. Since that time, HazyResearch has already updated their flash attention to work for Causal language modeling, so we can just use it directly from the source. Maybe in the future PyTorch will fix their version.\n\nAlso, all credit for the code goes to LAION who already implemented it for LLaMa:\n\n[https://github.com/LAION-AI/Open-Assistant/blob/main/model/model\\_training/models/patching\\_llama.py](https://github.com/LAION-AI/Open-Assistant/blob/main/model/model_training/models/patching_llama.py)\n\nI merely took the code there and adapted it to work without the rest of their scaffolding.\n\n**NOTE:**\n\n**This PR works, but it is not ready to be merged since flash attention requires either `pip install flash-attn` or compilation from the source by cloning HazyResearch's repo and following the source compiling instructions. I already had it compiled on my machine, so until it is confirmed that downstream users can compile on their machine or the pip install works for them (Personally it froze for me which is why I had to compile from source with ninja), I would not recommend pushing this through just yet.**\n\nBut once it is pushed I think everyone should use Flash Attention by default -- it exists for a reason, there is no catch!\n\n👍1jepjoo reacted with thumbs up emoji❤️4oobabooga, matatonic, Maykeye, and gr1336 reacted with heart emoji\n\n- 👍1 reaction\n- ❤️4 reactions\n\n`\n Add fixed Flash attention code\n`\n\n`\n c45cfe3\n`\n\n| |\n| --- |\n| Had to compile flash attention, which worked;
Started ooba with
CUDA\\_VISIBLE\\_DEVICE=0,1,2 python server.py --model tsumeone\\_llama-30b-supercot-4bit-128g-cuda --wbits 4 --groupsize 128 --flash-attention --model\\_type llama --pre\\_layer 16 36 60
(3 x 12gb rtx 3060)
Under Parameters, changed \"Truncate the prompt up to this length\" from 2048 to 2600
Fed it a short story, with the ending removed
2592 tokens in the input.
hit generate;
It wrote some stuff, but after ~200 tokens or so output becomes repetitive.
Output generated in 115.07 seconds (3.04 tokens/s, 350 tokens, context 2249, seed 1551651864)
Is this for specific models? |\n\nContributorAuthor\n\n| |\n| --- |\n| [@Fairfax-Mooresby](https://github.com/Fairfax-Mooresby)
No, it is for any model. The reason it is repetitive is well, most models aren't designed to go past 2048, since it is the base context limit of LLaMa. However, they can be trained to go beyond that, such as with Bluemoon 13B and 30B. You might want to try those as well. More models will come in the future to leverage higher context as well, I'm sure |\n\n| |\n| --- |\n| Flash attention always failed on my 3090 or P40/P6000. You found a way for it to work? Does it have to be compiled from source? |\n\n| |\n| --- |\n| [@kaiokendev](https://github.com/kaiokendev)
Unfortunately the increased token count for those models seems to break --pre\\_layer loading support, making 30b multi-gpu slow and use more Vram than it should. I'll give it a whirl on 13b and see how far i can push the context. Thanks for your work on this. |\n\nContributorAuthor\n\n| |\n| --- |\n| > Flash attention always failed on my 3090 or P40/P6000. You found a way for it to work? Does it have to be compiled from source?
I don't know about P40/P6000, it works on my 3090. This PR is using the direct implementation, and not via PyTorch since their implementation seems to be old |\n\n| |\n| --- |\n| I will try it on both. Less memory will definitely help training 30b |\n\n| |\n| --- |\n| This looks very promising. I am trying to get it running but I am probably doing something wrong. This command
```
pip install flash-attn
```
fails with error
```
File \"/tmp/pip-build-env-hmoohnwz/overlay/lib/python3.10/site-packages/setuptools/build_meta.py\", line 338, in run_setup
exec(code, locals())
File \"\", line 13, in
ModuleNotFoundError: No module named 'torch'
[end of output]
note: This error originates from a subprocess, and is likely not a problem with pip.
error: subprocess-exited-with-error
× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.
note: This error originates from a subprocess, and is likely not a problem with pip.
```
So I tried installing the previous version and it worked, although as you pointed out it took some 10 minutes to install:
```
pip install flash-attn==1.0.4
```
Starting the web UI with
```
python server.py --model llama-13b-4bit-128g --listen --flash-attention
```
works and I can generate with a 3700 tokens prompt on a RTX 3090, but it only outputs spaces.
For llama-30b I OOM without flash attention with a 1800 tokens context, and I also OOM with `--flash-attention` with the same context.
I'll try to install the latest version from source.
EDIT: same result installing flash-attn==1.0.6 from source :( |\n\n| |\n| --- |\n| I got it set up and installed fine. It doesn't compile a kernel for pascal so no go there. But I'm converting a 30b with groupsize 128 and see if it does help memory. From cursory looks at nvtop while generating, it used less. |\n\n| |\n| --- |\n| Tested all of these attention things last night on P40 and 3090. The only thing that gave slightly more tokens when generating at the edge of the 128g-30b was autogptq + quant\\_attn. Everything else went OOM about the same time. I got around 169X. The difference is small too. I got like 6 tokens generated before OOM vs 1. Goes to show there is no free lunch. That or something is wrong with the setup for everything including xformers but I don't think so. |\n\nContributorAuthor\n\n| |\n| --- |\n| > Tested all of these attention things last night on P40 and 3090. The only thing that gave slightly more tokens when generating at the edge of the 128g-30b was autogptq + quant\\_attn. Everything else went OOM about the same time. I got around 169X. The difference is small too. I got like 6 tokens generated before OOM vs 1. Goes to show there is no free lunch. That or something is wrong with the setup for everything including xformers but I don't think so.
Can you try with 13B and let me know your numbers with and without flash attention? |\n\n| |\n| --- |\n| Ok. I will try the 13b today. I have a bluemoonrp that does extra context in that size. |\n\n| |\n| --- |\n| Well, I checked with bluemoon RP and with wizard-storywriter-7b that has infinite context pretty much. I can get up to 8096ish for the former (it breaks) with or without flash attention. For the latter, somewhere around 6k (doesn't break).
I also notice GPU tends to go down to 30% utilization when I am using this. I compiled from source tho. |\n\n| |\n| --- |\n| hi~any progress about this commit? |\n\nContributorAuthor\n\n| |\n| --- |\n| [@laoda513](https://github.com/laoda513) I think with exllama I will close this PR since there is little to be gained for inference here. Personally I was never a fan of such a hacky solution to add flash attention especially when you cannot run the backward pass because LLaMA attn head dim size is 128. I also think the reason my result is different from others is due to memory fragmentation issues, and I don't have the time to debug that especially when exllama already fixed such issues. I have been using xformers as a replacement and it works perfectly, I even use it when training since it works on 3090 out of the box on CUDA 12.0, it is significantly improved and I don't have the headaches from flash attention. |\n\n👍1oobabooga reacted with thumbs up emoji\n\n- 👍1 reaction\n\n[on Jul 3, 2023Jul 3, 2023](https://github.com/oobabooga/textgen/pull/2427#event-9707996562)\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/oobabooga/textgen/pull/2427)"},{"resultType":"repo_readme","repo":"hao-ai-lab/flash-attention-fp4","url":"https://github.com/hao-ai-lab/flash-attention-fp4","readmeUrl":"https://raw.githubusercontent.com/hao-ai-lab/flash-attention-fp4/HEAD/README.md","snippet":"# FlashAttention (FP4 Fork)\n\n ## FP4 Flash Attention 4 on B200\n\n> [!NOTE]\n> **This is a fork of [flash-attention](https://github.com/Dao-AILab/flash-attention) with the addition of a CuTe DSL FP4/FP8 implementation of Flash Attention 4 on B200 for the [Attn-QAT](https://arxiv.org/abs/2603.00040) paper.** Supports NVFP4 and MXFP8 block-scaled QK with BF16 or FP8 PV, achieving up to **1.31x speedup** over BF16 FA4 (2018 vs 1545 TFLOPS). Peaks at **2018 TFLOPS** (NVFP4+FP8), **1948 TFLOPS** (MXFP8+FP8), and **1920 TFLOPS** (NVFP4+BF16) on B200. See **[flash_attn/cute/README.md](flash_attn/cute/README.md)** for results and usage.\n\n---\n\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\nimport flash_attn_interface\nflash_attn_interface.flash_attn_func()\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- `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). 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~~. 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. MI200 or MI300 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 the [Flash Attention v2](https://tridao.me/publications/flash2/flash2.pdf) is currently a work in progress.\n\nIt supports AMD's CDNA (MI200, MI300) and RDNA GPU's using fp16, bf16 and fp32 datatypes.\n\nThese features are supported in Fwd and Bwd\n1) Fwd and Bwd with causal masking\n2) Variable sequence lengths\n3) Arbitrary Q and KV sequence lengths\n4) Arbitrary head sizes\n5) Multi and grouped query attention\n6) Dropout\n7) Rotary embeddings\n8) ALiBi\n\nWe are working on the following things\n1) Paged Attention \n2) Sliding Window\n3) FP8\n4) Performance Improvements\n\n##### Getting Started\nTo get started with the triton backend for AMD, follow the steps below.\n\nFirst install the recommended Triton version \n\n```\npip install triton==3.2.0\n```\nThen install Flash Attention with the flag `FLASH_ATTENTION_TRITON_AMD_ENABLE` set to `\"TRUE\"`.\n\n```\ncd flash-attention\ngit checkout main_perf\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" python setup.py install\n```\n\nTo test that things are working, you can run our tests. These tests take hours so you don't need to run the full thing.\n```\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pytest tests/test_flash_attn_triton_amd.py\n```\n\nYou can use autotune for better performance by using this flag `FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=\"TRUE\"`\n```\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=\"TRUE\" python $PATH_TO_CODE\n```\n\n###### Docker\nYou can also use the Dockerfile below which does the above steps on top of the latest rocm/pytorch image.\n```\nFROM rocm/pytorch:latest\n\nWORKDIR /workspace\n\n# install triton\nRUN pip install triton==3.2.0\n\n# install flash attention\nENV FLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\"\n\nRUN git clone https://github.com/ROCm/flash-attention.git &&\\ \n cd flash-attention &&\\\n git checkout main_perf &&\\\n python setup.py install\n\n# set working dir\nWORKDIR /workspace/flash-attention\n```\n\nTo build the docker file\n```\ndocker build -t fa_triton .\n```\n\nTo run the docker image\n```\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 fa_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## 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```"}]}