{"id":"mamba","name":"mamba-architecture","summary":"O(n) の複雑さを持つ状態空間モデルとトランスフォーマーの O(n²) モデル。5×より高速な推論、百万トークンのシーケンス、KVキャッシュなし。","body":"# Mamba - Selective State Space Models\n\n## Quick start\n\nMamba is a state-space model architecture achieving O(n) linear complexity for sequence modeling.\n\n**Installation**:\n```bash\n# Install causal-conv1d (optional, for efficiency)\npip install causal-conv1d>=1.4.0\n\n# Install Mamba\npip install mamba-ssm\n# Or both together\npip install mamba-ssm[causal-conv1d]\n```\n\n**Prerequisites**: Linux, NVIDIA GPU, PyTorch 1.12+, CUDA 11.6+\n\n**Basic usage** (Mamba block):\n```python\nimport torch\nfrom mamba_ssm import Mamba\n\nbatch, length, dim = 2, 64, 16\nx = torch.randn(batch, length, dim).to(\"cuda\")\n\nmodel = Mamba(\n    d_model=dim,      # Model dimension\n    d_state=16,       # SSM state dimension\n    d_conv=4,         # Conv1d kernel size\n    expand=2          # Expansion factor\n).to(\"cuda\")\n\ny = model(x)  # O(n) complexity!\nassert y.shape == x.shape\n```\n\n## Common workflows\n\n### Workflow 1: Language model with Mamba-2\n\n**Complete LM with generation**:\n```python\nfrom mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\nfrom mamba_ssm.models.config_mamba import MambaConfig\nimport torch\n\n# Configure Mamba-2 LM\nconfig = MambaConfig(\n    d_model=1024,           # Hidden dimension\n    n_layer=24,             # Number of layers\n    vocab_size=50277,       # Vocabulary size\n    ssm_cfg=dict(\n        layer=\"Mamba2\",     # Use Mamba-2\n        d_state=128,        # Larger state for Mamba-2\n        headdim=64,         # Head dimension\n        ngroups=1           # Number of groups\n    )\n)\n\nmodel = MambaLMHeadModel(config, device=\"cuda\", dtype=torch.float16)\n\n# Generate text\ninput_ids = torch.randint(0, 1000, (1, 20), device=\"cuda\", dtype=torch.long)\noutput = model.generate(\n    input_ids=input_ids,\n    max_length=100,\n    temperature=0.7,\n    top_p=0.9\n)\n```\n\n### Workflow 2: Use pretrained Mamba models\n\n**Load from HuggingFace**:\n```python\nfrom transformers import AutoTokenizer\nfrom mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\n\n# Load pretrained model\nmodel_name = \"state-spaces/mamba-2.8b\"\ntokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neox-20b\")  # Use compatible tokenizer\nmodel = MambaLMHeadModel.from_pretrained(model_name, device=\"cuda\", dtype=torch.float16)\n\n# Generate\nprompt = \"The future of AI is\"\ninput_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids.to(\"cuda\")\noutput_ids = model.generate(\n    input_ids=input_ids,\n    max_length=200,\n    temperature=0.7,\n    top_p=0.9,\n    repetition_penalty=1.2\n)\ngenerated_text = tokenizer.decode(output_ids[0])\nprint(generated_text)\n```\n\n**Available models**:\n- `state-spaces/mamba-130m`\n- `state-spaces/mamba-370m`\n- `state-spaces/mamba-790m`\n- `state-spaces/mamba-1.4b`\n- `state-spaces/mamba-2.8b`\n\n### Workflow 3: Mamba-1 vs Mamba-2\n\n**Mamba-1** (smaller state):\n```python\nfrom mamba_ssm import Mamba\n\nmodel = Mamba(\n    d_model=256,\n    d_state=16,      # Smaller state dimension\n    d_conv=4,\n    expand=2\n).to(\"cuda\")\n```\n\n**Mamba-2** (multi-head, larger state):\n```python\nfrom mamba_ssm import Mamba2\n\nmodel = Mamba2(\n    d_model=256,\n    d_state=128,     # Larger state dimension\n    d_conv=4,\n    expand=2,\n    headdim=64,      # Head dimension for multi-head\n    ngroups=1        # Parallel groups\n).to(\"cuda\")\n```\n\n**Key differences**:\n- **State size**: Mamba-1 (d_state=16) vs Mamba-2 (d_state=128)\n- **Architecture**: Mamba-2 has multi-head structure\n- **Normalization**: Mamba-2 uses RMSNorm\n- **Distributed**: Mamba-2 supports tensor parallelism\n\n### Workflow 4: Benchmark vs Transformers\n\n**Generation speed comparison**:\n```bash\n# Benchmark Mamba\npython benchmarks/benchmark_generation_mamba_simple.py \\\n  --model-name \"state-spaces/mamba-2.8b\" \\\n  --prompt \"The future of machine learning is\" \\\n  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2\n\n# Benchmark Transformer\npython benchmarks/benchmark_generation_mamba_simple.py \\\n  --model-name \"EleutherAI/pythia-2.8b\" \\\n  --prompt \"The future of machine learning is\" \\\n  --topp 0.9 --temperature 0.7 --repetition-penalty 1.2\n```\n\n**Expected results**:\n- **Mamba**: 5× faster inference\n- **Memory**: No KV cache needed\n- **Scaling**: Linear with sequence length\n\n## When to use vs alternatives\n\n**Use Mamba when**:\n- Need long sequences (100K+ tokens)\n- Want faster inference than Transformers\n- Memory-constrained (no KV cache)\n- Building streaming applications\n- Linear scaling important\n\n**Advantages**:\n- **O(n) complexity**: Linear vs quadratic\n- **5× faster inference**: No attention overhead\n- **No KV cache**: Lower memory usage\n- **Million-token sequences**: Hardware-efficient\n- **Streaming**: Constant memory per token\n\n**Use alternatives instead**:\n- **Transformers**: Need best-in-class performance, have compute\n- **RWKV**: Want RNN+Transformer hybrid\n- **RetNet**: Need retention-based architecture\n- **Hyena**: Want convolution-based approach\n\n## Common issues\n\n**Issue: CUDA out of memory**\n\nReduce batch size or use gradient checkpointing:\n```python\nmodel = MambaLMHeadModel(config, device=\"cuda\", dtype=torch.float16)\nmodel.gradient_checkpointing_enable()  # Enable checkpointing\n```\n\n**Issue: Slow installation**\n\nInstall binary wheels (not source):\n```bash\npip install mamba-ssm --no-build-isolation\n```\n\n**Issue: Missing causal-conv1d**\n\nInstall separately:\n```bash\npip install causal-conv1d>=1.4.0\n```\n\n**Issue: Model not loading from HuggingFace**\n\nUse `MambaLMHeadModel.from_pretrained` (not `AutoModel`):\n```python\nfrom mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\nmodel = MambaLMHeadModel.from_pretrained(\"state-spaces/mamba-2.8b\")\n```\n\n## Advanced topics\n\n**Selective SSM**: See [references/selective-ssm.md](references/selective-ssm.md) for mathematical formulation, state-space equations, and how selectivity enables O(n) complexity.\n\n**Mamba-2 architecture**: See [references/mamba2-details.md](references/mamba2-details.md) for multi-head structure, tensor parallelism, and distributed training setup.\n\n**Performance optimization**: See [references/performance.md](references/performance.md) for hardware-aware design, CUDA kernels, and memory efficiency techniques.\n\n## Hardware requirements\n\n- **GPU**: NVIDIA with CUDA 11.6+\n- **VRAM**:\n  - 130M model: 2GB\n  - 370M model: 4GB\n  - 790M model: 8GB\n  - 1.4B model: 14GB\n  - 2.8B model: 28GB (FP16)\n- **Inference**: 5× faster than Transformers\n- **Memory**: No KV cache (lower than Transformers)\n\n**Performance** (vs Transformers):\n- **Speed**: 5× faster inference\n- **Memory**: 50% less (no KV cache)\n- **Scaling**: Linear vs quadratic\n\n## Resources\n\n- Paper (Mamba-1): https://arxiv.org/abs/2312.00752 (Dec 2023)\n- Paper (Mamba-2): https://arxiv.org/abs/2405.21060 (May 2024)\n- GitHub: https://github.com/state-spaces/mamba ⭐ 13,000+\n- Models: https://huggingface.co/state-spaces\n- Docs: Repository README and wiki","author":"@Orchestra-Research","ownerProfile":null,"authorContacts":null,"sourceUrl":"https://github.com/Orchestra-Research/AI-Research-SKILLs/tree/main/01-model-architecture/mamba","license":"MIT","category":"coding","lang":"en","tokens":1938,"stars":0,"calls30d":1,"claimed":false,"visibility":"public","origin":"crawler","version":"0.1.0","createdAt":"2026-08-22","updatedAt":"2026-08-22","files":[{"path":"references/architecture-details.md","size":5456,"sha256":"09cf810fbcc2949e889f502c910bca97be094f0efde759294d40300591dd77ea"},{"path":"references/benchmarks.md","size":8120,"sha256":"7f64a92015656c1ee318b23d11f8d3951a03c00be8dcf1614e01c6133f8a269f"},{"path":"references/training-guide.md","size":9012,"sha256":"e382cfc6e0b6af4a04862001c983ea746b4ad327a1eec43b8b2955035b3b0fb0"}],"requires":{"mcp":[],"tools":[]},"safety":{"flags":[],"scannedAt":"2026-08-22","hasScripts":false,"networkEndpoints":["arxiv.org","download.pytorch.org","huggingface.co"]}}