{"id":15600952,"url":"https://github.com/lucidrains/flash-attention-jax","last_synced_at":"2025-04-05T21:05:41.192Z","repository":{"id":44759694,"uuid":"512988377","full_name":"lucidrains/flash-attention-jax","owner":"lucidrains","description":"Implementation of Flash Attention in Jax","archived":false,"fork":false,"pushed_at":"2024-03-01T23:42:21.000Z","size":185,"stargazers_count":207,"open_issues_count":9,"forks_count":24,"subscribers_count":4,"default_branch":"main","last_synced_at":"2025-03-29T20:02:48.579Z","etag":null,"topics":["artificial-intelligence","attention-mechanisms","deep-learning","jax","long-context-attention"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-07-12T03:41:19.000Z","updated_at":"2025-03-29T00:49:40.000Z","dependencies_parsed_at":"2024-10-23T01:35:34.520Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/flash-attention-jax","commit_stats":{"total_commits":52,"total_committers":3,"mean_commits":"17.333333333333332","dds":"0.038461538461538436","last_synced_commit":"d73dcdd6cb4b4f8d786922e2ec6826036dae896a"},"previous_names":[],"tags_count":15,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fflash-attention-jax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fflash-attention-jax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fflash-attention-jax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fflash-attention-jax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/flash-attention-jax/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247399871,"owners_count":20932876,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","attention-mechanisms","deep-learning","jax","long-context-attention"],"created_at":"2024-10-03T02:10:09.229Z","updated_at":"2025-04-05T21:05:41.168Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"./flash-attention.png\" width=\"450px\"\u003e\u003c/img\u003e\n\n## Flash Attention - Jax\n\nImplementation of \u003ca href=\"https://arxiv.org/abs/2205.14135\"\u003eFlash Attention\u003c/a\u003e in Jax. It will likely not be as performant as with the \u003ca href=\"https://github.com/HazyResearch/flash-attention\"\u003eofficial CUDA version\u003c/a\u003e, given lack of ability for fine memory management. But just for educational purposes as well as to see how clever XLA compiler is (or is not).\n\n## Install\n\n```bash\n$ pip install flash-attention-jax\n```\n\n## Usage\n\n```python\nfrom jax import random\nfrom flash_attention_jax import flash_attention\n\nrng_key = random.PRNGKey(42)\n\nq = random.normal(rng_key, (1, 2, 131072, 512))  # (batch, heads, seq, dim)\nk = random.normal(rng_key, (1, 2, 131072, 512))\nv = random.normal(rng_key, (1, 2, 131072, 512))\nmask = random.randint(rng_key, (1, 131072,), 0, 2) # (batch, seq)\n\nout, _ = flash_attention(q, k, v, mask)\n\nout.shape  # (1, 2, 131072, 512) - (batch, heads, seq, dim)\n```\n\nQuick sanity check\n\n\n```python\nfrom flash_attention_jax import plain_attention, flash_attention, value_and_grad_difference\n\ndiff, (dq_diff, dk_diff, dv_diff) = value_and_grad_difference(\n    plain_attention,\n    flash_attention,\n    seed = 42\n)\n\nprint('shows differences between normal and flash attention for output, dq, dk, dv')\nprint(f'o: {diff}')       # \u003c 1e-4\nprint(f'dq: {dq_diff}')   # \u003c 1e-6\nprint(f'dk: {dk_diff}')   # \u003c 1e-6\nprint(f'dv: {dv_diff}')   # \u003c 1e-6\n```\n\nAutoregressive Flash Attention - GPT-like decoder attention\n\n```python\nfrom jax import random\nfrom flash_attention_jax import causal_flash_attention\n\nrng_key = random.PRNGKey(42)\n\nq = random.normal(rng_key, (131072, 512))\nk = random.normal(rng_key, (131072, 512))\nv = random.normal(rng_key, (131072, 512))\n\nout, _ = causal_flash_attention(q, k, v)\n\nout.shape  # (131072, 512)\n```\n\n## Todo\n\n- [x] leading dimensions for causal flash attention variant\n\n- [ ] figure out issue with jit and static argnums\n- [ ] comment with references to paper algorithms and explanations\n- [ ] make sure it can work one-headed key / values, as in PaLM\n\n## Citations\n\n```bibtex\n@article{Dao2022FlashAttentionFA,\n    title   = {FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness},\n    author  = {Tri Dao and Daniel Y. Fu and Stefano Ermon and Atri Rudra and Christopher R'e},\n    journal = {ArXiv},\n    year    = {2022},\n    volume  = {abs/2205.14135}\n}\n```\n\n```bibtex\n@article{Rabe2021SelfattentionDN,\n    title   = {Self-attention Does Not Need O(n2) Memory},\n    author  = {Markus N. Rabe and Charles Staats},\n    journal = {ArXiv},\n    year    = {2021},\n    volume  = {abs/2112.05682}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fflash-attention-jax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fflash-attention-jax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fflash-attention-jax/lists"}