{"id":15600943,"url":"https://github.com/lucidrains/memory-efficient-attention-pytorch","last_synced_at":"2025-04-04T09:09:23.467Z","repository":{"id":38395944,"uuid":"465595488","full_name":"lucidrains/memory-efficient-attention-pytorch","owner":"lucidrains","description":"Implementation of a memory efficient multi-head attention as proposed in the paper, \"Self-attention Does Not Need O(n²) Memory\"","archived":false,"fork":false,"pushed_at":"2023-07-18T14:24:24.000Z","size":35755,"stargazers_count":374,"open_issues_count":4,"forks_count":35,"subscribers_count":8,"default_branch":"main","last_synced_at":"2025-03-28T08:05:55.092Z","etag":null,"topics":["artificial-intelligence","attention-mechanism","deep-learning","memory-efficient"],"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-03-03T06:26:31.000Z","updated_at":"2025-03-26T09:17:23.000Z","dependencies_parsed_at":"2024-06-21T04:27:21.952Z","dependency_job_id":null,"html_url":"https://github.com/lucidrains/memory-efficient-attention-pytorch","commit_stats":{"total_commits":62,"total_committers":3,"mean_commits":"20.666666666666668","dds":"0.032258064516129004","last_synced_commit":"d54f391370ecbf843a871f0e260425d076995550"},"previous_names":[],"tags_count":31,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemory-efficient-attention-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemory-efficient-attention-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemory-efficient-attention-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fmemory-efficient-attention-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/memory-efficient-attention-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247149502,"owners_count":20891954,"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-mechanism","deep-learning","memory-efficient"],"created_at":"2024-10-03T02:09:54.797Z","updated_at":"2025-04-04T09:09:23.444Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## Memory Efficient Attention Pytorch (obsolete)\n\nImplementation of a memory efficient multi-head attention as proposed in the paper, \u003ca href=\"https://arxiv.org/abs/2112.05682\"\u003eSelf-attention Does Not Need O(n²) Memory\u003c/a\u003e. In addition, the module will take care of masking, causal masking, as well as cross attention.\n\nThis repository also contains a \u003ca href=\"https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py\"\u003enaive non-CUDA implementation\u003c/a\u003e of the improvements made by \u003ca href=\"https://tridao.me/\"\u003eTri Dao\u003c/a\u003e with his \u003ca href=\"https://github.com/HazyResearch/flash-attention\"\u003eFlash Attention 2\u003c/a\u003e paper, for educational purposes. It is a game changer for attention and building long-context transformers.\n\nUpdate: from now on, you should just be using the \u003ca href=\"https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html?highlight=scaled_dot_product#torch.nn.functional.scaled_dot_product_attention\"\u003e`F.scaled_dot_product_attention`\u003c/a\u003e function in Pytorch 2.0 for built-in Flash Attention v1 support - or use Flash Attention v2 at the \u003ca href=\"https://github.com/Dao-AILab/flash-attention\"\u003eofficial repository\u003c/a\u003e\n\n## Install\n\n```bash\n$ pip install memory-efficient-attention-pytorch\n```\n\n## Usage\n\nFor autoregressive language model\n\n```python\nimport torch\nfrom memory_efficient_attention_pytorch import Attention\n\nattn = Attention(\n    dim = 512,\n    dim_head = 64,                # dimension per head\n    heads = 8,                    # number of attention heads\n    causal = True,                # autoregressive or not\n    memory_efficient = True,      # whether to use memory efficient attention (can be turned off to test against normal attention)\n    q_bucket_size = 1024,         # bucket size along queries dimension\n    k_bucket_size = 2048          # bucket size along key / values dimension\n).cuda()\n\nx = torch.randn(1, 65536, 512).cuda()\nout = attn(x) # (1, 65536, 512)\n```\n\nCross attention\n\n```python\nimport torch\nfrom memory_efficient_attention_pytorch import Attention\n\ncross_attn = Attention(\n    dim = 512,\n    dim_head = 64,\n    heads = 8,\n    memory_efficient = True,\n    q_bucket_size = 1024,\n    k_bucket_size = 2048\n).cuda()\n\nx = torch.randn(1, 65536, 512).cuda()\ncontext = torch.randn(1, 65536, 512).cuda()\nmask = torch.ones(1, 65536).bool().cuda()\n\nout = cross_attn(x, context = context, mask = mask) # (1, 65536, 512)\n```\n\n## Citations\n\n```bibtex\n@misc{rabe2021selfattention,\n    title   = {Self-attention Does Not Need $O(n^2)$ Memory}, \n    author  = {Markus N. Rabe and Charles Staats},\n    year    = {2021},\n    eprint  = {2112.05682},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.LG}\n}\n```\n\n```bibtex\n@misc{liu2021swin,\n    title   = {Swin Transformer V2: Scaling Up Capacity and Resolution},\n    author  = {Ze Liu and Han Hu and Yutong Lin and Zhuliang Yao and Zhenda Xie and Yixuan Wei and Jia Ning and Yue Cao and Zheng Zhang and Li Dong and Furu Wei and Baining Guo},\n    year    = {2021},\n    eprint  = {2111.09883},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\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{dao2023flashattention2,\n  title     = {Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning,\n  author    = {Dao, Tri},\n  year      = {2023}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fmemory-efficient-attention-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fmemory-efficient-attention-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fmemory-efficient-attention-pytorch/lists"}