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

Mamba Architecture Details

Selective State Space Mechanism

Mamba's core innovation is the Selective SSM (S6) layer that makes state space model parameters input-dependent.

How S6 Works

Traditional SSMs (non-selective):

# Fixed A, B, C matrices for all inputs
h(t) = A * h(t-1) + B * x(t)  # State update
y(t) = C * h(t)                # Output

Mamba's Selective SSM:

# Input-dependent parameters
B(t) = Linear_B(x(t))  # Selection mechanism
C(t) = Linear_C(x(t))  # Output projection
Δ(t) = Linear_Δ(x(t))  # Discretization step

# Selective state update
h(t) = discretize(A, Δ(t)) * h(t-1) + Δ(t) * B(t) * x(t)
y(t) = C(t) * h(t)

Key Advantages

1. Content-based reasoning:

• Can selectively remember or forget based on input
• Addresses discrete modality weakness of traditional SSMs
• Example: Remembers important tokens, forgets padding

2. Input-dependent selection:

# Mamba decides per token what to remember
if is_important(x(t)):
    Δ(t) = large_value   # Keep in state
else:
    Δ(t) = small_value   # Forget quickly

3. No attention required:

• Replaces O(n²) attention with O(n) state updates
• State dimension is constant (typically 16)

Model Configuration

Core Parameters

from mamba_ssm import Mamba

model = Mamba(
    d_model=256,      # Hidden dimension (256, 512, 768, 1024, 2048)
    d_state=16,       # SSM state dimension (fixed at 16 is optimal)
    d_conv=4,         # Local convolution width (4 is standard)
    expand=2,         # Expansion factor (1.5-2.0)
    dt_rank="auto",   # Rank of Δ projection (auto = d_model / 16)
    dt_min=0.001,     # Min Δ init (controls forgetting rate)
    dt_max=0.1,       # Max Δ init
    dt_init="random", # Δ initialization (random, constant)
    dt_scale=1.0,     # Δ scaling factor
    conv_bias=True,   # Use bias in convolution
    bias=False        # Use bias in linear projections
)

Parameter Impact

d_state (SSM state dimension):

• Standard: 16 (optimal from ablations)
• Smaller (8): Faster but less capacity
• Larger (32, 64): Minimal improvement, 2× slower

expand (block expansion):

• Standard: 2.0
• Range: 1.5-2.0
• Controls inner dimension = expand * d_model

d_conv (convolution width):

• Standard: 4
• Local context window before SSM
• Helps with positional information

dt_rank (Δ projection rank):

• Auto: d_model / 16 (recommended)
• Controls Δ parameter efficiency
• Lower rank = more efficient but less expressive

Mamba Block Structure

# Mamba block (replaces Transformer block)
class MambaBlock(nn.Module):
    def __init__(self, d_model):
        self.norm = RMSNorm(d_model)
        self.mamba = Mamba(d_model, d_state=16, d_conv=4, expand=2)

    def forward(self, x):
        return x + self.mamba(self.norm(x))  # Residual

# Full model (stack of Mamba blocks)
model = nn.Sequential(
    Embedding(...),
    *[MambaBlock(d_model) for _ in range(n_layers)],
    RMSNorm(d_model),
    LMHead(...)
)

Key differences from Transformers:

• No multi-head attention (MHA)
• No feedforward network (FFN)
• Single Mamba layer per block
• 2× more layers than equivalent Transformer

Hardware-Aware Implementation

Parallel Algorithm

Mamba uses a scan-based parallel algorithm for training:

# Parallel mode (training)
# GPU kernel fuses operations
y = parallel_scan(A, B, C, x)  # O(n log n) parallel

# Sequential mode (inference)
# Constant memory RNN-style
h = 0
for x_t in sequence:
    h = A*h + B*x_t
    y_t = C*h

Memory Efficiency

Training:

• Recomputes activations in backward pass
• Similar to FlashAttention strategy
• Memory: O(batch_size seq_len d_model)

Inference:

• RNN-style sequential processing
• State size: O(d_model * d_state) = constant
• No KV cache needed (huge advantage!)

