Citation and evidence

Mamba 2

32 min full readUpdated 31 references

This article's verification

Report a problem with this article

Edit article
More

Use this article

Raw MarkdownExplore connections

Improve this page

Revision historyDiscuss this article

Topics

AI ModelsDeep LearningModel Architecture

Cite this article

Mamba 2 is a state space model architecture introduced in the paper "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality" by Tri Dao and Albert Gu, published at ICML 2024 (arXiv:2405.21060) on May 31, 2024.[1] It refines the selective SSM layer from the original Mamba model through a theoretical framework called State Space Duality (SSD), which establishes a formal equivalence between structured state space models and a class of attention mechanisms.[1] The new layer is 2 to 8 times faster than the Mamba 1 selective scan at the same state dimension, and allows the state dimension to grow from 16 in Mamba 1 to 64 or 128 by default, while remaining competitive with Transformer architectures on language modeling benchmarks.[1] As the authors put it, the SSD framework "allows us to design a new architecture (Mamba-2) whose core layer is a refinement of Mamba's selective SSM that is 2-8X faster, while continuing to be competitive with Transformers on language modeling."[1]

The SSD framework reveals that SSMs and attention are not competing families of sequence models but rather two computational views of the same underlying mathematical object: structured semiseparable matrices.[1] This theoretical result brought a wave of practical benefits, including compatibility with the tensor parallelism and sequence parallelism techniques developed for Transformers, and motivated a family of hybrid architectures that interleave SSM layers with attention layers to capture advantages of both. Since publication, Mamba 2 layers have been used in models such as Mistral Codestral Mamba,[8] Zyphra's Zamba2,[10] NVIDIA Nemotron-H,[26] IBM Bamba and Granite 4.0,[11][29] TII's Falcon-H1,[30] and Tencent Hunyuan-TurboS.[17] Some prominent hybrids chose the original Mamba layer instead: AI21's Jamba 1.5 and NVIDIA Hymba both use Mamba-1-style SSM layers.[25][27]

Who created Mamba 2?

Mamba 2 was developed by Tri Dao and Albert Gu.[1]

Tri Dao is an assistant professor at Princeton University. He completed his PhD at Stanford, where he worked on efficient deep learning systems. He is also the primary author of Flash Attention, the IO-aware attention algorithm that became a standard component in large language model training.[4] His research focuses on hardware-aware algorithm design, aiming to close the gap between the theoretical complexity of machine learning operations and the practical throughput achievable on real hardware.

Albert Gu is an assistant professor at Carnegie Mellon University. His doctoral work at Stanford, supervised by Christopher Re, introduced the S4 model and the broader family of structured state space models that preceded Mamba.[3] He is also Chief Scientist at Cartesia, which he co-founded to build real-time AI products on top of SSM architectures.[19] Gu was named to Time's 100 Most Influential People in AI.

The two researchers had previously co-authored the original Mamba paper (December 2023)[2] and have continued to collaborate on SSM research through both their academic positions and the Goomba Lab blog, where they published a four-part series explaining the theory, algorithms, and systems of Mamba 2.[21]

Background

What problem do state space models solve?

State space models originated in control theory, where they describe continuous-time dynamical systems through linear differential equations. In the sequence modeling context, they map an input sequence to an output sequence through a hidden state that evolves over time. The theoretical appeal is significant: SSMs process sequences in linear time relative to sequence length, whereas self-attention in Transformers scales quadratically. Inference is also efficient because the hidden state has fixed size, so generating each new token requires only an O(1) update rather than attending to a growing key-value cache.

For most of the 2010s, recurrent models such as LSTMs and GRUs were the dominant alternative to attention, but they struggled with long-range dependencies and were difficult to parallelize during training. Structured state space models, starting with S4 (Gu et al., 2021), showed that with the right parameterization a linear recurrence could be trained efficiently in parallel and could capture long-range dependencies competitive with Transformers on sequence tasks.[3] S4, S5, and related work demonstrated strong results on the Long Range Arena benchmark.[3]

The first Mamba model (Gu and Dao, 2023) added input-dependent selectivity to the SSM framework.[2] Prior SSMs used time-invariant transition matrices, meaning the model processed every token with the same dynamics. Mamba introduced selective state spaces, where the step size Delta and the B and C matrices are functions of the input (so the discretized transition depends on the input), allowing the model to gate what information to retain in the hidden state.[2] This brought SSMs to near-parity with Transformers on language modeling tasks while retaining linear complexity.[2] Training Mamba 1 relied on a parallel associative scan algorithm that runs efficiently on GPUs but does not use the tensor cores (matrix multiplication units) that dominate modern GPU throughput.[2]

Why does GPU hardware favor matrix multiplication?

Modern GPUs are designed around matrix multiplication. An A100 achieves 312 TFLOPS in BF16 matrix multiplication but only 19 TFLOPS in FP32 scalar arithmetic. An H100 achieves 989 TFLOPS in BF16 matrix multiplication but only 67 TFLOPS in scalar arithmetic. This is roughly a 15x gap. Mamba 1's parallel scan operates element-wise rather than through matrix multiplications, which means it uses the slower arithmetic units and leaves most GPU capacity idle. The SSD algorithm in Mamba 2 was designed specifically to express the SSM computation as matrix multiplications, closing this hardware gap.[22]

