Deploy AI21 Jamba2 on GPU Cloud: Long-Context Guide (2026)

AI21 shipped Jamba2 on January 8, 2026, under Apache 2.0, and it has had close to zero deploy-guide coverage since. That's a gap worth closing, because Jamba2 Mini isn't just another open-weights MoE model: 28 of its 32 layers are Mamba, not attention, which changes the VRAM and KV-cache math you'd normally run for a model this size. This guide covers what Jamba2 actually is, how much GPU you need for the 3B dense and Mini (52B total, 12B active) variants, and the exact vLLM setup and cost math for running it on rented hardware.

For background on how KV cache growth drives inference cost as context length increases, see the AI memory wall guide. If you've already deployed a pure state-space model, the Mamba-3 GPU deployment guide covers the fully-recurrent end of this design space, and the Liquid AI LFM2 guide covers a different hybrid ratio (18 linear layers to 6 GQA layers) worth comparing against Jamba2's 28-to-4 split.

What Jamba2 Is: Mamba-Transformer Hybrid Architecture Explained

Jamba2 Mini is not a transformer with a few tricks bolted on. Of its 32 total hidden layers, only 4 use full attention, at positions 4, 12, 20, and 28 (attn_layer_offset=4, attn_layer_period=8 in the published config). The other 28 layers are Mamba state-space blocks with mamba_d_state=16 and mamba_d_conv=4. That's a 12.5% attention ratio, which is the core fact that determines everything downstream: how much KV cache the model needs, how it scales to 256K context, and which GPU tier makes sense for you.

On top of that hybrid backbone, Jamba2 Mini routes through a sparse mixture of experts: 16 experts total, top-2 active per token, 4096 hidden size, and 8 KV heads using grouped-query attention on the layers that do attend. Total parameters land at 52B, with only 12B active on any given forward pass. The documented context window is 256K tokens (262,144 in max_position_embeddings), matching AI21's published model card. Don't confuse this with the 1M-token figures floating around for other 2026 releases; Jamba2's verified maximum is 256K.

This lineage traces back to the original Jamba paper from AI21, which introduced the interleaved Transformer-Mamba-MoE design and reported that the resulting model "fits in a single 80GB GPU" at up to 256K context. Jamba2 is the current production evolution of that same architectural bet, now released as two open-weights models.

AI21 didn't just release BF16 weights. They also shipped a native FP8 checkpoint, AI21-Jamba2-Mini-FP8, at launch, which matters a lot for GPU sizing, covered below.

The training recipe behind Jamba2 also differs from a typical open-weights release. AI21 mid-trained Jamba2 itself on 500B curated tokens with an elevated share of math and code, then ran SFT, DPO, and multiple rounds of on-policy RL combining verifiable and model-based rewards. The stated goal isn't to compete with reasoning models on chain-of-thought benchmarks. It's to give enterprise teams a model that's reliable and fast on grounded, instruction-following work without the token overhead of a "thinking" model. In human evaluation against Ministral 3 14B on enterprise tasks (QA, summarization), Jamba2 Mini showed a statistically significant win-rate advantage, and it leads on IFBench, IFEval, Collie, and FACTS, the instruction-following and grounding benchmarks AI21 cites as most relevant to that positioning.

Jamba2 vs Jamba Reasoning 3B: Two Different AI21 Releases

It's easy to conflate these because both carry the Jamba name and both are small enough to matter for self-hosting, but they're separate products with separate goals.

Jamba Reasoning 3B shipped in October 2025 as a math- and code-reasoning specialist, tuned to run at up to 250K context on consumer hardware, including a MacBook Pro. It's a distillation-style release aimed at local inference, not a rented-GPU workload.

Jamba2's 3B dense model, released alongside Jamba2 Mini in January 2026, is a different animal. It's tuned for instruction-following and grounding rather than reasoning distillation, and it ships under the same Apache 2.0 license as Jamba2 Mini. If your workload needs actual GPU capacity, the model this guide is about is Jamba2 Mini: the 52B-total, 12B-active MoE that's genuinely too large to run comfortably on a laptop and needs a rented H100 or H200 to serve at production concurrency.

Jamba2 Model Family: 3B Dense vs Mini (12B Active / 52B Total MoE)

AI21 published two sizes on January 8, 2026, both under Apache 2.0:

The 3B dense model is small enough to self-host on a single mid-range GPU and is the more practical option if your workload is lightweight instruction-following or you're prototyping before committing to Jamba2 Mini's larger footprint. Jamba2 Mini is the model this guide focuses on, because it's the size where GPU sizing decisions actually matter and where the hybrid architecture's VRAM savings pay off.

