mamba-architecture

Mamba - Selective State Space Models

Quick start

Mamba is a state-space model architecture achieving O(n) linear complexity for sequence modeling.

Installation:

# Install causal-conv1d (optional, for efficiency)
pip install causal-conv1d>=1.4.0
 
# Install Mamba
pip install mamba-ssm
# Or both together
pip install mamba-ssm[causal-conv1d]

Prerequisites: Linux, NVIDIA GPU, PyTorch 1.12+, CUDA 11.6+

Basic usage (Mamba block):

import

Common workflows

Workflow 1: Language model with Mamba-2

Complete LM with generation:

Workflow 2: Use pretrained Mamba models

Load from HuggingFace:

Available models:

• state-spaces/mamba-130m
• state-spaces/mamba-370m
• state-spaces/mamba-790m
• state-spaces/mamba-1.4b
• state-spaces/mamba-2.8b

Workflow 3: Mamba-1 vs Mamba-2

Mamba-1 (smaller state):

Mamba-2 (multi-head, larger state):

Key differences:

• State size: Mamba-1 (d_state=16) vs Mamba-2 (d_state=128)
• Architecture: Mamba-2 has multi-head structure
• Normalization: Mamba-2 uses RMSNorm
• Distributed: Mamba-2 supports tensor parallelism

Workflow 4: Benchmark vs Transformers

Generation speed comparison:

Expected results:

• Mamba: 5× faster inference
• Memory: No KV cache needed
• Scaling: Linear with sequence length

When to use vs alternatives

Use Mamba when:

• Need long sequences (100K+ tokens)
• Want faster inference than Transformers
• Memory-constrained (no KV cache)
• Building streaming applications
• Linear scaling important

Advantages:

• O(n) complexity: Linear vs quadratic
• 5× faster inference: No attention overhead
• No KV cache: Lower memory usage
• Million-token sequences: Hardware-efficient
• Streaming: Constant memory per token

Use alternatives instead:

• Transformers: Need best-in-class performance, have compute
• RWKV: Want RNN+Transformer hybrid
• RetNet: Need retention-based architecture
• Hyena: Want convolution-based approach

Common issues

Issue: CUDA out of memory

Reduce batch size or use gradient checkpointing:

Issue: Slow installation

Install binary wheels (not source):

Issue: Missing causal-conv1d

Install separately:

Issue: Model not loading from HuggingFace

Use MambaLMHeadModel.from_pretrained (not AutoModel):

Advanced topics

Selective SSM: See references/selective-ssm.md for mathematical formulation, state-space equations, and how selectivity enables O(n) complexity.

Mamba-2 architecture: See references/mamba2-details.md for multi-head structure, tensor parallelism, and distributed training setup.

Performance optimization: See references/performance.md for hardware-aware design, CUDA kernels, and memory efficiency techniques.

Hardware requirements

• GPU: NVIDIA with CUDA 11.6+
• VRAM:
  • 130M model: 2GB
  • 370M model: 4GB
  • 790M model: 8GB
  • 1.4B model: 14GB
  • 2.8B model: 28GB (FP16)
• Inference: 5× faster than Transformers
• Memory: No KV cache (lower than Transformers)

Performance (vs Transformers):

• Speed: 5× faster inference
• Memory: 50% less (no KV cache)
• Scaling: Linear vs quadratic

Resources

• Paper (Mamba-1): https://arxiv.org/abs/2312.00752 (Dec 2023)
• Paper (Mamba-2): https://arxiv.org/abs/2405.21060 (May 2024)
• GitHub: https://github.com/state-spaces/mamba ⭐ 13,000+
• Models: https://huggingface.co/state-spaces
• Docs: Repository README and wiki