Mamba architecture

--- name: mamba-architecture description: State-space model with O(n) complexity vs. Transformers' O(n²). 5× faster inference, million-token sequences, no KV cache. Selective SSM with hardware-aware design. Mamba-1 (dstate =16) and Mamba-2 (dstate =128, multi-head). Models ranging from 130M to 2.8B on HuggingFace. version: 1.0.0 author: Orchestra Research license: MIT tags: [Model Architecture, Mamba, State Space Models, SSM, Linear Complexity, Long Context, Efficient Inference, Hardware-Aware, Alternative to Transformers] dependencies: [mamba-ssm, torch, transformers, causal-conv1d] --- # Mamba - Selective State Space Models ## Quick Start Mamba is a state-space model architecture that achieves O(n) linear complexity for sequence modeling. Installation: ``bash # 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): python 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**: python 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**: python 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): python 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): python 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**: ``bash # Mamba Benchmark python benchmarks/benchmarkgenerationmamba_simple.py\ --m