The hardware gap had grown wider with successive GPU generations. NVIDIA's roadmap from Volta (2017) through Hopper (2022) to Blackwell (2024) emphasized BF16 and FP8 matrix throughput at the expense of scalar throughput, on the assumption that workloads would be dominated by Transformer-style dense matrix multiplication. Algorithms that did not fit this assumption, including most variants of linear attention and recurrent SSMs, were therefore disadvantaged at the kernel level. The SSD reformulation was as much a response to this hardware reality as it was a theoretical advance.

What is State Space Duality?

The central theoretical contribution of the Mamba 2 paper is the State Space Duality (SSD) framework. It connects structured SSMs and attention through the mathematics of semiseparable matrices.[1] The paper itself frames the result as a "rich framework of theoretical connections between SSMs and variants of attention, connected through various decompositions of a well-studied class of structured semiseparable matrices."[1]

What is a semiseparable matrix?

A semiseparable matrix is a structured lower-triangular matrix with the property that every submatrix contained on and below its diagonal is low-rank. The key insight of the SSD framework is that the output-to-input mapping of any causal linear SSM can be written as multiplication by a semiseparable matrix.[1] Entry (i, j) of this matrix, for i >= j, equals the product C_i times the chain A_{j+1} * A_{j+2} * ... * A_i times B_j, which describes how the contribution of input at position j propagates to output at position i through the recurrence. An SSM with a scalar-times-identity transition matrix (A_t = alpha_t * I) produces a 1-semiseparable matrix, meaning each off-diagonal submatrix has rank at most 1.[1]

This representation is not just a notational trick. It means that computing an SSM is equivalent to multiplying a vector by a semiseparable matrix, and any algorithm for that matrix multiplication corresponds to a valid algorithm for the SSM. The parallel associative scan used in Mamba 1 is one such algorithm; the SSD algorithm is another, chosen to exploit tensor cores. The duality is therefore not merely a conceptual bridge but an explicit menu of computational strategies, each with different memory and throughput tradeoffs, that practitioners can pick from based on hardware constraints and sequence length.

Structured Masked Attention

The SSD framework also shows that a certain class of attention mechanisms produces the same semiseparable matrix.[1] Standard self-attention computes output as the product of a softmax-normalized score matrix with a value matrix. Linear attention removes the softmax, making the score matrix a product of query and key matrices directly. Structured Masked Attention (SMA) generalizes this by replacing the causal mask with any structured matrix L that supports fast matrix-vector multiplication.[1]

When the mask L is chosen to be a 1-semiseparable matrix encoding the same scalar discount factors used in an SSM, SMA and the SSM compute identical results.[1] They are the same model expressed in two different computational forms. This is the duality: the SSM view naturally leads to recurrent algorithms, while the SMA view naturally leads to attention-like algorithms with matrix multiplications.

The paper also proves a converse result: any kernel attention method that has a fast recurrent form must be an SSM.[1] This establishes SSMs and kernel attention as the same class of models, not merely analogous families.

How does Mamba 2 relate to standard attention?

SSD differs from standard softmax self-attention in exactly two ways. First, it removes the softmax normalization. Second, it applies a multiplicative mask that encodes the input-dependent discount factors from the SSM recurrence.[1] These discount factors represent how much the influence of a token decays over distance: if the model assigns a low discount value at position t, earlier contributions are suppressed rapidly.

From this perspective, Mamba 2 can be understood as causal linear attention with a learnable, input-dependent positional mask.[1] The model controls this mask through the scalar A values, which determine how much context each position carries forward. This view clarifies why SSMs tend to be good at smooth, predictable patterns but can struggle with exact retrieval: the discount mechanism compresses distant tokens rather than preserving them precisely.

The duality also connects to the broader literature on efficient attention. Linear attention (Katharopoulos et al., 2020) removes the softmax from standard attention to achieve linear complexity, but lacks a mechanism for position-dependent weighting. Retention (Sun et al., 2023) adds a geometric decay factor to linear attention, bringing it closer to SSMs. The SSD framework shows that these are all instances of a single parametric family indexed by the choice of mask matrix L.[1] SSMs with learned input-dependent A values occupy a particularly expressive position in this family because their mask is data-driven at each step rather than fixed.

Relation to other linear attention variants

The SSD framework has clarified the relationships between several formerly disparate efficient attention proposals. RetNet's retention operator corresponds to SMA with a fixed, data-independent geometric decay mask (with a different decay rate per head). RWKV-4 and RWKV-5 use fixed decays, while RWKV-6 introduces data-dependent decay that brings it structurally closer to Mamba 2. Gated Linear Attention (GLA, Yang et al., 2024) uses a learnable input-dependent gate per channel that, in the SSD vocabulary, corresponds to a diagonal rather than scalar transition matrix.[23] DeltaNet adds an explicit delta-rule update that lets the model overwrite specific state slots, addressing one of the recall weaknesses of softmax-free recurrences.

