mamba-ssm

What it is and what it does

Mamba-ssm provides PyTorch implementations of Mamba state-space model architectures designed to match or exceed Transformer performance on language modeling while maintaining linear or near-linear time complexity. The package exposes multiple abstraction levels: the core selective SSM layer, the Mamba block (wrapping selective SSM), and complete language model backbones with repeating Mamba blocks. It includes three architecture variants—Mamba, Mamba-2, and Mamba-3—each with different parameter efficiency and design principles, along with pretrained models trained on the Pile dataset at scales from 130M to 2.8B parameters.

The package is built for GPU-accelerated inference and training on NVIDIA hardware, with hardware-aware optimizations inspired by FlashAttention. It requires ten runtime dependencies spanning PyTorch, Triton, specialized kernel libraries (quack-kernels, tilelang, apache-tvm-ffi), and standard ML tools (transformers, einops, packaging). Installation demands careful attention to build isolation and CUDA version compatibility, making it primarily suitable for researchers and practitioners already working in GPU-accelerated deep learning environments.

Use it for:

• Replace Transformer blocks in language models when linear-time scaling is critical for long sequences.
• Benchmark state-space models against Transformer baselines using pretrained Mamba models on standard NLP tasks.
• Integrate Mamba blocks into custom sequence models where selective state-space mechanisms offer efficiency gains.
• Fine-tune pretrained Mamba models (130M–2.8B parameters) on downstream language tasks via Hugging Face.
• Experiment with Mamba-3's MIMO and structured state-space duality for research on sequence modeling architectures.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Mamba-ssm implements state-space model architectures (Mamba, Mamba-2, Mamba-3) as alternatives to Transformers for sequence modeling, optimized for language tasks and information-dense data.

Yes, if you have GPU infrastructure (NVIDIA, CUDA 11.6+, Linux) and are researching or deploying sequence models where Transformer alternatives are worth exploring. The active maintenance, zero known vulnerabilities, and permissive license support this. No, if you lack GPU access, need CPU-only inference, or require a simpler install path—the high dependency friction and platform specificity make it unsuitable for casual experimentation or resource-constrained environments.

Install

mamba-ssm on PyPI

pip install mamba-ssm

uv add mamba-ssm

poetry add mamba-ssm

Installing mamba-ssm

Before you install

Installation requires high friction: ten runtime dependencies including torch, triton, and specialized kernels (quack-kernels, tilelang, apache-tvm-ffi), plus mandatory compilation flags and NVIDIA GPU with CUDA 11.6+. The package is actively maintained (last commit 2026-07-22, 97 days since latest release) but the build process is complex and platform-specific.

License in practice

Apache License 2.0 is permissive and allows commercial use, modification, and distribution with minimal restrictions. You must include license and copyright notices in derivative works, but no patent retaliation clause applies.

Quickstart

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

import torch
from mamba_ssm import Mamba

x = torch.randn(2, 64, 16).to("cuda")
model = Mamba(d_model=16, d_state=16, d_conv=4, expand=2).to("cuda")
y = model(x)

Requires Linux, NVIDIA GPU, PyTorch 1.12+, CUDA 11.6+, and --no-build-isolation flag to avoid CPU-only PyTorch being installed during build.

Verify before relying

• Whether pretrained model downloads (mamba-130m through mamba-2.8b) are automatically cached or require manual setup.
• Performance comparison to Transformers on specific tasks beyond the paper's evaluation set.
• AMD GPU support status and any additional prerequisites beyond those listed.
• Whether Mamba-3 requires source installation for all use cases or if PyPI packages support it.

Package facts

Evidence: mamba_ssm-2.3.2.post1.tar.gz

Keywords: cuda, pytorch, state-space model

Tags

More Artificial Intelligence packages