Both models are available on the ai21labs HuggingFace organization, including the base BF16 checkpoints and the FP8 variant covered next.

VRAM and GPU Sizing for Jamba2 Mini at 256K+ Context

Jamba2 Mini's 52B total parameters set the weight footprint, and that number alone determines your GPU tier before you even get to KV cache.

At BF16 (2 bytes per parameter): 52B x 2 = ~104GB. That's larger than a single 80GB H100, so BF16 Jamba2 Mini needs either a 141GB H200 (with headroom left for KV cache) or a 2-GPU tensor-parallel setup across smaller cards.

BF16 vs FP8: Why the Quantized Build Changes Your GPU Choice

AI21 published AI21-Jamba2-Mini-FP8 as a native checkpoint, not a community quantization. At FP8 (1 byte per parameter): 52B x 1 = ~52GB. That fits comfortably on a single H100 (80GB), leaving roughly 28GB for KV cache, activations, and batching.

That's the practical decision point for most teams: BF16 forces you onto an H200 or a multi-GPU H100 setup, while the native FP8 build runs on a single H100. Unless you have a specific reason to need full BF16 precision, the FP8 checkpoint is the more cost-efficient starting point, and it's officially supported rather than a third-party quant.

Because only 4 of Jamba2 Mini's 32 layers generate a KV cache, that headroom goes further than it would on a comparable pure transformer, covered next.

Why Hybrid SSM Models Change the Long-Context Cost Curve vs Pure Transformers

Here's the part that actually matters for GPU economics: KV cache is what makes long-context transformer serving expensive, and Jamba2 Mini mostly doesn't have one.

A standard transformer's KV cache grows with every layer, every KV head, and every token in context. At 256K tokens, a 32-layer dense transformer with full attention on all layers has to hold a cache proportional to all 32 layers' worth of keys and values. That's the reason long-context serving on pure transformers pushes you toward higher-VRAM GPUs even when the model weights themselves would fit on something smaller: the cache, not the weights, becomes the binding constraint.

Jamba2's 4-in-32 Attention Layer Ratio and What It Means for KV Cache

Jamba2 Mini only accumulates a KV cache from its 4 attention layers (positions 4, 12, 20, and 28). The other 28 layers are Mamba blocks, which carry a fixed-size recurrent state (governed by mamba_d_state=16 and mamba_d_conv=4) that does not grow with sequence length at all. Whether you're processing 2K tokens or the full 256K context window, the Mamba layers' memory footprint is constant.

The practical result: Jamba2 Mini's cache overhead at 256K context is a fraction of what a 32-layer dense transformer of comparable size would need, because 87.5% of its layers contribute nothing to that growth. This is the same directional pattern other hybrid releases have shown in production. IBM's Bamba, which cuts KV cache overhead the same way, reported up to 2.5x the throughput and 2x the latency improvement of standard transformers in vLLM. NVIDIA's Nemotron-H tested multiple sizes on long-context workloads: its 8B variant hit up to 3x faster throughput than Llama 3.1 8B, while the 47B variant, which replaces most attention layers with Mamba2 blocks, reached up to 2.9x faster throughput than similarly sized Qwen-2.5-72B and Llama-3.1-70B baselines. Together AI's research on Mamba-Llama distillation notes that Mamba architectures can offer up to 5x higher inference throughput than transformers on general workloads. Those are separate models, not Jamba2 benchmarks, but they establish that Jamba2's design pattern (mostly-Mamba layers, a thin slice of attention) is a proven way to buy back GPU headroom at long context, not a novel or unproven bet.

For the transformer-side techniques that mitigate this same problem (PagedAttention, prefix caching, quantized cache dtypes), see the KV cache optimization guide and the context engineering guide for production agents. Both apply directly to Jamba2's 4 attention layers, just against a much smaller base.

Step-by-Step: Serving Jamba2 Mini on Rented GPUs with vLLM

AI21's model card specifies vLLM 0.12.0 or higher as the supported serving path. Here's the full deployment sequence.

1. Provision the instance. Log in to app.spheron.ai and select a GPU matching your checkpoint choice: a single H100 SXM5 for the FP8 build, or an H200 for BF16. Deploy with Ubuntu 22.04 and the NVIDIA Docker runtime template, then confirm the card is visible with nvidia-smi.

2. Install vLLM.

pip install "vllm>=0.12.0"
vllm --version