When these models are placed in the SSD coordinate system, they differ primarily in three axes: the form of the transition matrix (scalar versus diagonal versus more general), the dependence of the transition on input (fixed versus data-dependent), and the presence or absence of corrective update rules. Mamba 2 occupies the data-dependent scalar transition corner, accepting a constraint on the transition shape in exchange for an algorithm that maps cleanly onto matrix multiplication units. Several research groups have followed up with hybrids that pair Mamba 2 blocks with DeltaNet-style update rules or GLA-style diagonal transitions to recover expressiveness without sacrificing throughput.

Computational comparison

The following table compares the core computational properties of standard attention, linear attention, and SSD:

PropertySoftmax attentionLinear attentionSSD (Mamba 2)
Training complexityO(T^2 d)O(T d N)O(T d N)
Inference (per token)O(T d) (grows with T)O(d N)O(d N)
Memory (inference)O(T d) KV cacheO(d N) fixed stateO(d N) fixed state
Mask typeCausal (uniform)NoneInput-dependent 1-semiseparable
Exact token retrievalYesApproximateApproximate
Tensor core friendlyYesPartialYes
Supports tensor parallelismYesPartialYes (Mamba 2)

Expanded article table

The training complexity for SSD is O(T d N) rather than O(T^2 d) because the SSD algorithm operates on chunks of size Q, running the quadratic form within each chunk and the linear recurrence across chunks.[1]

How does the Mamba 2 architecture differ from Mamba 1?

Constraints on the transition matrix

The core architectural change from Mamba 1 to Mamba 2 is a constraint on the recurrent transition matrix A. In Mamba 1, A is a diagonal matrix (each element of the state vector has its own decay rate). In Mamba 2, A is further restricted to scalar-times-identity: A_t = alpha_t * I, where alpha_t is a single scalar. All dimensions of the hidden state share the same decay rate at each time step.[1]

This restriction looks like it should reduce expressiveness. In practice, the authors found that the SSD models with scalar A remain highly competitive on language modeling, and the constraint enables the entire block decomposition strategy that makes the SSD algorithm efficient.[1] The expressiveness lost from the constraint on A is recovered by expanding the state dimension N from 16 (Mamba 1 default) to 64 or 128 in Mamba 2, which the faster algorithm makes affordable.[1]

Parallel input projections

In Mamba 1, the SSM parameters B, C, and delta (the discretization step) were computed sequentially, with delta applied to A and B before use. Mamba 2 moves all data-dependent projections to a single parallel step at the beginning of the block. The inputs X, A, B, and C are all projected from the input in parallel, with no sequential dependency between them.[1] This simplifies the data flow, reduces the number of passes through GPU memory, and makes the block more amenable to tensor parallelism.

Multi-head structure

Mamba 2 introduces a multi-head structure analogous to multi-head attention (MHA) in Transformers.[1] The input X is split into multiple heads (head dimension P, typically 64 or 128), while the B and C projections have a single head shared across all X heads. In attention terms B and C play the roles of keys and queries and X plays the role of values, so the paper calls this the multi-input SSM (MIS) or multi-value attention (MVA) pattern.

The paper also defines other sharing patterns by analogy with attention: a multi-contract SSM (X and B shared, C per head, analogous to multi-query attention) and a multi-expand SSM (X and C shared, B per head, analogous to multi-key attention). In ablations the multi-value pattern performed best. For tensor parallelism, Mamba 2 can use a grouped version, grouped-value attention (GVA), with several independent B and C heads, analogous to grouped-query attention in Transformers.[1] These sharing schemes allow larger effective model capacity without proportional increases in parameters or computation.

The multi-head structure also enables tensor parallelism across GPUs. In Mamba 1, the custom selective scan kernel required two all-reduce operations per layer during distributed training, compared to one for Transformer layers. Mamba 2's parallel projection structure and head-based organization reduces this to one all-reduce per layer, matching Transformers and enabling more efficient large-scale training.[20]

Normalization

Mamba 2 adds a group normalization (or RMS normalization) layer applied to the output of the SSM computation before the final output projection.[1] The authors relate it to NormFormer; Falcon Mamba later added extra RMS normalization layers to the original Mamba for a similar stability reason.[1][13]

How does the SSD algorithm work?

The SSD algorithm is a four-step procedure that computes the SSM output by decomposing the semiseparable matrix into blocks and using matrix multiplication for most of the computation.[22]

Given a sequence of length T, the algorithm selects a chunk size Q (typically between 64 and 256 tokens). The sequence is divided into T/Q chunks. The four steps are:

  1. Intra-chunk outputs: For each chunk, compute the contribution to the output from tokens within the same chunk using the attention-like (quadratic) form of SSD. This is a Q x Q matrix multiplication per chunk and runs fully in parallel across all chunks.

  2. Chunk states: For each chunk, compute the SSM state at the end of the chunk assuming a zero initial state. This summarizes what each chunk contributes to future chunks.

  3. Pass states: Run a sequential SSM scan over the T/Q chunk-level states computed in step 2. This is the only sequential step, but it operates on T/Q elements rather than T, reducing its cost by a factor of Q. For Q = 64, the sequential work is 64 times smaller than a full sequential scan.

  4. Output states: Compute the additional contribution to each token's output from the true initial state that was propagated through the inter-chunk scan in step 3. This is again a matrix multiplication.

