Jamba — Hybrid SSM-Transformer

Attention is quadratic in sequence length, and state space models are linear. The gap widens fast: at 256k tokens a single attention map has about 65B entries per head, while an SSM's recurrent state stays the same size at any sequence length.

Pure-SSM models (Mamba, Mamba-2) match Transformer perplexity at small scale but lag on state-tracking tasks and on some kinds of in-context retrieval. My read is that an SSM compresses history into a fixed state, so when the history is long, some of it gets lost. Attention keeps the keys and values of past tokens around and pays quadratic cost for the privilege.

So use both. Put Transformer layers where exact recall matters, use SSM layers for the rest, and tune the ratio. Jamba is a production-scale model that ships this hybrid recipe (52B total, 12B active, 256k context, single 80GB GPU). Jamba 1.5 extends the family to 398B total and 94B active. Mamba-3 (ICLR 2026) is a recent pure-SSM design that a future hybrid could be rebuilt around.

This lesson reads the three papers and builds the mental model for picking the right ratio.

An SSM in one page

A state space model processes a sequence x_1, ..., x_N via a fixed-size state h:

h_t = A h_{t-1} + B x_t
y_t = C h_t

At each step the state evolves via a linear dynamics A, takes input B x_t, and emits output C h_t. A, B, C can be learned. The property that matters: computing y_t needs only h_{t-1} and x_t, not any earlier x. Memory is constant and inference is O(1) per token.

Modeling quality comes down to the structure of A. S4 (Gu 2021) used a highly structured matrix that could be evaluated as a long convolution during training. Mamba (Gu, Dao 2023) replaced the fixed A, B, C with data-dependent ones, which is the "selective" part. Mamba-2 (2024) simplified the structure further, and Mamba-3 (2026) adds complexity back in specific places.

For a decoder LLM, an SSM layer can take the place of an attention layer, with a fixed-size per-layer state instead of a growing KV cache.

The Jamba block

A Jamba block interleaves layers according to two numbers:

• l: the attention-to-Mamba ratio. Jamba uses l = 8, meaning 1 Transformer layer for every 7 Mamba layers (7 Mamba + 1 Attention = 8 layers per group).
• e: the MoE frequency. Jamba uses e = 2, meaning every other layer applies MoE.

The layer sequence within a block:

M  M  M  M  M  M  M  A    (7 Mamba + 1 Attention)
|  M  |  M  |  M  |  M    (where | marks MoE applied)

Each Jamba block is 8 layers. Stack 4 blocks (32 layers) and you get 28 Mamba layers and 4 attention layers, with MoE on 16 of the 32.

Why the 1:7 ratio

AI21 ran ablations to find the attention-to-Mamba ratio with the best perplexity per parameter and the best in-context recall on their long-context evals.

• Too much attention (1:1): quality improves, but memory and speed degrade.
• Too little attention (1:15): memory looks great, but in-context retrieval suffers.
• Sweet spot: 1:7 or 1:8.

The intuition is that the Transformer layers handle exact recall and state tracking while the Mamba layers do the cheap bulk of the processing.

Positional encoding

Mamba layers are position-aware through the recurrence, so the attention layers in the original Mamba-based hybrids did not use RoPE. The SSM layers supplied position information. Jamba 1.5 adds RoPE to the attention layers to generalize to longer context, a refinement based on long-context evaluation.

The memory budget

For a Jamba-1 shape (32 layers: 28 Mamba + 4 Attention, hidden 4096, 32 attention heads):

• KV cache (attention layers only): 2 * 4 * 32 * 128 * 256k * 2 = 16 GB at 256k BF16. Only the 4 attention layers contribute.
• SSM state: 28 * hidden * state_size, fixed per layer and independent of sequence length. With a state of 16 per feature and hidden 4096, that is 28 * 4096 * 16 * 2 = 3.7 MB total.

Compare to a pure Transformer at 32 layers, same hidden, full MHA at 32 heads: 2 * 32 * 32 * 128 * 256k * 2 = 128 GB at 256k BF16. That is an 8x reduction in KV cache. Even against a GQA(8) baseline (2 * 32 * 8 * 128 * 256k * 2 = 32 GB), the 1:7 hybrid at 16 GB is 2x smaller.

This is the arithmetic behind "256k context on a single 80GB GPU." The full-MHA KV cache of a pure Transformer would not fit, and a GQA baseline leaves little room for weights and activations. The hybrid's cache is small enough to leave that room.

Mamba-3: the pure-SSM baseline in 2026

Mamba-3 (ICLR 2026, arXiv:2603.15569) introduces three innovations on the pure-SSM side:

1. Exponential-trapezoidal discretization. Replaces the Euler-method discretization in Mamba-2 with a more expressive recurrence. A convolution-like operation is applied to the state-input inside the core recurrence, instead of as an outer convolution on x_t.

2. Complex-valued state update. Earlier Mambas simplified the state matrix from complex (S4) to real diagonal (Mamba) to scaled identity (Mamba-2). Mamba-3 brings complex values back, which is equivalent to a data-dependent rotary embedding on the state. That restores the state-tracking ability the real-valued simplifications cost.