CUDA Kernel Optimizations

# Fused kernel operations
- Discretization (continuous → discrete A, B)
- SSM recurrence (parallel scan)
- Convolution (efficient 1D conv)
- All in single GPU kernel

Layer Count Scaling

Mamba models use 2× layers compared to Transformers:

Why 2× layers?

• Mamba blocks are simpler (no MHA, no FFN)
• ~50% fewer parameters per layer
• Doubling layers matches compute budget

Initialization Strategy

# Δ (discretization step) initialization
dt_init_floor = 1e-4
dt = torch.exp(
    torch.rand(d_inner) * (math.log(dt_max) - math.log(dt_min))
    + math.log(dt_min)
).clamp(min=dt_init_floor)

# A (state transition) initialization
A = -torch.exp(torch.rand(d_inner, d_state))  # Negative for stability

# B, C (input/output) initialization
B = torch.randn(d_inner, d_state)
C = torch.randn(d_inner, d_state)

Critical for stability:

• A must be negative (exponential decay)
• Δ in range [dt_min, dt_max]
• Random initialization helps diversity

Resources

• Paper: https://arxiv.org/abs/2312.00752 (Mamba-1)
• Paper: https://arxiv.org/abs/2405.21060 (Mamba-2)
• GitHub: https://github.com/state-spaces/mamba
• Models: https://huggingface.co/state-spaces
• CUDA kernels: https://github.com/state-spaces/mamba/tree/main/csrc

Mamba Performance Benchmarks

Inference Speed Comparison

Throughput (tokens/sec)

Mamba-1.4B vs Transformer-1.3B on single A100 80GB:

Key insight: Speedup grows with sequence length (Mamba O(n) vs Transformer O(n²))

Latency (ms per token)

Generation latency (batch size 1, autoregressive):

Mamba advantage: Constant per-token latency regardless of context length

Memory Usage

Training Memory (BF16, per GPU)

Mamba-1.4B training memory breakdown:

Note: Transformer OOMs at 8K sequence length on 40GB A100

Inference Memory (FP16, batch size 1)

Mamba stores no KV cache - constant memory per token!

Actual Mamba state size:

• 130M: ~3 MB (d_model × d_state × n_layers = 768 × 16 × 24)
• 2.8B: ~13 MB (2560 × 16 × 64)

Language Modeling Benchmarks

Perplexity on Common Datasets

Models trained on The Pile (300B tokens):

Mamba consistently outperforms Transformers of similar size by 10-20%

Zero-Shot Task Performance

Mamba-2.8B vs Transformer-2.7B on common benchmarks:

Average improvement: +2.2 points across benchmarks

Audio Modeling Benchmarks

SC09 (Speech Commands)

Task: Audio classification (10 classes)

LJSpeech (Speech Generation)

Task: Text-to-speech quality (MOS score)

RTF (Real-Time Factor): Lower is better (0.21 = 5× faster than real-time)

Genomics Benchmarks

Human Reference Genome (HG38)

Task: Next nucleotide prediction

Mamba handles million-length sequences efficiently

Scaling Laws

Compute-Optimal Training

FLOPs vs perplexity (The Pile validation):

Scaling coefficient: Mamba achieves same perplexity as Transformer with 0.8× compute

Parameter Efficiency

Perplexity 10.0 target on The Pile:

Mamba needs ~30% fewer parameters for same performance

Long-Range Arena (LRA)

Task: Long-context understanding benchmarks

Average: Mamba 85.0%, S4 83.4%, Transformer 54.4%

Training Throughput

Tokens/sec During Training

8× A100 80GB cluster, BF16, different sequence lengths:

Mamba scales to longer sequences without OOM

Hardware Utilization

GPU Memory Bandwidth

Mamba-1.4B inference on different GPUs:

High efficiency: Mamba is memory-bandwidth bound (good!)

Multi-GPU Scaling

Mamba-2.8B training throughput:

Near-linear scaling up to 16 GPUs

Cost Analysis