Steps 1, 2, and 4 are fully parallel and use matrix multiplications, which run on tensor cores. Only step 3 is sequential, and it operates on a much shorter sequence. The reference implementation of this algorithm requires approximately 25 lines of code, compared to the more complex selective scan in Mamba 1.[22]

This design is similar in spirit to Flash Attention, which uses block-level chunking and tiling to keep computation within fast on-chip SRAM rather than reading from slower HBM.[4] The SSD algorithm likewise works block by block, combining results across chunks without materializing the full semiseparable matrix.[22]

How much faster is Mamba 2?

Training throughput

The primary benchmarked improvement is training speed. Measured on the same state dimension, the SSD algorithm is 2 to 8 times faster than the Mamba 1 parallel scan.[1] The paper attributes the gap mainly to state size: the Mamba 1 scan slows down roughly linearly as the state dimension grows, while SSD, which runs on tensor cores, handles much larger state dimensions with little slowdown.[1]

Because the SSD algorithm is faster, Mamba 2 can afford to run with a state dimension 4 to 8 times larger than Mamba 1 at comparable or lower training cost. The larger state improves performance on associative recall tasks, where the model needs to store many key-value associations in its hidden state. On the synthetic Multi-Query Associative Recall (MQAR) benchmark used by the authors, Mamba 2 with N=64 substantially outperforms Mamba 1 with N=16.[1]

Language modeling results

The paper trains Mamba 2 on 300 billion tokens from the Pile dataset and evaluates against Mamba 1 and the Transformer++ baseline (a Transformer with RoPE, SwiGLU, and other modern improvements).[1] Mamba 2 matches or slightly exceeds Mamba 1 and sits close to the Transformer++ at equivalent parameter counts and token budgets.[1] At 2.7B parameters, Mamba 2 reaches a Pile validation perplexity of 6.09, compared to 6.13 for a Transformer++ trained with the same data and hyperparameters.[1]

Sequence length scaling

In the paper's kernel benchmarks on an A100 GPU, the SSD layer's time grows linearly with sequence length and is faster than FlashAttention-2 from 2K tokens upward. The authors caution that a whole Mamba 2 model may still train less efficiently than a same-size Transformer at short lengths such as 2K, because a Transformer's MLP layers are very hardware-efficient.[1]

NVIDIA's 8B-parameter study predicted that a Mamba-2-Hybrid model would be up to 8 times faster than a comparable Transformer when generating tokens.[14] This asymptotic advantage reflects the linear vs. quadratic complexity difference: attention compute grows as O(T^2) while the SSM computation grows as O(T).

Systems-level improvements

Beyond the core algorithm, Mamba 2's parallel projection structure enables several systems optimizations that were not available for Mamba 1.[20]

Tensor parallelism splits model parameters across GPUs. Mamba 2 achieves this by dividing the input and output projection matrices across the tensor parallel degree (2, 4, or 8 shards), applying per-GPU group normalization, and using a single all-reduce per layer.[20] This matches the communication pattern of Transformer layers and allows Mamba 2 to scale across multiple GPUs with no additional overhead compared to Transformers.

Sequence parallelism splits long sequences across multiple GPUs. Mamba 2 supports this through the block decomposition property of SSD: each GPU computes local outputs for its portion of the sequence, then passes final chunk states to subsequent GPUs via point-to-point communication.[20]

Variable-length batches, where sequences within a batch have different lengths, normally require padding to a uniform length, which wastes computation on padding tokens. Mamba 2 handles this by treating the entire batch as a single concatenated sequence and setting the transition scalar A_t to zero at sequence boundaries, preventing state from leaking between sequences.[20]

Why do hybrid models combine Mamba 2 with attention?

A widely adopted finding from both the Mamba 1 and Mamba 2 literature is that hybrid architectures, which alternate SSM layers with occasional attention layers, outperform either architecture used alone. SSM layers provide efficient long-range context compression; attention layers provide precise token-level retrieval. The combination addresses the primary weakness of pure SSMs (limited in-context recall) while retaining most of their efficiency advantage.

The Mamba 2 paper itself tested a hybrid of 6 attention layers and 58 SSD blocks, which outperformed both the pure Mamba 2 model and the pure Transformer baseline on the 300B-token Pile training run.[1] This result has been replicated across a wide range of subsequent work, and by 2025 essentially every major SSM deployment uses a hybrid configuration rather than a pure SSM stack.

Codestral Mamba (Mistral AI)

Mistral released Codestral Mamba 7B (also marketed as Mamba-Codestral-7B-v0.1) on July 16, 2024, one of the first openly released code models built on the Mamba 2 architecture.[8] The model is released under an Apache 2.0 license and was developed with input from Albert Gu and Tri Dao.[8] Mistral describes it as an instructed code model with 7,285,403,648 parameters, trained with advanced code and reasoning capabilities.[8]

