VKVlad KulikovPublic / lectures
Modern LLMs8 / 20

Module 8. Beyond the Transformer: SSMs, Mamba, and Hybrid Sequence Models

Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.

Download course files (.zip) 2588.1 KB

The archive preserves the source folder structure, both languages, and supporting files.

Notebook 8. SSMs, selective scans, and linear memory — from formula to contract

Main module notebook

Download .ipynb module_08_notebook_ssm.ipynb · 61.6 KB

Show notebook text

Notebook 8. SSMs, selective scans, and linear memory — from formula to contract

Module 8 — hands-on notebook

This is engineering-disciplined educational code, not a fused Mamba kernel. The mandatory path carries one state-space idea through seven testable levels:

  1. exact diagonal zero-order-hold discretization;

  2. a selective recurrent scan and the cost of the Δ·B approximation;

  3. associative affine composition and a chunked scan;

  4. exact equivalence of sequential and chunked linear attention;

  5. the delta rule with a normalized key and a decay gate;

  6. finite-state interference against explicit softmax addressing;

  7. SSM/attention/hybrid memory and arithmetic under strict units.

Shape contract: time, state, input, and value axes are validated explicitly. Empty sequences, incompatible shapes, zero keys, and non-finite values do not pass silently.

Unit contract: KiB/MiB/GiB use powers of two; FLOPs are an arithmetic scenario. Recurrent state is counted per sequence and scales with batch size, but not with context length.

Colab contract: Runtime → Run all completes on CPU in a clean session. The notebook reads no repository files and needs no upload, Google Drive, API key, or network access. Installing a missing NumPy requires network access; the experiments then run offline. The frozen M8 numerical contract is embedded and verified by SHA-256.

0. Embedded numerical contract

The reference data belong to M8 and are independent of mutable lecture prose and visualization files. Compact JSON is embedded in the notebook; SHA-256 prevents accidental drift among the builder, answer sheet, and published .ipynb.

Python
import hashlib
import json
import math
import subprocess
import sys
from dataclasses import dataclass
from numbers import Integral, Real
from typing import Any, Iterable

try:
    import numpy as np
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy>=1.24,<3"])
    import numpy as np

np.set_printoptions(precision=6, suppress=True)
REFERENCE_RAW = '{"affine_scan":{"chunk_size":16,"max_abs_difference":8.881784197001252e-16,"output_checksum":2.5500288962870172,"seed":11,"sequence_length":53,"state_width":6},"decode_flops":{"attention_core_flops":[16777216,536870912,17179869184],"inner_width":8192,"model_width":4096,"positions":[1024,32768,1048576],"ratios":[25.6,819.2,26214.4],"scope":"leading-order attention matmuls versus diagonal state update plus readout, with prepared coefficients; not end-to-end latency","ssm_core_flops":[655360,655360,655360],"state_width":16},"delta":{"additive_retrieval":[4.0,6.0],"full_retrieval":[3.0,4.0],"gated_half_retrieval":[1.9000000000000001,2.8],"half_retrieval":[2.0,3.0]},"linear_attention":{"chunk_size":16,"key_width":8,"max_abs_difference":2.1316282072803006e-14,"output_checksum":111.6891373299517,"seed":19,"sequence_length":50,"value_width":6},"memory":{"full_attention_256k_bytes":34359738368,"hybrid_256k":{"attention_layers":4,"conv_bytes":1835008,"kv_bytes":4294967296,"recurrent_bytes":9175040,"ssm_bytes":7340032,"ssm_layers":28,"total_bytes":4304142336},"kv_bytes_per_token_full_attention":131072,"mamba1":{"conv_bytes":2097152,"crossover_first_integer_token":80,"crossover_tokens_exact":80.0,"first_strictly_larger_token":81,"ssm_bytes":8388608,"total_bytes":10485760},"mamba2_like":{"conv_bytes":2162688,"crossover_first_integer_token":529,"crossover_tokens_exact":528.5,"first_strictly_larger_token":529,"ssm_bytes":67108864,"total_bytes":69271552}},"recall":{"expected_additive_mse":[0.09375,0.21875,0.46875,0.96875,1.96875,3.96875],"key_width":32,"mse":[0.13978625740853395,0.21200348269044997,0.5995392976261167,0.9892555844637765,1.8757385532834472,4.141977097591836],"paired_softmax_mse":[3.7460446307862835e-10,4.0983742994017256e-12,9.826283934707642e-10,1.7655760849531266e-09,1.1566029996764256e-09,1.0724543396718937e-08],"seed":7,"softmax_inverse_temperature":20.0,"softmax_mse_at_128":1.648958588765511e-08,"softmax_seed":107,"stored_pairs":[4,8,16,32,64,128],"value_width":32},"schema_version":2,"zoh":{"approximate_outputs":[0.2,-1.9496785275591946,1.0528816219675536],"approximate_states":[[0.1,0.1],[-0.9632120558828557,-0.9864664716763387],[0.4157823563022096,0.637099265665344]],"continuous_a":[-1.0,-2.0],"continuous_b":[[1.0],[1.0]],"deltas":[0.1,1.0,0.5],"exact_input_matrices":[[[0.09516258196404043],[0.09063462346100906]],[[0.6321205588285577],[0.43233235838169365]],[[0.39346934028736663],[0.31606027941427883]]],"exact_outputs":[0.1857972054250495,-1.0171784972997493,0.9023586277858214],"exact_states":[[0.09516258196404043,0.09063462346100906],[-0.5971122013551949,-0.42006629594455425],[0.42477182316430406,0.47758680462151737]],"inputs":[[1.0],[-1.0],[2.0]],"output_c":[[1.0,1.0]],"relative_input_error_first_mode":[0.05083319447750502,0.5819767068693265,0.270747041268399],"ten_percent_relative_error_root_a_minus_1":0.1937475579949906,"transitions":[[0.9048374180359595,0.8187307530779818],[0.36787944117144233,0.1353352832366127],[0.6065306597126334,0.36787944117144233]],"zero_mode_input_matrix":[[0.25],[0.22119921692859512]]}}'
EXPECTED_SHA256 = '28cbfd26cb5163eff23dd9ac6fd3084a325ee3a5efcb96192edd9562ae9b7387'
actual_sha256 = hashlib.sha256(REFERENCE_RAW.encode("utf-8")).hexdigest()
assert actual_sha256 == EXPECTED_SHA256
REFERENCE = json.loads(REFERENCE_RAW)
assert REFERENCE["schema_version"] == 2
print('The embedded M8 contract is verified', actual_sha256[:12])


"""Minimal M8 operators; used by the validated reference and embedded in notebooks."""
import numpy as np

def zoh_diagonal(a, b, delta):
    """Diagonal A: [n], B: [n,p], positive scalar step; caller validates shapes."""
    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)
    transition = np.exp(delta * a)
    factor = np.full_like(a, float(delta))
    np.divide(np.expm1(delta * a), a, out=factor, where=a != 0)
    return transition, factor[:, None] * b

def delta_step(state, key, value, beta, decay=1.0):
    """State [dv,dk], unit key [dk], value [dv]; decay before correction."""
    base = decay * state
    prediction = base @ key
    return base + beta * np.outer(value - prediction, key)

def compose_affine_pair(later, earlier):
    """Elementwise transforms; returns later(earlier(h)), not the reverse."""
    a2, b2 = later
    a1, b1 = earlier
    return a2 * a1, a2 * b1 + b2
Python
def _real(name: str, value: Real) -> float:
    """Return one finite real scalar, rejecting booleans and non-finite values."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    number = float(value)
    if not math.isfinite(number):
        raise ValueError(f"{name} must be finite")
    return number


def _positive_real(name: str, value: Real) -> float:
    """Return one finite strictly positive scalar."""

    number = _real(name, value)
    if number <= 0.0:
        raise ValueError(f"{name} must be positive")
    return number


def _unit_interval(name: str, value: Real) -> float:
    """Return a finite scalar in the closed unit interval."""

    number = _real(name, value)
    if number < 0.0 or number > 1.0:
        raise ValueError(f"{name} must lie in [0, 1]")
    return number


def _positive_int(name: str, value: Integral) -> int:
    """Return one strictly positive integer, rejecting booleans."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    number = int(value)
    if number <= 0:
        raise ValueError(f"{name} must be positive")
    return number


