mirror of
https://github.com/rohitg00/ai-engineering-from-scratch.git
synced 2026-10-02 01:54:39 +08:00
feat(phase-10/21): Jamba hybrid SSM Transformer
This commit is contained in:
Rohit Ghumare
@@ -0,0 +1,114 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 960 560" font-family="Georgia, 'Times New Roman', serif">
|
||||
<defs>
|
||||
<marker id="arr" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="6" markerHeight="6" orient="auto">
|
||||
<path d="M0,0 L10,5 L0,10 z" fill="#1a1a1a"/>
|
||||
</marker>
|
||||
<style>
|
||||
.box { fill: #faf6ef; stroke: #1a1a1a; stroke-width: 1.5; }
|
||||
.attn { fill: #fff1d6; stroke: #c0392b; stroke-width: 1.5; }
|
||||
.mamba { fill: #e6f4ea; stroke: #2e7d32; stroke-width: 1.5; }
|
||||
.moe { fill: #e8eaf6; stroke: #3949ab; stroke-width: 1.2; }
|
||||
.old { fill: #eeeeee; stroke: #888; stroke-width: 1.2; }
|
||||
.label { font-size: 13px; font-weight: 600; fill: #1a1a1a; }
|
||||
.mono { font-size: 11px; font-family: 'Menlo', monospace; fill: #222; }
|
||||
.small { font-size: 10px; font-family: 'Menlo', monospace; fill: #555; }
|
||||
.caption { font-size: 11px; fill: #555; font-style: italic; }
|
||||
.title { font-size: 16px; font-weight: 700; fill: #1a1a1a; }
|
||||
.head { font-size: 12px; font-weight: 700; fill: #1a1a1a; }
|
||||
.hi { font-size: 22px; font-weight: 700; fill: #c0392b; }
|
||||
</style>
|
||||
</defs>
|
||||
|
||||
<text x="480" y="28" text-anchor="middle" class="title">jamba — one attention per seven mamba, moe every other layer</text>
|
||||
|
||||
<!-- Layer stack: show 8 layers (one block) with attention on layer 1 and 9, moe every even layer -->
|
||||
<text x="30" y="70" class="head">one jamba block (l=8)</text>
|
||||
|
||||
<!-- Layer row: 8 cells. Layer 1 = attn. Layers 2..8 = mamba. Moe on even layers (2,4,6,8). -->
|
||||
<g transform="translate(30,82)">
|
||||
<!-- Layer 1: attn + dense -->
|
||||
<rect x="0" y="0" width="100" height="36" class="attn"/>
|
||||
<text x="50" y="22" text-anchor="middle" class="mono">L1 attn + mlp</text>
|
||||
<!-- Layer 2: mamba + moe -->
|
||||
<rect x="110" y="0" width="100" height="36" class="mamba"/>
|
||||
<text x="160" y="22" text-anchor="middle" class="mono">L2 mamba + moe</text>
|
||||
<!-- Layer 3: mamba + dense -->
|
||||
<rect x="220" y="0" width="100" height="36" class="mamba"/>
|
||||
<text x="270" y="22" text-anchor="middle" class="mono">L3 mamba + mlp</text>
|
||||
<!-- Layer 4: mamba + moe -->
|
||||
<rect x="330" y="0" width="100" height="36" class="mamba"/>
|
||||
<text x="380" y="22" text-anchor="middle" class="mono">L4 mamba + moe</text>
|
||||
<!-- Layer 5: mamba + dense -->
|
||||
<rect x="440" y="0" width="100" height="36" class="mamba"/>
|
||||
<text x="490" y="22" text-anchor="middle" class="mono">L5 mamba + mlp</text>
|
||||
<!-- Layer 6: mamba + moe -->
|
||||
<rect x="550" y="0" width="100" height="36" class="mamba"/>
|
||||
<text x="600" y="22" text-anchor="middle" class="mono">L6 mamba + moe</text>
|
||||
<!-- Layer 7: mamba + dense -->
|
||||
<rect x="660" y="0" width="100" height="36" class="mamba"/>
|
||||
<text x="710" y="22" text-anchor="middle" class="mono">L7 mamba + mlp</text>
|
||||
<!-- Layer 8: mamba + moe -->
|
||||
<rect x="770" y="0" width="100" height="36" class="mamba"/>
|
||||
<text x="820" y="22" text-anchor="middle" class="mono">L8 mamba + moe</text>
|
||||
</g>
|
||||
|
||||
<text x="30" y="142" class="small">ratio 1:7 attention : mamba · moe on every 2nd layer (16 experts, top-2)</text>
|
||||
<text x="30" y="158" class="small">block repeats 4 times -> 32 layers · 4 attention · 28 mamba · 16 moe mlps</text>
|
||||
|
||||
<!-- KV cache budget comparison -->
|
||||
<rect x="30" y="185" width="900" height="160" class="box"/>
|
||||
<text x="480" y="208" text-anchor="middle" class="label">memory at 128k context · BF16 · same hidden dim</text>
|
||||
|
||||
<!-- Row 1: pure transformer -->
|
||||
<text x="60" y="240" class="mono">pure transformer</text>
|
||||
<text x="60" y="256" class="small">32 attention layers · GQA 32/8</text>
|
||||
<rect x="300" y="224" width="530" height="28" class="old"/>
|
||||
<rect x="300" y="224" width="530" height="28" fill="#c0392b" fill-opacity="0.25"/>
|
||||
<text x="565" y="244" text-anchor="middle" class="mono">KV cache · 16.0 GB</text>
|
||||
|
||||
<!-- Row 2: jamba -->
|
||||
<text x="60" y="292" class="mono">jamba 1:7</text>
|
||||
<text x="60" y="308" class="small">4 attention layers · 28 mamba · GQA 32/8</text>
|
||||
<rect x="300" y="276" width="66" height="28" class="old"/>
|
||||
<rect x="300" y="276" width="66" height="28" fill="#2e7d32" fill-opacity="0.30"/>
|
||||
<text x="380" y="296" class="mono">KV 2.0 GB</text>
|
||||
<rect x="470" y="276" width="14" height="28" class="old"/>
|
||||
<text x="495" y="296" class="mono">+ SSM state 7 MB (fixed)</text>
|
||||
|
||||
<text x="60" y="332" class="hi">8× smaller KV</text>
|
||||
|
||||
<!-- Mamba block internals -->
|
||||
<text x="30" y="380" class="head">inside one mamba layer (Gu & Dao 2023)</text>
|
||||
<rect x="30" y="395" width="900" height="100" class="box"/>
|
||||
|
||||
<rect x="55" y="420" width="100" height="52" class="box"/>
|
||||
<text x="105" y="443" text-anchor="middle" class="mono">input proj</text>
|
||||
<text x="105" y="460" text-anchor="middle" class="small">h -> 2·d_inner</text>
|
||||
|
||||
<line x1="155" y1="446" x2="185" y2="446" stroke="#1a1a1a" stroke-width="1.5" marker-end="url(#arr)"/>
|
||||
|
||||
<rect x="185" y="420" width="100" height="52" class="box"/>
|
||||
<text x="235" y="443" text-anchor="middle" class="mono">conv1d · SiLU</text>
|
||||
<text x="235" y="460" text-anchor="middle" class="small">local mixing</text>
|
||||
|
||||
<line x1="285" y1="446" x2="315" y2="446" stroke="#1a1a1a" stroke-width="1.5" marker-end="url(#arr)"/>
|
||||
|
||||
<rect x="315" y="420" width="220" height="52" class="mamba"/>
|
||||
<text x="425" y="441" text-anchor="middle" class="mono">selective SSM</text>
|
||||
<text x="425" y="458" text-anchor="middle" class="small">(B, C, Δ) = f(x_t) · parallel scan</text>
|
||||
|
||||
<line x1="535" y1="446" x2="565" y2="446" stroke="#1a1a1a" stroke-width="1.5" marker-end="url(#arr)"/>
|
||||
|
||||
<rect x="565" y="420" width="100" height="52" class="box"/>
|
||||
<text x="615" y="443" text-anchor="middle" class="mono">gate · SiLU</text>
|
||||
<text x="615" y="460" text-anchor="middle" class="small">output gate</text>
|
||||
|
||||
<line x1="665" y1="446" x2="695" y2="446" stroke="#1a1a1a" stroke-width="1.5" marker-end="url(#arr)"/>
|
||||
|
||||
<rect x="695" y="420" width="100" height="52" class="box"/>
|
||||
<text x="745" y="443" text-anchor="middle" class="mono">out proj</text>
|
||||
<text x="745" y="460" text-anchor="middle" class="small">d_inner -> h</text>
|
||||
|
||||
<text x="30" y="520" class="caption">O(state_dim) hidden per layer · no kv cache · selective = content-aware</text>
|
||||
<text x="480" y="548" text-anchor="middle" class="caption">jamba keeps attention only where it matters for recall · mamba carries the rest</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 6.3 KiB |
@@ -0,0 +1,259 @@
|
||||
"""Jamba hybrid SSM-Transformer calculator.
|
||||
|
||||
Simulate the Jamba block schedule (attention / Mamba / MoE interleave) and
|
||||
print parameter counts, KV cache, SSM state memory, and the ratio versus an
|
||||
equivalent pure-Transformer stack. Stdlib only. No tensors, no training.
|
||||
|
||||
Reference config: Jamba-v0.1 (Lieber et al. 2024) -- 52B total, 12B active,
|
||||
32 layers, 1 attention per 7 Mamba, MoE every 2 layers, 16 experts top-2.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
SHARED = {
|
||||
"vocab_size": 65536,
|
||||
"max_position_embeddings": 131072,
|
||||
"mamba_state_dim": 16,
|
||||
"mamba_conv_kernel": 4,
|
||||
"mamba_expand": 2,
|
||||
"attn_to_mamba_ratio": (1, 7),
|
||||
"moe_every": 2,
|
||||
"num_experts": 16,
|
||||
"experts_per_token": 2,
|
||||
}
|
||||
|
||||
|
||||
def make_config(**overrides) -> dict:
|
||||
return {**SHARED, **overrides}
|
||||
|
||||
|
||||
CONFIGS = {
|
||||
"jamba-v0.1": make_config(
|
||||
hidden_size=4096, intermediate_size=14336,
|
||||
num_layers=32, num_attention_heads=32, num_key_value_heads=8,
|
||||
),
|
||||
"jamba-1.5-mini": make_config(
|
||||
hidden_size=4096, intermediate_size=14336,
|
||||
num_layers=32, num_attention_heads=32, num_key_value_heads=8,
|
||||
),
|
||||
"jamba-1.5-large": make_config(
|
||||
hidden_size=8192, intermediate_size=28672,
|
||||
num_layers=72, num_attention_heads=64, num_key_value_heads=8,
|
||||
),
|
||||
"pure-mamba-7b": make_config(
|
||||
hidden_size=4096, intermediate_size=14336,
|
||||
num_layers=32, num_attention_heads=32, num_key_value_heads=8,
|
||||
attn_to_mamba_ratio=(0, 1), moe_every=0,
|
||||
num_experts=1, experts_per_token=1,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Breakdown:
|
||||
name: str
|
||||
schedule: list[tuple[str, str]]
|
||||
total_params: int
|
||||
active_params: int
|
||||
attn_layers: int
|
||||
mamba_layers: int
|
||||
moe_layers: int
|
||||
kv_cache_bytes: int
|
||||
ssm_state_bytes: int
|
||||
pure_transformer_kv_bytes: int
|
||||
|
||||
|
||||
def build_schedule(cfg: dict) -> list[tuple[str, str]]:
|
||||
a, m = cfg["attn_to_mamba_ratio"]
|
||||
period = a + m if (a + m) > 0 else 1
|
||||
moe_every = cfg["moe_every"]
|
||||
schedule: list[tuple[str, str]] = []
|
||||
for i in range(cfg["num_layers"]):
|
||||
slot = i % period
|
||||
attn_type = "attn" if slot < a else "mamba"
|
||||
if moe_every and ((i + 1) % moe_every == 0):
|
||||
mlp_type = "moe"
|
||||
else:
|
||||
mlp_type = "mlp"
|
||||
schedule.append((attn_type, mlp_type))
|
||||
return schedule
|
||||
|
||||
|
||||
def attention_params(cfg: dict) -> int:
|
||||
h = cfg["hidden_size"]
|
||||
q_heads = cfg["num_attention_heads"]
|
||||
kv_heads = cfg["num_key_value_heads"]
|
||||
head_dim = h // q_heads
|
||||
q = h * h
|
||||
kv = 2 * h * (kv_heads * head_dim)
|
||||
o = h * h
|
||||
return q + kv + o
|
||||
|
||||
|
||||
def mamba_params(cfg: dict) -> int:
|
||||
h = cfg["hidden_size"]
|
||||
expand = cfg["mamba_expand"]
|
||||
n = cfg["mamba_state_dim"]
|
||||
d_inner = expand * h
|
||||
k = cfg["mamba_conv_kernel"]
|
||||
in_proj = h * (2 * d_inner)
|
||||
conv = d_inner * k + d_inner
|
||||
x_proj = d_inner * (2 * n + 1)
|
||||
dt_proj = d_inner
|
||||
a_log = d_inner * n
|
||||
d_param = d_inner
|
||||
out_proj = d_inner * h
|
||||
return in_proj + conv + x_proj + dt_proj + a_log + d_param + out_proj
|
||||
|
||||
|
||||
def swiglu_params(h: int, ff: int) -> int:
|
||||
return 2 * h * ff + ff * h
|
||||
|
||||
|
||||
def rmsnorm_params(h: int) -> int:
|
||||
return h
|
||||
|
||||
|
||||
def analyze(name: str, cfg: dict) -> Breakdown:
|
||||
h = cfg["hidden_size"]
|
||||
ff = cfg["intermediate_size"]
|
||||
vocab = cfg["vocab_size"]
|
||||
n_layers = cfg["num_layers"]
|
||||
k_experts = cfg["experts_per_token"]
|
||||
num_experts = cfg["num_experts"]
|
||||
|
||||
schedule = build_schedule(cfg)
|
||||
attn_count = sum(1 for a, _ in schedule if a == "attn")
|
||||
mamba_count = sum(1 for a, _ in schedule if a == "mamba")
|
||||
moe_count = sum(1 for _, m in schedule if m == "moe")
|
||||
|
||||
emb = vocab * h
|
||||
attn_p = attention_params(cfg)
|
||||
mamba_p = mamba_params(cfg)
|
||||
mlp_p = swiglu_params(h, ff)
|
||||
norms_per_layer = 2 * rmsnorm_params(h)
|
||||
final_norm = rmsnorm_params(h)
|
||||
router_p = h * num_experts if num_experts > 1 else 0
|
||||
|
||||
total = emb + final_norm
|
||||
active = emb + final_norm
|
||||
for attn_type, mlp_type in schedule:
|
||||
mixer = attn_p if attn_type == "attn" else mamba_p
|
||||
if mlp_type == "moe":
|
||||
layer_total = mixer + mlp_p * num_experts + router_p + norms_per_layer
|
||||
layer_active = mixer + mlp_p * k_experts + router_p + norms_per_layer
|
||||
else:
|
||||
layer_total = mixer + mlp_p + norms_per_layer
|
||||
layer_active = layer_total
|
||||
total += layer_total
|
||||
active += layer_active
|
||||
|
||||
head_dim = h // cfg["num_attention_heads"]
|
||||
kv_heads = cfg["num_key_value_heads"]
|
||||
seq = cfg["max_position_embeddings"]
|
||||
kv_cache_bytes = 2 * attn_count * kv_heads * head_dim * seq * 2
|
||||
|
||||
d_inner = cfg["mamba_expand"] * h
|
||||
ssm_state_bytes = mamba_count * d_inner * cfg["mamba_state_dim"] * 2
|
||||
|
||||
pure_kv = 2 * n_layers * kv_heads * head_dim * seq * 2
|
||||
|
||||
return Breakdown(
|
||||
name=name,
|
||||
schedule=schedule,
|
||||
total_params=total,
|
||||
active_params=active,
|
||||
attn_layers=attn_count,
|
||||
mamba_layers=mamba_count,
|
||||
moe_layers=moe_count,
|
||||
kv_cache_bytes=kv_cache_bytes,
|
||||
ssm_state_bytes=ssm_state_bytes,
|
||||
pure_transformer_kv_bytes=pure_kv,
|
||||
)
|
||||
|
||||
|
||||
def fmt_params(x: int) -> str:
|
||||
if x >= 1_000_000_000:
|
||||
return f"{x / 1e9:.1f}B"
|
||||
if x >= 1_000_000:
|
||||
return f"{x / 1e6:.1f}M"
|
||||
return f"{x:,}"
|
||||
|
||||
|
||||
def fmt_bytes(b: float) -> str:
|
||||
for unit in ("B", "KB", "MB", "GB", "TB"):
|
||||
if b < 1024:
|
||||
return f"{b:.1f} {unit}"
|
||||
b /= 1024
|
||||
return f"{b:.1f} PB"
|
||||
|
||||
|
||||
def render_schedule(schedule: list[tuple[str, str]]) -> str:
|
||||
cells = []
|
||||
for attn_type, mlp_type in schedule:
|
||||
a = "A" if attn_type == "attn" else "M"
|
||||
m = "e" if mlp_type == "moe" else "d"
|
||||
cells.append(f"{a}{m}")
|
||||
return " ".join(cells)
|
||||
|
||||
|
||||
def print_breakdown(b: Breakdown, cfg: dict) -> None:
|
||||
a, m = cfg["attn_to_mamba_ratio"]
|
||||
pure = b.pure_transformer_kv_bytes
|
||||
ratio = f"{pure / b.kv_cache_bytes:.1f}x" if b.kv_cache_bytes else "inf"
|
||||
active_frac = b.active_params / max(b.total_params, 1)
|
||||
print(f"\n{b.name}")
|
||||
print("-" * 72)
|
||||
print(f" layers : {cfg['num_layers']} "
|
||||
f"(attn={b.attn_layers}, mamba={b.mamba_layers}, moe={b.moe_layers})")
|
||||
print(f" interleave : {a}:{m} "
|
||||
f"MoE every {cfg['moe_every']}, {cfg['num_experts']} experts "
|
||||
f"top-{cfg['experts_per_token']}")
|
||||
print(f" total / active: {fmt_params(b.total_params)} / "
|
||||
f"{fmt_params(b.active_params)} ({active_frac:.1%})")
|
||||
print(f" context : {cfg['max_position_embeddings']:,} tokens")
|
||||
print(f" KV cache : {fmt_bytes(b.kv_cache_bytes)} "
|
||||
f"(vs {fmt_bytes(pure)} pure transformer -> {ratio} reduction)")
|
||||
print(f" SSM state : {fmt_bytes(b.ssm_state_bytes)} "
|
||||
f"(constant vs seq_len)")
|
||||
print(f" schedule (Ad=attn+dense, Me=mamba+moe, Md=mamba+dense):")
|
||||
sched = render_schedule(b.schedule)
|
||||
for i in range(0, len(sched), 48):
|
||||
print(f" {sched[i:i + 48]}")
|
||||
|
||||
|
||||
def print_summary(results: list[Breakdown]) -> None:
|
||||
print()
|
||||
print("=" * 72)
|
||||
print("HEADLINE: SSM vs TRANSFORMER AT 128K")
|
||||
print("=" * 72)
|
||||
print(f" {'model':18s} {'total':>8s} {'active':>8s} "
|
||||
f"{'KV@ctx':>10s} {'pure KV':>10s} {'ratio':>6s}")
|
||||
for b in results:
|
||||
ratio = (f"{b.pure_transformer_kv_bytes / b.kv_cache_bytes:5.1f}x"
|
||||
if b.kv_cache_bytes else " inf")
|
||||
print(f" {b.name:18s} "
|
||||
f"{fmt_params(b.total_params):>8s} "
|
||||
f"{fmt_params(b.active_params):>8s} "
|
||||
f"{fmt_bytes(b.kv_cache_bytes):>10s} "
|
||||
f"{fmt_bytes(b.pure_transformer_kv_bytes):>10s} "
|
||||
f"{ratio:>6s}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
print("=" * 72)
|
||||
print("JAMBA HYBRID SSM-TRANSFORMER CALCULATOR")
|
||||
print("=" * 72)
|
||||
results = []
|
||||
for name, cfg in CONFIGS.items():
|
||||
b = analyze(name, cfg)
|
||||
print_breakdown(b, cfg)
|
||||
results.append(b)
|
||||
print_summary(results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,232 @@
|
||||
# Jamba: Hybrid SSM-Transformer
|
||||
|
||||
> At 128k context a pure Transformer spends more memory on its KV cache than on its weights. State-space models have O(1) cache but lose at in-context recall. Jamba interleaves one attention layer for every seven Mamba layers and adds MoE on top — recurrence for throughput, attention for recall, experts for capacity. It is the first hybrid that actually matches dense Transformers at scale.
|
||||
|
||||
**Type:** Learn
|
||||
**Languages:** Python (stdlib)
|
||||
**Prerequisites:** Phase 10, Lessons 04, 12, 14 (Pre-training, Inference optimization, Open model architectures)
|
||||
**Time:** ~45 minutes
|
||||
|
||||
## Learning Objectives
|
||||
|
||||
- Explain why the Transformer KV cache dominates long-context memory and how an SSM removes that cost
|
||||
- Describe Mamba's selective state space update and why it is content-aware, unlike S4
|
||||
- Reconstruct Jamba's 1:7 attention-to-Mamba interleave with MoE every other layer and justify each ratio
|
||||
- Compute parameter and memory footprints for Jamba at 128k context and compare against a dense Transformer of equivalent active params
|
||||
- Name the exact failure modes of pure SSM models and why the attention layers are there
|
||||
|
||||
## The Problem
|
||||
|
||||
You trained a GPT in Lesson 04 and read five different Transformer dialects in Lesson 14. Every one of them pays the same tax: the KV cache. For each token generated, every layer stores a key vector and a value vector for every previous token. That cache grows linearly with sequence length and layer count, so 128k-context serving is KV cache serving. For Llama 3 8B at 128k, the KV cache is roughly 17 GB — larger than the 16 GB of weights themselves.
|
||||
|
||||
Recurrent networks do not have this problem. A classical RNN has O(1) state regardless of sequence length — one hidden vector that summarizes everything so far. The catch, until 2021, was that RNNs were impossibly slow to train (sequential, no parallelism) and structurally bad at in-context recall (information is crushed into a fixed-size state).
|
||||
|
||||
The state-space family fixed the training problem. S4 (Gu, Goel, Re 2021) showed you could parameterize a linear SSM so it has an equivalent convolutional form, which trains in parallel, while still running as a recurrence at inference. Mamba (Gu and Dao 2023) added input-dependent selectivity so the model could decide to remember or forget each token, closing most of the language-quality gap with Transformers while keeping linear time and constant cache.
|
||||
|
||||
Mamba alone still loses on tasks that need precise retrieval from far back in the context — phone-book lookups, in-context learning from many examples, long-range copy. Those tasks want exact content-addressed recall, which attention does natively and SSMs do not. Jamba (Lieber et al. 2024) is the architectural answer: keep Mamba for the bulk of the layers so the KV cache stays tiny, sprinkle in one attention layer every seven layers for the recall, and add Mixture of Experts to pack in capacity without paying for it per token.
|
||||
|
||||
## The Concept
|
||||
|
||||
### Why the KV cache hurts
|
||||
|
||||
For a Transformer decoder, generating token T+1 needs the keys and values of tokens 1..T at every layer. The cache size per sequence is:
|
||||
|
||||
```
|
||||
kv_cache_bytes = 2 * num_layers * num_kv_heads * head_dim * seq_len * bytes
|
||||
```
|
||||
|
||||
Every factor except `bytes` is fixed by the model. Only `seq_len` is yours. Doubling context doubles the cache. GQA and MLA (Lesson 14) shrink `num_kv_heads` — they do not change the linear growth.
|
||||
|
||||
An SSM layer, by contrast, carries one fixed-size hidden state `h` of shape (state_dim,) regardless of how long the sequence is. If `state_dim` is 16, every SSM layer stores 16 numbers total, not 16 * seq_len. That is the "O(1) cache" advantage.
|
||||
|
||||
### Linear state-space models (S4)
|
||||
|
||||
A linear SSM maps an input sequence `u` to an output sequence `y` via a hidden state `h`:
|
||||
|
||||
```
|
||||
h'(t) = A h(t) + B u(t)
|
||||
y(t) = C h(t) + D u(t)
|
||||
```
|
||||
|
||||
Discretize with step size Δ and it becomes the classic RNN update `h_t = A_bar h_{t-1} + B_bar u_t`. Three properties make this useful:
|
||||
|
||||
1. **Recurrent form.** Inference is one matmul per step. Constant cache. Linear in sequence length.
|
||||
2. **Convolutional form.** The same computation can be expressed as a global convolution `y = K * u` where `K` is a kernel derived from (A, B, C). Trains in parallel across the whole sequence.
|
||||
3. **Structured A.** S4 uses a HiPPO-based parameterization that gives the continuous system well-behaved long-range memory. Mamba uses a simpler diagonal form but keeps the structure.
|
||||
|
||||
S4 ships these two forms as a duality: train as a convolution, infer as a recurrence.
|
||||
|
||||
### Mamba's selection mechanism
|
||||
|
||||
S4's parameters (A, B, C, Δ) are fixed per channel. Every token hits the same filter. That is why pure S4 struggles on text: it cannot choose which tokens to pay attention to.
|
||||
|
||||
Mamba makes (B, C, Δ) input-dependent. For each token `u_t`, the model computes:
|
||||
|
||||
```
|
||||
B_t = Linear_B(u_t)
|
||||
C_t = Linear_C(u_t)
|
||||
Δ_t = softplus(Linear_Δ(u_t))
|
||||
```
|
||||
|
||||
The SSM's input matrix, output matrix, and step size are now functions of the current token. If `Δ_t` is large, more of `u_t` is written into `h_t`. If `Δ_t` is small, the state passes through. This is the RNN analog of attention: the model can selectively propagate or ignore information.
|
||||
|
||||
The price: the convolutional duality breaks. With time-varying (B, C, Δ) there is no fixed kernel. Mamba replaces the parallel conv with a hardware-aware parallel scan — a custom CUDA kernel that computes the recurrence in parallel using the associativity of the state update. Linear time, parallel training, selective recurrence.
|
||||
|
||||
The Mamba block in its entirety:
|
||||
|
||||
```
|
||||
x -> proj -> [conv1d -> SiLU -> selective_ssm] -> gate * proj -> y
|
||||
```
|
||||
|
||||
One projection in, a 1D causal conv (cheap local mixing), a SiLU, the selective SSM, and a gated output projection. No attention anywhere.
|
||||
|
||||
### What Mamba loses
|
||||
|
||||
Benchmarks where Mamba still underperforms Transformers of equivalent size:
|
||||
|
||||
- **Exact recall from far back.** "The password is MAGENTA. ...10,000 tokens of filler... What was the password?" Transformers nail this with attention. Mamba's finite-dim state has to compress the password along with everything else.
|
||||
- **Induction heads and in-context learning.** Mechanistic interpretability work shows Transformers form induction heads that do pattern copying. SSMs can approximate these but need more parameters to do so.
|
||||
- **Many-shot in-context examples.** The whole point of 128k context is stuffing examples into the prompt. SSMs summarize examples into state; attention indexes them directly.
|
||||
|
||||
Those are exactly the workloads that benefit most from big context. So the answer is not "replace attention with Mamba," it is "keep enough attention to cover the recall tasks, use Mamba for everything else."
|
||||
|
||||
### The Jamba block
|
||||
|
||||
Jamba's core unit is a block of `l` layers. Each layer is either an attention layer or a Mamba layer, each followed by an MLP. The layer types interleave at a ratio `a : m` (attention : Mamba). MLPs every `e` layers can be swapped for Mixture of Experts.
|
||||
|
||||
The released Jamba-v0.1 configuration:
|
||||
|
||||
- `l = 8` layers per block
|
||||
- `a : m = 1 : 7` — one attention layer per seven Mamba layers
|
||||
- `e = 2` — MoE replaces the MLP on every other layer
|
||||
- 16 total experts, top-2 active per token
|
||||
- Block repeats 4 times (32 layers total)
|
||||
|
||||
So out of 32 layers: 4 are attention, 28 are Mamba, and 16 use MoE MLPs. Total parameters: 52B. Active per token: 12B. Fits on a single 80GB GPU with 256k context.
|
||||
|
||||
### Why 1:7
|
||||
|
||||
The Jamba ablations swept attention ratios from pure-Mamba (0:8) up to 1:1. The findings:
|
||||
|
||||
- Pure Mamba loses perplexity to a matched Transformer by a few tenths of nats. It specifically fails needle-in-a-haystack.
|
||||
- Adding one attention layer per block (1:7) closes the gap and passes needle-in-a-haystack at 256k.
|
||||
- Going denser (1:3, 1:1) gives little additional quality while reimposing the KV cache cost.
|
||||
|
||||
The 1:7 ratio is the Pareto point. It is the minimum attention that keeps the recall tasks honest.
|
||||
|
||||
### Why no positional embeddings
|
||||
|
||||
Attention is permutation-equivariant on its own, which is why Transformers need RoPE or learned positions. Mamba is a recurrence — the order is baked into the scan. In a hybrid stack, the Mamba layers carry the positional signal into the attention layers. Jamba confirms this empirically and ships with no explicit positional encoding at all.
|
||||
|
||||
### KV cache at 128k
|
||||
|
||||
With 4 attention layers out of 32 at 128k context, 8 KV heads, head_dim 128, BF16:
|
||||
|
||||
```
|
||||
kv_cache = 2 * 4 * 8 * 128 * 131072 * 2 = 2.15 GB
|
||||
```
|
||||
|
||||
Compare to a pure Transformer with the same dims but all 32 attention layers:
|
||||
|
||||
```
|
||||
kv_cache = 2 * 32 * 8 * 128 * 131072 * 2 = 17.2 GB
|
||||
```
|
||||
|
||||
Eight times smaller. That is the whole product thesis.
|
||||
|
||||
### Jamba-1.5
|
||||
|
||||
Jamba-1.5-Mini is the post-trained version of the original (12B active, 52B total). Jamba-1.5-Large scales the same architecture to 94B active, 398B total, and ships with ExpertsInt8 — an INT8 quantization scheme for the expert weights that fits the whole model on 8x80GB GPUs at 256k context. Same 1:7 interleave. Same MoE every other layer. Different scale.
|
||||
|
||||
### Mamba-3 (2026)
|
||||
|
||||
Mamba-3 (ICLR 2026 Oral) keeps the hybrid-ready SSM primitive but improves three things inside the SSM block itself: a trapezoidal discretization for a second-order state update, a complex-valued recurrence equivalent to data-dependent RoPE (which restores the state-tracking capability Mamba-2 had regressed on), and a MIMO formulation that increases arithmetic intensity during decode. At 1.5B scale Mamba-3 (MIMO) matches Mamba-2 with half the state size. The Jamba-shaped assembly is unchanged.
|
||||
|
||||
### Building block comparison
|
||||
|
||||
| Block | KV cache | Recall | FLOPs/token | Notes |
|
||||
|-------|----------|--------|-------------|-------|
|
||||
| Transformer (MHA) | O(L * H * d) per layer | excellent | O(L * d) | quadratic attention, full cache |
|
||||
| Transformer (GQA) | O(L * H/G * d) per layer | excellent | O(L * d) | Llama 3 default |
|
||||
| Mamba | O(state_dim) per layer | lossy on long recall | O(d * state_dim) | parallel scan, no attention |
|
||||
| Jamba block | KV only on attention layers | excellent | dominated by Mamba | 1 attention per 7 Mamba |
|
||||
|
||||
That table is the whole lesson in five rows.
|
||||
|
||||
## Build It
|
||||
|
||||
The code for this lesson simulates a Jamba-shaped stack as pure Python. It does not train anything — the point is to compute the cache and parameter budget that justify the architecture. See `code/main.py`.
|
||||
|
||||
### Step 1: Represent the interleave pattern
|
||||
|
||||
Given `a : m`, a block of `l` layers with MoE every `e` layers, `main.py` produces a layer-by-layer schedule:
|
||||
|
||||
```
|
||||
layers = ["mamba", "mamba", "mamba", "attn", "mamba", "mamba", "mamba", "mamba"]
|
||||
mlp = ["moe", "mlp", "moe", "mlp", "moe", "mlp", "moe", "mlp"]
|
||||
```
|
||||
|
||||
Repeat the block until the target layer count is reached. This mirrors the Jamba-v0.1 config precisely.
|
||||
|
||||
### Step 2: Count parameters
|
||||
|
||||
Attention layers are standard GQA blocks (Q/K/V/output projections and one MLP per layer). Mamba layers follow Gu and Dao 2023: an input projection, a 1D causal conv, the SSM projections for (B, C, Δ), and an output gate. MoE MLPs multiply the MLP parameter count by `num_experts` for total, but charge only `top_k` experts for active params.
|
||||
|
||||
### Step 3: Compute the KV cache at 128k
|
||||
|
||||
Only the attention layers contribute. The Mamba state is counted separately as `num_layers * state_dim * hidden_size * bytes`, typically under 100 MB for a 50B-parameter model.
|
||||
|
||||
### Step 4: Compare to a pure Transformer
|
||||
|
||||
Same hidden dim, same depth, but every layer is attention. The calculator prints side-by-side KV cache numbers and the ratio. For the Jamba-v0.1 config at 128k, the number is an 8x reduction. At 256k it is larger.
|
||||
|
||||
### Step 5: Active vs total params
|
||||
|
||||
Sum all experts into total params. Sum `top_k` experts per MoE layer plus dense MLPs plus all attention/Mamba params into active params. Print the ratio — you should see the Jamba-v0.1 numbers (52B total, 12B active, ~23% active).
|
||||
|
||||
See `code/main.py` for the implementation. Running it on the bundled configs reproduces the published numbers within rounding.
|
||||
|
||||
## Use It
|
||||
|
||||
Run `python code/main.py` to print the layer schedule, the parameter breakdown, and the KV cache for Jamba-v0.1 alongside an equivalent dense Transformer. Change `attn_to_mamba_ratio` to see what happens at 1:3 or 2:6. Drop `num_experts` to 1 to see the non-MoE variant.
|
||||
|
||||
The implementation is intentionally schematic — no tensor ops, no training. Once you believe the budget, you would reach for the official AI21 implementation to actually run weights. The point of this lesson is that the architecture is specified by a handful of ratios and the memory budget falls out of them.
|
||||
|
||||
## Ship It
|
||||
|
||||
This lesson produces `outputs/skill-jamba.md`, a skill that takes a deployment target (context length, GPU VRAM, latency budget) and decides whether a hybrid SSM-Transformer is the right answer over a pure dense Transformer. It checks the KV cache at the target context, the expected decode throughput from the SSM majority, and the recall-task requirements of the workload.
|
||||
|
||||
## Exercises
|
||||
|
||||
1. Compute the Jamba-v0.1 KV cache at 32k, 128k, and 1M context. Then compute the Mamba state size at the same three numbers. Plot both in your head. At which context length does the attention KV cache overtake the Mamba state for the first time?
|
||||
|
||||
2. Change the interleave from 1:7 to 1:3 at 32 layers. How many attention layers do you now have? Recompute the 128k KV cache. How much of the Jamba advantage remains?
|
||||
|
||||
3. The original Jamba paper reported that a pure Mamba model of matched size failed a needle-in-a-haystack test at 256k. Design the minimum experiment (prompt format, what to vary, what to measure) that would convince you the attention layers were the fix, not something else in the stack.
|
||||
|
||||
4. Jamba-1.5-Large has 94B active and 398B total parameters. From the 1:7 interleave and MoE-every-other-layer structure, back-calculate a plausible layer count, hidden dim, and number of experts. You should land within 20% of the published spec.
|
||||
|
||||
5. Mamba-3 introduces a complex-valued recurrence equivalent to data-dependent RoPE. In a Jamba hybrid, the Mamba layers already carry positional information implicitly. Does swapping Mamba-1 for Mamba-3 change whether you need explicit RoPE on the attention layers? State your prediction and the experiment that would falsify it.
|
||||
|
||||
## Key Terms
|
||||
|
||||
| Term | What people say | What it actually means |
|
||||
|------|-----------------|-----------------------|
|
||||
| SSM | "The new RNN" | State-space model: linear recurrence h' = Ah + Bu, y = Ch + Du, with a structured A matrix |
|
||||
| S4 | "Structured SSM" | Gu/Goel/Re 2021: HiPPO-based A, runs as either a convolution (train) or a recurrence (infer) |
|
||||
| Selective SSM | "Content-aware Mamba" | (B, C, Δ) become functions of the current input, so the SSM can choose what to remember |
|
||||
| Mamba block | "SSM layer" | proj -> conv1d -> SiLU -> selective_ssm -> gated_output — one full Mamba layer |
|
||||
| Parallel scan | "How Mamba trains fast" | Hardware-aware custom kernel that computes a time-varying recurrence in parallel using associativity |
|
||||
| Jamba block | "The hybrid unit" | l layers mixing attention and Mamba at ratio a:m, with MLP or MoE every e layers |
|
||||
| 1:7 interleave | "One attention per seven Mamba" | Jamba's canonical ratio: minimum attention that preserves long-range recall |
|
||||
| MoE in Jamba | "Sparse experts every other layer" | 16 experts, top-2, applied on every second MLP — 52B total params, 12B active |
|
||||
| ExpertsInt8 | "Jamba-1.5 quantization" | INT8 storage for expert weights, dequantized to BF16 at compute, fits Jamba-1.5-Large on 8x80GB |
|
||||
| Needle-in-a-haystack | "Long-context recall benchmark" | Hide a sentence inside a long document, ask the model to retrieve it verbatim — the test pure Mamba fails |
|
||||
|
||||
## Further Reading
|
||||
|
||||
- [Lieber et al., 2024 -- "Jamba: A Hybrid Transformer-Mamba Language Model"](https://arxiv.org/abs/2403.19887) -- the original architecture paper, ablations on interleave ratio and MoE placement
|
||||
- [Jamba Team, 2024 -- "Jamba-1.5: Hybrid Transformer-Mamba Models at Scale"](https://arxiv.org/abs/2408.12570) -- 94B active / 398B total, ExpertsInt8 quantization, 256k context at serving
|
||||
- [Gu and Dao, 2023 -- "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"](https://arxiv.org/abs/2312.00752) -- selective SSMs, parallel scan, 5x inference throughput over Transformers
|
||||
- [Gu, Goel, Re, 2021 -- "Efficiently Modeling Long Sequences with Structured State Spaces"](https://arxiv.org/abs/2111.00396) -- S4, the structured SSM that made state-space models competitive
|
||||
- [Lahoti et al., 2026 -- "Mamba-3: Improved Sequence Modeling using State Space Principles"](https://openreview.net/forum?id=HwCvaJOiCj) -- ICLR 2026 Oral: trapezoidal discretization, complex-valued recurrence, MIMO decode
|
||||
@@ -0,0 +1,38 @@
|
||||
---
|
||||
name: jamba-hybrid-picker
|
||||
description: Decide whether a hybrid SSM-Transformer (Jamba family) beats a dense Transformer for a given deployment target, with explicit KV-cache and recall-task reasoning.
|
||||
version: 1.0.0
|
||||
phase: 10
|
||||
lesson: 21
|
||||
tags: [jamba, mamba, ssm, state-space, hybrid, moe, long-context, kv-cache]
|
||||
---
|
||||
|
||||
Given a deployment target (GPU type, VRAM per GPU, number of GPUs, target context length, target p50/p99 latency, peak concurrent requests) and a task profile (chat, code, reasoning, long-context RAG, many-shot in-context learning, precise recall), decide whether the right architecture is a hybrid SSM-Transformer (Jamba-1.5-Mini, Jamba-1.5-Large) or a dense Transformer (Llama 3, Qwen 2.5, DeepSeek V3) from Lesson 14. Justify the call against the SSM vs attention tradeoff taught in Lesson 21.
|
||||
|
||||
Produce:
|
||||
|
||||
1. **Architecture verdict.** Hybrid (Jamba-family) or dense (Lesson-14 family). State the single dominant reason — KV cache at target context, recall-task sensitivity, MoE capacity, or serving throughput.
|
||||
|
||||
2. **KV cache budget at target context.** Compute the KV cache for the top candidate using the formula from `code/main.py`. For Jamba, only the attention layers count; the SSM state is constant vs sequence length. Compare against the dense alternative side-by-side.
|
||||
|
||||
3. **Recall-task check.** Score the task profile against Mamba's known failure modes: exact verbatim recall from far back, induction-head-style pattern copying, many-shot in-context examples. If any scores high, confirm that the Jamba candidate has at least 1 attention per 8 layers (Jamba-1.5 does; pure Mamba does not). Reject pure SSM for recall-heavy workloads.
|
||||
|
||||
4. **Throughput sanity check.** Jamba's headline advantage is decode throughput from the Mamba majority. Estimate tokens/sec from the active-param count (11.7B or 94B) and the GPU memory bandwidth. Compare to the same-active-param dense Transformer.
|
||||
|
||||
5. **Quantization choice.** For Jamba-1.5-Large, default to ExpertsInt8 (INT8 storage on expert weights, BF16 compute). For Jamba-1.5-Mini on a single 80GB GPU, BF16 is fine. For dense alternatives, reuse the Lesson-11 matrix (GPTQ-4bit, AWQ-4bit, FP8, BF16).
|
||||
|
||||
6. **Fallback.** Name a second choice. If the recommendation is Jamba-1.5-Mini and the workload is recall-dominant with no long-context need, fall back to Llama 3 8B dense. If the recommendation is Jamba-1.5-Large and expert-parallel support is missing in the serving stack, fall back to Mixtral 8x22B (standard MoE tooling).
|
||||
|
||||
Hard rejects:
|
||||
- Pure Mamba or pure SSM for workloads that require exact retrieval from >32k context.
|
||||
- Jamba on a serving stack that does not implement the selective-scan kernel (throughput collapses to Python loop).
|
||||
- Dense Transformer above 70B at 128k+ context on a single 80GB GPU (KV cache alone exceeds VRAM).
|
||||
- MoE models on a stack without expert-parallel support.
|
||||
- Any recommendation that does not name the specific revision (e.g., "Jamba-1.5-Mini 2024-08", not "Jamba").
|
||||
|
||||
Worth reconsidering if:
|
||||
- The workload does not need long context (below 16k). A dense 7B-13B Transformer ties or beats Jamba on throughput there with simpler tooling.
|
||||
- The task is purely code or math with strong per-step reasoning. Dense reasoning-tuned models (Qwen, DeepSeek) lead on these benchmarks as of 2026.
|
||||
- The serving stack lacks a fused selective-scan kernel. Jamba's advantage depends on it; without it, use a dense alternative.
|
||||
|
||||
Output: a one-page verdict naming the model revision, serving stack, quantization, and context budget, with numbered evidence for each decision. End with a "change my mind if" paragraph that names the specific workload parameter that would flip the call.
|
||||
Reference in New Issue
Block a user