Mamba

Part I: The mental model

input x : (B, L, d_model)
      |
      v
in_proj            d_model -> 2 * d_inner     (d_inner = expand * d_model)
      |  split into two branches
      +--> x branch (B, L, d_inner)      +--> z branch = the gate
      |
      v
causal conv1d (depthwise, width d_conv) + SiLU     short-range mixing
      |
      v
x_proj             d_inner -> dt_rank + 2 * d_state
      |  split three ways, all input-dependent
      +--> dt -> dt_proj -> softplus = Delta   (B, L, d_inner)
      +--> B  (B, L, d_state)
      +--> C  (B, L, d_state)
      |
      v
selective_scan (S6)    A from A_log, discretize with Delta, scan over L
      |                custom CUDA kernel, state (B, d_inner, d_state) in SRAM
      v
y = y + D * x          per-channel skip
y = y * SiLU(z)        gated by the z branch
      |
      v
out_proj           d_inner -> d_model
      |
      v
output : (B, L, d_model)

The one-sentence identity: Mamba is a state-space model made selective, the recurrence parameters that control what the hidden state keeps and forgets are computed from the current token instead of being fixed, and that content-dependence is what lets a linear-time recurrence match attention on language. A classical state-space model reads an input sequence through a small hidden state with fixed matrices A, B, and C. It is elegant and cheap but it treats every token the same, so it cannot decide that one token matters and another is filler. Mamba lets B, C, and the step size Delta depend on the input, and suddenly the same hidden state can be told to reset, hold, or overwrite based on content. That is the whole architectural bet.

The second load-bearing idea is the price of the first. Once the recurrence is input-dependent it is no longer time-invariant, and a time-varying linear recurrence cannot be written as a single fixed convolution, so the FFT-based training path that made earlier models like S4 fast is gone. Mamba pays that bill with a hardware-aware parallel scan, a custom CUDA kernel that computes the recurrence directly while keeping the large expanded state in fast on-chip memory and never writing it to HBM. The design instinct is the same one behind FlashAttention, fuse the operation, keep the big intermediate on chip, and recompute in the backward pass instead of storing it. Everything below is described against the public repository as it stands in 2026. The project moves, so where a file path is likely to have shifted I name the component by its role and stay at concept level.

Part II: Using it

Mamba is a Linux-and-NVIDIA-GPU project. The scan and the depthwise convolution are compiled CUDA extensions, so a working install needs a CUDA toolkit whose version matches the one your PyTorch was built against, plus a compiler. The common path is the prebuilt wheel:

pip install mamba-ssm
# the fast depthwise causal conv1d lives in a companion package
pip install causal-conv1d>=1.4.0
# or build from source when a matching wheel is not published
git clone https://github.com/state-spaces/mamba
cd mamba
pip install -e .

If the CUDA extensions fail to build, Mamba still runs. The block has a slower pure-PyTorch reference path that is used when the compiled kernels or causal-conv1d are absent, which is enough for correctness checks on a small model though not for serious throughput. On CPU or macOS you can read the code and run the reference path on tiny inputs, but the point of the project is the GPU kernel.

The smallest useful program is a single Mamba block. It is a drop in replacement for a Transformer layer, same input and output shape, no attention mask, no positions:

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 1024, 512
x = torch.randn(batch, length, dim, device="cuda")

layer = Mamba(
    d_model=dim,   # the model width, unchanged through the block
    d_state=16,    # N, the SSM state expansion per channel
    d_conv=4,      # width of the local causal convolution
    expand=2,      # inner expansion, d_inner = expand * d_model
).to("cuda")

y = layer(x)              # (2, 1024, 512), same shape in and out
assert y.shape == x.shape

The second thing worth doing is running a real pretrained language model. The state-spaces organization publishes Mamba models trained on the Pile at 130M, 370M, 790M, 1.4B, and 2.8B parameters, plus Mamba-2 checkpoints. They use the GPT-NeoX tokenizer, so pair them with that vocabulary:

import torch
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
model = MambaLMHeadModel.from_pretrained(
    "state-spaces/mamba-130m", device="cuda", dtype=torch.float16
)

ids = tok("The key idea behind Mamba is", return_tensors="pt").input_ids.to("cuda")
out = model.generate(input_ids=ids, max_length=100, cg=True, temperature=0.7)
print(tok.decode(out[0]))

The cg=True flag captures the decode step into a CUDA graph, which matters because Mamba generation is a tight loop of tiny per-token launches where kernel launch overhead dominates. Under the hood generate keeps an InferenceParams object, a cache that holds the recurrent SSM state and the small convolution window for every layer. This is the payoff of the recurrent form, decoding token t+1 costs the same as token t regardless of how long the context already is, and memory does not grow with sequence length the way a Transformer KV cache does.