def _non_negative_int(name: str, value: Integral) -> int:
    """Return one non-negative integer, rejecting booleans."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    number = int(value)
    if number < 0:
        raise ValueError(f"{name} must be non-negative")
    return number


def _finite_array(
    name: str,
    value: Any,
    *,
    ndim: int | None = None,
    non_empty: bool = True,
) -> np.ndarray:
    """Return a float64 array with explicit dimensionality and finiteness checks."""

    array = np.asarray(value, dtype=np.float64)
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have ndim={ndim}, got shape {array.shape}")
    if non_empty and array.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def bytes_to_kib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary KiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**10


def bytes_to_mib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary MiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**20


def bytes_to_gib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary GiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**30

1. Exact ZOH discretization and the selective scan

For the diagonal continuous system

h˙=Ah+Bx,\dot h = Ah + Bx,

the exact constant-input step is

Aˉ=exp⁡(ΔA),Bˉi=exp⁡(Δai)−1aiBi.\bar A=\exp(\Delta A),\qquad \bar B_i=\frac{\exp(\Delta a_i)-1}{a_i}B_i.

At a_i = 0 the implementation uses the limit Δ rather than a hidden 0/0. The same scan is then run with exact B̄ and with the educational Δ·B approximation so that the error is visible in the state, not only in an isolated equation.

Python
@dataclass(frozen=True)
class ZOHDiscretization:
    """Exact diagonal zero-order-hold discretization for one time step.

    ``transition`` has shape ``[d_state]`` and represents the diagonal of
    ``exp(delta*A)``. ``input_matrix`` has shape ``[d_state, d_input]``.
    """

    transition: np.ndarray
    input_matrix: np.ndarray


@dataclass(frozen=True)
class SSMScanResult:
    """States, outputs, and per-step discrete parameters of a diagonal SSM scan."""

    states: np.ndarray
    outputs: np.ndarray
    transitions: np.ndarray
    input_matrices: np.ndarray


def zoh_discretize_diagonal(
    continuous_a: Any,
    continuous_b: Any,
    delta: Real,
) -> ZOHDiscretization:
    """Discretize ``dh/dt = A h + B x`` for diagonal stable ``A``.

    Parameters
    ----------
    continuous_a:
        Diagonal entries of ``A`` with shape ``[d_state]``. Non-positive values are
        accepted; positive values are rejected because this lesson models stable modes.
    continuous_b:
        Input matrix with shape ``[d_state, d_input]``. A vector is treated as one input
        channel.
    delta:
        Strictly positive time step.

    Notes
    -----
    The exact input factor is ``expm1(delta*a)/a``. At ``a=0`` the continuous limit is
    exactly ``delta``; handling that limit explicitly avoids a hidden ``0/0``.
    """

    a = _finite_array("continuous_a", continuous_a, ndim=1)
    if np.any(a > 0.0):
        raise ValueError("continuous_a must be non-positive for the stable lesson model")
    b = _finite_array("continuous_b", continuous_b)
    if b.ndim == 1:
        b = b[:, None]
    if b.ndim != 2 or b.shape[0] != a.size:
        raise ValueError("continuous_b must have shape [d_state, d_input]")
    step = _positive_real("delta", delta)

    transition, input_matrix = zoh_diagonal(a, b, step)
    if not np.all(np.isfinite(input_matrix)):
        raise FloatingPointError("ZOH discretization produced a non-finite input matrix")
    return ZOHDiscretization(transition=transition, input_matrix=input_matrix)


def scan_diagonal_ssm(
    continuous_a: Any,
    continuous_b: Any,
    output_c: Any,
    inputs: Any,
    deltas: Any,
    *,
    initial_state: Any | None = None,
    exact_input_discretization: bool = True,
) -> SSMScanResult:
    """Run a selective diagonal SSM over a non-empty sequence.

    Shapes are ``A:[d_state]``, ``B:[d_state,d_input]``, ``C:[d_output,d_state]``,
    ``inputs:[T,d_input]``, and ``deltas:[T]``. A one-dimensional input sequence is
    interpreted as ``d_input=1``. The output includes every post-update state.

    Setting ``exact_input_discretization=False`` keeps the exact transition but replaces
    the ZOH input matrix by the common small-step approximation ``delta*B``.
    """

    a = _finite_array("continuous_a", continuous_a, ndim=1)
    b = _finite_array("continuous_b", continuous_b)
    if b.ndim == 1:
        b = b[:, None]
    c = _finite_array("output_c", output_c)
    if c.ndim == 1:
        c = c[None, :]
    x = _finite_array("inputs", inputs)
    if x.ndim == 1:
        x = x[:, None]
    dt = _finite_array("deltas", deltas, ndim=1)

    if b.ndim != 2 or b.shape[0] != a.size:
        raise ValueError("continuous_b must have shape [d_state, d_input]")
    if c.ndim != 2 or c.shape[1] != a.size:
        raise ValueError("output_c must have shape [d_output, d_state]")
    if x.ndim != 2 or x.shape[1] != b.shape[1]:
        raise ValueError("inputs must have shape [T, d_input]")
    if x.shape[0] != dt.size:
        raise ValueError("inputs and deltas must have the same sequence length")
    if np.any(dt <= 0.0):
        raise ValueError("every delta must be positive")
    if np.any(a > 0.0):
        raise ValueError("continuous_a must be non-positive for the stable lesson model")

    state = (
        np.zeros(a.size, dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != a.shape:
        raise ValueError("initial_state must have shape [d_state]")

    states = np.empty((x.shape[0], a.size), dtype=np.float64)
    outputs = np.empty((x.shape[0], c.shape[0]), dtype=np.float64)
    transitions = np.empty_like(states)
    input_matrices = np.empty((x.shape[0], a.size, b.shape[1]), dtype=np.float64)

    for index, (token, step) in enumerate(zip(x, dt)):
        discrete = zoh_discretize_diagonal(a, b, float(step))
        input_matrix = (
            discrete.input_matrix
            if exact_input_discretization
            else float(step) * b
        )
        state = discrete.transition * state + input_matrix @ token
        states[index] = state
        outputs[index] = c @ state
        transitions[index] = discrete.transition
        input_matrices[index] = input_matrix

    if not np.all(np.isfinite(states)) or not np.all(np.isfinite(outputs)):
        raise FloatingPointError("SSM scan produced non-finite states or outputs")
    return SSMScanResult(states, outputs, transitions, input_matrices)


z = REFERENCE["zoh"]
a = np.asarray(z["continuous_a"], dtype=np.float64)
b = np.asarray(z["continuous_b"], dtype=np.float64)
c = np.asarray(z["output_c"], dtype=np.float64)
x = np.asarray(z["inputs"], dtype=np.float64)
delta = np.asarray(z["deltas"], dtype=np.float64)

exact = scan_diagonal_ssm(a, b, c, x, delta)
approx = scan_diagonal_ssm(a, b, c, x, delta, exact_input_discretization=False)
assert exact.states.shape == (3, 2)
assert exact.outputs.shape == (3, 1)
assert np.allclose(exact.transitions, z["transitions"], rtol=1e-12, atol=1e-12)
assert np.allclose(exact.input_matrices, z["exact_input_matrices"], rtol=1e-12, atol=1e-12)
assert np.allclose(exact.states, z["exact_states"], rtol=1e-12, atol=1e-12)
assert np.allclose(exact.outputs[:, 0], z["exact_outputs"], rtol=1e-12, atol=1e-12)
assert np.allclose(approx.states, z["approximate_states"], rtol=1e-12, atol=1e-12)

relative_errors = []
for step in delta:
    discrete = zoh_discretize_diagonal(a, b, float(step))
    relative_errors.append(abs(step * b[0, 0] - discrete.input_matrix[0, 0]) / discrete.input_matrix[0, 0])
assert np.allclose(relative_errors, z["relative_input_error_first_mode"], rtol=1e-12, atol=1e-12)

# The zero continuous mode uses the analytical limit phi(0)=delta.
zero_mode = zoh_discretize_diagonal([0.0, -1.0], [[1.0], [1.0]], 0.25)
assert np.allclose(zero_mode.input_matrix, z["zero_mode_input_matrix"], atol=1e-12)
assert zero_mode.input_matrix[0, 0] == 0.25
# A small a is not a zero mode when delta*a is order one.
with np.errstate(divide="raise", invalid="raise"):
    long_step = zoh_discretize_diagonal([-1e-15, 0.0], [[1.0], [1.0]], 1e15)
np.testing.assert_allclose(long_step.input_matrix[:, 0], [(1-np.exp(-1))*1e15, 1e15], rtol=1e-14)


root = z["ten_percent_relative_error_root_a_minus_1"]
root_error = root / (-math.expm1(-root)) - 1.0
assert math.isclose(root, 0.1937475579949906, rel_tol=0.0, abs_tol=1e-13)
assert math.isclose(root_error, 0.1, rel_tol=0.0, abs_tol=1e-12)

for index, step in enumerate(delta):
    print(
        f"t={index+1} delta={step:.1f} | Abar={exact.transitions[index]} | "
        f"relative ΔB error={relative_errors[index]:.1%} | y={exact.outputs[index,0]:+.6f}"
    )
print('ZOH, the selective scan, and the a=0 limit are verified')

2. Associativity and a chunked selective scan

Each diagonal step is an affine transformation

Tt(h)=at⊙h+bt.T_t(h)=a_t\odot h+b_t.

Composition is associative:

(a2,b2)∘(a1,b1)=(a2⊙a1,  a2⊙b1+b2).(a_2,b_2)\circ(a_1,b_1) =(a_2\odot a_1,\;a_2\odot b_1+b_2).

That algebra is why a scan can be parallelized. The NumPy code below does not imitate a fused GPU kernel: it exposes the within-chunk contribution and the boundary state, then checks an uneven final chunk.

Python
@dataclass(frozen=True)
class AffineTransform:
    """Elementwise affine state transform ``h_out = multiplier*h_in + bias``."""

    multiplier: np.ndarray
    bias: np.ndarray


def compose_affine(later: AffineTransform, earlier: AffineTransform) -> AffineTransform:
    """Compose two diagonal affine transforms as ``later(earlier(h))``."""

    a2 = _finite_array("later.multiplier", later.multiplier, ndim=1)
    b2 = _finite_array("later.bias", later.bias, ndim=1)
    a1 = _finite_array("earlier.multiplier", earlier.multiplier, ndim=1)
    b1 = _finite_array("earlier.bias", earlier.bias, ndim=1)
    if a1.shape != b1.shape or a2.shape != b2.shape or a1.shape != a2.shape:
        raise ValueError("all affine components must share shape [d_state]")
    return AffineTransform(multiplier=a2 * a1, bias=a2 * b1 + b2)


def affine_scan_sequential(
    multipliers: Any,
    biases: Any,
    *,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate ``h_t = a_t*h_{t-1} + b_t`` sequentially.

    ``multipliers`` and ``biases`` must both have shape ``[T, d_state]``. The returned
    array contains post-update states with the same shape.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")
    result = np.empty_like(b)
    for index in range(a.shape[0]):
        state = a[index] * state + b[index]
        result[index] = state
    return result


def affine_scan_chunked(
    multipliers: Any,
    biases: Any,
    *,
    chunk_size: Integral,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate the same affine recurrence in non-overlapping chunks.

    Inside each chunk, the local zero-state recurrence and cumulative multipliers can be
    computed in parallel by a scan kernel. Only one boundary state crosses each chunk.
    This NumPy version keeps the mechanics transparent; it is not a fused GPU scan.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    chunk = _positive_int("chunk_size", chunk_size)
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")

    result = np.empty_like(b)
    for start in range(0, a.shape[0], chunk):
        end = min(start + chunk, a.shape[0])
        local_a = a[start:end]
        local_b = b[start:end]

        # Zero-start local states carry all within-chunk contributions.
        zero_state = np.zeros(a.shape[1], dtype=np.float64)
        local_zero = np.empty_like(local_b)
        for offset in range(local_a.shape[0]):
            zero_state = local_a[offset] * zero_state + local_b[offset]
            local_zero[offset] = zero_state

        # The chunk-start state is multiplied by the cumulative transition to each row.
        prefix_multiplier = np.cumprod(local_a, axis=0)
        result[start:end] = local_zero + prefix_multiplier * state
        state = result[end - 1].copy()
    return result


ref = REFERENCE["affine_scan"]
rng = np.random.default_rng(ref["seed"])
a_steps = rng.uniform(0.70, 0.99, size=(ref["sequence_length"], ref["state_width"]))
b_steps = rng.standard_normal((ref["sequence_length"], ref["state_width"]))
h0 = rng.standard_normal(ref["state_width"])
seq = affine_scan_sequential(a_steps, b_steps, initial_state=h0)
chunked = affine_scan_chunked(a_steps, b_steps, chunk_size=ref["chunk_size"], initial_state=h0)
difference = float(np.max(np.abs(seq - chunked)))
assert difference < 1e-12
assert math.isclose(float(np.sum(seq)), ref["output_checksum"], rel_tol=1e-11, abs_tol=1e-12)

# Composition itself is associative; check it on three transforms.
T1 = AffineTransform(np.array([0.9, 0.8]), np.array([1.0, -1.0]))
T2 = AffineTransform(np.array([0.7, 0.6]), np.array([0.5, 2.0]))
T3 = AffineTransform(np.array([0.5, 0.4]), np.array([-0.5, 1.0]))
left = compose_affine(T3, compose_affine(T2, T1))
right = compose_affine(compose_affine(T3, T2), T1)
assert np.allclose(left.multiplier, right.multiplier, atol=1e-15)
assert np.allclose(left.bias, right.bias, atol=1e-15)
print(f"max|sequential - chunked| = {difference:.3e}; uneven tail = {ref['sequence_length'] % ref['chunk_size']}")
print('Affine associativity and the chunked scan are verified')

3. A special case: sequential and chunked linear attention

For additive memory

St=St−1+vtkt⊤,yt=Stqt,S_t=S_{t-1}+v_tk_t^\top,\qquad y_t=S_tq_t,

the chunked form adds an inter-chunk q @ Sᵀ term and an intra-chunk lower-triangular matrix product. This is a reordering of computation, not an approximation. We check chunk sizes 1, an uneven tail, and a chunk longer than the sequence.

Python
def _validate_qkv(q: Any, k: Any, v: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Validate a non-empty causal linear-attention triplet."""

    query = _finite_array("q", q, ndim=2)
    key = _finite_array("k", k, ndim=2)
    value = _finite_array("v", v, ndim=2)
    if query.shape[0] != key.shape[0] or query.shape[0] != value.shape[0]:
        raise ValueError("q, k, and v must share sequence length T")
    if query.shape[1] != key.shape[1]:
        raise ValueError("q and k must share feature width d_key")
    return query, key, value