The released checkpoint is a pure Mamba 2 model with no attention layers: its configuration lists 64 Mamba 2 layers with a state size of 128.[24] Mistral says it tested Codestral Mamba on in-context retrieval up to 256K tokens, far beyond the 8K context of the original Mistral 7B, and expected it to be a great local code assistant.[8] Inference benefits from the SSM's fixed-state behavior: token throughput remains roughly constant as the prompt grows, in contrast to Transformer code models whose KV cache memory and bandwidth requirements scale linearly with context length.

On the HumanEval benchmark, Codestral Mamba achieves results competitive with similarly sized Transformer-based code models, including DeepSeek-Coder and CodeLlama variants.[8] Mistral released Mathstral 7B, a math-specialist Transformer model, on the same day.

Jamba 1.5 and 1.6 (AI21 Labs)

Jamba, developed by AI21 Labs, was among the first commercially released hybrid models to combine Mamba layers with attention.[5] The Jamba 1.5 series, released in 2024, interleaves Mamba and attention layers with mixture-of-experts (MoE) routing. Jamba 1.5 Large has 72 layers (9 blocks of 8) with 94B active of 398B total parameters, and Jamba 1.5 Mini has 12B active of 52B total.[27] AI21 tested Mamba 2 for this release but found that a Mamba-1-attention hybrid worked better than a Mamba-2-attention hybrid, so Jamba 1.5 uses Mamba-1 layers.[27] Jamba 1.5 Mini and Jamba 1.5 Large support 256K-token context windows and achieve faster inference than comparable Transformer-only models on long contexts.[6]

On a 262,144-token input, Jamba 1.5 Mini generates approximately 62 tokens per second, compared to 41 for LLaMA 3.1 8B and 39 for Mixtral under the same conditions.[28] The model fits in a smaller GPU memory footprint than Transformer equivalents because SSM layers do not require a KV cache, and the original Jamba release fit up to 140K tokens of context on a single 80 GB GPU.[5]

AI21 released Jamba 1.6 in early 2025 as a successor focused on enterprise retrieval-augmented generation use cases.[7] Jamba 1.6 keeps the hybrid SSM-attention-MoE backbone and 256K-token context window but retunes the model for grounded question answering and data classification.[7] AI21 cites customer results: the retailer Fnac reported a 26% improvement in data-classification output quality with Jamba 1.6 Mini, which let it move from Jamba 1.5 Large to Jamba 1.6 Mini with a roughly 40% latency improvement, and the education provider Educa Edtech reported more than 90% retrieval accuracy and citation reliability in grounded question answering.[7]

Zamba 2 (Zyphra)

Zamba2, developed by Zyphra, replaces the Mamba 1 blocks from the original Zamba architecture with Mamba 2 blocks.[10] The architecture uses a backbone of Mamba 2 layers interleaved with two shared attention blocks in an alternating pattern. The shared attention blocks use a low-rank adaptation (LoRA) projector at each invocation to allow the shared MLP weights to specialize per layer at minimal parameter cost.[10]

Zamba2-7B was pre-trained on a 3-trillion-token dataset and then annealed on 100 billion high-quality tokens.[10] The model has extremely low inference latency due to the Mamba 2 blocks, which the Zyphra team reports have roughly four times the throughput of an equal-parameter Transformer block.[10] Zamba2 models are available in 1.2B, 2.7B, and 7B parameter sizes.[10]

Bamba (IBM Research)

Bamba is an open-source hybrid Mamba 2 model developed by IBM Research in collaboration with Princeton, CMU, and UIUC, with direct involvement from Albert Gu and Tri Dao.[11] Bamba-9B was trained on 2.2 trillion tokens and achieves benchmark performance comparable to LLaMA 3.1 8B despite having been trained on roughly one-seventh the data.[11]

At inference time on vLLM, Bamba-9B delivers 2.5 times the throughput and 2 times lower latency compared to a standard Transformer of similar size.[12] IBM released the full training recipe, data mixtures, and a quantization framework alongside the model weights, making Bamba one of the most fully open hybrid SSM models available.[11] The architecture later informed IBM's Granite 4.0 series of enterprise models, released in October 2025 with a majority of Mamba 2 layers.[11][29]

Falcon-H1 (Falcon family at TII)

The Technology Innovation Institute (TII) in Abu Dhabi released Falcon Mamba 7B in 2024 as a pure SSM model based on Mamba 1 with additional RMS normalization layers for training stability.[13] This was followed by the Falcon-H1 series in 2025, which uses a hybrid architecture where attention heads and Mamba 2 SSM components run in parallel within each layer. Falcon-H1 models range from 0.5B to 34B parameters and support up to 256K-token context windows across 18 languages.[30]

NVIDIA Hymba and Nemotron-H

NVIDIA developed the Nemotron-H family of hybrid Mamba-Transformer models as part of its NeMo framework.[15] Nemotron-H comes in 8B, 47B, and 56B parameter sizes and is designed to offer Transformer-quality accuracy at lower inference cost. NVIDIA reports that Nemotron-H models match or beat similarly sized Transformers such as Qwen-2.5 and Llama-3.1 while being up to 3 times faster at inference.[26] The NeMo framework provides pre-training and fine-tuning recipes for Mamba 2, Mamba 2 hybrid, and Nemotron-H models.[15]