Training Cost (USD)

Training to The Pile perplexity 10.0 on cloud GPUs:

Savings: 36% cost reduction vs Transformer

Inference Cost (USD/million tokens)

API-style inference (batch size 1, 2K context):

Mamba provides: 2.6× faster, 57% cheaper, better quality

Resources

• Benchmarks code: https://github.com/state-spaces/mamba/tree/main/benchmarks
• Paper (Mamba-1): https://arxiv.org/abs/2312.00752 (Section 4: Experiments)
• Paper (Mamba-2): https://arxiv.org/abs/2405.21060 (Section 5: Experiments)
• Pretrained models: https://huggingface.co/state-spaces

Mamba Training Guide

Training from Scratch

Setup Environment

# Install dependencies
pip install torch>=1.12.0 --extra-index-url https://download.pytorch.org/whl/cu116
pip install packaging ninja
pip install causal-conv1d>=1.1.0
pip install mamba-ssm

# Verify CUDA
python -c "import torch; print(torch.cuda.is_available())"

Basic Training Loop

import torch
from mamba_ssm import Mamba
from torch.utils.data import DataLoader

# Model setup
model = Mamba(
    d_model=512,
    d_state=16,
    d_conv=4,
    expand=2
).cuda()

# Optimizer (same as GPT)
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=6e-4,
    betas=(0.9, 0.95),
    weight_decay=0.1
)

# Training loop
for batch in dataloader:
    inputs, targets = batch
    inputs, targets = inputs.cuda(), targets.cuda()

    # Forward
    logits = model(inputs)
    loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))

    # Backward
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()

Distributed Training

Single-Node Multi-GPU (DDP)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# Initialize process group
dist.init_process_group("nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# Wrap model
model = Mamba(...).cuda()
model = DDP(model, device_ids=[local_rank])

# Train
optimizer = torch.optim.AdamW(model.parameters(), lr=6e-4)
for batch in dataloader:
    loss = compute_loss(model, batch)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Launch:

torchrun --nproc_per_node=8 train.py

Multi-Node Training

# Node 0 (master)
torchrun --nproc_per_node=8 \
  --nnodes=4 --node_rank=0 \
  --master_addr=$MASTER_ADDR --master_port=29500 \
  train.py

# Node 1-3 (workers)
torchrun --nproc_per_node=8 \
  --nnodes=4 --node_rank=$NODE_RANK \
  --master_addr=$MASTER_ADDR --master_port=29500 \
  train.py

Mixed Precision Training

BF16 (Recommended)

from torch.cuda.amp import autocast, GradScaler

# BF16 (no scaler needed on A100/H100)
for batch in dataloader:
    with autocast(dtype=torch.bfloat16):
        logits = model(inputs)
        loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

FP16 (with gradient scaling)

scaler = GradScaler()

for batch in dataloader:
    with autocast(dtype=torch.float16):
        logits = model(inputs)
        loss = F.cross_entropy(logits.view(-1, vocab_size), targets.view(-1))

    optimizer.zero_grad()
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

Hyperparameter Recommendations

Learning Rate Schedule

# Cosine decay with warmup (GPT-3 style)
def get_lr(it, warmup_iters=2000, lr_decay_iters=600000):
    max_lr = 6e-4
    min_lr = 6e-5

    # Warmup
    if it < warmup_iters:
        return max_lr * it / warmup_iters

    # Decay
    if it > lr_decay_iters:
        return min_lr

    # Cosine
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (max_lr - min_lr)

# Apply in training loop
for it, batch in enumerate(dataloader):
    lr = get_lr(it)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

Batch Size Recommendations

# Gradient accumulation
accumulation_steps = 8
optimizer.zero_grad()

for i, batch in enumerate(dataloader):
    loss = compute_loss(model, batch) / accumulation_steps
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()

Optimizer Configuration

# AdamW (recommended)
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=6e-4,           # Peak learning rate
    betas=(0.9, 0.95), # Standard for LLMs
    eps=1e-8,
    weight_decay=0.1   # Important for generalization
)