def linear_attention_sequential(q: Any, k: Any, v: Any) -> np.ndarray:
    """Causal additive linear attention with state shape ``[d_value, d_key]``."""

    query, key, value = _validate_qkv(q, k, v)
    state = np.zeros((value.shape[1], key.shape[1]), dtype=np.float64)
    output = np.empty((query.shape[0], value.shape[1]), dtype=np.float64)
    for index in range(query.shape[0]):
        state += np.outer(value[index], key[index])
        output[index] = state @ query[index]
    return output


def linear_attention_chunked(
    q: Any,
    k: Any,
    v: Any,
    *,
    chunk_size: Integral,
) -> np.ndarray:
    """Exact chunked form of the same causal additive linear attention.

    The inter-chunk term reads the carried state; the intra-chunk term is a masked
    matrix multiplication. The final partial chunk is supported explicitly.
    """

    query, key, value = _validate_qkv(q, k, v)
    chunk = _positive_int("chunk_size", chunk_size)
    state = np.zeros((value.shape[1], key.shape[1]), dtype=np.float64)
    output = np.empty((query.shape[0], value.shape[1]), dtype=np.float64)
    for start in range(0, query.shape[0], chunk):
        end = min(start + chunk, query.shape[0])
        q_chunk = query[start:end]
        k_chunk = key[start:end]
        v_chunk = value[start:end]
        inter = q_chunk @ state.T
        causal_scores = np.tril(q_chunk @ k_chunk.T)
        output[start:end] = inter + causal_scores @ v_chunk
        state += v_chunk.T @ k_chunk
    return output


