Beyond attention with Mamba, state space models and hybrids

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import json
import math
import time

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoTokenizer

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

t = np.linspace(0, 10, 1001)
splash = (np.abs(t - 1) < 0.05) * 10.0  # a short burst whose total is 1
stream = ((t > 4) & (t < 6)) * 1.0
dt = t[1] - t[0]

fig, ax = plt.subplots(figsize=(7.4, 3.4))
colors = [COLORS["enc"], COLORS["att"], COLORS["data"], COLORS["dec"]]
for a, color in zip([-0.125, -0.25, -0.5, -1.0], colors):  # four holes, small to large
    level = np.zeros_like(t)
    for i in range(1, len(t)):  # exact step of h' = a h + x for constant x over dt
        x = splash[i - 1] + stream[i - 1]
        level[i] = level[i - 1] * math.exp(a * dt) + (math.exp(a * dt) - 1) / a * x
    ax.plot(t, level, color=color, linewidth=2, label=f"a = {a:g}")
ax.fill_between(t, 0, stream * 0.15, color=COLORS["line"], alpha=0.4, label="input")
ax.set_xlabel("time")
ax.set_ylabel("state")
ax.set_title("Four leaky buckets: the same input, four memories")
ax.legend(frameon=False, fontsize=9)
plt.show()

for product in [0.001, 0.01, 0.1, 0.5, 1.0, 2.0, 5.0]:
    exact_over_euler = (1 - math.exp(-product)) / product
    print(f"|Δa| = {product:<5}  exact b̄ / Euler b̄ = {exact_over_euler:.3f}")

REPO = "state-spaces/mamba-130m-hf"
mamba_config = json.load(open(hf_hub_download(REPO, "config.json")))
for key in [
    "hidden_size",
    "intermediate_size",
    "state_size",
    "time_step_rank",
    "conv_kernel",
    "num_hidden_layers",
    "vocab_size",
]:
    print(f"{key:18} {mamba_config[key]}")

def scan_loop(decay, write, c):
    # decay, write: (batch, channels, L, N); c: (batch, L, N) -> y: (batch, channels, L)
    h = torch.zeros(decay.shape[0], decay.shape[1], decay.shape[3], device=decay.device)
    ys = []
    for t in range(decay.shape[2]):
        h = decay[:, :, t] * h + write[:, :, t]
        ys.append((h * c[:, None, t]).sum(-1))
    return torch.stack(ys, dim=-1)