Now the mistakes people make. First, forgetting causal-conv1d, which does not error loudly but silently drops onto the slow reference conv and halves your throughput. Second, running in float32 out of habit, the kernels are tuned for bfloat16 and float16 and the SSM state is kept in float32 internally for stability, so passing fp16 or bf16 inputs is both correct and much faster. Third, treating d_state as free, the expanded state is d_inner * d_state per position and the naive path materializes it across the whole sequence, so a large d_state without the fused kernel is where people run out of memory. Fourth, and most conceptual, reaching for an attention mask. Mamba has no attention and no quadratic score matrix, causality is intrinsic to the left-to-right recurrence and the causal convolution, so there is nothing to mask and nothing that sees the future.

Part III: When it is the right tool

Mamba is the right tool when sequence length is the binding constraint and you want subquadratic cost without giving up too much quality. Long documents, audio and genomics where sequences run to hundreds of thousands of tokens, and high-throughput generation where the constant-memory recurrent decode is a real operational win. It is also the cleanest codebase for learning how a modern state-space model is actually built and made fast, which is a reason to read it even if you deploy something else.

The honest competition. Plain Transformers with FlashAttention remain the default and beat Mamba on any task that needs precise recall or copying from context, because attention can look directly at every earlier token while a fixed-size state must compress the past and will drop details. Among subquadratic alternatives, S4 and its diagonal simplifications S4D and DSS are the time-invariant predecessors Mamba grew out of, faster to train through the convolutional view but without selectivity. RWKV and RetNet reach similar linear-time recurrent goals from the linear-attention direction. Hyena uses long implicit convolutions rather than a recurrence. Gated linear attention and its relatives sit in the same design space, all honest choices with different tradeoffs.

The most important practical finding is that the strongest systems are usually hybrids. Because a pure state-space model struggles with exact in-context retrieval, interleaving a few attention layers among many Mamba layers recovers that ability while keeping most of the efficiency, which is the recipe behind models like Jamba and the gated-recurrence-plus-local-attention family. So the realistic framing is not Mamba versus Transformers but how many attention layers you can remove before recall breaks. Selectivity buys efficiency by compressing history into a fixed state, and the tasks Mamba is worst at are exactly the ones where that compression throws away the token you needed. The full comparison, with the linear-attention and RWKV and RetNet branches drawn out, lives on the state-space models class page.

Part IV: The full life of one forward pass

The specimen is a single training forward pass of one Mamba block over an input of shape (B, L, d_model). I follow the data down through the block, into the selective scan kernel, and back out, then note what changes in the backward pass and in single-step decoding. The block lives in mamba_ssm/modules/mamba_simple.py and the scan in mamba_ssm/ops/selective_scan_interface.py over a CUDA kernel under csrc/selective_scan/.

Stage 1: input projection and the two branches

The block first widens the input. A single linear layer in_proj maps d_model to 2 * d_inner, where d_inner = expand * d_model and expand defaults to 2. The result is split into two equal branches. One branch, call it x, is the signal that will actually pass through the state-space model. The other branch, z, is set aside to gate the output at the very end. This split is why a Mamba block folds what would be a separate mixing layer and a separate gated feed-forward into one unit.

Stage 2: the causal convolution

The x branch is transposed to channel-first and run through a depthwise causal 1D convolution of width d_conv (default 4), followed by a SiLU activation. Depthwise means each of the d_inner channels has its own tiny filter and channels do not mix here. Causal means the convolution is padded on the left only, so position t sees positions t-3 through t and never the future. This short convolution is doing local token mixing, the kind of neighbor interaction a state-space recurrence is not naturally shaped for, and it is the fast path that causal-conv1d accelerates.

Stage 3: producing Delta, B, and C from the input

Here is where the model becomes selective. The convolved x is fed through a linear layer x_proj that produces dt_rank + 2 * d_state numbers per position. Those are split into three input-dependent quantities. The first, dt_rank wide, is passed through a second linear layer dt_proj and a softplus to become Delta, the step size, of shape (B, L, d_inner), one positive timescale per channel per position. The other two are B and C, each of shape (B, L, d_state). In a classical state-space model B and C would be fixed weight matrices. Here they are activations, recomputed for every token, and that is precisely the S6 selectivity. The state matrix A is the one piece that stays a learned parameter, stored as A_log and read back as A = -exp(A_log) so it is always negative and the recurrence decays rather than explodes.

Stage 4: discretization

The continuous state-space equations are h'(t) = A h(t) + B x(t) and y(t) = C h(t). To run them on a discrete sequence each step is discretized with the current Delta using a zero-order hold. The A term becomes Abar = exp(Delta * A), an elementwise exponential because A is diagonal, and the B term is discretized in proportion to Delta, which the code implements with the simplified rule Bbar = Delta * B that the paper notes is a fine Euler approximation because A dominates the dynamics. The upshot is a per-position, per-channel linear recurrence h_t = Abar_t * h_{t-1} + Bbar_t * x_t with a readout y_t = C_t * h_t, where every coefficient carries a time subscript because everything is input-dependent.