ref = REFERENCE["linear_attention"]
rng = np.random.default_rng(ref["seed"])
q = rng.standard_normal((ref["sequence_length"], ref["key_width"]))
k = rng.standard_normal((ref["sequence_length"], ref["key_width"]))
v = rng.standard_normal((ref["sequence_length"], ref["value_width"]))
sequential = linear_attention_sequential(q, k, v)
for chunk_size in (1, 7, ref["chunk_size"], 64):
    candidate = linear_attention_chunked(q, k, v, chunk_size=chunk_size)
    assert candidate.shape == sequential.shape
    assert np.allclose(candidate, sequential, rtol=1e-12, atol=1e-12)
max_difference = float(np.max(np.abs(linear_attention_chunked(q, k, v, chunk_size=16) - sequential)))
assert max_difference < 1e-12
assert math.isclose(float(np.sum(sequential)), ref["output_checksum"], rel_tol=1e-11, abs_tol=1e-12)
# Perturbing future tokens must not change earlier causal outputs.
changed = v.copy(); changed[-1] += 100
np.testing.assert_allclose(linear_attention_chunked(q, k, changed, chunk_size=7)[:-1], sequential[:-1])
print(f"output shape={sequential.shape}, max difference={max_difference:.3e}")
print('Sequential and chunked linear attention agree')

4. The delta rule: address-specific overwrite requires a normalized key

For a unit key k:

S′=gS+β (v−gSk)k⊤.S' = gS + \beta\,(v-gSk)k^\top.

β=1 overwrites retrieval at that key exactly; β=0 leaves only the global decay g. Key normalization is not cosmetic: without it the exact-overwrite claim is false, so the contract either normalizes the key or explicitly requires unit norm.

Python
@dataclass(frozen=True)
class DeltaUpdateResult:
    """State and retrieval trace for one normalized-key delta-rule update."""

    state: np.ndarray
    normalized_key: np.ndarray
    retrieval_before_write: np.ndarray
    retrieval_after_write: np.ndarray


def normalize_key(key: Any) -> np.ndarray:
    """Return a unit-L2 key, rejecting the zero vector."""

    vector = _finite_array("key", key, ndim=1)
    scale = float(np.max(np.abs(vector)))
    if scale == 0.0:
        raise ValueError("key must have non-zero norm")
    scaled = vector / scale
    return scaled / np.linalg.norm(scaled)


def delta_update(
    state: Any,
    key: Any,
    value: Any,
    *,
    beta: Real,
    decay: Real = 1.0,
    normalize: bool = True,
) -> DeltaUpdateResult:
    """Apply a gated delta-rule write to associative state ``[d_value, d_key]``.

    With a unit key, the update is

    ``base = decay * state``
    ``new = base + beta * outer(value - base@key, key)``.

    Thus ``beta=1`` overwrites retrieval at that key exactly, while ``beta=0`` only
    applies the global decay. If ``normalize=False``, the caller must supply a unit key;
    the explicit check prevents a misleading "exact overwrite" claim for arbitrary norm.
    """

    matrix = _finite_array("state", state, ndim=2)
    raw_key = _finite_array("key", key, ndim=1)
    target = _finite_array("value", value, ndim=1)
    if matrix.shape != (target.size, raw_key.size):
        raise ValueError("state must have shape [d_value, d_key]")
    write_rate = _unit_interval("beta", beta)
    forget_rate = _unit_interval("decay", decay)
    unit_key = normalize_key(raw_key) if normalize else raw_key.copy()
    if not normalize and not np.isclose(np.linalg.norm(unit_key), 1.0, rtol=0.0, atol=1e-10):
        raise ValueError("key must have unit norm when normalize=False")

    base = forget_rate * matrix
    before = base @ unit_key
    updated = delta_step(matrix, unit_key, target, write_rate, forget_rate)
    after = updated @ unit_key
    return DeltaUpdateResult(updated, unit_key, before, after)


ref = REFERENCE["delta"]
k = np.array([1.0, 0.0])
v1 = np.array([1.0, 2.0])
v2 = np.array([3.0, 4.0])
state = np.outer(v1, k)
additive = state + np.outer(v2, k)
half = delta_update(state, k, v2, beta=0.5)
full = delta_update(state, k, v2, beta=1.0)
gated = delta_update(state, k, v2, beta=0.5, decay=0.8)
no_write = delta_update(state, k, v2, beta=0.0, decay=0.5)

assert np.allclose(additive @ k, ref["additive_retrieval"])
assert np.allclose(half.retrieval_after_write, ref["half_retrieval"])
assert np.allclose(full.retrieval_after_write, ref["full_retrieval"])
assert np.allclose(gated.retrieval_after_write, ref["gated_half_retrieval"])
assert np.allclose(full.retrieval_after_write, v2)  # exact overwrite for a unit key
assert np.allclose(no_write.state, 0.5 * state)

# normalize=True makes a scaled copy of the same key semantically identical.
scaled_key = delta_update(state, 7.0 * k, v2, beta=1.0, normalize=True)
assert np.allclose(scaled_key.state, full.state)
for magnitude in (1e300, 1e-300):
    np.testing.assert_allclose(normalize_key([magnitude, magnitude]), np.ones(2)/np.sqrt(2))
# Targeted overwrite can change another, non-orthogonal key.
other = np.array([1., 1.]) / np.sqrt(2)
assert not np.allclose(full.state @ other, state @ other)

print("additive", additive @ k, "| beta=.5", half.retrieval_after_write, "| beta=1", full.retrieval_after_write)
print('The delta rule, decay gate, and key normalization are verified')

5. Capacity: a compressed state versus an explicit table

Toy additive S = Σ v kᵀ compresses all pairs into a fixed number of coordinates, so non-orthogonal keys interfere. This is not a Mamba benchmark. The control block lets softmax address an explicit key/value table and, at temperature 0.05 (logit multiplier 20), produces near one-hot retrieval. It does not claim mathematically zero error at every finite temperature.

Python
@dataclass(frozen=True)
class RecallCurve:
    """Point-retrieval MSE for additive associative memory at several loads."""

    stored_pairs: tuple[int, ...]
    mse: tuple[float, ...]
    key_width: int
    value_width: int
    seed: int


def associative_recall_curve(
    stored_pairs: Iterable[Integral],
    *,
    key_width: Integral,
    value_width: Integral | None = None,
    seed: Integral = 7,
) -> RecallCurve:
    """Measure interference in ``S = sum(v k^T)`` for normalized random keys."""

    widths = _positive_int("key_width", key_width)
    values = widths if value_width is None else _positive_int("value_width", value_width)
    random_seed = _non_negative_int("seed", seed)
    loads = tuple(_positive_int("stored_pairs item", item) for item in stored_pairs)
    if not loads:
        raise ValueError("stored_pairs must be non-empty")

    rng = np.random.default_rng(random_seed)
    errors: list[float] = []
    for count in loads:
        keys = rng.standard_normal((count, widths))
        keys /= np.linalg.norm(keys, axis=1, keepdims=True)
        targets = rng.standard_normal((count, values))
        state = targets.T @ keys
        retrieved = (state @ keys.T).T
        errors.append(float(np.mean((retrieved - targets) ** 2)))
    return RecallCurve(loads, tuple(errors), widths, values, random_seed)