In November 2024, months before the Nemotron-H paper appeared in April 2025, NVIDIA released Hymba, a 1.5B-parameter small language model that takes the integration one step further.[16][25][26] Rather than alternating Mamba and attention layers in series, Hymba runs them in parallel within a single hybrid head: each Hymba block contains a normalization layer, a hybrid head that applies attention and SSM in parallel and then fuses their outputs, another normalization, and a feedforward network.[16] The block stack uses full attention only in the first, middle, and final blocks; the rest use sliding window attention to bound KV cache memory.[16] NVIDIA reports that over half of the work formerly done by attention can be replaced by SSM processing with negligible accuracy loss, and Hymba's cross-layer KV cache sharing reduces cache memory by up to ten times relative to a same-size Transformer.[16] In NVIDIA's comparison of sub-2B models, Hymba 1.5B has the highest average across MMLU, ARC-E, ARC-C, PIQA, Winogrande, HellaSwag, and SQuAD-C, ahead of Llama 3.2 1B, OpenELM 1B, and Qwen 2.5 1.5B, while using far less cache memory, although Qwen 2.5 1.5B scores higher on MMLU (60.9 vs 51.2).[16] Hymba's SSM heads use the original Mamba formulation rather than Mamba 2.[25]

Hunyuan-TurboS (Tencent)

In May 2025 Tencent published a technical report on Hunyuan-TurboS, which it describes as the first industry-deployed large-scale Mamba-based model.[17] Hunyuan-TurboS uses a mixture-of-experts backbone with 56B activated parameters out of 560B total, organized into 128 layers that alternate Mamba 2, attention, and feedforward sublayers using a structured AMF/MF block pattern.[17] In the AMF configuration, each block sequentially applies an attention layer, a Mamba 2 layer (a faster variant the authors call "Faster Mamba2"), and a feedforward MoE layer.[17] Grouped-query attention is used to bound the KV cache size of the attention sublayers.[17]

The model was pretrained on 16 trillion tokens, supports a 256K context window, and uses an adaptive long-short chain-of-thought scheme that lets the model switch between a fast direct-answer mode and a deliberative reasoning mode depending on prompt difficulty.[17] Tencent reports that its inference optimizations for the hybrid architecture deliver a 1.8 times speedup over Hunyuan-Turbo, its earlier pure-Transformer MoE model.[17] On the LMSYS Chatbot Arena leaderboard at release, Hunyuan-TurboS reached an Elo of 1356, placing it in the top tier of contemporary models and ahead of Gemini 2.0 Flash and the OpenAI o4-mini-2025-04-16 snapshot at the time.[17] The deployment is significant because it confirms that hybrid SSM architectures can be operated economically at frontier scale in real consumer-facing products.

Cartesia

Cartesia is an AI company co-founded by Albert Gu (as Chief Scientist) and other researchers to build real-time AI applications on top of state space model technology.[19] The company, based in San Francisco, announced $22 million in new funding led by Index Ventures in December 2024, bringing its total to $27 million,[19] and raised a $64 million Series A led by Kleiner Perkins in March 2025.[31]

Cartesia's primary product is Sonic, a text-to-speech model based on a state space model architecture.[18] Sonic was released in May 2024 with a reported model latency of 135 milliseconds, which Cartesia called the fastest for a model of its class.[18] The latency advantage comes directly from the SSM architecture: because SSMs process tokens through a fixed-size state with O(1) cost per token, there is no growing KV cache to read during generation, which reduces memory bandwidth pressure and enables consistent low-latency output regardless of prompt length.

Cartesia also released Mamba-3B-SlimPJ, an SSM trained on the SlimPajama dataset that demonstrated competitive performance against Transformer models of the same size, and has continued developing multi-modal real-time models. The company has argued that SSMs are particularly well-suited to always-on, streaming applications because their fixed memory footprint and constant per-token cost make resource consumption predictable.

What are the limitations of Mamba 2?

In-context recall

The main documented limitation of Mamba 2, shared with Mamba 1 and most SSMs, is imprecise in-context retrieval. Transformers store all past tokens in a KV cache and can attend directly to any of them during generation, making exact lookup tasks straightforward. SSMs compress past context into a fixed-size hidden state, which means tokens that were not retained strongly during compression cannot be precisely retrieved later.

NVIDIA's empirical study of Mamba-based language models found that both Mamba 1 and Mamba 2 begin to fail at phonebook lookup tasks once input sequences exceed approximately 500 tokens.[14] A Transformer-based 8B model maintained near-perfect accuracy on the same task up to its full pretraining context length of 4,096 tokens.[14] This gap persisted even after training Mamba 2 on 3.5 trillion tokens, suggesting it is architectural rather than a matter of training scale.[14]

Five-shot in-context learning, where the model reads a few examples in the prompt and generalizes to new inputs, also lags behind Transformers.[14] Because few-shot learning relies on the model attending back to the demonstration examples during generation, the SSM's compressed state representation is at a structural disadvantage.

Fixed state size