# Weight decay exemptions (optional)
decay = set()
no_decay = set()
for name, param in model.named_parameters():
    if 'norm' in name or 'bias' in name:
        no_decay.add(param)
    else:
        decay.add(param)

optimizer = torch.optim.AdamW([
    {'params': list(decay), 'weight_decay': 0.1},
    {'params': list(no_decay), 'weight_decay': 0.0}
], lr=6e-4, betas=(0.9, 0.95))

Memory Optimization

Gradient Checkpointing

from torch.utils.checkpoint import checkpoint

class MambaBlock(nn.Module):
    def __init__(self, d_model, use_checkpoint=False):
        super().__init__()
        self.use_checkpoint = use_checkpoint
        self.norm = RMSNorm(d_model)
        self.mamba = Mamba(d_model)

    def forward(self, x):
        if self.use_checkpoint and self.training:
            return x + checkpoint(self._forward, x, use_reentrant=False)
        return x + self._forward(x)

    def _forward(self, x):
        return self.mamba(self.norm(x))

# Enable for training
model = MambaLM(use_checkpoint=True)

Memory savings: ~30-40% with minimal speed impact

Flash Attention Integration

Mamba's CUDA kernels already use flash-attention-style optimizations:

• Fused operations in single kernel
• Recomputation in backward pass
• No intermediate activation storage

Long Context Training

Sequence Length Progression

# Start short, increase gradually
training_stages = [
    {'seq_len': 512,  'iters': 50000},
    {'seq_len': 1024, 'iters': 100000},
    {'seq_len': 2048, 'iters': 150000},
    {'seq_len': 4096, 'iters': 200000},
]

for stage in training_stages:
    dataloader = create_dataloader(seq_len=stage['seq_len'])
    train(model, dataloader, max_iters=stage['iters'])

Memory Requirements (Batch Size 1)

Mamba advantage: Memory grows linearly, Transformers grow quadratically

Common Training Issues

Issue: OOM during training

Solution 1: Reduce batch size

per_gpu_batch = 8  # Reduce from 16
gradient_accumulation = 8  # Increase from 4

Solution 2: Enable gradient checkpointing

model = MambaLM(use_checkpoint=True)

Solution 3: Use smaller sequence length

seq_len = 1024  # Reduce from 2048

Issue: Training unstable (loss spikes)

Solution 1: Check gradient norm

grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
print(f"Grad norm: {grad_norm}")  # Should be < 10

Solution 2: Lower learning rate

max_lr = 3e-4  # Reduce from 6e-4

Solution 3: Check Δ initialization

# Ensure dt_min, dt_max are reasonable
model = Mamba(
    d_model=512,
    dt_min=0.001,  # Not too small
    dt_max=0.1     # Not too large
)

Issue: Slow training speed

Solution 1: Verify CUDA kernels installed

import mamba_ssm
print(mamba_ssm.__version__)  # Should have CUDA kernels

Solution 2: Use BF16 on A100/H100

with autocast(dtype=torch.bfloat16):  # Faster than FP16
    loss = model(inputs)

Solution 3: Increase batch size if possible

per_gpu_batch = 16  # Increase from 8 (better GPU utilization)

Checkpointing

Save/Load Model

# Save
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'iter': iteration,
    'config': model_config
}
torch.save(checkpoint, f'checkpoint_{iteration}.pt')

# Load
checkpoint = torch.load('checkpoint_100000.pt')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
iteration = checkpoint['iter']

Best Practices

# Save every N iterations
if iteration % save_interval == 0:
    save_checkpoint(model, optimizer, iteration)

# Keep only last K checkpoints
checkpoints = sorted(glob.glob('checkpoint_*.pt'))
if len(checkpoints) > keep_last:
    for ckpt in checkpoints[:-keep_last]:
        os.remove(ckpt)

Resources

• Training code: https://github.com/state-spaces/mamba/tree/main/benchmarks
• Pretrained models: https://huggingface.co/state-spaces
• CUDA installation: https://github.com/state-spaces/mamba#installation