def stable_softmax_rows(logits: Any) -> np.ndarray:
    """Compute row-wise softmax for a finite non-empty matrix."""

    scores = _finite_array("logits", logits, ndim=2)
    shifted = scores - np.max(scores, axis=1, keepdims=True)
    weights = np.exp(shifted)
    denominator = np.sum(weights, axis=1, keepdims=True)
    if np.any(denominator <= 0.0) or not np.all(np.isfinite(denominator)):
        raise FloatingPointError("softmax denominator must be positive and finite")
    return weights / denominator


def softmax_addressing_mse(
    *,
    stored_pairs: Integral,
    key_width: Integral,
    value_width: Integral | None = None,
    inverse_temperature: Real = 20.0,
    seed: Integral = 107,
) -> float:
    """Return finite-temperature self-retrieval MSE from an explicit key/value table.

    This is a near-one-hot addressing contrast, not a claim that finite-temperature
    softmax is mathematically exact for arbitrary random keys.
    """

    count = _positive_int("stored_pairs", stored_pairs)
    width = _positive_int("key_width", key_width)
    values = width if value_width is None else _positive_int("value_width", value_width)
    temperature = _positive_real("inverse_temperature", inverse_temperature)
    random_seed = _non_negative_int("seed", seed)
    rng = np.random.default_rng(random_seed)
    keys = rng.standard_normal((count, width))
    keys /= np.linalg.norm(keys, axis=1, keepdims=True)
    targets = rng.standard_normal((count, values))
    weights = stable_softmax_rows(temperature * (keys @ keys.T))
    return float(np.mean((weights @ targets - targets) ** 2))


ref = REFERENCE["recall"]
curve = associative_recall_curve(
    ref["stored_pairs"],
    key_width=ref["key_width"],
    value_width=ref["value_width"],
    seed=ref["seed"],
)
assert np.allclose(curve.mse, ref["mse"], rtol=1e-11, atol=1e-12)
assert all(later > earlier for earlier, later in zip(curve.mse[2:], curve.mse[3:]))
softmax_mse = softmax_addressing_mse(
    stored_pairs=128,
    key_width=32,
    value_width=32,
    inverse_temperature=ref["softmax_inverse_temperature"],
    seed=ref["softmax_seed"],
)
assert math.isclose(softmax_mse, ref["softmax_mse_at_128"], rel_tol=1e-10, abs_tol=1e-12)
assert 0.0 < softmax_mse < 1e-6
# The plotted control uses the same keys and values at every load.
paired_rng = np.random.default_rng(ref["seed"])
paired_errors = []
for count in ref["stored_pairs"]:
    keys = paired_rng.standard_normal((count, 32))
    keys /= np.linalg.norm(keys, axis=1, keepdims=True)
    values = paired_rng.standard_normal((count, 32))
    prediction = stable_softmax_rows(20 * keys @ keys.T) @ values
    paired_errors.append(float(np.mean((prediction-values)**2)))
np.testing.assert_allclose(paired_errors, ref["paired_softmax_mse"], rtol=1e-10, atol=1e-15)
print("linear-state MSE:", [round(value, 4) for value in curve.mse])
print(f"finite-temperature softmax MSE at 128 pairs: {softmax_mse:.3e}")
print('Recall interference and finite-temperature softmax are verified')

6. Complete recurrent cache and central-operator arithmetic

Persistent state includes both the main SSM tensor and the causal-convolution tail:

L dinner dstate s+L dconv_channels dconv s.L\,d_{inner}\,d_{state}\,s +L\,d_{conv\_channels}\,d_{conv}\,s.

Full-attention KV is

2LHkvdhsbytes per token.2L H_{kv}d_hs\quad\text{bytes per token}.

We compare attention, Mamba-1-/Mamba-2-like caches, and a 1:8 hybrid. The ratio of 4·N·d_model to 5·d_inner·d_state compares the two attention matmuls with state update plus readout. Use d_model=4096, d_inner=8192, d_state=16 for this separate experiment. SSM coefficients are already available; preparing them, softmax, and other block operations are excluded. The leading count has three operations per coordinate for update and two for readout.

Python
@dataclass(frozen=True)
class SSMStateGeometry:
    """Per-sequence inference-cache geometry for an SSM stack.

    A recurrent Mamba-style layer can retain both the main SSM state and the
    short causal-convolution state. ``conv_channels`` is explicit because
    Mamba-2 may cache x plus B/C channels rather than only ``inner_width``.
    """

    layers: int
    inner_width: int
    state_width: int
    conv_width: int = 0
    conv_channels: int = 0
    bytes_per_value: int = 2

    def __post_init__(self) -> None:
        _positive_int("layers", self.layers)
        _positive_int("inner_width", self.inner_width)
        _positive_int("state_width", self.state_width)
        _positive_int("bytes_per_value", self.bytes_per_value)
        _non_negative_int("conv_width", self.conv_width)
        _non_negative_int("conv_channels", self.conv_channels)
        if (self.conv_width == 0) != (self.conv_channels == 0):
            raise ValueError("conv_width and conv_channels must both be zero or both be positive")

    def ssm_bytes_per_sequence(self, *, batch_size: Integral = 1) -> int:
        """Return the main recurrent SSM-state payload."""

        batch = _positive_int("batch_size", batch_size)
        return batch * self.layers * self.inner_width * self.state_width * self.bytes_per_value

    def conv_bytes_per_sequence(self, *, batch_size: Integral = 1) -> int:
        """Return the short causal-convolution cache payload."""

        batch = _positive_int("batch_size", batch_size)
        return batch * self.layers * self.conv_channels * self.conv_width * self.bytes_per_value

    def bytes_per_sequence(self, *, batch_size: Integral = 1) -> int:
        """Return the complete fixed recurrent cache for one or more sequences."""

        return self.ssm_bytes_per_sequence(batch_size=batch_size) + self.conv_bytes_per_sequence(batch_size=batch_size)


@dataclass(frozen=True)
class KVCacheGeometry:
    """Per-token K/V geometry for a stack of attention layers."""

    layers: int
    kv_heads: int
    head_width: int
    bytes_per_value: int = 2

    def __post_init__(self) -> None:
        _positive_int("layers", self.layers)
        _positive_int("kv_heads", self.kv_heads)
        _positive_int("head_width", self.head_width)
        _positive_int("bytes_per_value", self.bytes_per_value)

    @property
    def bytes_per_token(self) -> int:
        return 2 * self.layers * self.kv_heads * self.head_width * self.bytes_per_value

    def total_bytes(self, context_tokens: Integral, *, batch_size: Integral = 1) -> int:
        tokens = _non_negative_int("context_tokens", context_tokens)
        batch = _positive_int("batch_size", batch_size)
        return batch * tokens * self.bytes_per_token


@dataclass(frozen=True)
class HybridMemoryEstimate:
    """Per-sequence KV and recurrent cache of a hybrid stack."""

    total_layers: int
    attention_layers: int
    ssm_layers: int
    context_tokens: int
    kv_bytes: int
    ssm_bytes: int
    conv_bytes: int
    recurrent_bytes: int
    total_bytes: int