3. Download the weights. For the FP8 build (recommended starting point):

huggingface-cli download ai21labs/AI21-Jamba2-Mini-FP8 --local-dir ./jamba2-mini-fp8

For BF16:

huggingface-cli download ai21labs/AI21-Jamba2-Mini --local-dir ./jamba2-mini

4. Launch the server with the Mamba-aware flags. This is where Jamba2 differs from a standard vLLM launch command:

vllm serve ./jamba2-mini-fp8 \
  --mamba-ssm-cache-dtype float32 \
  --enable-prefix-caching \
  --max-model-len 262144 \
  --host 0.0.0.0 \
  --port 8000

--mamba-ssm-cache-dtype float32 keeps the recurrent state in full precision even when the weights are quantized. Skipping this flag on the FP8 build degrades output quality, since the Mamba state is more sensitive to precision loss than the attention KV cache is. --enable-prefix-caching is worth keeping on for the 4 attention layers, same as any transformer deployment. If you're instead serving a separate Int8-experts checkpoint rather than this native FP8 build, add --quantization experts_int8 to that command; don't attach it to the native FP8 checkpoint shown above.

If you'd rather use the HuggingFace Transformers path instead of vLLM, you'll need flash-attn (installed with --no-build-isolation), causal-conv1d>=1.2.0, mamba-ssm, and transformers>=4.54.0. vLLM is the faster path for production serving; Transformers is useful for fine-tuning or debugging.

5. Validate.

curl http://localhost:8000/v1/models

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "jamba2-mini-fp8", "messages": [{"role": "user", "content": "Summarize this document."}]}'

For tool-calling workloads, AI21's recommended command also adds --enable-auto-tool-choice --tool-call-parser hermes, which is worth including if your application makes function calls.

Cost Analysis: Running Jamba2 Mini on Spheron GPUs

The checkpoint choice is also the cost choice, since it determines whether you need one GPU or two.

Pricing fluctuates based on GPU availability. The prices above are based on 14 Jul 2026 and may have changed. Check current GPU pricing → for live rates.

The gap between the FP8 and BF16 rows is smaller than you'd expect given the VRAM difference, since H100 and H200 on-demand rates currently sit close together. The real saving isn't a cheaper GPU, it's that FP8 lets you run on a single card at all rather than needing tensor parallelism across two H100s to hold 104GB of BF16 weights. That's one fewer GPU in your bill, plus one less thing to configure.

For teams that only need the 3B dense model, an A100 80GB on-demand covers it with plenty of headroom, and spot pricing on either tier is worth using for batch or non-latency-sensitive workloads where a reclaim mid-job is tolerable.

When Jamba2 Beats a Transformer Model for Your Use Case (and When It Doesn't)

Jamba2 Mini is the stronger choice when:

• Your workload genuinely needs 100K+ token context and you want to avoid the KV-cache-driven GPU tax that comes with running a same-size pure transformer at that length
• You're running grounded, instruction-following enterprise work (QA, summarization, and document analysis) rather than open-ended chain-of-thought reasoning
• You want the reliability profile AI21 is targeting: consistent, non-"thinking" outputs without the extra latency and token cost of a reasoning model
• You need an Apache 2.0 license with no usage restrictions

A pure transformer may still be the better call when:

• Your context lengths are consistently short (under 8K tokens), where the hybrid architecture's KV-cache savings don't matter and mature FlashAttention-tuned transformer tooling wins on raw throughput
• You need deep fine-tuning ecosystem support. LoRA and PEFT tooling for dense and standard MoE transformers is far more mature than tooling for hybrid Mamba-MoE architectures like Jamba2
• Your task depends on complex multi-step chain-of-thought reasoning, which is explicitly not what Jamba2 is trained for; a dedicated reasoning model will outperform it there
• You need a context window beyond 256K tokens, which is Jamba2's documented ceiling

Jamba2 Mini is a genuinely different bet than the other MoE models we've covered on this blog, including GLM-5.2. Most large open-weights MoEs are pure transformers with sparse routing bolted on for compute efficiency; Jamba2 changes the layer composition itself, and the VRAM and KV-cache math follows from that. If your workload sits in the 100K-256K context range on grounded enterprise tasks, it's worth a real evaluation against whatever dense transformer you're running today. For a broader VRAM-tier view across open-weights models generally, see the best open-source LLMs to self-host guide.

Jamba2 Mini's FP8 checkpoint fits a single H100 with room for a 256K-token KV cache, no multi-GPU setup required.

View live GPU pricing →