Stage 5: the selective scan itself

This recurrence is the operation the whole repository exists to make fast. Conceptually it is a loop over the sequence carrying a hidden state of shape (B, d_inner, d_state), which is d_state times larger than the input. Materializing that expanded state for all L positions is what a naive implementation does and what runs out of memory. The kernel instead fuses discretization, the scan, the C readout, the D skip, and the z gate into one pass, and it never writes the expanded state to HBM. It streams over the sequence keeping the state in registers and shared memory, and it parallelizes across the sequence with a work-efficient parallel scan rather than a literal serial loop, because a linear recurrence is associative. The interface function makes the shapes explicit:

from mamba_ssm.ops.selective_scan_interface import selective_scan_fn

# u:     (B, D, L)   the conv+SiLU output, D = d_inner
# delta: (B, D, L)   per-channel, per-step timescale
# A:     (D, N)      diagonal state matrix, negative real, N = d_state
# B, C:  (B, N, L)   input-dependent, this is what makes it S6 rather than S4
# D:     (D,)        per-channel skip connection
# z:     (B, D, L)   gate branch, applies y * silu(z) inside the kernel
y = selective_scan_fn(u, delta, A, B, C, D=D, z=z,
                      delta_bias=dt_bias, delta_softplus=True)

Passing delta_softplus=True lets the kernel apply the softplus internally rather than in Python, and giving it z folds the final gate into the same kernel. In training the block often calls a further fused entry point, mamba_inner_fn, which pulls the input projection, the convolution, the x_proj split, the scan, and the output projection into one fused path so the intermediate tensors between them never hit memory.

Stage 6: skip, gate, and output projection

The scan output has a per-channel skip added, y = y + D * u, which lets each channel pass a direct copy of its input around the state, and then it is gated by the branch set aside in stage 1, y = y * SiLU(z). The gate is what turns the block from a bare sequence mixer into something with the expressive shape of a gated multiplicative unit. Finally out_proj maps d_inner back down to d_model, and the block returns a tensor the same shape as it received. A full model stacks these blocks with RMSNorm and residual connections in mamba_ssm/modules/block.py, wrapped by MixerModel and topped with a tied language-model head in mamba_ssm/models/mixer_seq_simple.py.

Stage 7: backward and decoding

The backward pass mirrors FlashAttention's memory trick. Rather than having stored the expanded per-position states during the forward pass, the kernel recomputes them on the way back while it accumulates gradients for u, Delta, A, B, C, D, and z. That recomputation trades a little extra arithmetic for a large memory saving, which is the entire reason the fused kernel exists instead of an autograd graph over primitive ops. Single-step decoding takes a different route entirely. There is no sequence to scan, just one new token, so the block runs a step path that advances the small convolution window and applies one recurrence update through a Triton selective_state_update kernel, reading and writing the cached state in InferenceParams. That closes the loop, one token in, a fixed-size state updated in place, one token out, at a cost independent of context length.

Part V: Internals deep dives

Deep dive: S6, and why selectivity forbids the convolution

The clearest way to see Mamba's central idea is to compare the two equivalent views of a time-invariant state-space model and then watch selectivity destroy one of them. When A, B, C, and Delta are fixed for all time, the recurrence h_t = Abar * h_{t-1} + Bbar * x_t, y_t = C * h_t can be unrolled. The output at position t is a weighted sum of all past inputs with weights C * Abar^k * Bbar for the input k steps back. Those weights do not depend on t, only on the gap k, so the whole mapping from x to y is a convolution with one fixed kernel:

time-invariant (S4):   y = K * x     with kernel
   K = ( C Bbar,  C Abar Bbar,  C Abar^2 Bbar,  ... )

  two equivalent views of the SAME operation:
    recurrent   ->  O(1) memory per step, great for inference
    convolution ->  one big FFT over the sequence, great for training

That equivalence is the engine of S4. Train with the convolution using an FFT so the whole sequence is processed in parallel, then switch to the recurrence for cheap autoregressive inference. It works only because the kernel K is the same at every position.

Now make it selective. Mamba computes B, C, and Delta from the input, so they carry a time subscript, and the unrolled weight from position t back to position t-k is C_t * Abar_t * Abar_{t-1} * ... * Bbar_{t-k}, a product of factors that are all different at every t. There is no longer a single kernel K, there is a different effective kernel at every position, so the operation is no longer a convolution and the FFT training path simply does not apply. Selectivity and the convolutional view are mutually exclusive. You can have fixed parameters and a cheap convolution, or input-dependent parameters and content selection, but not both. Mamba chooses selection and is left holding a time-varying linear recurrence that it must compute directly. The recurrence is still linear and therefore associative, which is the crack the parallel scan gets its lever into.

