Crate burn_mamba 

burn-mamba — Mamba-1/2/3 selective state space models on Burn

A minimal, readable reference implementation of the Mamba-1, Mamba-2, and Mamba-3 SSM architectures on top of the Burn deep learning framework.

The goal is clarity: the official CUDA/Triton kernels are ported down to standard, portable Burn tensor operations, so the same code runs on every backend (CPU, WGPU, CUDA, Metal, LibTorch, …). There are no custom kernels.

Module families

Each family lives in its own module and follows the same composition (Network → Layers → Layer → Block):

• mamba1 — the original selective SSM (conv1d + sequential selective scan).
• mamba2 — Structured State Space Duality (SSD): the recurrence is recast as a chunkwise, GEMM-friendly algorithm.
• mamba3 — SSD extended with trapezoidal discretisation, data-dependent RoPE on B/C, and MIMO rank expansion.

Shared infrastructure lives in modules (the family-generic layer/network composition, plus activations, norms and losses) and utils (virtual-layer scheduling, class tokens, and the custom-backward plumbing).

Two execution modes

Every block, layer, and network exposes both a parallel forward() (used for training and prompt prefill) and a recurrent step() (used for token-by-token decoding). The two are mathematically equivalent: a forward() over a sequence equals unrolling step() token by token from the same initial cache — a parity property the test suites assert on outputs, final cache, and gradients.

Modules

mamba1

Mamba-1: the original selective state space model.

mamba2

Mamba-2: Structured State Space Duality (SSD).

mamba3

Mamba-3: trapezoidal SSD with data-dependent RoPE and MIMO.

modules

Family-generic composition (Layer/Layers/networks/bidi/caches) plus the shared neural modules (activations, norms, losses, tensor helpers).

prelude

Convenience re-exports: use burn_mamba::prelude::*; brings the enabled model families and their public types into scope.

utils

Virtual-layer/LR scheduling, class tokens, and custom-backward plumbing.

Macros

decl_ssd_autodiff_backend_ext

Declare a marker trait $autodiff_trait: Backend + $ext_trait + AutodiffBackend and a blanket impl for burn::backend::Autodiff<B> whenever B: $ext_trait.

impl_ssd_backend_ext_for_burn_backends

Emit impl $trait_name for <backend> {} blocks for every burn backend supported by this crate, opting in to the trait’s default body.

Constants

DENY_INF

DENY_NAN