mamba-architecture

1# Mamba - Selective State Space Models
2 
3## Quick start
4 
5Mamba is a state-space model architecture achieving O(n) linear complexity for sequence modeling.
6 
7**Installation**:
8```bash
9# Install causal-conv1d (optional, for efficiency)
10pip install causal-conv1d>=1.4.0
11 
12# Install Mamba
13pip install mamba-ssm
14# Or both together
15pip install mamba-ssm[causal-conv1d]
16```
17 
18**Prerequisites**: Linux, NVIDIA GPU, PyTorch 1.12+, CUDA 11.6+
19 
20**Basic usage** (Mamba block):
21```python
22import torch
23from mamba_ssm import Mamba
24 
25batch, length, dim = 2, 64, 16
26x = torch.randn(batch, length, dim).to("cuda")
27 
28model = Mamba(
29    d_model=dim,      # Model dimension
30    d_state=16,       # SSM state dimension
31    d_conv=4,         # Conv1d kernel size
32    expand=2          # Expansion factor
33).to("cuda")
34 
35y = model(x)  # O(n) complexity!
36assert y.shape == x.shape
37```
38 
39## Common workflows
40 
41### Workflow 1: Language model with Mamba-2
42 
43**Complete LM with generation**:
44```python
45from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
46from mamba_ssm.models.config_mamba import MambaConfig
47import torch
48 
49# Configure Mamba-2 LM
50config = MambaConfig(
51    d_model=1024,           # Hidden dimension
52    n_layer=24,             # Number of layers
53    vocab_size=50277,       # Vocabulary size
54    ssm_cfg=dict(
55        layer="Mamba2",     # Use Mamba-2
56        d_state=128,        # Larger state for Mamba-2
57        headdim=64,         # Head dimension
58        ngroups=1           # Number of groups
59    )
60)
61…