3. Multi-input multi-output (MIMO) projections. Replaces per-feature scalar projections with matrix-valued ones. This improves modeling power and hardware utilization at inference without increasing decode latency.

At 1.5B parameters, Mamba-3 improves average downstream accuracy by 0.6 points over Gated DeltaNet; the MIMO variant adds 1.2 more for a total 1.8-point gain. At the same state size, Mamba-3 matches Mamba-2 with half the state.

Mamba-3 is not yet in a production hybrid at scale, but it is a natural candidate for the SSM side of the next Jamba-class model.

When to reach for a hybrid

Hybrids win when:

• Context is long enough that a pure-Transformer KV cache becomes painful (64k and up).
• Tasks mix short-range structure (good for SSM) with long-range recall (needs attention).
• You deploy on a single-GPU memory budget where the Transformer KV cache alone would not fit.

Hybrids lose when:

• Context is short (under 16k). The SSM machinery buys little, and a pure Transformer is fine.
• Tasks need dense attention across the whole input, such as deep reasoning or multi-document cross-reference. Having few attention layers in the hybrid hurts there.
• You are scaling toward frontier models. Pure-Transformer + MoE + MLA (DeepSeek-V3 style) is the recipe I see doing best on capability right now.

The competitive landscape

In this landscape, pure-Transformer MoE leads at the frontier and hybrids look strongest in the 256k-plus context niche. Mamba-3's state-tracking gains may push hybrid ratios toward more SSM and less attention in the next generation.

code/main.py is a memory calculator for hybrid architectures. Given the number of attention layers, a hidden size, a head config, and an SSM state size, it computes:

• KV cache at a target context length.
• SSM state memory.
• The sum of the two at 8k, 65k, 131k, and 262k context (8192, 65536, 131072, and 262144 tokens), for each configuration.

It covers these configurations:

• A pure-Transformer baseline (full MHA and GQA 8), where the KV cache grows with context.
• Jamba-style hybrids at 1:7, 1:3, and 1:15.
• A pure SSM with no KV cache.

The numbers are illustrative. They count KV cache and SSM state only, with no weights or activations, so treat them as a way to compare shapes and not as a deployment budget. The closing takeaway is computed from the same functions as the table, so it prints the same numbers: 128 GB for the full-MHA Transformer and 16 GB for the 1:7 hybrid at 256k.

Integration considerations for a real deployment:

• Production inference servers such as vLLM and SGLang support Jamba and Mamba. Check the specific version.
• At 256k context, Jamba's memory advantage shows up in concurrent-request throughput, because the same VRAM holds more Jamba sequences than Transformer sequences.
• Mamba-3 as a standalone model is a research preview at 1.5B.

This lesson produces outputs/skill-hybrid-picker.md. Given a workload specification (context length profile, task mix, memory budget), it recommends between a pure Transformer, a Jamba-style hybrid, and a pure SSM, with explicit reasoning about the memory and quality tradeoffs.

1. Run code/main.py to compute KV cache at 256k context for a 32-layer pure Transformer (hidden 4096, 32 heads) and for a Jamba-1 hybrid of the same shape. Verify the ~8x memory reduction the AI21 paper claims.

2. Modify the calculator to model a 1:3 hybrid (3 Mamba : 1 Attention) and a 1:15 hybrid (15 Mamba : 1 Attention). Plot KV cache vs ratio. At what ratio does the KV cache equal the SSM state memory?

3. Read Section 3 of the Jamba paper (arXiv:2403.19887). Explain why AI21 uses Mamba-1 rather than Mamba-2 despite Mamba-2 being faster. Hint: the hybrid ablation section documents this.

4. Compute the parameter overhead of MoE-every-other-layer in Jamba 1.5 Large (398B total, 94B active). Compare the active ratio to DeepSeek-V3 (37B/671B) and explain why Jamba's architecture pushes the active ratio higher.

5. Read Section 3 of the Mamba-3 paper (arXiv:2603.15569). Explain in three sentences why a complex-valued state update is equivalent to a data-dependent rotary embedding. Tie the answer to Phase 7 · Lesson 04's RoPE derivation.

• Lieber et al. - Jamba: A Hybrid Transformer-Mamba Language Model (arXiv:2403.19887): the original Jamba paper, with ratio ablations and the 256k context claim
• AI21 - Jamba 1.5: Hybrid Transformer-Mamba at Scale (arXiv:2408.12570): the scaled-up family, 398B/94B and 12B/52B public releases
• Gu, Dao - Mamba: Linear-Time Sequence Modeling with Selective State Spaces (arXiv:2312.00752): the selective SSM paper Jamba builds on
• Dao, Gu - Mamba-2 (arXiv:2405.21060): the simplified structured-state-space successor
• Lahoti et al. - Mamba-3 (arXiv:2603.15569, ICLR 2026): complex-valued state and MIMO, a recent pure-SSM design
• Gu et al. - Efficiently Modeling Long Sequences with Structured State Spaces (arXiv:2111.00396): the S4 paper, where the SSM line for LLMs starts