# the recurrence is a linear scan:  h_t = a_t * h_{t-1} + b_t
# it composes with an ASSOCIATIVE operator over (a, b) pairs:
def combine(left, right):
    a_l, b_l = left
    a_r, b_r = right
    return (a_r * a_l, a_r * b_l + b_r)
# associativity is exactly what a Blelloch prefix scan needs, so all L
# states can be produced in O(log L) parallel depth instead of a serial loop

Deep dive: the hardware-aware scan kernel

The kernel under csrc/selective_scan/ is where the performance lives, and its design is dominated by one number, the size of the expanded state. For each of d_inner channels the model carries a state of width d_state, so the working set per position is d_inner * d_state values, and across a length-L sequence a naive implementation would write a tensor of shape (B, L, d_inner, d_state) to HBM. That tensor is N times larger than the activations themselves and moving it dominates runtime, exactly the memory-bandwidth wall described in the FlashAttention chapter. The kernel refuses to create it. A conceptual sketch of what one thread's slice does, with the real kernel doing the across-thread reduction as a parallel scan rather than this serial loop:

// one CUDA block owns a (batch, channel) slice and streams over L.
// the state h[N] lives in registers / shared memory, never in HBM.
float h[N] = {0.f};
for (int t = 0; t < L; ++t) {
    float dt = softplus(delta[t] + dt_bias);   // per-step timescale
    float acc = 0.f;
    #pragma unroll
    for (int n = 0; n < N; ++n) {
        float dA = expf(dt * A[n]);             // Abar = exp(dt * A)
        float dB = dt * B[t][n];                // Bbar ~= dt * B  (Euler)
        h[n] = dA * h[n] + dB * u[t];           // the recurrence
        acc += C[t][n] * h[n];                  // readout with C
    }
    y[t] = acc + D * u[t];                      // fused skip
}
// backward recomputes h[] rather than reloading a stored (B,L,D,N) tensor

Three ideas make this fast. Kernel fusion, the discretization, the scan, the C readout, the D skip, and the z gate are one launch, so no intermediate leaves the chip. State in fast memory, the expanded state never touches HBM. Recomputation, the backward pass regenerates the states it needs instead of having stored them. The kernel is not a clever algorithm so much as a disciplined memory schedule, it wins by keeping the one tensor that is too big to move on the chip, and by paying flops to avoid bytes. This is also why a pure-PyTorch version exists but is only a correctness reference, an eager graph over these primitives would materialize precisely the tensor the kernel is built to avoid.

Deep dive: the Mamba block as an architecture

Step back from the kernel and the block itself is a deliberate fusion of two earlier ideas. Prior state-space language models put an SSM mixing layer and a separate gated feed-forward layer in sequence. Mamba collapses them. The in_proj that splits into an x branch and a z gate is the gated-unit half, and the selective scan is the mixing half, so one Mamba block does the work of two conventional layers, which is part of why a Mamba language model can be all identical blocks with no attention and no explicit feed-forward.

The small pieces each earn their place. The depthwise causal convolution supplies short-range mixing that a diagonal recurrence handles poorly, so the recurrence can specialize in long-range structure. The SiLU gate gives the block a multiplicative, data-dependent nonlinearity in the spirit of gated linear units. The per-channel D skip lets a channel bypass the state entirely when the state is not helping. And keeping A diagonal, stored as A_log and negated, makes the exponential in discretization an elementwise operation and guarantees stable decay. None of these are incidental, together they turn a bare recurrence into a block that trains stably and stacks deep. The configuration knobs that matter are d_state for how much the state can hold, expand for the inner width, d_conv for the local window, and dt_rank, which defaults to about d_model / 16 and controls how expressive the input-dependent Delta is.

Deep dive: Mamba-2 and state-space duality

Mamba-2 comes from the paper on structured state-space duality, or SSD, and it starts from an observation. If you restrict the state matrix A even further, from a general diagonal down to a scalar times the identity per head, then the whole selective SSM can be written as a matrix operation that looks like a form of attention with a decay mask. The two views, a linear recurrence and a masked quadratic attention, are dual descriptions of the same structured matrix, one that is semiseparable. That duality is not just theory, it changes how the operation is computed.

Mamba-1 (S6)      diagonal A,   scan kernel,   state N ~ 16
Mamba-2 (SSD)     scalar A,     matmul-heavy,  state N ~ 64..128

SSD chunked algorithm over the sequence:
   split into chunks -> inside a chunk: quadratic attention-like matmul
                     -> between chunks: pass one recurrent state forward
   most work is now big matmuls that saturate the tensor cores