The state dimension in Mamba 2 is a fixed hyperparameter chosen at training time. While the SSD algorithm allows this dimension to be set much larger than in Mamba 1 without prohibitive cost, there is still a ceiling: once the model's hidden state is full, new information can only be accommodated by discarding or overwriting old information. Transformers have no equivalent constraint; their KV cache grows proportionally with context length.

This tradeoff means Mamba 2 is well-suited to tasks where the relevant information is distributed throughout a long sequence (dense context), but can be unreliable on tasks that require finding and retrieving a specific fact buried in a very long document (sparse retrieval).

Training instability at scale

Several groups training large SSM models reported that plain Mamba architectures can exhibit instability at scale without additional normalization. Falcon Mamba 7B addressed this by adding RMS normalization layers.[13] The Mamba 2 architecture includes group normalization on the SSM output, which partially addresses this, but training SSMs at the scale of the largest Transformer models (70B+ parameters, trillions of tokens) remains less well-characterized than Transformer training at those scales.[1] The 2025 deployment of Hunyuan-TurboS at the 560B-parameter MoE scale has begun to fill in this gap, but the recipes for stable training at frontier scale remain less openly documented than their Transformer analogs.[17]

Ecosystem maturity

As of mid-2024, the Mamba 2 ecosystem was less mature than the Transformer ecosystem. Transformer models benefit from years of optimized inference kernels, serving frameworks, and deployment tooling. Mamba 2 required custom CUDA kernels for efficient training, which initially limited availability to researchers with access to appropriate GPU hardware running Linux. Integration into mainstream frameworks like HuggingFace Transformers arrived later in 2024, and vLLM added first-class support for Mamba 2 inference in early 2025, which substantially closed the deployment gap.

Pretrained models

The reference implementation at the state-spaces/mamba GitHub repository provides five Mamba 2 models trained on 300 billion tokens from the Pile:[1]

ModelLayersModel dimensionParameters
mamba2-130m24768130M
mamba2-370m481,024370M
mamba2-780m481,536780M
mamba2-1.3b482,0481.3B
mamba2-2.7b642,5602.7B

Expanded article table

All models use a default state dimension of 64 or 128, compared to 16 in the corresponding Mamba 1 models. A hybrid variant, mamba2attn-2.7b, adds attention layers to the 2.7B model.[1]

Beyond the reference checkpoints, the following table summarizes representative production models built on Mamba-family SSM layers (Mamba 2 unless noted):

ModelYearOrganizationParametersArchitecture patternContext
Codestral Mamba2024Mistral AI7BPure Mamba 2 (no attention layers)256K
Jamba 1.5 Mini / Large2024AI21 Labs12B active of 52B / 94B active of 398B, MoEInterleaved Mamba-1 (not Mamba 2) + attention + MoE256K
Jamba 1.6 Mini / Large2025AI21 Labs12B active / 94B active, MoERefined Jamba 1.5 hybrid (Mamba-1 layers)256K
Zamba22024Zyphra1.2B / 2.7B / 7BMamba 2 backbone, shared attention with LoRA4K-16K
Bamba2024IBM Research9BHybrid Mamba 2 + attention4K (training)
Falcon-H12025TII0.5B-34BParallel Mamba 2 and attention heads256K
Nemotron-H2025NVIDIA8B / 47B / 56BMamba 2 + attention8K (pre-training)
Hymba2024NVIDIA1.5BParallel hybrid head, SWA + Mamba (Mamba-1-style SSM heads)8K
Hunyuan-TurboS2025Tencent56B active / 560B MoEAMF block pattern256K

Expanded article table

Continue reading

More from See also

More to read

Topics:AI ModelsDeep LearningModel Architecture

