{"id":14436175,"url":"https://github.com/tommyip/mamba2-minimal","last_synced_at":"2025-07-13T01:33:41.276Z","repository":{"id":244394944,"uuid":"815102867","full_name":"tommyip/mamba2-minimal","owner":"tommyip","description":"Minimal Mamba-2 implementation in PyTorch","archived":false,"fork":false,"pushed_at":"2024-06-17T11:21:52.000Z","size":31,"stargazers_count":181,"open_issues_count":3,"forks_count":15,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-04-11T20:13:00.408Z","etag":null,"topics":["llm","machine-learning","mamba","pytorch","ssm"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/tommyip.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":"2024-06-14T11:04:57.000Z","updated_at":"2025-04-10T09:12:52.000Z","dependencies_parsed_at":"2024-12-22T12:15:25.599Z","dependency_job_id":"1d85a0fa-9cb2-4add-8464-fccf37343fe9","html_url":"https://github.com/tommyip/mamba2-minimal","commit_stats":{"total_commits":8,"total_committers":1,"mean_commits":8.0,"dds":0.0,"last_synced_commit":"4f4413731f24303b4e046962bce6ca5b6a3b13d6"},"previous_names":["tommyip/mamba2-minimal"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/tommyip/mamba2-minimal","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tommyip%2Fmamba2-minimal","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tommyip%2Fmamba2-minimal/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tommyip%2Fmamba2-minimal/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tommyip%2Fmamba2-minimal/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/tommyip","download_url":"https://codeload.github.com/tommyip/mamba2-minimal/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/tommyip%2Fmamba2-minimal/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":265080004,"owners_count":23708098,"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":["llm","machine-learning","mamba","pytorch","ssm"],"created_at":"2024-08-31T05:01:09.664Z","updated_at":"2025-07-13T01:33:41.173Z","avatar_url":"https://github.com/tommyip.png","language":"Python","funding_links":[],"categories":["Tutorials"],"sub_categories":["Codes"],"readme":"# mamba2-minimal\n\nA minimal, single-file implementation of the Mamba-2 model in PyTorch.\n\n![Mamba-2](https://github.com/state-spaces/mamba/blob/f9dbb4fdb2705d71282e0db184d177c6375623f0/assets/ssd_algorithm.png)\n\u003e **Transformers are SSMs: Generalized Models and Efficient Algorithms**\\\n\u003e     **Through Structured State Space Duality**\\\n\u003e Tri Dao*, Albert Gu*\\\n\u003e Paper: https://arxiv.org/abs/2405.21060\n\nMamba is a new class of foundation models, most notable for _not_ being based on the Transformer architecture. Instead it is in the family of State Space Models (SSMs) that maps a sequence through a hidden state in the fashion of RNNs. This approach enables linear scaling in computation and memory with respect to sequence length during training (unlike transformer's quadratic complexity), as well as constant time per step during inference. Mamba-2 builds upon Mamba-1 by imposing additional constraints on certain SSM parameters, allowing it to have much larger state dimensions and significantly improved training speed.\n\nThis implementation is device agnostic and have been tested to work on the CPU and MPS (Metal Performance Shaders) backends. The model's output logits follow the same distribution as the reference implementation but are not numerically equivalent. \n\n## Usage\n\nInstall dependencies (`torch`, `einops` and `transformers`):\n\n```\npip install -r requirements.txt\n```\n\n**See [demo.ipynb](./demo.ipynb) for using Mamba-2 as part of an end-to-end language model with pretrained weights for text generation.**\n\nThe core Mamba-2 model can be used as follows:\n\n```py\nimport torch\n\nfrom mamba2 import Mamba2, Mamba2Config\n\nconfig = Mamba2Config(d_model=768)\nmodel = Mamba2(config)\n\nx = torch.randn(2, 64, 768)  # (batch, seqlen, d_model)\ny = model(x)  # same shape as x\n```\n\n## TODOs\n\n- [x] Constant time (wrt sequence length) autoregressive inference\n- [ ] Remove dependency on `einops` (depends on whether resulting code is still readable)\n\n## Credits\n\n* [Albert Gu], [Tri Dao] - authors of the Mamba-2 architecture\n* [John Ma] - author of [johnma2006/mamba-minimal], who inspired this repo\n\n## Resources\n\nSome resources to understand Mamba and SSMs.\n\n* [Mamba-1/2 reference implementation]\n* [Mamba-1 paper]\n* [Mamba-2 paper]\n* [The Annotated S4] (literate programming for the S4 model)\n* [Mamba-2 blog post]\n\n[Albert Gu]: https://github.com/albertfgu\n[Tri Dao]: https://github.com/tridao\n[John Ma]: https://github.com/johnma2006\n[johnma2006/mamba-minimal]: https://github.com/johnma2006/mamba-minimal\n[Mamba-1 paper]: https://arxiv.org/abs/2312.00752\n[Mamba-2 paper]: https://arxiv.org/abs/2405.21060\n[The Annotated S4]: https://srush.github.io/annotated-s4/\n[Mamba-2 blog post]: https://tridao.me/blog/2024/mamba2-part1-model/\n[Mamba-1/2 reference implementation]: https://github.com/state-spaces/mamba\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftommyip%2Fmamba2-minimal","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ftommyip%2Fmamba2-minimal","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ftommyip%2Fmamba2-minimal/lists"}