hybrid attention
**Hybrid attention-SSM architectures** interleave a small number of full quadratic-attention layers with a majority of linear-time state-space (SSM/Mamba) layers in a single model — capturing attention's perfect recall for rare or distant tokens while keeping Mamba's O(1)-per-token inference cost for the bulk of context processing. The result: models that match pure-Transformer quality on language benchmarks at 2–5× lower decode latency and dramatically lower KV-cache memory, especially at long context (32k–1M tokens). Jamba (AI21, 2024), Zamba (Zyphra, 2024), StripedHyena (Together AI, 2023), Griffin (Google DeepMind, 2024), and RecurrentGemma are the leading examples.
**Why pure attention and pure SSM each leave something on the table.** A pure Transformer with $L$ layers has $L$ KV-cache entries per token — the cache grows linearly with both sequence length and layer count, so a 70B model at 128k context can consume >100 GB of HBM just for KV. A pure SSM (Mamba) replaces the cache with a fixed-size recurrent state ($d_{\text{state}} \times d_{\text{model}}$ per layer, independent of sequence length), so memory is constant — but the state has finite capacity, and empirically pure-SSM models underperform attention on recall-intensive tasks (multi-hop reasoning, exact copying over very long distances, retrieval from arbitrary positions).
The hybrid insight: **a few attention layers placed strategically give the model a "scratchpad" for exact recall, while the SSM layers handle the bulk of sequential reasoning at constant memory cost.**
**Architecture patterns.** The ratio of attention-to-SSM layers, their placement, and whether they share KV-cache or use grouped-query attention (GQA) vary across designs:
| Model | Total layers | Attention layers | SSM layers | Ratio (attn:ssm) | MoE? | Context | Key design choice |
|---|---|---|---|---|---|---|---|
| Jamba (AI21, 52B) | 32 | 8 (every 4th) | 24 | 1:3 | Yes (16 experts, top-2) | 256k | Attention + Mamba + MoE in same block |
| Zamba-7B (Zyphra) | 36 | 6 (shared KV) | 30 | 1:5 | No | 4k+ | Shared attention KV across all attn layers |
| StripedHyena-7B | 32 | 8 (interleaved) | 24 | 1:3 | No | 32k–128k | Hyena (long-conv) + attention |
| Griffin (DeepMind) | varies | ~25% | ~75% | 1:3 | No | ∞ (recurrent) | RG-LRU (gated linear recurrence) + local attn |
| RecurrentGemma-9B | 26 | 6 | 20 | ~1:3 | No | 8k (local) | Griffin-based, local sliding-window attn |
| Mamba-2-Hybrid (Nvidia) | 56 | 8 | 48 | 1:6 | No | 8k | SSD (structured state-space duality) + attn |
**The 1:3 to 1:6 sweet spot.** Empirically, placing one attention layer for every 3–6 SSM layers recovers 95–100% of pure-Transformer quality while cutting KV-cache by 70–85%. The attention layers act as "information highways" — positions where the model can perform exact copying, attend to any arbitrary position in context, and aggregate information that the SSM layers' finite state can't perfectly retain.
**Memory and latency analysis at inference.** For a model with $L$ total layers, $L_a$ attention layers, $L_s$ SSM layers, sequence length $S$, hidden dim $D$, and KV-head dim $d_k$:
$$\text{KV cache} = 2 \cdot L_a \cdot S \cdot n_{\text{kv\_heads}} \cdot d_k \cdot \text{bytes}$$
$$\text{SSM state} = L_s \cdot d_{\text{state}} \cdot D \cdot \text{bytes}$$
For a Jamba-52B-class model ($L_a = 8$, $L_s = 24$, $S = 128\text{k}$, GQA with 8 KV heads, $d_k = 128$, fp16): KV cache ≈ 2 × 8 × 128k × 8 × 128 × 2 = **4 GB** (vs ~50 GB for a pure 32-layer Transformer at 128k). SSM state ≈ 24 × 64 × 8192 × 2 = **24 MB** — negligible. Total memory for sequence state: **~4 GB vs ~50 GB** for equivalent-quality pure attention.
**Decode latency** scales with the number of attention layers (each requires a KV-cache read across all past positions), while SSM layers are O(1) — just a matrix multiply on the fixed state vector. With 8 attention layers instead of 32, decode self-attention cost drops 4×; the SSM layers add negligible latency (a small matmul per layer).
**Mamba-2 and Structured State-Space Duality (SSD).** Mamba-2 (Dao & Gu, 2024) reframes the selective state-space model as a structured-masked attention operation — showing that the SSM recurrence is mathematically equivalent to a specific (block-diagonal + causal) attention pattern. This "duality" means SSM layers can be implemented using the same hardware-efficient tiled matmul kernels as FlashAttention, achieving near-attention-level hardware utilization on modern GPUs/TPUs while preserving O(1) recurrent inference. Mamba-2-Hybrid stacks these SSD layers with a few conventional attention layers for exact recall.
```svg
```
**Training considerations.** Hybrid models train with the same parallelism strategies as pure Transformers (tensor-parallel, pipeline-parallel, FSDP), because the SSM layers have identical per-layer parameter counts. The key training difference: SSM layers can process prefill in both recurrent mode (sequential, O(S) total) or parallel-scan mode (log-depth, O(S log S) total). Mamba-2's SSD formulation enables a chunked parallel-scan that processes prefill as matmuls — matching FlashAttention's hardware efficiency during training while preserving O(1) recurrent inference.
**When to choose hybrid over pure Transformer.** Hybrid attention-SSM is the strongest fit when: (1) inference context is routinely long (32k–1M tokens) and KV-cache memory dominates serving cost; (2) streaming / real-time decode is needed (chatbots, code completion) where per-token latency matters; (3) the task requires some exact recall (so pure SSM underperforms) but not maximum recall across every position; (4) cost-per-token must be minimized at scale. At short context (<4k) with small batch, pure Transformers are simpler and equally fast — the hybrid advantage emerges at scale.
**Hardware implications.** For chip architects, hybrid models shift the inference bottleneck: the few attention layers remain memory-bandwidth-bound (reading the KV-cache), while the majority SSM layers are compute-bound (state-update matmuls). This means an ideal hybrid-model accelerator should be balanced — high FLOPS for SSM layers but also high HBM bandwidth for the periodic attention layers — which is exactly what the CFS Inference Simulator at /infer models (roofline analysis showing compute-vs-memory bottleneck per layer type).