def hybrid_memory_estimate(
    *,
    total_layers: Integral,
    attention_layers: Integral,
    context_tokens: Integral,
    inner_width: Integral,
    state_width: Integral,
    conv_width: Integral,
    conv_channels: Integral,
    kv_heads: Integral,
    head_width: Integral,
    bytes_per_value: Integral = 2,
) -> HybridMemoryEstimate:
    """Return KV plus complete recurrent-cache payload for one hybrid sequence."""

    total = _positive_int("total_layers", total_layers)
    if isinstance(attention_layers, (bool, np.bool_)) or not isinstance(attention_layers, Integral):
        raise TypeError("attention_layers must be an integer")
    attention = int(attention_layers)
    if attention < 0 or attention > total:
        raise ValueError("attention_layers must lie in [0, total_layers]")
    tokens = _non_negative_int("context_tokens", context_tokens)
    inner = _positive_int("inner_width", inner_width)
    state = _positive_int("state_width", state_width)
    conv = _non_negative_int("conv_width", conv_width)
    channels = _non_negative_int("conv_channels", conv_channels)
    if (conv == 0) != (channels == 0):
        raise ValueError("conv_width and conv_channels must both be zero or both be positive")
    heads = _positive_int("kv_heads", kv_heads)
    width = _positive_int("head_width", head_width)
    item_bytes = _positive_int("bytes_per_value", bytes_per_value)
    ssm_layers = total - attention

    kv_bytes = 0 if attention == 0 else KVCacheGeometry(attention, heads, width, item_bytes).total_bytes(tokens)
    if ssm_layers == 0:
        ssm_bytes = conv_bytes = recurrent_bytes = 0
    else:
        geometry = SSMStateGeometry(
            ssm_layers, inner, state, conv, channels, item_bytes
        )
        ssm_bytes = geometry.ssm_bytes_per_sequence()
        conv_bytes = geometry.conv_bytes_per_sequence()
        recurrent_bytes = geometry.bytes_per_sequence()
    return HybridMemoryEstimate(
        total_layers=total,
        attention_layers=attention,
        ssm_layers=ssm_layers,
        context_tokens=tokens,
        kv_bytes=kv_bytes,
        ssm_bytes=ssm_bytes,
        conv_bytes=conv_bytes,
        recurrent_bytes=recurrent_bytes,
        total_bytes=kv_bytes + recurrent_bytes,
    )


@dataclass(frozen=True)
class DecodeFlopEstimate:
    """Arithmetic of the central token-mixing operators only.

    The numbers exclude projections, gating, convolution, FFN, weight traffic,
    communication and kernel efficiency. They must not be read as latency or
    end-to-end model speedups.
    """

    context_tokens: int
    attention_core_flops: int
    ssm_core_flops: int
    ratio: float


def decode_flop_estimate(
    *,
    context_tokens: Integral,
    model_width: Integral,
    inner_width: Integral,
    state_width: Integral,
) -> DecodeFlopEstimate:
    """Return the declared core-only scenario ``4*N*d`` versus ``5*d_inner*d_state``."""

    tokens = _positive_int("context_tokens", context_tokens)
    model = _positive_int("model_width", model_width)
    inner = _positive_int("inner_width", inner_width)
    state = _positive_int("state_width", state_width)
    attention = 4 * tokens * model
    # State update: two multiplies + one add; readout: multiply + add.
    # Leading-order counts omit one readout addition per channel.
    ssm = 5 * inner * state
    return DecodeFlopEstimate(tokens, attention, ssm, attention / ssm)


ref = REFERENCE["memory"]
kv = KVCacheGeometry(layers=32, kv_heads=8, head_width=128, bytes_per_value=2)
mamba1 = SSMStateGeometry(
    layers=32, inner_width=8192, state_width=16,
    conv_width=4, conv_channels=8192, bytes_per_value=2,
)
mamba2 = SSMStateGeometry(
    layers=32, inner_width=8192, state_width=128,
    conv_width=4, conv_channels=8192 + 2 * 128, bytes_per_value=2,
)
assert kv.bytes_per_token == ref["kv_bytes_per_token_full_attention"] == 128 * 2**10
assert mamba1.ssm_bytes_per_sequence() == ref["mamba1"]["ssm_bytes"] == 8 * 2**20
assert mamba1.conv_bytes_per_sequence() == ref["mamba1"]["conv_bytes"] == 2 * 2**20
assert mamba1.bytes_per_sequence() == ref["mamba1"]["total_bytes"] == 10 * 2**20
assert ref["mamba1"]["crossover_tokens_exact"] == 80
assert ref["mamba1"]["crossover_first_integer_token"] == 80  # first >=
assert ref["mamba1"]["first_strictly_larger_token"] == 81
assert kv.total_bytes(80) == mamba1.bytes_per_sequence()
assert kv.total_bytes(81) > mamba1.bytes_per_sequence()
assert mamba2.ssm_bytes_per_sequence() == ref["mamba2_like"]["ssm_bytes"] == 64 * 2**20
assert math.isclose(bytes_to_mib(mamba2.conv_bytes_per_sequence()), 2.0625)
assert mamba2.bytes_per_sequence() == ref["mamba2_like"]["total_bytes"]
assert ref["mamba2_like"]["crossover_tokens_exact"] == 528.5
assert ref["mamba2_like"]["crossover_first_integer_token"] == 529

full_256k = kv.total_bytes(256 * 1024)
hybrid = hybrid_memory_estimate(
    total_layers=32, attention_layers=4, context_tokens=256 * 1024,
    inner_width=8192, state_width=16, conv_width=4, conv_channels=8192,
    kv_heads=8, head_width=128, bytes_per_value=2,
)
assert full_256k == ref["full_attention_256k_bytes"] == 32 * 2**30
assert hybrid.kv_bytes == ref["hybrid_256k"]["kv_bytes"] == 4 * 2**30
assert hybrid.ssm_bytes == ref["hybrid_256k"]["ssm_bytes"] == 7 * 2**20
assert hybrid.conv_bytes == ref["hybrid_256k"]["conv_bytes"] == int(1.75 * 2**20)
assert hybrid.recurrent_bytes == ref["hybrid_256k"]["recurrent_bytes"] == int(8.75 * 2**20)
assert hybrid.total_bytes == ref["hybrid_256k"]["total_bytes"]

assert mamba1.bytes_per_sequence(batch_size=3) == 3 * mamba1.bytes_per_sequence()
assert kv.total_bytes(8192, batch_size=3) == 3 * kv.total_bytes(8192)

flop_ref = REFERENCE["decode_flops"]
flop_rows = [decode_flop_estimate(context_tokens=position, model_width=4096, inner_width=8192, state_width=16) for position in flop_ref["positions"]]
assert [row.attention_core_flops for row in flop_rows] == flop_ref["attention_core_flops"]
assert [row.ssm_core_flops for row in flop_rows] == flop_ref["ssm_core_flops"]
assert np.allclose([row.ratio for row in flop_rows], flop_ref["ratios"], rtol=1e-12)
assert "not end-to-end latency" in flop_ref["scope"]
# Independent arithmetic: two attention dot-product stages; SSM write + read.
for position,row in zip(flop_ref["positions"],flop_rows):
    assert row.attention_core_flops == 2 * (2 * position * 4096)
    update = 8192 * 16 * (2 + 1)
    readout = 8192 * (16 + 16)  # leading estimate of the reduction
    assert row.ssm_core_flops == update + readout
assert np.allclose([row.ratio for row in flop_rows], [25.6,819.2,26214.4])

print(f"KV/token={bytes_to_kib(kv.bytes_per_token):.0f} KiB")
print(f"Mamba-1 cache={bytes_to_mib(mamba1.bytes_per_sequence()):.0f} MiB; crossover=80")
print(f"Mamba-2-like cache={bytes_to_mib(mamba2.bytes_per_sequence()):.4f} MiB; crossover=528.5 / first integer 529")
print(f"256K hybrid={bytes_to_gib(hybrid.total_bytes):.6f} GiB (KV 4 GiB + recurrent 8.75 MiB)")
print("core-only ratios:", [round(row.ratio) for row in flop_rows])
print('SSM/KV/hybrid memory and the FLOP scenario are verified')

7. Negative tests

Negative examples pin the contract: an unstable positive mode, zero Δ, an empty sequence, incompatible shapes, a zero key, a non-unit key when automatic normalization is disabled, an invalid chunk size, and an impossible hybrid configuration.

Python
def expect_failure(exception_type, function, *args, **kwargs):
    """Assert that one invalid educational input fails with the declared exception."""

    try:
        function(*args, **kwargs)
    except exception_type:
        return
    raise AssertionError(f"expected {exception_type.__name__} from {function.__name__}")