The practical consequence is speed. Mamba-1's scan is memory-bound and does not use the tensor cores that GPUs devote most of their flops to. Mamba-2 reformulates the computation as a chunked algorithm, quadratic within each chunk as a dense matmul and recurrent between chunks by passing a single state across chunk boundaries, so the bulk of the work becomes large matrix multiplications that saturate the tensor cores. The paper reports that this makes the core operation several times faster than Mamba-1's scan, and because it is cheaper it can afford a much larger state dimension, which improves quality. Mamba-2 also adopts a multi-head structure with grouped B and C shared across heads, closely echoing multi-query and grouped-query attention. The SSD kernels are written in Triton and live under mamba_ssm/ops/triton/, with the block in mamba_ssm/modules/mamba2.py. Mamba-2's lesson is that giving up a little expressiveness in A to gain a matmul-shaped algorithm is a good trade on hardware that is built for matmuls, and that the recurrence-attention duality is the bridge that lets you choose the shape you compute. The theory of that duality, and how it connects to linear attention, is worked through on the state-space models class page.

Part VI: Reading the repository

The tree is small and readable in an afternoon. Paths below are organized by role, and a few may have moved as the project evolves.

Stage 0, orientation. Read the README.md, which has the block example, the pretrained models, and a pointer to the two papers. Get the vocabulary straight before touching code, SSM, S4, S6, selective, and SSD. The question to hold, what exactly is input-dependent in Mamba and what is not.

Stage 1, the block. Read mamba_ssm/modules/mamba_simple.py top to bottom, the Mamba class. Trace the forward from in_proj through the conv, x_proj, dt_proj, the selective scan, the gate, and out_proj, matching each line to a stage of Part IV. Then read the step method and see how single-token decoding avoids the scan entirely. Questions, where does A become -exp(A_log), why is there both a fused mamba_inner_fn path and an unfused path, and what does the z branch do.

Stage 2, the scan interface. mamba_ssm/ops/selective_scan_interface.py. This is the autograd boundary, a torch.autograd.Function whose forward calls the CUDA kernel and whose backward calls the CUDA backward. Read the docstrings for the exact shapes of u, delta, A, B, C, D, and z, and notice how the reference PyTorch implementation in the same neighborhood spells out the math the kernel encodes. Questions, which arguments are input-dependent, and why the backward needs to recompute states.

Stage 3, the CUDA kernel. Under csrc/selective_scan/, the forward and backward kernels and the common header. You do not need to follow every index to get the point, look for where the state array is declared, confirm it stays on chip, and see the discretization and the scan fused in one pass. Read it against the FlashAttention story, the shapes of the wins are the same.

Stage 4, the model. mamba_ssm/modules/block.py for the norm-and-residual wrapper, then mamba_ssm/models/mixer_seq_simple.py for MixerModel, MambaLMHeadModel, and the config. This is a full language model in remarkably little code, an embedding, a stack of identical blocks, a final norm, and a tied head. Then mamba_ssm/utils/generation.py for the InferenceParams cache and the CUDA-graph decode loop.

Stage 5, Mamba-2 and SSD. mamba_ssm/modules/mamba2.py and the Triton kernels in mamba_ssm/ops/triton/, especially the combined SSD chunk-scan entry point. This is denser reading, the chunked algorithm and the multi-head bookkeeping, so leave it until the Mamba-1 story is solid. The benchmarks/ and tests/ directories are the place to confirm your understanding against numbers and against the reference implementations the tests compare the kernels to.

Where not to start, the CUDA kernel before the block, and Mamba-2 before Mamba-1. Selectivity in mamba_simple.py is the one idea everything else depends on, so read it until it is obvious before descending into either the kernel or the duality.

Part VII: Hands-on labs

Labs 1 through 4 need one GPU, lab 5 is CPU-only reading and arithmetic. Log formats and exact numbers vary with the model and the hardware.

Lab 1: a block, forward and backward. Concept: shapes and the two branches.

import torch
from mamba_ssm import Mamba

layer = Mamba(d_model=256, d_state=16, d_conv=4, expand=2).cuda()
x = torch.randn(1, 64, 256, device="cuda", requires_grad=True)
y = layer(x)
y.sum().backward()
print(y.shape, x.grad.shape)          # both (1, 64, 256)
print(sum(p.numel() for p in layer.parameters()))

Confirm the output matches the input shape and that gradients flow. Then print the shapes inside the block by temporarily logging after in_proj and after x_proj, and match them to 2 * d_inner and dt_rank + 2 * d_state.

Lab 2: selectivity forbids convolution, empirically. Concept: time-varying weights.

import torch
from mamba_ssm.ops.selective_scan_interface import selective_scan_fn

