[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"project-2306":3},{"id":4,"name":5,"fullName":6,"owner":7,"repo":5,"description":8,"homepage":9,"htmlUrl":10,"language":11,"languages":10,"totalLinesOfCode":10,"stars":12,"forks":13,"watchers":14,"openIssues":15,"contributorsCount":16,"subscribersCount":16,"size":16,"stars1d":17,"stars7d":18,"stars30d":19,"stars90d":16,"forks30d":16,"starsTrendScore":20,"compositeScore":21,"rankGlobal":10,"rankLanguage":10,"license":22,"archived":23,"fork":23,"defaultBranch":24,"hasWiki":25,"hasPages":23,"topics":26,"createdAt":10,"pushedAt":10,"updatedAt":27,"readmeContent":28,"aiSummary":29,"trendingCount":16,"starSnapshotCount":16,"syncStatus":30,"lastSyncTime":31,"discoverSource":32},2306,"flash-attention","Dao-AILab\u002Fflash-attention","Dao-AILab","Fast and memory-efficient exact attention","",null,"Python",24115,2823,150,1009,0,5,76,381,40,110,"BSD 3-Clause \"New\" or \"Revised\" License",false,"main",true,[],"2026-06-12 04:00:14","# FlashAttention\nThis repository provides the official implementation of FlashAttention and\nFlashAttention-2 from the\nfollowing papers.\n\n**FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness**  \nTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré  \nPaper: https:\u002F\u002Farxiv.org\u002Fabs\u002F2205.14135  \nIEEE Spectrum [article](https:\u002F\u002Fspectrum.ieee.org\u002Fmlperf-rankings-2022) about our submission to the MLPerf 2.0 benchmark using FlashAttention.\n![FlashAttention](assets\u002Fflashattn_banner.jpg)\n\n**FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning**  \nTri Dao\n\nPaper: https:\u002F\u002Ftridao.me\u002Fpublications\u002Fflash2\u002Fflash2.pdf\n\n![FlashAttention-2](assets\u002Fflashattention_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:\u002F\u002Fgithub.com\u002FDao-AILab\u002Fflash-attention\u002Fblob\u002Fmain\u002Fusage.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:\u002F\u002Ftridao.me\u002Fblog\u002F2024\u002Fflash3\u002F\n\nPaper: https:\u002F\u002Ftridao.me\u002Fpublications\u002Fflash3\u002Fflash3.pdf\n\n![FlashAttention-3 speedup on H100 80GB SXM5 with FP16](assets\u002Fflash3_fp16_fwd.png)\n\nThis is a beta release for testing \u002F benchmarking before we integrate that with\nthe rest of the repo.\n\nCurrently released:\n- FP16 \u002F BF16 forward and backward, FP8 forward\n\nRequirements: H100 \u002F 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## FlashAttention-4 (CuTeDSL)\n\nFlashAttention-4 is written in CuTeDSL and optimized for Hopper and Blackwell GPUs (e.g. H100, B200).\n\nTo install:\n```sh\npip install flash-attn-4\n```\n\nIf you're on CUDA 13, we recommend installing with the `cu13` extra for best performance:\n```sh\npip install \"flash-attn-4[cu13]\"\n```\n\nOnce installed, you can use it as follows:\n```python\nfrom flash_attn.cute import flash_attn_func\n\nout = flash_attn_func(q, k, v, causal=True)\n```\n\n## Installation and features\n**Requirements:**\n- CUDA toolkit or ROCm toolkit\n- PyTorch 2.2 and above.\n- `packaging` Python package (`pip install packaging`)\n- `psutil` Python package (`pip install psutil`)\n- `ninja` Python package (`pip install ninja`) *\n- Linux. Might work for Windows starting v2.3.2 (we've seen a few positive [reports](https:\u002F\u002Fgithub.com\u002FDao-AILab\u002Fflash-attention\u002Fissues\u002F595)) 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\u002Fflash_attention_interface.py`\n\n### NVIDIA CUDA Support\n**Requirements:**\n- CUDA 12.0 and above.\n\nWe recommend the\n[Pytorch](https:\u002F\u002Fcatalog.ngc.nvidia.com\u002Forgs\u002Fnvidia\u002Fcontainers\u002Fpytorch)\ncontainer from Nvidia, which has all the required tools to install FlashAttention.\n\nFlashAttention-2 with CUDA currently supports:\n1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). For Turing GPUs (T4, RTX 2080), see the separate [flash-attention-turing](https:\u002F\u002Fgithub.com\u002Fssiu\u002Fflash-attention-turing) repo, which supports a core subset of FlashAttention features on Turing.\n2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).\n3. All head dimensions up to 256. ~~Head dim > 192 backward requires A100\u002FA800 or H100\u002FH800~~. 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:\u002F\u002Fgithub.com\u002FROCm\u002Fcomposable_kernel) (ck) which is the default backend and a [Triton](https:\u002F\u002Fgithub.com\u002Ftriton-lang\u002Ftriton) backend. They provide an implementation of FlashAttention-2.\n\n**Requirements:**\n- ROCm 6.0 and above.\n\nWe recommend the\n[Pytorch](https:\u002F\u002Fhub.docker.com\u002Fr\u002Frocm\u002Fpytorch)\ncontainer from ROCm, which has all the required tools to install FlashAttention.\n\n#### Composable Kernel Backend\nFlashAttention-2 ROCm CK backend currently supports:\n1. MI200x, MI250x, MI300x, MI355x, and RDNA 3\u002F4 GPUs.\n2. Datatype fp16 and bf16\n3. Both forward's and backward's head dimensions up to 256.\n4. RDNA 3 GPUs do not currently support backward, and RDNA 4 GPUs support backward only with deterministic=False\n\n#### Triton Backend\nThe Triton implementation of [Flash Attention](https:\u002F\u002Ftridao.me\u002Fpublications\u002Fflash2\u002Fflash2.pdf) supports AMD's CDNA (MI200, MI300) and RDNA GPUs using fp16, bf16, and fp32 datatypes. It provides forward and backward passes with causal masking, variable sequence lengths, arbitrary Q\u002FKV sequence lengths and head sizes, MQA\u002FGQA, dropout, rotary embeddings, ALiBi, paged attention, and FP8 (via the Flash Attention v3 interface). Sliding window attention is currently a work in progress.\n\nThe Triton backend kernels are provided by the [aiter](https:\u002F\u002Fgithub.com\u002FROCm\u002Faiter) package, included as a git submodule at `third_party\u002Faiter` and automatically installed during setup.\n\nTo install, first get PyTorch for ROCm from https:\u002F\u002Fpytorch.org\u002Fget-started\u002Flocally\u002F, then install Flash Attention:\n```sh\ncd flash-attention\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pip install --no-build-isolation .\n```\n\nTo use a specific aiter commit (e.g., for testing or development):\n```sh\ncd flash-attention\ncd third_party\u002Faiter && git fetch origin && git checkout \u003Ccommit-sha> && cd ..\u002F..\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pip install --no-build-isolation .\n```\n\nTo run the tests (note: full suite takes hours):\n```sh\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pytest tests\u002Ftest_flash_attn_triton_amd.py\n```\n\nThe Triton backend uses a default kernel configuration optimized for determinism and reasonable performance across workloads. For peak throughput, enable `FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=\"TRUE\"` to search for optimal settings, which incurs a one-time warmup cost.\n\nAlternativly, if _not_ autotuning, `FLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON` may be used to set a single triton config overriding the hardcoded defaults for `attn_fwd`. E.g.\n```sh\nFLASH_ATTENTION_FWD_TRITON_AMD_CONFIG_JSON='{\"BLOCK_M\":128,\"BLOCK_N\":64,\"waves_per_eu\":1,\"PRE_LOAD_V\":false,\"num_stages\":1,\"num_warps\":8}'\n```\n\nFor a quick start with Docker:\n```dockerfile\nFROM rocm\u002Fpytorch:latest\n\nWORKDIR \u002Fworkspace\n\n# build flash attention with triton backend\nRUN git clone https:\u002F\u002Fgithub.com\u002FDao-AILab\u002Fflash-attention &&\\ \n    cd flash-attention &&\\\n    FLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pip install --no-build-isolation .\n\n# set working dir\nWORKDIR \u002Fworkspace\u002Fflash-attention\n\n# set env variable to use triton backend\nENV FLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\"\n```\n\nBuild and run:\n```sh\ndocker build -t flash-attn-triton .\ndocker run -it --network=host --user root --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --shm-size 16G --device=\u002Fdev\u002Fkfd --device=\u002Fdev\u002Fdri flash-attn-triton\n```\n\n## How to use FlashAttention\n\nThe main functions implement scaled dot product attention (softmax(Q @ K^T *\nsoftmax_scale) @ V):\n```python\nfrom flash_attn import flash_attn_qkvpacked_func, flash_attn_func\n```\n\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False,\n                          window_size=(-1, -1), alibi_slopes=None, deterministic=False):\n\"\"\"dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\nArguments:\n    qkv: (batch_size, seqlen, 3, nheads, headdim)\n    dropout_p: float. Dropout probability.\n    softmax_scale: float. The scaling of QK^T before applying softmax.\n        Default to 1 \u002F 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\u002FGQA) 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 \u002F 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\u002Fvalues from\n    the previous step, and update them with the new keys\u002Fvalues from the current step, and do\n    attention with the updated cache, all in 1 kernel.\n\n    If you pass in k \u002F 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\u002Ftest_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\u002FGQA) 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 \u002F 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 \u002F 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 \u002F 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 \u002F 2, 1 & rotary_dim \u002F 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:\u002F\u002Fgithub.com\u002FDao-AILab\u002Fflash-attention\u002Fblob\u002Fmain\u002Fflash_attn\u002Fmodules\u002Fmha.py).\n\n### Using with 🤗 Kernels\n\nIf your hardware environment belongs to any of the above-mentioned, you can also use the [`kernels` library](https:\u002F\u002Fgithub.com\u002Fhuggingface\u002Fkernels)\nto use Flash Attention 2 and 3 right away.\n\n```py\n# pip install kernels\n\nfrom kernels import get_kernel\n\n# FA2\nfa_module = get_kernel(\"kernels-community\u002Fflash-attn2\", version=1)\nflash_attn_func = fa_module.flash_attn_func\n\n# FA3\nfa3_module = get_kernel(\"kernels-community\u002Fflash-attn3\", version=1)\nflash_attn_func = fa3_module.flash_attn_func\n```\n\n## Changelog\n\n### 2.0: Complete rewrite, 2x faster\nUpgrading from FlashAttention (1.x) to FlashAttention-2\n\nThese functions have been renamed:\n- `flash_attn_unpadded_func` -> `flash_attn_varlen_func`\n- `flash_attn_unpadded_qkvpacked_func` -> `flash_attn_varlen_qkvpacked_func`\n- `flash_attn_unpadded_kvpacked_func` -> `flash_attn_varlen_kvpacked_func`\n\nIf the inputs have the same sequence lengths in the same batch, it is simpler\nand faster to use these functions:\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)\n```\n```python\nflash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)\n```\n### 2.1: Change behavior of causal flag\n\nIf seqlen_q != seqlen_k and causal=True, the causal mask is aligned to the\nbottom right corner of the attention matrix, instead of the top-left corner.\n\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 =\nmasked out) is:  \nv2.0:  \n    1 0 0 0 0  \n    1 1 0 0 0  \nv2.1:  \n    1 1 1 1 0  \n    1 1 1 1 1  \n\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is:  \nv2.0:  \n    1 0  \n    1 1  \n    1 1  \n    1 1  \n    1 1  \nv2.1:  \n    0 0  \n    0 0  \n    0 0  \n    1 0  \n    1 1  \nIf the row of the mask is all zero, the output will be zero.\n\n### 2.2: Optimize for inference\n\nOptimize for inference (iterative decoding) when query has very small sequence\nlength (e.g., query sequence length = 1). The bottleneck here is to load KV\ncache as fast as possible, and we split the loading across different thread\nblocks, with a separate kernel to combine results.\n\nSee the function `flash_attn_with_kvcache` with more features for inference\n(perform rotary embedding, updating KV cache inplace).\n\nThanks to the xformers team, and in particular Daniel Haziza, for this\ncollaboration.\n\n### 2.3: Local (i.e., sliding window) attention\n\nImplement sliding window attention (i.e., local attention). Thanks to [Mistral\nAI](https:\u002F\u002Fmistral.ai\u002F) and in particular Timothée Lacroix for this\ncontribution. Sliding window was used in the [Mistral 7B](https:\u002F\u002Fmistral.ai\u002Fnews\u002Fannouncing-mistral-7b\u002F) 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:\u002F\u002Farxiv.org\u002Fabs\u002F2309.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\u003C!-- * [RTX 3090](#rtx-3090) -->\n\u003C!-- * [T4](#t4) -->\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 \u002F seqlen.\n\n#### Speedup\n\n![FlashAttention speedup on A100 80GB SXM5 with FP16\u002FBF16](assets\u002Fflash2_a100_fwd_bwd_benchmark.png)\n\n#### Memory\n\n![FlashAttention memory](assets\u002Fflashattn_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\u002FBF16](assets\u002Fflash2_h100_fwd_bwd_benchmark.png)\n\n## Full model code and training script\n\nWe have released the full GPT model\n[implementation](https:\u002F\u002Fgithub.com\u002FDao-AILab\u002Fflash-attention\u002Fblob\u002Fmain\u002Fflash_attn\u002Fmodels\u002Fgpt.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\u002Fsec per A100, equivalent to 72% model FLOPs utilization (we don't need\nany activation checkpointing).\n\nWe also include a training\n[script](https:\u002F\u002Fgithub.com\u002FDao-AILab\u002Fflash-attention\u002Ftree\u002Fmain\u002Ftraining) 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:\u002F\u002Fgithub.com\u002Fopenai\u002Ftriton\u002Fblob\u002Fmaster\u002Fpython\u002Ftutorials\u002F06-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:\u002F\u002Fgithub.com\u002FDao-AILab\u002Fflash-attention\u002Fblob\u002Fmain\u002Fflash_attn\u002Fflash_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 \u002F non-causal).\n\nTo run the tests:\n```sh\npytest -q -s tests\u002Ftest_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\u002Ftest_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```\n","FlashAttention 是一个快速且内存高效的精确注意力机制实现项目。它通过IO感知技术优化了传统注意力机制在计算速度和内存使用上的表现，特别适用于需要高效处理大规模数据集的深度学习模型中。FlashAttention-2 进一步改进了并行性和工作分区策略以提高性能；而 FlashAttention-3 和 4 则针对最新的 NVIDIA Hopper 和 Blackwell 架构进行了专门优化，支持 FP16、BF16 及 FP8 数据类型，在这些硬件上实现了显著的速度提升。该项目适合于对计算资源敏感的应用场景，如自然语言处理任务中的长序列建模等。",2,"2026-06-11 02:49:25","top_language"]