References

  1. ^1 ^2 ^3 ^4 ^5 ^6 ^7 ^8 ^9 ^10 ^11 ^12 ^13 ^14 ^15 ^16 ^17 ^18 ^19 ^20 ^21 ^22 ^23 ^24 ^25 ^26 ^27 ^28 ^29 ^30 ^31 ^32 ^33 ^34 ^35 ^36 ^37Dao, T., & Gu, A. (2024). "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality." *ICML 2024*. arXiv:2405.21060.
  2. ^1 ^2 ^3 ^4 ^5Gu, A., & Dao, T. (2023). "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv:2312.00752.
  3. ^1 ^2 ^3Gu, A., Goel, K., & Re, C. (2021). "Efficiently Modeling Long Sequences with Structured State Spaces." *ICLR 2022*. arXiv:2111.00396.
  4. ^1 ^2Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Re, C. (2022). "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." *NeurIPS 2022*. arXiv:2205.14135.
  5. ^1 ^2Lieber, O., et al. (2024). "Jamba: A Hybrid Transformer-Mamba Language Model." AI21 Labs. arXiv:2403.19887.
  6. ^AI21 Labs. (2024). "Jamba 1.5: Hybrid Transformer-Mamba Models at Scale." ai21.com/...brid-transformer-mamba-models-at-scale
  7. ^1 ^2 ^3AI21 Labs. (2025). "Introducing Jamba 1.6." ai21.com/...introducing-jamba-1-6
  8. ^1 ^2 ^3 ^4 ^5 ^6Mistral AI. (2024). "Codestral Mamba." mistral.ai/...codestral-mamba
  9. Mistral AI. (2024). "Mamba-Codestral-7B-v0.1 Model Card." huggingface.co/...Mamba-Codestral-7B-v0.1
  10. ^1 ^2 ^3 ^4 ^5 ^6Zyphra. (2024). "Zamba2-7B." zyphra.com/...zamba2-7b
  11. ^1 ^2 ^3 ^4 ^5IBM Research. (2025). "Meet Bamba, IBM's new attention-state space model." research.ibm.com/...bamba-ssm-transformer-model
  12. ^Hugging Face. (2024). "Bamba: Inference-Efficient Hybrid Mamba2 Model." huggingface.co/...bamba
  13. ^1 ^2 ^3Technology Innovation Institute. (2024). "Welcome Falcon Mamba: The first strong attention-free 7B model." huggingface.co/...falconmamba
  14. ^1 ^2 ^3 ^4 ^5NVIDIA. (2024). "An Empirical Study of Mamba-based Language Models." arXiv:2406.07887.
  15. ^1 ^2NVIDIA. (2024). "Mamba 2 and Hybrid Models - NeMo Framework User Guide." Archived copy: web.archive.org/...mamba
  16. ^1 ^2 ^3 ^4 ^5NVIDIA. (2024). "Hymba: A Hybrid-Head Architecture for Small Language Models." NVIDIA Technical Blog. developer.nvidia.com/...language-model-performance
  17. ^1 ^2 ^3 ^4 ^5 ^6 ^7 ^8 ^9Liu, Y., et al. (2025). "Hunyuan-TurboS: Advancing Large Language Models through Mamba-Transformer Synergy and Adaptive Chain-of-Thought." Tencent. arXiv:2505.15431.
  18. ^1 ^2Cartesia. (2024). "Announcing Sonic: a low-latency voice model for lifelike speech." cartesia.ai/...sonic
  19. ^1 ^2 ^3Index Ventures. (2024). "Building the Next Generation of Real-Time AI Models." indexventures.com/...ls-our-investment-in-cartesia
  20. ^1 ^2 ^3 ^4 ^5Dao, T. (2024). "State Space Duality (Mamba-2) Part IV - The Systems." tridao.me/...mamba2-part4-systems
  21. ^Gu, A. (2024). "State Space Duality (Mamba-2) Part I - The Model." goombalab.github.io/...mamba2-part1-model
  22. ^1 ^2 ^3 ^4Goomba Lab. (2024). "State Space Duality (Mamba-2) Part III - The Algorithm." goombalab.github.io/...mamba2-part3-algorithm
  23. ^Yang, S., et al. (2024). "Gated Linear Attention Transformers with Hardware-Efficient Training." arXiv:2312.06635.
  24. ^Mistral AI. "mistralai/Mamba-Codestral-7B-v0.1" config.json (model_type mamba2, num_hidden_layers 64, state_size 128). Hugging Face. huggingface.co/...config.json
  25. ^1 ^2 ^3Dong, X., et al. (2024). "Hymba: A Hybrid-head Architecture for Small Language Models." arXiv:2411.13676, submitted November 20, 2024. arxiv.org/...2411.13676
  26. ^1 ^2 ^3NVIDIA. (2025). "Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models." arXiv:2504.03624, submitted April 4, 2025. arxiv.org/...2504.03624
  27. ^1 ^2 ^3AI21 Labs. (2024). "Jamba-1.5: Hybrid Transformer-Mamba Models at Scale." arXiv:2408.12570. arxiv.org/...2408.12570
  28. ^DeepLearning.AI. (2024). "AI21 Labs' Jamba 1.5 Outpaces Transformers in Long-Text Processing." The Batch. charonhub.deeplearning.ai/...-long-text-processing
  29. ^1 ^2IBM. (2025). "IBM Granite 4.0: hyper-efficient, high performance hybrid models for enterprise." October 2, 2025. ibm.com/...fficient-high-performance-hybrid-models
  30. ^1 ^2Zuo, J., et al. (2025). "Falcon-H1: A Family of Hybrid-Head Language Models Redefining Efficiency and Performance." arXiv:2507.22448. arxiv.org/...2507.22448
  31. ^Garfinkle, A. (2025). "Exclusive: Cartesia, voice AI startup, raises $64 million Series A." Fortune, March 11, 2025. fortune.com/...-startup-raises-64-million-series-a

Improve this article

Add missing citations, update stale details, or write a clearer explanation. Anyone can edit: changes go live when you publish, and every edit is checked by an editor.

8 revisions · v8 · 6,409 words · full history

Fact-checks are independent of edits: a reviewer re-verifies the article against its sources and stamps the date. How we verify

Research and drafting on this wiki are AI-assisted, under named human editorial standards. How AI is used here

Reviewer note: Independent fact-check of the full article against its cited sources (170 claims, 34 sources fetched); 11 material + 28 minor fixed

Cite this page: AI Wiki. "Mamba 2." aiwiki.ai, updated 31 Oct 2026, fact-checked 2 Nov 2026. CC BY 4.0. https://aiwiki.ai/wiki/mamba_2

Edit article