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 torch from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16 x = torch.randn(batch, length, dim).to("cuda")

model = Mamba( d_model=dim, # Model dimension d_state=16, # SSM state dimension d_conv=4, # Conv1d kernel size expand=2 # Expansion factor ).to("cuda")

y = model(x) # O(n) complexity! assert y.shape == x.shape

Common workflows

Workflow 1: Language model with Mamba-2

Complete LM with generation:

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel from mamba_ssm.models.config_mamba import MambaConfig import torch

Configure Mamba-2 LM

config = MambaConfig( d_model=1024, # Hidden dimension n_layer=24, # Number of layers vocab_size=50277, # Vocabulary size ssm_cfg=dict( layer="Mamba2", # Use Mamba-2 d_state=128, # Larger state for Mamba-2 headdim=64, # Head dimension ngroups=1 # Number of groups ) )

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16)

Generate text

input_ids = torch.randint(0, 1000, (1, 20), device="cuda", dtype=torch.long) output = model.generate( input_ids=input_ids, max_length=100, temperature=0.7, top_p=0.9 )

Workflow 2: Use pretrained Mamba models

Load from HuggingFace:

from transformers import AutoTokenizer from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

Load pretrained model

model_name = "state-spaces/mamba-2.8b" tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") # Use compatible tokenizer model = MambaLMHeadModel.from_pretrained(model_name, device="cuda", dtype=torch.float16)

Generate

prompt = "The future of AI is" input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda") output_ids = model.generate( input_ids=input_ids, max_length=200, temperature=0.7, top_p=0.9, repetition_penalty=1.2 ) generated_text = tokenizer.decode(output_ids[0]) print(generated_text)

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):

from mamba_ssm import Mamba

model = Mamba( d_model=256, d_state=16, # Smaller state dimension d_conv=4, expand=2 ).to("cuda")

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

from mamba_ssm import Mamba2

model = Mamba2( d_model=256, d_state=128, # Larger state dimension d_conv=4, expand=2, headdim=64, # Head dimension for multi-head ngroups=1 # Parallel groups ).to("cuda")

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:

Benchmark Mamba

python benchmarks/benchmark_generation_mamba_simple.py
--model-name "state-spaces/mamba-2.8b"
--prompt "The future of machine learning is"
--topp 0.9 --temperature 0.7 --repetition-penalty 1.2

Benchmark Transformer

python benchmarks/benchmark_generation_mamba_simple.py
--model-name "EleutherAI/pythia-2.8b"
--prompt "The future of machine learning is"
--topp 0.9 --temperature 0.7 --repetition-penalty 1.2

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:

model = MambaLMHeadModel(config, device="cuda", dtype=torch.float16) model.gradient_checkpointing_enable() # Enable checkpointing

Issue: Slow installation

Install binary wheels (not source):

pip install mamba-ssm --no-build-isolation

Issue: Missing causal-conv1d

Install separately:

pip install causal-conv1d>=1.4.0

Issue: Model not loading from HuggingFace

Use MambaLMHeadModel.from_pretrained (not AutoModel ):

from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel model = MambaLMHeadModel.from_pretrained("state-spaces/mamba-2.8b")

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