class RMSNorm(nn.Module):
    def __init__(self, size, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.eps = eps

    def forward(self, x):
        x32 = x.float()
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x32.to(x.dtype)

class Mixer(nn.Module):
    def __init__(self, c):
        super().__init__()
        d, e = c["hidden_size"], c["intermediate_size"]
        self.n, self.rank = c["state_size"], c["time_step_rank"]
        k = c["conv_kernel"]
        self.in_proj = nn.Linear(d, 2 * e, bias=c["use_bias"])
        self.conv1d = nn.Conv1d(e, e, k, groups=e, padding=k - 1, bias=c["use_conv_bias"])
        self.x_proj = nn.Linear(e, self.rank + 2 * self.n, bias=False)
        self.dt_proj = nn.Linear(self.rank, e, bias=True)
        self.A_log = nn.Parameter(torch.zeros(e, self.n))
        self.D = nn.Parameter(torch.ones(e))
        self.out_proj = nn.Linear(e, d, bias=c["use_bias"])
        self.scan = scan_loop
        self.keep = False  # section 5 reads the steps
        self.steps = None

    def forward(self, x):
        L = x.shape[1]
        u, gate = self.in_proj(x).transpose(1, 2).chunk(2, dim=1)  # (batch, E, L) each
        u = F.silu(self.conv1d(u)[..., :L])
        seed, b, c = self.x_proj(u.transpose(1, 2)).split(
            [self.rank, self.n, self.n], dim=-1
        )
        step = F.softplus(self.dt_proj(seed)).transpose(1, 2)  # (batch, E, L)
        if self.keep:
            self.steps = step.detach()
        a = -torch.exp(self.A_log.float())  # (E, N), every entry negative
        decay = torch.exp(step[..., None] * a[:, None, :])  # (batch, E, L, N)
        write = step[..., None] * b[:, None].float() * u[..., None].float()
        y = self.scan(decay, write, c.float()).to(u.dtype)
        y = y + u * self.D[None, :, None]
        y = y * F.silu(gate)
        return self.out_proj(y.transpose(1, 2))

class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.norm = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
        self.mixer = Mixer(c)

    def forward(self, x):
        return x + self.mixer(self.norm(x))


class Mamba(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.embeddings = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList([Block(c) for _ in range(c["num_hidden_layers"])])
        self.norm_f = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)

    def forward(self, ids):
        x = self.embeddings(ids)
        for layer in self.layers:
            x = layer(x)
        return self.lm_head(self.norm_f(x)).float()

block = Block(mamba_config)
rows = []
for name, p in block.mixer.named_parameters():
    rows.append(
        {
            "part": name,
            "shape": " × ".join(str(s) for s in p.shape),
            "parameters": p.numel(),
        }
    )
parts = pd.DataFrame(rows)
parts["share"] = parts.parameters / parts.parameters.sum()
table(parts, digits=3)

from transformers import MambaConfig, MambaForCausalLM

if SMOKE:
    small = MambaConfig(
        vocab_size=mamba_config["vocab_size"],
        hidden_size=64,
        intermediate_size=128,
        state_size=8,
        time_step_rank=4,
        num_hidden_layers=2,
    )
    torch.manual_seed(0)
    reference = MambaForCausalLM(small).eval()
    mamba_config = small.to_dict()
    state = {k.removeprefix("backbone."): v for k, v in reference.state_dict().items()}
else:
    reference = MambaForCausalLM.from_pretrained(REPO, dtype=torch.float32).eval()
    state = load_file(hf_hub_download(REPO, "model.safetensors"))
    state = {k.removeprefix("backbone."): v for k, v in state.items()}
state.setdefault("lm_head.weight", state["embeddings.weight"])  # tied to the embeddings
mamba = Mamba(mamba_config)
mamba.load_state_dict(state, strict=True)
mamba.eval()

tok = AutoTokenizer.from_pretrained(REPO)


def tokens_of(ids):
    return [tok.decode([i]) for i in ids[0]]


sentence = (
    "In December 1911 Roald Amundsen reached the South Pole with four men and fifty-two "
    "dogs; Robert Falcon Scott arrived thirty-four days later and found a Norwegian tent."
)
ids = tok(sentence, return_tensors="pt").input_ids
with torch.no_grad():
    theirs = reference(ids).logits
    ours = mamba(ids)
print("tokens:", ids.shape[1])
print("largest logit:", f"{theirs.abs().max().item():.1f}")
print("max difference from Hugging Face:", f"{(ours - theirs).abs().max().item():.1e}")

m = mamba.layers[0].mixer
with torch.no_grad():
    x = mamba.layers[0].norm(mamba.embeddings(ids))
    print("tokens in, normalised:          ", tuple(x.shape), "  (batch, L, D)")
    both = m.in_proj(x).transpose(1, 2)
    print("in_proj, channels first:        ", tuple(both.shape), "  (batch, 2E, L)")
    u, gate = both.chunk(2, dim=1)
    print("stream u and gate z:            ", tuple(u.shape), "each")
    L = ids.shape[1]
    u = F.silu(m.conv1d(u)[..., :L])
    print("after the convolution and SiLU: ", tuple(u.shape))
    seed, b, c = m.x_proj(u.transpose(1, 2)).split([m.rank, m.n, m.n], dim=-1)
    print(
        "x_proj → step seed, B, C:       ",
        tuple(seed.shape),
        tuple(b.shape),
        tuple(c.shape),
    )
    step = F.softplus(m.dt_proj(seed)).transpose(1, 2)
    print(
        "dt_proj, softplus → Δ:          ",
        tuple(step.shape),
        "  one step per channel and token",
    )
    a = -torch.exp(m.A_log.float())
    print("A, all negative:                ", tuple(a.shape))
    decay = torch.exp(step[..., None] * a[:, None, :])
    write = step[..., None] * b[:, None] * u[..., None]
    print("decay e^(Δa) and write ΔBu:     ", tuple(decay.shape), "  (batch, E, L, N)")
    y = scan_loop(decay, write, c)
    print("the scan, read out by C:        ", tuple(y.shape))
    y = (y + u * m.D[None, :, None]) * F.silu(gate)
    out = m.out_proj(y.transpose(1, 2))
    print(
        "skip D, gate, out_proj:         ",
        tuple(out.shape),
        "  back to the width of the model",
    )
    print("same as the block's own forward:", f"{(out - m(x)).abs().max():.1e}")

print()
word = 4
print(f"channel 0 at token {word} ({tokens_of(ids)[word]!r}):")
print("  step Δ:", f"{step[0, 0, word]:.4f}", "  share of the state kept:", end=" ")
print(" ".join(f"{v:.2f}" for v in decay[0, 0, word, :6].tolist()), "…")

for layer in mamba.layers:
    layer.mixer.keep = True
with torch.no_grad():
    mamba(ids)
tokens = tokens_of(ids)
replaced = []
for layer in mamba.layers:
    m = layer.mixer
    a = -torch.exp(m.A_log.detach().float())  # (E, N)
    kept = torch.exp(m.steps[0][:, :, None] * a[:, None, :])  # (E, L, N)
    replaced.append((1 - kept).mean(dim=(0, 2)))
    m.keep = False
replaced = torch.stack(replaced)  # (layers, L)
relative = (
    replaced / replaced.mean(dim=1, keepdim=True)
).numpy()  # 1 = the layer's average

show = list(range(0, len(mamba.layers), max(1, len(mamba.layers) // 6)))
fig, ax = plt.subplots(figsize=(9.6, 3.0))
shades = LinearSegmentedColormap.from_list("shades", [COLORS["card"], COLORS["att"]])
ax.imshow(relative[show], cmap=shades, aspect="auto", vmin=0.4, vmax=1.6)
ax.set_xticks(
    range(len(tokens)), [s.strip() or "␣" for s in tokens], rotation=70, fontsize=8
)
ax.set_yticks(range(len(show)), [f"layer {i}" for i in show])
ax.set_title("How much of the state each token replaces, against the layer's average")
plt.show()

middle = relative[len(relative) // 12 : len(relative) // 2 + 2].mean(axis=0)
order = np.argsort(-middle)
print("replace most: ", " · ".join(tokens[i].strip() for i in order[:8]))
print("replace least:", " · ".join(tokens[i].strip() for i in order[-8:]))

page = tok(
    "The expedition left the Bay of Whales on 20 October with four sledges. The weather "
    "was clear, the dogs were strong, and by the end of the first week they had covered "
    "more ground than Scott's party would in a month. Amundsen had planned the depots "
    "with care: every one was marked with flags across the line of march, so that a "
    "party returning in poor light could not miss it.",
    return_tensors="pt",
).input_ids
for layer in mamba.layers:
    layer.mixer.keep = True
with torch.no_grad():
    mamba(page)
horizons = []
for layer in mamba.layers:
    m = layer.mixer
    a = torch.exp(m.A_log.detach().float())  # |a|, (E, N)
    typical_step = m.steps[0].median(dim=-1).values  # (E,)
    horizons.append((1 / (typical_step[:, None] * a)).flatten())
    m.keep = False
horizons = torch.cat(horizons).clamp(max=1e6).numpy()

fig, ax = plt.subplots(figsize=(7.2, 3.2))
bins = np.logspace(-2, 6, 50)
ax.hist(horizons, bins=bins, color=COLORS["att"])
ax.set_xscale("log")
ax.set_xlabel("memory horizon, tokens")
ax.set_ylabel("channel × state pairs")
ax.set_title("How long the channels of mamba-130m remember")
plt.show()
for q in [0.1, 0.5, 0.9, 0.99]:
    print(f"{q:.0%} of horizons are under {np.quantile(horizons, q):,.2g} tokens")

def scan_parallel(decay, write, c):
    a, h = decay, write
    shift = 1
    while shift < a.shape[2]:
        a_left = F.pad(
            a[:, :, :-shift], (0, 0, shift, 0), value=1.0
        )  # identity beyond the start
        h_left = F.pad(h[:, :, :-shift], (0, 0, shift, 0), value=0.0)
        h = a * h_left + h
        a = a * a_left
        shift *= 2
    return (h * c[:, None]).sum(-1)


torch.manual_seed(SEED)
decay = torch.rand(1, 3, 37, 4)
write = torch.randn(1, 3, 37, 4)
c = torch.randn(1, 37, 4)
gap = (scan_parallel(decay, write, c) - scan_loop(decay, write, c)).abs().max()
print("parallel against loop:", f"{gap:.1e}")

sizes = [256, 512] if SMOKE else [256, 1024, 4096]
timing = []
for L in sizes:
    torch.manual_seed(SEED)
    decay = torch.rand(1, 1536, L, 16, device=device)
    write = torch.randn(1, 1536, L, 16, device=device)
    c = torch.randn(1, L, 16, device=device)
    for name, scan in [("loop", scan_loop), ("parallel", scan_parallel)]:
        scan(decay, write, c)
        if device == "cuda":
            torch.cuda.synchronize()
        start = time.time()
        scan(decay, write, c)
        if device == "cuda":
            torch.cuda.synchronize()
        timing.append({"tokens": L, "scan": name, "ms": (time.time() - start) * 1000})
ms = pd.DataFrame(timing).pivot(index="tokens", columns="scan", values="ms")
table(ms.reset_index(), digits=1)

def mixer_step(m, x, state):
    window, h = (
        state  # the convolution's last inputs (batch, E, k), the state (batch, E, N)
    )
    u, gate = m.in_proj(x).chunk(2, dim=-1)  # (batch, E) each
    window = torch.cat([window[:, :, 1:], u[:, :, None]], dim=-1)
    u = F.silu((window * m.conv1d.weight[:, 0]).sum(-1) + m.conv1d.bias)
    seed, b, c = m.x_proj(u).split([m.rank, m.n, m.n], dim=-1)
    step = F.softplus(m.dt_proj(seed))  # (batch, E)
    a = -torch.exp(m.A_log.float())
    h = torch.exp(step[..., None] * a) * h + step[..., None] * b[:, None] * u[..., None]
    y = (h * c[:, None]).sum(-1) + u * m.D
    return m.out_proj(y * F.silu(gate)), (window, h)


def model_step(model, token, states):
    x = model.embeddings(token)
    new_states = []
    for layer, state in zip(model.layers, states):
        out, state = mixer_step(layer.mixer, layer.norm(x), state)
        x = x + out
        new_states.append(state)
    return model.lm_head(model.norm_f(x)).float(), new_states


def empty_states(model, batch=1):
    states = []
    for layer in model.layers:
        m = layer.mixer
        e, k = m.D.shape[0], m.conv1d.weight.shape[-1]
        dev = m.D.device
        states.append(
            (torch.zeros(batch, e, k, device=dev), torch.zeros(batch, e, m.n, device=dev))
        )
    return states


with torch.no_grad():
    full = mamba(ids)
    states = empty_states(mamba)
    stepped = []
    for t in range(ids.shape[1]):
        logits, states = model_step(mamba, ids[:, t], states)
        stepped.append(logits)
stepped = torch.stack(stepped, dim=1)
state_bytes = sum(w.numel() + h.numel() for w, h in states) * 2  # in 16-bit floats
print("step by step against the whole sequence:", f"{(stepped - full).abs().max():.1e}")
print(f"the whole state of the model: {state_bytes / 1e6:.2f} MB, at any length")

pythia_config = json.load(open(hf_hub_download("EleutherAI/pythia-160m", "config.json")))
per_token = (
    2 * pythia_config["num_hidden_layers"] * pythia_config["hidden_size"] * 2
)  # K and V, fp16
memory = []
for n in [2048, 32768, 131072, 1_000_000]:
    memory.append(
        {
            "tokens read": n,
            "Pythia-160m cache, MB": per_token * n / 1e6,
            "mamba-130m state, MB": state_bytes / 1e6,
        }
    )
table(pd.DataFrame(memory), digits=2)

def segment_sums(a):
    # a (..., T) -> (..., T, T): out[t, s] = a[s+1] + ... + a[t] for s <= t,
    # and -inf above the diagonal
    T = a.shape[-1]
    below = torch.ones(T, T, dtype=torch.bool, device=a.device).tril(-1)
    x = a[..., None].expand(*a.shape, T).masked_fill(~below, 0)
    x = torch.cumsum(x, dim=-2)
    on_or_below = torch.ones(T, T, dtype=torch.bool, device=a.device).tril()
    return x.masked_fill(~on_or_below, float("-inf"))


def ssd(x, a, b, c, chunk=64):
    # x (batch, L, heads, P), already times Δ; a (batch, L, heads) = Δ·A;
    # b and c (batch, L, heads, N)
    batch, L, heads, p = x.shape
    pad = (chunk - L % chunk) % chunk
    x, b, c = [F.pad(t, (0, 0, 0, 0, 0, pad)) for t in (x, b, c)]
    a = F.pad(a, (0, 0, 0, pad))
    n = x.shape[1] // chunk
    x, b, c = [t.reshape(batch, n, chunk, heads, -1) for t in (x, b, c)]
    a = a.reshape(batch, n, chunk, heads).permute(
        0, 3, 1, 2
    )  # (batch, heads, chunks, chunk)
    a_cum = torch.cumsum(a, dim=-1)
    # 1. inside each chunk: attention with the decay as its mask
    mask = torch.exp(segment_sums(a))
    y_inside = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", c, b, mask, x)
    # 2. the state each chunk leaves behind
    to_end = torch.exp(a_cum[..., -1:] - a_cum)
    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", b, to_end, x)
    # 3. the recurrence between chunks
    states = torch.cat([torch.zeros_like(states[:, :1]), states], dim=1)
    between = torch.exp(segment_sums(F.pad(a_cum[..., -1], (1, 0))))
    states = torch.einsum("bhzc,bchpn->bzhpn", between, states)[:, :-1]
    # 4. what each position reads from the state that entered its chunk
    y_entering = torch.einsum("bclhn,bchpn,bhcl->bclhp", c, states, torch.exp(a_cum))
    return (y_inside + y_entering).reshape(batch, n * chunk, heads, p)[:, :L]


def ssd_loop(x, a, b, c):
    h = torch.zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1], device=x.device)
    ys = []
    for t in range(x.shape[1]):
        h = (
            torch.exp(a[:, t])[..., None, None] * h
            + x[:, t, :, :, None] * b[:, t, :, None, :]
        )
        ys.append((h * c[:, t, :, None, :]).sum(-1))
    return torch.stack(ys, dim=1)


torch.manual_seed(SEED)
x = torch.randn(1, 150, 3, 8)
a = -torch.rand(1, 150, 3) * 0.3
b = torch.randn(1, 150, 3, 16)
c = torch.randn(1, 150, 3, 16)
print(
    "chunks of 32 against the step-by-step loop:",
    f"{(ssd(x, a, b, c, 32) - ssd_loop(x, a, b, c)).abs().max():.1e}",
)

class GatedRMSNorm(nn.Module):
    def __init__(self, size, eps):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))
        self.eps = eps

    def forward(self, x, gate):
        x32 = x.float() * F.silu(gate.float())
        x32 = x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * x32.to(x.dtype)


class Mixer2(nn.Module):
    def __init__(self, c):
        super().__init__()
        d = c["hidden_size"]
        self.heads, self.p, self.n = c["num_heads"], c["head_dim"], c["state_size"]
        self.groups = c["n_groups"]
        self.e = self.heads * self.p
        conv = self.e + 2 * self.groups * self.n
        k = c["conv_kernel"]
        self.in_proj = nn.Linear(
            d, self.e + conv + self.heads, bias=c.get("use_bias", False)
        )
        self.conv1d = nn.Conv1d(conv, conv, k, groups=conv, padding=k - 1, bias=True)
        self.dt_bias = nn.Parameter(torch.zeros(self.heads))
        self.A_log = nn.Parameter(torch.zeros(self.heads))
        self.D = nn.Parameter(torch.ones(self.heads))
        self.norm = GatedRMSNorm(self.e, c["layer_norm_epsilon"])
        self.out_proj = nn.Linear(self.e, d, bias=c.get("use_bias", False))
        self.chunk = c.get("chunk_size", 64)
        self.keep = False  # the next cell reads the matrix of one head
        self.parts = None

    def select(self, dt, b, c):
        # the step, B and C of every token: what makes the layer selective (section 11)
        return F.softplus((dt + self.dt_bias).float()), b, c  # the step in 32 bits

    def forward(self, x):
        batch, L, _ = x.shape
        sizes = [self.e, self.e + 2 * self.groups * self.n, self.heads]
        gate, xbc, dt = self.in_proj(x).split(sizes, dim=-1)
        xbc = F.silu(self.conv1d(xbc.transpose(1, 2))[..., :L].transpose(1, 2))
        u, b, c = xbc.split([self.e, self.groups * self.n, self.groups * self.n], dim=-1)
        step, b, c = self.select(dt, b, c)
        a = -torch.exp(self.A_log.float())  # one number per head
        u = u.reshape(batch, L, self.heads, self.p)
        share = self.heads // self.groups
        b = b.reshape(batch, L, self.groups, self.n).repeat_interleave(share, dim=2)
        c = c.reshape(batch, L, self.groups, self.n).repeat_interleave(share, dim=2)
        if self.keep:
            self.parts = (u.detach(), step.detach(), a.detach(), b.detach(), c.detach())
        y = ssd(u.float() * step[..., None], step * a, b.float(), c.float(), self.chunk)
        y = y + u * self.D[:, None]
        y = self.norm(y.reshape(batch, L, self.e).to(x.dtype), gate)
        return self.out_proj(y)


class Mamba2(nn.Module):
    def __init__(self, c, mixer=Mixer2):
        super().__init__()
        self.embeddings = nn.Embedding(c["vocab_size"], c["hidden_size"])
        self.layers = nn.ModuleList()
        for _ in range(c["num_hidden_layers"]):
            block = nn.Module()
            block.norm = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
            block.mixer = mixer(c)
            self.layers.append(block)
        self.norm_f = RMSNorm(c["hidden_size"], c["layer_norm_epsilon"])
        self.lm_head = nn.Linear(c["hidden_size"], c["vocab_size"], bias=False)

    def hidden(self, ids):
        x = self.embeddings(ids)
        for block in self.layers:
            x = x + block.mixer(block.norm(x))
        return self.norm_f(x)

    def forward(self, ids):
        return self.lm_head(self.hidden(ids)).float()

from transformers import Mamba2Config, Mamba2ForCausalLM

REPO2 = "AntonV/mamba2-130m-hf"
if SMOKE:
    small2 = Mamba2Config(
        vocab_size=50288,
        hidden_size=64,
        num_heads=4,
        head_dim=32,
        state_size=16,
        n_groups=1,
        num_hidden_layers=2,
        chunk_size=16,
    )
    torch.manual_seed(0)
    reference2 = Mamba2ForCausalLM(small2).eval()
    mamba2_config = small2.to_dict()
    state = {k.removeprefix("backbone."): v for k, v in reference2.state_dict().items()}
else:
    reference2 = Mamba2ForCausalLM.from_pretrained(REPO2, dtype=torch.float32).eval()
    mamba2_config = json.load(open(hf_hub_download(REPO2, "config.json")))
    state = load_file(hf_hub_download(REPO2, "model.safetensors"))
    state = {k.removeprefix("backbone."): v for k, v in state.items()}
state.setdefault("lm_head.weight", state["embeddings.weight"])
mamba2 = Mamba2(mamba2_config)
mamba2.load_state_dict(state, strict=True)
mamba2.eval()
with torch.no_grad():
    gap = (mamba2(ids) - reference2(ids).logits).abs().max()
print("Mamba-2, max difference from Hugging Face:", f"{gap:.1e}")
for key in ["num_heads", "head_dim", "state_size", "chunk_size"]:
    print(f"{key:12} {mamba2_config[key]}")
del reference2

layer2 = len(mamba2.layers) // 2
mixer2 = mamba2.layers[layer2].mixer
mixer2.keep = True
with torch.no_grad():
    mamba2(ids)
mixer2.keep = False
u, step, a, b, c = mixer2.parts
head = 0
L = ids.shape[1]
decay = torch.exp(segment_sums(step[0, :, head] * a[head]))  # (L, L)
M = (c[0, :, head] @ b[0, :, head].T) * decay * step[0, None, :, head]
y_matrix = M @ u[0, :, head]
y_scan = ssd(u * step[..., None], step * a, b, c, mixer2.chunk)[0, :, head]
print("matrix form against the chunked scan:", f"{(y_matrix - y_scan).abs().max():.1e}")

from transformers import AutoModelForCausalLM

PYTHIA = (
    "EleutherAI/pythia-14m" if SMOKE else "EleutherAI/pythia-160m"
)  # 14m: same tokenizer, tiny
pythia = AutoModelForCausalLM.from_pretrained(
    PYTHIA, dtype=torch.float32, attn_implementation="eager"
).eval()
with torch.no_grad():
    attention = pythia(
        ids, output_attentions=True
    ).attentions  # one (1, heads, L, L) per layer
pythia_head = attention[len(attention) // 2][0, 0].numpy()

labels = [s.strip() or "␣" for s in tokens_of(ids)]
fig, axes = plt.subplots(1, 2, figsize=(10.4, 5.0))
weights = M.abs().numpy()
weights = weights / weights.sum(
    axis=1, keepdims=True
)  # each row as a share, like attention
for ax, matrix, title in [
    (axes[0], weights, f"Mamba-2, layer {layer2}, head {head}: |M|, rows as shares"),
    (axes[1], pythia_head, f"Pythia-160m, layer {len(attention) // 2}, head 0: softmax"),
]:
    ax.imshow(np.sqrt(matrix), cmap=shades, vmin=0)
    ax.set_xticks(range(L), labels, rotation=90, fontsize=6)
    ax.set_yticks(range(L), labels, fontsize=6)
    ax.set_title(title, fontsize=9)
plt.show()

import pandas as pd

pg19 = pd.read_parquet(
    "hf://datasets/emozilla/pg19-test/data/test-00000-of-00001-29a571947c0b5ccc.parquet"
)
LENGTH = 512 if SMOKE else 8192
N_BOOKS = 2 if SMOKE else 6
book_ids = []
for text in pg19.text:
    t = tok(text[:120_000], return_tensors="pt").input_ids
    if t.shape[1] > LENGTH:
        book_ids.append(t[:, : LENGTH + 1])
    if len(book_ids) == N_BOOKS:
        break

for layer in mamba.layers:
    layer.mixer.scan = scan_parallel
mamba.to(device)
pythia.to(device)


@torch.no_grad()
def position_losses(model_logits, ids):
    logits = model_logits(ids[:, :-1].to(device))
    return F.cross_entropy(logits[0], ids[0, 1:].to(device), reduction="none").cpu()


curves = {"mamba-130m": torch.zeros(LENGTH), "Pythia-160m": torch.zeros(LENGTH)}
for ids_book in book_ids:
    curves["mamba-130m"] += position_losses(mamba, ids_book) / len(book_ids)
    curves["Pythia-160m"] += position_losses(
        lambda x: pythia(x).logits.float(), ids_book
    ) / len(book_ids)

fig, ax = plt.subplots(figsize=(7.4, 3.4))
width = max(4, LENGTH // 64)
for name, color in [("Pythia-160m", COLORS["muted"]), ("mamba-130m", COLORS["att"])]:
    smooth = np.convolve(curves[name].numpy(), np.ones(width) / width, mode="valid")
    ax.plot(smooth, color=color, linewidth=2, label=name)
ax.axvline(2048, color=COLORS["line"], linestyle="--", linewidth=1)
ax.set_xlabel("position in the book")
ax.set_ylabel("loss, nats per token")
ax.set_title(f"{len(book_ids)} PG-19 books, both trained on 2,048 tokens")
ax.legend(frameon=False)
plt.show()
for name in curves:
    inside = curves[name][1024:2048].mean().item()
    beyond = curves[name][-2048:].mean().item()
    print(f"{name:12} loss at 1–2k: {inside:.2f}, at the last 2k: {beyond:.2f}")

import random

rng = random.Random(SEED)
words = sorted(
    {w.lower() for w in pg19.text.iloc[0].split() if w.isalpha() and 3 <= len(w) <= 7}
)


def copy_prompt(n_words):
    example = " ".join(rng.sample(words, 6))
    target = " ".join(rng.sample(words, n_words))
    text = f"Repeat the words exactly.\nWords: {example}\nRepeat: {example}\n\n"
    text += f"Words: {target}\nRepeat:"
    return text, target


@torch.no_grad()
def mamba_continue(model, ids, steps):
    states = empty_states(model)
    for t in range(ids.shape[1]):
        logits, states = model_step(model, ids[:, t].to(device), states)
    out = []
    for _ in range(steps):
        token = logits.argmax(-1)
        out.append(token.item())
        logits, states = model_step(model, token, states)
    return out


@torch.no_grad()
def pythia_continue(ids, steps):
    out = pythia.generate(
        ids.to(device), max_new_tokens=steps, do_sample=False, pad_token_id=tok.eos_token_id
    )
    return out[0, ids.shape[1] :].tolist()


copy_lengths = [5, 10] if SMOKE else [10, 25, 50, 100, 200]
COPY_TRIALS = 2 if SMOKE else 10
copied = {"mamba-130m": [], "Pythia-160m": []}
for n in copy_lengths:
    hits = {"mamba-130m": 0, "Pythia-160m": 0}
    for _ in range(COPY_TRIALS):
        text, target = copy_prompt(n)
        prompt = tok(text, return_tensors="pt").input_ids
        steps = len(tok(" " + target).input_ids) + 2
        said = {
            "mamba-130m": tok.decode(mamba_continue(mamba, prompt, steps)),
            "Pythia-160m": tok.decode(pythia_continue(prompt, steps)),
        }
        for name in hits:
            hits[name] += said[name].strip().startswith(target)
    for name in hits:
        copied[name].append(hits[name] / COPY_TRIALS)

fig, ax = plt.subplots(figsize=(7.2, 3.2))
for name, color in [("Pythia-160m", COLORS["muted"]), ("mamba-130m", COLORS["att"])]:
    ax.plot(copy_lengths, copied[name], color=color, marker="o", linewidth=2, label=name)
ax.set_xscale("log")
ax.set_xticks(copy_lengths, [str(n) for n in copy_lengths])
ax.set_ylim(-0.05, 1.05)
ax.set_xlabel("words to repeat")
ax.set_ylabel("lists repeated exactly")
ax.set_title(f"Repeating a list of random words, {COPY_TRIALS} lists per length")
ax.legend(frameon=False)
plt.show()

def layout(repo):
    c = read_hybrid_config(repo)
    n = c.get("num_hidden_layers") or c.get("n_layer")
    if "hybrid_override_pattern" in c:  # Nemotron-H: M Mamba, * attention, - an MLP layer
        names = {"M": "ssm", "*": "attention", "-": "mlp"}
        return [names[ch] for ch in c["hybrid_override_pattern"]]
    if "layer_types" in c:  # Granite 4.0
        return ["attention" if "attention" in t else "ssm" for t in c["layer_types"]]
    if "layers_block_type" in c:  # Zamba2: "hybrid" layers add a shared attention block
        return ["attention" if t == "hybrid" else "ssm" for t in c["layers_block_type"]]
    if "attn_layer_period" in c:  # Jamba
        period, offset = c["attn_layer_period"], c["attn_layer_offset"]
        return ["attention" if i % period == offset else "ssm" for i in range(n)]
    if "full_attention_interval" in c:  # Qwen3-Next: Gated DeltaNet, then attention
        k = c["full_attention_interval"]
        return ["attention" if (i + 1) % k == 0 else "linear" for i in range(n)]
    if "linear_attn_config" in c:  # Kimi Linear: its attention layers are counted from 1
        full = c["linear_attn_config"]["full_attn_layers"]
        return ["attention" if i + 1 in full else "linear" for i in range(n)]
    if "attn_layer_idx" in c:  # Mamba-2 with attention, from the Mamba-2 paper
        return ["attention" if i in c["attn_layer_idx"] else "ssm" for i in range(n)]
    return ["ssm"] * n


def read_hybrid_config(repo):
    with open(hf_hub_download(repo, "config.json")) as f:
        c = json.load(f)
    return c.get("text_config", c)


HYBRIDS = [
    ("state-spaces/mamba2attn-2.7b", "Mamba-2 + attention, 2024"),
    ("ai21labs/Jamba-v0.1", "Jamba, AI21, 2024"),
    ("Zyphra/Zamba2-1.2B", "Zamba2, Zyphra, 2024"),
    ("nvidia/Nemotron-H-8B-Base-8K", "Nemotron-H, NVIDIA, 2025"),
    ("ibm-granite/granite-4.0-h-tiny", "Granite 4.0 H, IBM, 2025"),
    ("Qwen/Qwen3-Next-80B-A3B-Instruct", "Qwen3-Next, 2025"),
    ("moonshotai/Kimi-Linear-48B-A3B-Instruct", "Kimi Linear, 2025"),
]
layouts = {name: layout(repo) for repo, name in HYBRIDS}
colors = {
    "attention": COLORS["dec"],
    "ssm": COLORS["att"],
    "linear": COLORS["enc"],
    "mlp": COLORS["line"],
}
fig, ax = plt.subplots(figsize=(10.4, 3.6))
for row, kinds in enumerate(layouts.values()):
    for col, kind in enumerate(kinds):
        ax.add_patch(plt.Rectangle((col, -row - 0.4), 0.9, 0.8, color=colors[kind]))
ax.set_xlim(0, max(len(k) for k in layouts.values()))
ax.set_ylim(-len(layouts) + 0.4, 0.6)
ax.set_yticks([-i for i in range(len(layouts))], list(layouts))
ax.set_xlabel("layer")
for kind, label in [
    ("attention", "attention"),
    ("ssm", "Mamba"),
    ("linear", "linear attention"),
    ("mlp", "MLP only"),
]:
    ax.plot([], [], color=colors[kind], linewidth=8, label=label)
ax.legend(frameon=False, ncol=4, loc="upper center", bbox_to_anchor=(0.5, 1.18))
plt.show()

rows = []
for name, kinds in layouts.items():
    mixing = [k for k in kinds if k != "mlp"]
    rows.append(
        {
            "model": name,
            "layers that mix tokens": len(mixing),
            "attention": mixing.count("attention"),
            "share of attention": mixing.count("attention") / len(mixing),
        }
    )
table(pd.DataFrame(rows), digits=2)