B, D, L, N = 1, 4, 8, 16
u = torch.randn(B, D, L, device="cuda")
delta = torch.rand(B, D, L, device="cuda")
A = -torch.rand(D, N, device="cuda")
# make B_mat and C the SAME across time -> effectively S4-like, one kernel exists
Bt = torch.randn(B, N, 1, device="cuda").expand(B, N, L).contiguous()
Ct = torch.randn(B, N, 1, device="cuda").expand(B, N, L).contiguous()
y_ti = selective_scan_fn(u, delta, A, Bt, Ct)
# now let them vary in time -> no single convolution kernel reproduces this
Bt2 = torch.randn(B, N, L, device="cuda")
Ct2 = torch.randn(B, N, L, device="cuda")
y_tv = selective_scan_fn(u, delta, A, Bt2, Ct2)
print((y_ti[..., -1] - y_tv[..., -1]).abs().mean())

The point is conceptual, not the number. With B and C constant in time the operation is a fixed-kernel convolution and could be done with an FFT. Once you let them vary per position, as Mamba does, there is no fixed kernel and the scan is the only way. Sit with why the second case cannot be a convolution.

Lab 3: constant-memory decoding. Concept: the recurrent inference cache.

import torch, time
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

model = MambaLMHeadModel.from_pretrained(
    "state-spaces/mamba-130m", device="cuda", dtype=torch.float16)