expect_failure(ValueError, zoh_discretize_diagonal, [0.1], [[1.0]], 0.1)
expect_failure(ValueError, zoh_discretize_diagonal, [-1.0], [[1.0]], 0.0)
expect_failure(ValueError, scan_diagonal_ssm, [-1.0], [[1.0]], [[1.0]], [], [])
expect_failure(ValueError, scan_diagonal_ssm, [-1.0], [[1.0]], [[1.0]], [[1.0], [2.0]], [0.1])
expect_failure(ValueError, affine_scan_chunked, [[0.9]], [[1.0]], chunk_size=0)
expect_failure(ValueError, linear_attention_sequential, [[1.0]], [[1.0, 2.0]], [[1.0]])
expect_failure(ValueError, normalize_key, [0.0, 0.0])
expect_failure(ValueError, delta_update, np.zeros((2, 2)), [2.0, 0.0], [1.0, 2.0], beta=1.0, normalize=False)
expect_failure(ValueError, stable_softmax_rows, [[0.0, np.nan]])
expect_failure(ValueError, hybrid_memory_estimate, total_layers=8, attention_layers=9, context_tokens=128, inner_width=64, state_width=16, conv_width=4, conv_channels=64, kv_heads=2, head_width=8)
expect_failure(TypeError, KVCacheGeometry, layers=True, kv_heads=8, head_width=128)
print('Negative tests passed')

Exercises

  1. [auto] Find the Δ at which the relative Δ·B error for a=-1 first reaches 10%, then verify it by substitution.

  2. [auto] Implement a parallel-prefix affine scan and compare it with affine_scan_sequential at lengths that are not powers of two.

  3. [auto] With two non-orthogonal keys, measure how overwriting the first key changes retrieval at the second key.

  4. [auto] Derive the crossover for a common batch size B. Show that B cancels because both caches scale proportionally with it, leaving the length threshold unchanged. Verify this for several B with equal-length sequences.

  5. [project] Compare the educational chunked linear attention with a PyTorch version and profile where the transparent NumPy implementation stops being a useful performance model.

The mandatory path is fully offline. It validates equations, shapes, units, and failure modes; it does not replace a fused scan kernel, profiler, or production serving runtime.

Module 8 project-exercise demo — Chunk boundaries and memory updates

Project demonstration

Download .ipynb module_08_project_demo.ipynb · 20.5 KB

Show notebook text

Module 8 project-exercise demo — Chunk boundaries and memory updates

Implementation demo, not a completed full-scale project. The first experiment tests state transfer between chunks; the second rewrites non-orthogonal keys. These are fixed algorithms with no training. All inputs are small, deterministic, offline, and CPU-friendly. The goal is to validate interfaces, metrics, invariants, and report structure before spending GPU or API budget.

Python
from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

plt.rcParams.update({"font.family":"DejaVu Sans", "font.size":12, "axes.titlesize":15, "axes.labelsize":13, "xtick.labelsize":11, "ytick.labelsize":11, "figure.dpi":150})

Demo 1. Carrying state across chunk boundaries

Use 53 affine transitions and a nonzero initial state. Chunk sizes 1, 7, 16, and 64 cover partial final chunks and a chunk larger than the entire sequence. Compare every intermediate state with sequential evaluation. A deliberately incorrect reset control makes the lost cross-boundary information visible.

Python
import math
from dataclasses import dataclass
from numbers import Integral, Real
from typing import Any


def _real(name: str, value: Real) -> float:
    """Return one finite real scalar, rejecting booleans and non-finite values."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real):
        raise TypeError(f"{name} must be a real scalar")
    number = float(value)
    if not math.isfinite(number):
        raise ValueError(f"{name} must be finite")
    return number


def _positive_real(name: str, value: Real) -> float:
    """Return one finite strictly positive scalar."""

    number = _real(name, value)
    if number <= 0.0:
        raise ValueError(f"{name} must be positive")
    return number


def _unit_interval(name: str, value: Real) -> float:
    """Return a finite scalar in the closed unit interval."""

    number = _real(name, value)
    if number < 0.0 or number > 1.0:
        raise ValueError(f"{name} must lie in [0, 1]")
    return number


def _positive_int(name: str, value: Integral) -> int:
    """Return one strictly positive integer, rejecting booleans."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    number = int(value)
    if number <= 0:
        raise ValueError(f"{name} must be positive")
    return number


def _non_negative_int(name: str, value: Integral) -> int:
    """Return one non-negative integer, rejecting booleans."""

    if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral):
        raise TypeError(f"{name} must be an integer")
    number = int(value)
    if number < 0:
        raise ValueError(f"{name} must be non-negative")
    return number


def _finite_array(
    name: str,
    value: Any,
    *,
    ndim: int | None = None,
    non_empty: bool = True,
) -> np.ndarray:
    """Return a float64 array with explicit dimensionality and finiteness checks."""

    array = np.asarray(value, dtype=np.float64)
    if ndim is not None and array.ndim != ndim:
        raise ValueError(f"{name} must have ndim={ndim}, got shape {array.shape}")
    if non_empty and array.size == 0:
        raise ValueError(f"{name} must be non-empty")
    if not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must contain only finite values")
    return array


def bytes_to_kib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary KiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**10


def bytes_to_mib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary MiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**20


def bytes_to_gib(byte_count: Real) -> float:
    """Convert a finite non-negative byte count to binary GiB."""

    number = _real("byte_count", byte_count)
    if number < 0.0:
        raise ValueError("byte_count must be non-negative")
    return number / 2**30



@dataclass(frozen=True)
class AffineTransform:
    """Elementwise affine state transform ``h_out = multiplier*h_in + bias``."""

    multiplier: np.ndarray
    bias: np.ndarray


def compose_affine(later: AffineTransform, earlier: AffineTransform) -> AffineTransform:
    """Compose two diagonal affine transforms as ``later(earlier(h))``."""

    a2 = _finite_array("later.multiplier", later.multiplier, ndim=1)
    b2 = _finite_array("later.bias", later.bias, ndim=1)
    a1 = _finite_array("earlier.multiplier", earlier.multiplier, ndim=1)
    b1 = _finite_array("earlier.bias", earlier.bias, ndim=1)
    if a1.shape != b1.shape or a2.shape != b2.shape or a1.shape != a2.shape:
        raise ValueError("all affine components must share shape [d_state]")
    return AffineTransform(multiplier=a2 * a1, bias=a2 * b1 + b2)


