Mamba-3 public release

.gitattributes

@@ -0,0 +1,5 @@
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,105 @@
+---
+license: apache-2.0
+library_name: mamba-ssm
+tags:
+- mamba
+- mamba-3
+- mimo
+- state-space-model
+- causal-lm
+- language-modeling
+---
+
+# Mamba-3 MIMO 895M
+
+`mamba3-mimo-894m` is a pretrained causal language model with 895M
+parameters. It is built from stacked blocks, each containing a Mamba-3
+MIMO mixer followed by a gated MLP. It contains no attention layers and
+is released with BF16 weights in the public `mamba_ssm` checkpoint format.
+
+## Model architecture
+
+| Property | Value |
+|---|---:|
+| Parameters | 895M |
+| Layers | 24 |
+| Model dimension | 1,536 |
+| SSM state size | 128 |
+| SSM head dimension | 64 |
+| SSM heads | 48 |
+| SSM groups | 1 |
+| MIMO rank | 4 |
+| Chunk size | 16 |
+| Context length | 2,048 |
+
+The model was pretrained on 100B tokens from
+[FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu).
+It uses tied input and output embeddings and the `meta-llama/Llama-3.1-8B` tokenizer.
+
+## Installation
+
+Install CUDA-enabled PyTorch first, followed by the latest Mamba source:
+
+```bash
+pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation
+```
+
+While this repository is private, authenticate with Hugging Face:
+
+```bash
+hf auth login
+```
+
+## Usage
+
+```python
+import torch
+from transformers import AutoTokenizer
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+
+model_id = "state-spaces/mamba3-mimo-894m"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "meta-llama/Llama-3.1-8B",
+)
+
+model = MambaLMHeadModel.from_pretrained(
+    model_id,
+    device="cuda",
+    dtype=torch.bfloat16,
+)
+model.eval()
+
+input_ids = tokenizer(
+    "Mamba-3 is",
+    return_tensors="pt",
+).input_ids.cuda()
+
+with torch.inference_mode():
+    logits = model(input_ids).logits
+
+print(logits.shape)
+```
+
+## References
+
+- [Mamba-3: Improved Sequence Modeling using State Space Principles](https://arxiv.org/abs/2603.15569)
+- [Official Mamba implementation](https://github.com/state-spaces/mamba)
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{lahoti2026mamba3improvedsequencemodeling,
+  title         = {Mamba-3: Improved Sequence Modeling using State Space Principles},
+  author        = {Aakash Lahoti and Kevin Y. Li and Berlin Chen and
+                   Caitlin Wang and Aviv Bick and J. Zico Kolter and
+                   Tri Dao and Albert Gu},
+  year          = {2026},
+  eprint        = {2603.15569},
+  archivePrefix = {arXiv},
+  primaryClass  = {cs.LG},
+  url           = {https://arxiv.org/abs/2603.15569}
+}
+```

config.json

@@ -0,0 +1,29 @@
+{
+    "d_model": 1536,
+    "d_intermediate": 2800,
+    "n_layer": 24,
+    "vocab_size": 128256,
+    "ssm_cfg": {
+        "layer": "Mamba3",
+        "d_state": 128,
+        "expand": 2,
+        "headdim": 64,
+        "ngroups": 1,
+        "rope_fraction": 0.5,
+        "dt_min": 0.001,
+        "dt_max": 0.1,
+        "dt_init_floor": 0.0001,
+        "A_floor": 0.0001,
+        "chunk_size": 16,
+        "is_mimo": true,
+        "mimo_rank": 4,
+        "is_outproj_norm": false
+    },
+    "attn_layer_idx": [],
+    "attn_cfg": {},
+    "rms_norm": true,
+    "residual_in_fp32": true,
+    "fused_add_norm": false,
+    "pad_vocab_size_multiple": 16,
+    "tie_embeddings": true
+}

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4463f5fc22589bed79e3984fdfa56237ab29ae21ddc9f1eb3ee36b4b4fe28aa9
+size 1789217227