Jamba’s Context Window: A Hybrid Mamba-Transformer Architecture

AI21 documents a 256K-token context window for Jamba. The number on its own is not the interesting part — several transformers claim as much. The interesting part is that Jamba is documented as fitting a large fraction of that window on a single 80GB GPU, and the reason is structural: most of its layers do not keep a KV cache at all.

The documented figure

Jamba was released by AI21 Labs in March 2024, with the Jamba 1.5 family — Mini and Large — following in August 2024. Across those releases the context window stated on AI21’s model cards is 256K tokens, and the original release announcement stated that a single 80GB A100 could hold roughly 140K tokens of context for the base model at half-precision. Both figures come from AI21’s own materials; the architecture is described in the Jamba paper, published on arXiv.

Parameter counts for the initial release: approximately 52B total with about 12B active per token. For Jamba 1.5 Large, approximately 398B total with about 94B active. The gap between those two numbers is mixture-of-experts routing and is discussed below; it is a separate mechanism from the one that produces the context window, and conflating them is the usual mistake.

What the layer stack actually looks like

A conventional decoder-only transformer repeats one block: attention, then a feed-forward network. Jamba repeats a larger unit built from three ingredients — attention layers, Mamba state-space layers, and MoE feed-forward layers — in a fixed proportion. In the published configuration the repeating block contains eight layers with a single attention layer among them and Mamba layers filling the rest, giving a ratio of roughly one attention layer to seven state-space layers. MoE is applied to the feed-forward position in every other layer, with 16 experts of which the router selects 2 per token.

The design question that produces this is not “can we replace attention”. Pure state-space models are cheap over long sequences but weaker at the precise recall that attention does well — pulling one specific earlier token back out of a long context. The hybrid keeps enough attention layers to retain that ability, and pays the attention cost on one layer in eight instead of on all of them.

It is worth being clear about what a Mamba layer does in place of attention, because “state-space layer” is opaque and the idea is not. Attention, at each new token, looks back over every previous token and computes a weighted blend. A state-space layer instead carries a fixed-size summary of everything seen so far, updates it with the new token, and reads from the updated summary. It is a recurrence, of the kind that predates transformers, made trainable at scale by a particular parameterisation. The trade is exactly what you would expect from that description: constant memory and constant work per token regardless of how much history there is, at the price of having compressed that history into a fixed budget rather than keeping it verbatim. Attention keeps everything and pays for it. Mamba summarises and cannot go back.

Why the KV cache is the real constraint

The thing that stops you serving a long context is usually not arithmetic, it is memory. During generation a transformer keeps, for every layer and every previous token, the key and value vectors that layer computed — the KV cache — because recomputing them at every step would be quadratic work. That cache grows linearly with sequence length and linearly with the number of attention layers, and it is per concurrent request, so it is what bounds your batch size.

The arithmetic is simple enough to do in a line. Cache bytes are approximately:

kv_bytes ≈ 2 (K and V)
         × layers_with_attention
         × kv_heads × head_dim
         × sequence_length
         × bytes_per_element

Every factor there is fixed by the architecture except sequence length, which is the user’s, and layers_with_attention, which the hybrid changes. Cutting the number of attention layers by roughly eight-fold cuts the KV cache by roughly eight-fold at the same sequence length. That is the whole mechanism behind the claim about fitting long contexts on one GPU. A Mamba layer does not keep a growing cache; it carries a fixed-size recurrent state whose memory cost is the same at token 200,000 as at token 20.

Note what this does not change: prefill still has to process every token of your prompt, so time-to-first-token still grows with prompt length, and you are still billed for those input tokens. The hybrid buys memory headroom and long-sequence throughput, not free input.

The batch-size consequence is the one with commercial teeth. A serving cluster’s economics are set by how many concurrent requests fit on a GPU, because the weights are loaded once and shared while the cache is per request. Cut the cache per request by a large factor and you fit proportionally more requests, which lowers the cost of every one of them. That is why long-context pricing from providers running hybrid models can look structurally different from long-context pricing on a dense transformer, and it is a better reason to care about the architecture than the headline window is.

The MoE part is a separate saving

Jamba is sparse as well as hybrid, and the two are easy to conflate because both are described as making a big model cheap. They make different things cheap. Mixture-of-experts routing reduces the arithmetic per token by using a fraction of the feed-forward weights — 12B active out of 52B total in the original release. That affects latency and per-token price. It does nothing for the KV cache, and it makes the memory requirement worse rather than better, because every expert has to be resident even though most are idle on any given token.

So the two mechanisms pull in opposite directions on memory and the same direction on compute. If you want the sparse-activation story on its own, with none of the state-space complication, the clearest example is DBRX’s 36B-active-of-132B routing, and the general form is in why a 400B model can cost like a 40B one.

What this changes when you call it

• Throughput at long context degrades more gently. A full-attention model’s tokens-per-second falls off as the context fills, because each generated token attends over everything before it. With most layers recurrent, that fall-off applies to a fraction of the stack.
• Long-context recall is still worth testing. A documented window is a capacity, not a guarantee of retrieval quality across it. The hybrid keeps attention layers precisely because recall needs them, but how well it holds at 200K is an empirical question about your task, not a property of the architecture.
• Self-hosting needs a runtime that knows the architecture. Mamba layers are not something a generic transformer kernel handles. Check that your serving stack lists Jamba support explicitly rather than assuming a Hugging Face config is enough.
• Quantisation behaves differently. The memory profile is dominated by expert weights rather than by cache, so the wins from weight quantisation are proportionally larger than they would be on a dense long-context transformer.