def affine_scan_sequential(
    multipliers: Any,
    biases: Any,
    *,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate ``h_t = a_t*h_{t-1} + b_t`` sequentially.

    ``multipliers`` and ``biases`` must both have shape ``[T, d_state]``. The returned
    array contains post-update states with the same shape.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")
    result = np.empty_like(b)
    for index in range(a.shape[0]):
        state = a[index] * state + b[index]
        result[index] = state
    return result


def affine_scan_chunked(
    multipliers: Any,
    biases: Any,
    *,
    chunk_size: Integral,
    initial_state: Any | None = None,
) -> np.ndarray:
    """Evaluate the same affine recurrence in non-overlapping chunks.

    Inside each chunk, the local zero-state recurrence and cumulative multipliers can be
    computed in parallel by a scan kernel. Only one boundary state crosses each chunk.
    This NumPy version keeps the mechanics transparent; it is not a fused GPU scan.
    """

    a = _finite_array("multipliers", multipliers, ndim=2)
    b = _finite_array("biases", biases, ndim=2)
    if a.shape != b.shape:
        raise ValueError("multipliers and biases must have identical shape")
    chunk = _positive_int("chunk_size", chunk_size)
    state = (
        np.zeros(a.shape[1], dtype=np.float64)
        if initial_state is None
        else _finite_array("initial_state", initial_state, ndim=1).copy()
    )
    if state.shape != (a.shape[1],):
        raise ValueError("initial_state must have shape [d_state]")

    result = np.empty_like(b)
    for start in range(0, a.shape[0], chunk):
        end = min(start + chunk, a.shape[0])
        local_a = a[start:end]
        local_b = b[start:end]

        # Zero-start local states carry all within-chunk contributions.
        zero_state = np.zeros(a.shape[1], dtype=np.float64)
        local_zero = np.empty_like(local_b)
        for offset in range(local_a.shape[0]):
            zero_state = local_a[offset] * zero_state + local_b[offset]
            local_zero[offset] = zero_state

        # The chunk-start state is multiplied by the cumulative transition to each row.
        prefix_multiplier = np.cumprod(local_a, axis=0)
        result[start:end] = local_zero + prefix_multiplier * state
        state = result[end - 1].copy()
    return result


rng = np.random.default_rng(SEED)
a = rng.uniform(.3,.99,(53,6)); b = rng.normal(size=(53,6)); h0=rng.normal(size=6)
expected=affine_scan_sequential(a,b,initial_state=h0)
sizes=[1,7,16,64]
errors=[]
for chunk in sizes:
    actual=affine_scan_chunked(a,b,chunk_size=chunk,initial_state=h0)
    errors.append(np.max(np.abs(actual-expected)))
    np.testing.assert_allclose(actual,expected,rtol=1e-12,atol=1e-12)
# A state reset at each chunk is an intentionally wrong control.
reset=np.concatenate([affine_scan_sequential(a[j:j+16],b[j:j+16]) for j in range(0,53,16)])
assert np.max(np.abs(reset-expected)) > .1
print("chunks:",sizes,"max errors:",errors,"reset error:",np.max(np.abs(reset-expected)))
fig,ax=plt.subplots(figsize=(11,4),layout="constrained")
ax.plot(expected[:,0],color="#5E81AC",label="Эталон" if DEMO_LANGUAGE=="ru" else "Reference")
ax.plot(reset[:,0],"--",color="#D08770",label="Контроль со сбросом" if DEMO_LANGUAGE=="ru" else "Reset control")
for x in (16,32,48):ax.axvline(x,color="#D8DEE9")
ax.set_xlabel("Шаг" if DEMO_LANGUAGE=="ru" else "Step")
ax.set_ylabel("h[0]");ax.legend();plt.show()

Small discrepancies between valid chunk sizes reflect floating-point ordering. The reset control loses initial and carried state. Regrouping transitions preserves the mathematical operator; reordering them generally does not.

Demo 2. Rewriting non-orthogonal keys

Four unit keys receive repeated updates in the same 3×8 memory. All methods see identical inputs. Additive writes accumulate values; delta updates correct the current direction; gated delta first multiplies the entire state by 0.9. At β=1, the current key must read back exactly, while other key directions can still change.

Python
"""Minimal M8 operators; used by the validated reference and embedded in notebooks."""
import numpy as np

def zoh_diagonal(a, b, delta):
    """Diagonal A: [n], B: [n,p], positive scalar step; caller validates shapes."""
    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)
    transition = np.exp(delta * a)
    factor = np.full_like(a, float(delta))
    np.divide(np.expm1(delta * a), a, out=factor, where=a != 0)
    return transition, factor[:, None] * b

def delta_step(state, key, value, beta, decay=1.0):
    """State [dv,dk], unit key [dk], value [dv]; decay before correction."""
    base = decay * state
    prediction = base @ key
    return base + beta * np.outer(value - prediction, key)

def compose_affine_pair(later, earlier):
    """Elementwise transforms; returns later(earlier(h)), not the reverse."""
    a2, b2 = later
    a1, b1 = earlier
    return a2 * a1, a2 * b1 + b2



@dataclass(frozen=True)
class DeltaUpdateResult:
    """State and retrieval trace for one normalized-key delta-rule update."""

    state: np.ndarray
    normalized_key: np.ndarray
    retrieval_before_write: np.ndarray
    retrieval_after_write: np.ndarray


def normalize_key(key: Any) -> np.ndarray:
    """Return a unit-L2 key, rejecting the zero vector."""

    vector = _finite_array("key", key, ndim=1)
    scale = float(np.max(np.abs(vector)))
    if scale == 0.0:
        raise ValueError("key must have non-zero norm")
    scaled = vector / scale
    return scaled / np.linalg.norm(scaled)


def delta_update(
    state: Any,
    key: Any,
    value: Any,
    *,
    beta: Real,
    decay: Real = 1.0,
    normalize: bool = True,
) -> DeltaUpdateResult:
    """Apply a gated delta-rule write to associative state ``[d_value, d_key]``.

    With a unit key, the update is

    ``base = decay * state``
    ``new = base + beta * outer(value - base@key, key)``.

    Thus ``beta=1`` overwrites retrieval at that key exactly, while ``beta=0`` only
    applies the global decay. If ``normalize=False``, the caller must supply a unit key;
    the explicit check prevents a misleading "exact overwrite" claim for arbitrary norm.
    """

    matrix = _finite_array("state", state, ndim=2)
    raw_key = _finite_array("key", key, ndim=1)
    target = _finite_array("value", value, ndim=1)
    if matrix.shape != (target.size, raw_key.size):
        raise ValueError("state must have shape [d_value, d_key]")
    write_rate = _unit_interval("beta", beta)
    forget_rate = _unit_interval("decay", decay)
    unit_key = normalize_key(raw_key) if normalize else raw_key.copy()
    if not normalize and not np.isclose(np.linalg.norm(unit_key), 1.0, rtol=0.0, atol=1e-10):
        raise ValueError("key must have unit norm when normalize=False")

    base = forget_rate * matrix
    before = base @ unit_key
    updated = delta_step(matrix, unit_key, target, write_rate, forget_rate)
    after = updated @ unit_key
    return DeltaUpdateResult(updated, unit_key, before, after)


# Repeated writes to the same four unit keys; compare on identical inputs.
rng=np.random.default_rng(SEED+1)
keys=rng.normal(size=(4,8));keys/=np.linalg.norm(keys,axis=1,keepdims=True)
values=rng.normal(size=(48,3));ids=np.arange(48)%4
states={name:np.zeros((3,8)) for name in ("additive","delta","gated")}
latest=np.zeros((4,3));history={name:[] for name in states}
for t,(i,v) in enumerate(zip(ids,values)):
    k=keys[i];latest[i]=v
    states["additive"] += np.outer(v,k)
    states["delta"]=delta_update(states["delta"],k,v,beta=1).state
    states["gated"]=delta_update(states["gated"],k,v,beta=1,decay=.9).state
    np.testing.assert_allclose(states["delta"]@k,v,atol=1e-12)
    np.testing.assert_allclose(states["gated"]@k,v,atol=1e-12)
    for name,state in states.items():
        history[name].append(np.mean((keys@state.T-latest)**2))
fig,ax=plt.subplots(figsize=(11,4.8),layout="constrained")
labels = {"additive":"Аддитивная запись", "delta":"Дельта-правило", "gated":"Дельта-правило с затуханием"}
for (name,errors),color in zip(history.items(),("#5E81AC","#A3BE8C","#D08770")):
    ax.plot(errors,label=labels[name] if DEMO_LANGUAGE=="ru" else name,color=color)
ax.set_xlabel("Запись" if DEMO_LANGUAGE=="ru" else "Write")
ax.set_ylabel("MSE относительно последних значений" if DEMO_LANGUAGE=="ru" else "MSE relative to latest values")
ax.legend();plt.show()
print({name:float(errors[-1]) for name,errors in history.items()})
print("Exact at the current key does not imply exact at all previous keys.")

The curve measures error against each key’s latest value, taking zero as the target before its first write. Separate assertions check the current key exactly. Errors at other keys expose interference, while decay changes the tradeoff between recent and older associations. A single final error does not establish a universal winner.

Scaling to the full project

Implement a parallel prefix scan, then repeat the memory study across dimensions and seeds. Separate current-key error from error over all previously written keys.

Minimum report

  1. Data, model, and code revisions, plus the random seed.

  2. The primary metric, plus quality, cost, latency, and memory where applicable.

  3. A baseline, a matched budget, and at least one negative control.

  4. Failed runs and the limits of the supported conclusions.