ids = torch.randint(0, 50277, (1, 16), device="cuda")
for max_len in (64, 256, 1024):
    torch.cuda.synchronize(); t0 = time.time()
    model.generate(input_ids=ids, max_length=max_len, cg=True)
    torch.cuda.synchronize()
    print(max_len, round(time.time() - t0, 3), "s",
          torch.cuda.max_memory_allocated() // 1024**2, "MB")

Watch peak memory stay roughly flat as the generated length grows, because the recurrent state is fixed size and there is no KV cache that scales with context. Compare that mentally to a Transformer of the same size, whose cache grows linearly with tokens.

Lab 4: Mamba-1 versus Mamba-2 throughput. Concept: scan versus matmul-shaped SSD.

import torch, time
from mamba_ssm import Mamba, Mamba2

x = torch.randn(4, 4096, 1024, device="cuda", dtype=torch.bfloat16)
for name, layer in [("mamba1", Mamba(d_model=1024, d_state=16)),
                    ("mamba2", Mamba2(d_model=1024, d_state=128))]:
    layer = layer.to("cuda", torch.bfloat16)
    torch.cuda.synchronize(); t0 = time.time()
    for _ in range(20):
        layer(x)
    torch.cuda.synchronize()
    print(name, round(time.time() - t0, 3), "s")

Note that Mamba-2 uses a far larger d_state yet stays competitive or faster, because its work is shaped as tensor-core matmuls rather than a bandwidth-bound scan. Try growing the sequence length and watch both stay linear in L while a same-width attention layer would grow quadratically.

Lab 5: derive dt_rank and the projection widths. Concept: the block's parameter shapes on paper.

# for d_model = 768, expand = 2, d_state = 16, d_conv = 4:
#   d_inner  = expand * d_model            = 1536
#   in_proj  : 768 -> 2 * 1536             = 3072 out
#   dt_rank  = ceil(d_model / 16)          = 48
#   x_proj   : 1536 -> dt_rank + 2*d_state = 48 + 32 = 80 out
#   dt_proj  : 48 -> 1536
#   out_proj : 1536 -> 768
# check these against a printed layer's named parameters

Instantiate Mamba(d_model=768) and print {n: p.shape for n, p in layer.named_parameters()}, then match every shape to the arithmetic above. When the numbers line up you understand the block.

Part VIII: Questions and model answers

Understanding checks. Answer aloud before reading.

1. What is Mamba, in one sentence?

A selective state-space sequence model that makes the recurrence parameters B, C, and the step size Delta functions of the input, so a fixed-size hidden state can decide what to keep and forget, giving linear-time training and constant-memory inference that rivals attention.

2. What does selectivity mean here, concretely?

That B, C, and Delta are computed from the current token by linear projections rather than being fixed weights. Because they change per position, the model can gate its state on content, holding information across some tokens and overwriting it at others, which a time-invariant SSM cannot do.

3. Why can a selective SSM not be computed as a convolution?

A convolutional view exists only when the unrolled input-to-output weights depend just on the gap between positions, which requires fixed A, B, C, and Delta. Once those are input-dependent the effective weight from one position to another differs at every position, so there is no single kernel and no FFT, only a direct time-varying recurrence.

4. If the convolution is gone, how is training still efficient?

The recurrence is linear and therefore associative, so a parallel prefix scan produces all states in logarithmic parallel depth. A custom CUDA kernel runs that scan while keeping the large expanded state in on-chip memory and fusing discretization, readout, skip, and gate, so it never writes the state to HBM.

5. Why is the scan a custom kernel and not plain PyTorch ops?

The hidden state is d_state times wider than the activations, so an eager graph over primitive ops would materialize a tensor of shape roughly (B, L, d_inner, d_state) in HBM and be dominated by moving it. The kernel exists precisely to keep that tensor on chip and recompute it in the backward pass.

6. Walk through the Mamba block in order.

Input projection widens and splits into an x branch and a z gate. The x branch gets a depthwise causal convolution and SiLU. A projection produces input-dependent Delta, B, and C. The selective scan runs the discretized recurrence with A read from A_log. A per-channel D skip is added, the result is gated by SiLU of z, and an output projection returns to the model width.

7. Why keep A diagonal and store it as A_log?

Diagonal A makes the discretization exponential elementwise and keeps the state channels independent, which is cheap and stable. Storing A_log and using A = -exp(A_log) forces A negative so the recurrence always decays rather than growing without bound.

8. What is in the inference cache and why does memory stay flat?

An InferenceParams object holds each layer's fixed-size recurrent SSM state plus the short convolution window. Decoding one token updates that state in place, so cost and memory are independent of how long the context already is, unlike a Transformer KV cache that grows with every token.

9. How does the backward pass save memory?

It does not store the per-position expanded states from the forward pass. Instead it recomputes them while accumulating gradients, trading extra arithmetic for a large reduction in memory traffic, the same recomputation strategy FlashAttention uses for its score matrix.

10. What changes in Mamba-2 relative to Mamba-1?

A is restricted from a general diagonal to a scalar times the identity per head, which exposes a duality between the SSM and a masked attention over a semiseparable matrix. That lets the computation be reshaped into a chunked, matmul-heavy algorithm that saturates tensor cores, so it is several times faster and can afford a much larger state.

11. When should you not use Mamba?

When the task needs exact recall or copying from far back in the context, because a fixed-size state compresses history and can drop the specific token you need. Attention looks at every past token directly and wins there, which is why strong systems often interleave a few attention layers among the Mamba layers.

12. Why does the block need a convolution at all if it has a recurrence?

A diagonal state-space recurrence is good at long-range structure but weak at tight local interactions between neighboring tokens. The short depthwise causal convolution supplies that local mixing cheaply, letting the recurrence specialize in the long range.

13. Name three honest alternatives and their tradeoffs.

Transformers with FlashAttention, best quality on recall but quadratic in length. S4 and its diagonal variants, time-invariant so they keep the fast convolution but lack selectivity. RWKV, RetNet, and gated linear attention, other linear-time recurrent designs from the linear-attention side with their own tradeoffs.

14. What is the single biggest performance mistake a new user makes?

Not installing causal-conv1d and running in float32, which silently drops onto slow reference paths. The kernels are built for bf16 or fp16 with the compiled convolution and scan, and without them you are measuring the reference implementation, not Mamba.

Part IX: Design lessons

Make the parameters depend on the data. The whole leap from S4 to Mamba is letting B, C, and Delta be functions of the input. Content-dependence is the difference between a filter that treats all tokens alike and a model that can select, and it is worth reorganizing an entire training strategy to get. The same instinct shows up wherever a system gains power by conditioning its behavior on its input rather than fixing it in advance.

Know exactly which optimization your change breaks. Selectivity was not free, it forfeited the convolutional training path that made the predecessor fast. Mamba's authors understood that precisely and had the replacement, a parallel scan, ready. The lesson is to trace which existing efficiency an architectural change invalidates before you make it, because that is where the real engineering cost lands.

Optimize for bytes moved, not flops. The scan kernel wins by keeping the oversized state on chip and recomputing in the backward pass, spending arithmetic to avoid memory traffic. On modern accelerators the bottleneck is almost always bandwidth, and the same fuse-and-recompute pattern powers FlashAttention and most fast kernels. The scan itself is a classic parallel prefix, the oldest trick in the parallel-algorithms book applied to a new recurrence.

Reshape the computation to fit the hardware. Mamba-2 gives up a little expressiveness in A so the operation becomes matmuls, because GPUs spend most of their silicon on matmul. Choosing an algorithm because the hardware is fast at its shape, rather than choosing the most general algorithm and hoping, is a repeatable source of large speedups.

Ship a slow reference beside the fast kernel. The repository keeps a pure-PyTorch implementation of the scan that the tests compare against the CUDA kernel. That reference is what makes the fast path trustworthy and readable, and it is the difference between a kernel you believe and one you merely run.

Part X: Memorization framework

The one-sentence summary: Mamba makes a state-space model's B, C, and Delta depend on the input, which forbids the S4 convolution and forces a hardware-aware parallel scan, wraps that scan in a block with a causal convolution and a SiLU gate, and in Mamba-2 restricts A to a scalar so the whole thing becomes matmul-shaped through state-space duality.

in_proj -> (x branch, z gate)
  -> causal conv1d + SiLU  (local mixing)
  -> x_proj -> input-dependent Delta, B, C   (this is S6 selectivity)
  -> selective_scan: h_t = exp(dt A) h_{t-1} + dt B x_t,  y_t = C h_t
  -> + D skip, * SiLU(z) gate
  -> out_proj

The chain mapped to source:

block          mamba_ssm/modules/mamba_simple.py  (class Mamba)
scan interface mamba_ssm/ops/selective_scan_interface.py  (selective_scan_fn)
cuda kernel    csrc/selective_scan/  (fwd + bwd)
model          mamba_ssm/models/mixer_seq_simple.py  (MambaLMHeadModel)
inference      mamba_ssm/utils/generation.py  (InferenceParams)
mamba-2 / ssd  mamba_ssm/modules/mamba2.py + mamba_ssm/ops/triton/

Memorize these blocks:

• S4 vs S6: S4 has fixed A, B, C, Delta and a convolutional view. S6 makes B, C, Delta input-dependent, which kills the convolution and demands a scan.
• Why a kernel: the state is d_state times wider than the activations, so it must stay on chip, and the backward recomputes it rather than storing a (B, L, d_inner, d_state) tensor.
• The block: in_proj splits into x and z, conv1d plus SiLU, x_proj gives Delta, B, C, selective scan, D skip, SiLU(z) gate, out_proj.
• Discretization: A = -exp(A_log), Abar = exp(Delta A), Bbar approximately Delta B, so h_t = Abar h_{t-1} + Bbar x_t and y_t = C h_t.
• Mamba-2: scalar A gives a duality with masked attention, a chunked matmul algorithm, tensor-core speed, and a much larger state.

Part XI: Papers and further reading

This walkthrough leans on a short chain of papers, and each one is readable on its own. Where this site derives the same idea in depth, the companion link points there.

1. Gu and Dao, Mamba, Linear-Time Sequence Modeling with Selective State Spaces, 2023. The paper this repository implements, the S6 selective scan and the hardware-aware kernel. The theory is derived on the state-space models class page, and the scan itself is a classic parallel prefix, treated in the parallel computing class.
2. Dao and Gu, Transformers are SSMs, Generalized Models and Efficient Algorithms Through Structured State Space Duality, 2024. The Mamba-2 paper, restricting A to a scalar per head so the whole computation reshapes into chunked matmuls.
3. Gu, Goel, and Ré, Efficiently Modeling Long Sequences with Structured State Spaces, 2021. S4, the time-invariant predecessor whose convolutional training path selectivity gives up.
4. Gu et al., HiPPO, Recurrent Memory with Optimal Polynomial Projections, 2020. The theory of how a small state can remember a long history, the origin of the A matrix that S4 inherits.
5. Gu et al., On the Parameterization and Initialization of Diagonal State Space Models, 2022. S4D, the diagonal simplification whose parameterization Mamba's A_log descends from.
6. Gupta et al., Diagonal State Spaces are as Effective as Structured State Spaces, 2022. DSS, the companion finding that diagonal state matrices lose little against the full S4 structure.
7. Dao et al., FlashAttention, Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022. The kernel discipline the selective scan follows, fuse the pass, keep the big tensor on chip, and recompute in the backward. Covered in the FlashAttention walkthrough.
8. Katharopoulos et al., Transformers are RNNs, Fast Autoregressive Transformers with Linear Attention, 2020. The linear-attention view that state-space duality later connects to the SSM side.
9. Peng et al., RWKV, Reinventing RNNs for the Transformer Era, 2023. A parallel route to linear-time language models from the RNN direction.
10. Sun et al., Retentive Network, A Successor to Transformer for Large Language Models, 2023. RetNet, another recurrence with a parallel training form and a useful contrast to selectivity.
11. Poli et al., Hyena Hierarchy, Towards Larger Convolutional Language Models, 2023. The long-convolution alternative that keeps the FFT path Mamba abandons.
12. Lieber et al., Jamba, A Hybrid Transformer-Mamba Language Model, 2024. The hybrid recipe, a few attention layers among many Mamba layers to recover exact recall.

Part XII: Final takeaway

If the continuous state-space equations, the zero-order-hold discretization, HiPPO and S4, the parallel associative scan, and the duality with linear attention are the gaps, the state-space models class page derives all of them from first principles, and the memory story behind the fused kernel is worked out in the FlashAttention chapter. Then come back and read mamba_simple.py once more. The forward pass will read like an ordinary block with one strange and powerful twist, the parameters of its recurrence are computed from the very sequence it is processing, and everything else is the machinery required to make that twist run fast.