Module 5. Modern Attention Mechanisms

# Install only missing dependencies; ordinary Colab already has NumPy.
import importlib.util
import subprocess
import sys

missing = [name for name in ("numpy",) if importlib.util.find_spec(name) is None]
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])
print('Environment ready.')

import gc
import hashlib
import json
import math
import tracemalloc
from typing import Literal

import numpy as np

_REFERENCE_JSON = '{\n  "schema_version": 2,\n  "context_tokens": 131072,\n  "kv_rows": {\n    "gpt_oss_gqa_all_full_equivalent": 73728,\n    "gpt_oss_mha_hypothetical": 589824,\n    "gpt_oss_mqa_hypothetical": 9216\n  },\n  "gpt_oss_hybrid": {\n    "layers": 36,\n    "full_layers": 18,\n    "sliding_layers": 18,\n    "window": 128,\n    "per_layer_bytes_per_token": 2048,\n    "full_bytes": 4831838208,\n    "sliding_bytes": 4718592,\n    "hybrid_total_bytes": 4836556800,\n    "all_full_bytes": 9663676416\n  },\n  "online_softmax_trace": {\n    "l1": 1.3678794411714423,\n    "n1": 1.7357588823428847,\n    "alpha": 0.36787944117144233,\n    "l2": 1.5852997230319539,\n    "n2": 3.966890002140263,\n    "output": 2.502296533903012\n  },\n  "mla": {\n    "bytes_per_token": 70272,\n    "mha_equivalent_bytes_per_token": 3997696,\n    "raw_element_ratio": 56.888888888888886\n  },\n  "score_matrix": {\n    "n_4096_fp32_bytes": 67108864,\n    "n_8192_fp32_bytes": 268435456,\n    "tile_128_fp32_bytes": 65536\n  }\n}\n'
_REFERENCE_SHA256 = '720d34f7da83628df36153a92218644fa6e643db4a345ba5b7b41d9c0206fc34'
assert hashlib.sha256(_REFERENCE_JSON.encode("utf-8")).hexdigest() == _REFERENCE_SHA256
REFERENCE = json.loads(_REFERENCE_JSON)
_CONFIG_JSON = '{"num_hidden_layers": 36, "num_experts": 128, "experts_per_token": 4, "vocab_size": 201088, "hidden_size": 2880, "intermediate_size": 2880, "swiglu_limit": 7.0, "head_dim": 64, "num_attention_heads": 64, "num_key_value_heads": 8, "sliding_window": 128, "initial_context_length": 4096, "rope_theta": 150000, "rope_scaling_factor": 32.0, "rope_ntk_alpha": 1, "rope_ntk_beta": 32}'
assert hashlib.sha256(_CONFIG_JSON.encode('utf-8')).hexdigest() == '063aaa98e012cac330d58d18471700f0b0eb90bb15c58fb227809d45090f3e65'
CONFIG = json.loads(_CONFIG_JSON)
assert (CONFIG['num_hidden_layers'], CONFIG['num_attention_heads'],
        CONFIG['num_key_value_heads'], CONFIG['head_dim'], CONFIG['sliding_window']) == (36, 64, 8, 64, 128)

rng = np.random.default_rng(4)
np.set_printoptions(precision=5, suppress=True)
print('Module-5 reference metrics unpacked and verified', "· NumPy", np.__version__)

EmptyRowPolicy = Literal["error", "zero"]

def make_mask(n: int, window: int | None = None, n_sink: int = 0) -> np.ndarray:
    """Return an additive causal mask of shape [n, n].

    Allowed entries are 0 and forbidden entries are -inf. `window` keeps only the
    latest `window` causal keys; `n_sink` preserves initial real positions as StreamingLLM-style anchors.
    """
    if isinstance(n, bool) or not isinstance(n, int) or n <= 0:
        raise ValueError("n must be positive")
    if window is not None and (isinstance(window, bool) or not isinstance(window, int) or window <= 0):
        raise ValueError("window must be positive")
    if isinstance(n_sink, bool) or not isinstance(n_sink, int) or not 0 <= n_sink <= n:
        raise ValueError("n_sink must satisfy 0 <= n_sink <= n")
    if n_sink and window is None:
        raise ValueError("n_sink is meaningful only together with a finite window")

    query_pos = np.arange(n)[:, None]
    key_pos = np.arange(n)[None, :]
    allowed = key_pos <= query_pos
    if window is not None:
        allowed &= (query_pos - key_pos) < window
        allowed |= (key_pos < n_sink) & (key_pos <= query_pos)
    return np.where(allowed, 0.0, -np.inf)

def sdpa(
    Q: np.ndarray,
    K: np.ndarray,
    V: np.ndarray,
    mask: np.ndarray | None = None,
    *,
    sink: float | np.ndarray | None = None,
    empty_row: EmptyRowPolicy = "error",
    return_weights: bool = False,
) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
    """Compute transparent scaled dot-product attention.

    Shapes:
        Q [..., T_q, d_k], K [..., T_k, d_k], V [..., T_k, d_v].
        `mask` must broadcast to [..., T_q, T_k]. A finite entry is an additive
        bias; -inf forbids the key. A per-head sink should be reshaped so it
        broadcasts to [..., T_q, 1].

    `empty_row="error"` rejects a row with no key and no sink. `"zero"` returns
    zero weights/output for that row, which is useful for explicit padding rows.
    """
    if empty_row not in {"error", "zero"}:
        raise ValueError("empty_row must be 'error' or 'zero'")

    Q, K, V = (np.asarray(x) for x in (Q, K, V))
    if min(Q.ndim, K.ndim, V.ndim) < 2 or any(0 in x.shape for x in (Q, K, V)):
        raise ValueError("Q, K, and V must each have at least two dimensions")
    if Q.shape[-1] != K.shape[-1]:
        raise ValueError("Q and K must have the same d_k")
    if K.shape[-2] != V.shape[-2]:
        raise ValueError("K and V must have the same T_k")
    try:
        np.broadcast_shapes(Q.shape[:-2], K.shape[:-2], V.shape[:-2])
    except ValueError as exc:
        raise ValueError("Q/K/V batch dimensions are not broadcast-compatible") from exc
    if not all(np.isfinite(x).all() for x in (Q, K, V)):
        raise ValueError("Q, K, and V must contain only finite values")

    # Keep float32/float64 inputs in their precision; integer examples become float64.
    dtype = np.result_type(Q.dtype, K.dtype, V.dtype, np.float32)
    Q, K, V = (x.astype(dtype, copy=False) for x in (Q, K, V))
    d_k = Q.shape[-1]
    scale = np.asarray(1.0 / math.sqrt(d_k), dtype=dtype)
    with np.errstate(over="ignore", invalid="ignore"):
        scores = (Q @ np.swapaxes(K, -1, -2)) * scale
    if not np.isfinite(scores).all():
        raise FloatingPointError("Q @ K.T overflowed; rescale the inputs")

    if mask is not None:
        mask = np.asarray(mask, dtype=dtype)
        if np.isnan(mask).any() or np.isposinf(mask).any():
            raise ValueError("mask may contain finite biases and -inf, but not NaN/+inf")
        try:
            mask = np.broadcast_to(mask, scores.shape)
        except ValueError as exc:
            raise ValueError(f"mask shape {mask.shape} does not broadcast to {scores.shape}") from exc
        finite_bias = np.isfinite(mask)
        with np.errstate(over="ignore", invalid="ignore"):
            scores = scores + mask
        if np.any(finite_bias & ~np.isfinite(scores)):
            raise FloatingPointError(
                "adding the finite mask bias overflowed; reduce the bias magnitude"
            )

    has_token = np.isfinite(scores).any(axis=-1, keepdims=True)
    token_max = np.max(scores, axis=-1, keepdims=True)

    sink_logits = None
    if sink is not None:
        sink_logits = np.asarray(sink, dtype=dtype)
        if not np.isfinite(sink_logits).all():
            raise ValueError("sink logits must be finite")
        try:
            sink_logits = np.broadcast_to(sink_logits, token_max.shape)
        except ValueError as exc:
            raise ValueError(
                f"sink shape {sink_logits.shape} does not broadcast to {token_max.shape}"
            ) from exc
        row_max = np.maximum(token_max, sink_logits)
    else:
        if empty_row == "error" and not has_token.all():
            bad = np.argwhere(~has_token.squeeze(-1))
            raise ValueError('Attention contains fully masked rows' + f": {bad[:8].tolist()}")
        # A zero-policy row gets a harmless max of 0; all its exponentials stay 0.
        row_max = np.where(has_token, token_max, np.zeros_like(token_max))

    token_exp = np.where(
        np.isfinite(scores), np.exp(scores - row_max), np.zeros_like(scores)
    )
    denominator = token_exp.sum(axis=-1, keepdims=True)
    if sink_logits is not None:
        denominator += np.exp(sink_logits - row_max)

    weights = np.divide(
        token_exp,
        denominator,
        out=np.zeros_like(token_exp),
        where=denominator > 0,
    )
    output = weights @ V
    return (output, weights) if return_weights else output

n, d = 24, 16
M_causal = make_mask(n)
M_window = make_mask(n, window=8)
M_sink = make_mask(n, window=8, n_sink=2)
assert (M_causal == 0).sum() == n * (n + 1) // 2
assert np.isneginf(M_window[23, 10]) and M_window[23, 16] == 0
assert M_sink[23, 0] == 0 and M_sink[23, 1] == 0 and np.isneginf(M_sink[23, 2])

Q = rng.standard_normal((n, d))
K = rng.standard_normal((n, d))
V = rng.standard_normal((n, d))
Z, A = sdpa(Q, K, V, M_causal, return_weights=True)
assert Z.shape == (n, d) and A.shape == (n, n)
assert np.isfinite(Z).all() and np.allclose(A.sum(-1), 1.0)

# Deliberately remove all support from one query row.
M_empty = M_causal.copy()
M_empty[5, :] = -np.inf
try:
    sdpa(Q, K, V, M_empty)
    raise AssertionError("the default policy must reject a fully masked row")
except ValueError as exc:
    assert str(exc)

Z_zero, A_zero = sdpa(Q, K, V, M_empty, empty_row="zero", return_weights=True)
assert np.all(A_zero[5] == 0) and np.all(Z_zero[5] == 0)
assert np.allclose(A_zero[np.arange(n) != 5].sum(-1), 1.0)
print('SDPA shapes, three masks, and the empty-row policy verified')

def gqa(
    Q: np.ndarray,
    K: np.ndarray,
    V: np.ndarray,
    n_heads: int,
    n_kv_heads: int,
    mask: np.ndarray | None = None,
) -> np.ndarray:
    """Educational grouped-query attention for 2D projected sequences.

    Q [T, H*d_k], K [T, H_kv*d_k], V [T, H_kv*d_v].
    The output has shape [T, H*d_v].
    """
    Q, K, V = (np.asarray(x) for x in (Q, K, V))
    if any(x.ndim != 2 or 0 in x.shape for x in (Q, K, V)):
        raise ValueError("GQA expects 2D projected sequences")
    if not (Q.shape[0] == K.shape[0] == V.shape[0]):
        raise ValueError("Q, K, and V must have the same sequence length")
    if any(isinstance(h, bool) or not isinstance(h, int) or h <= 0 for h in (n_heads, n_kv_heads)) or n_heads % n_kv_heads:
        raise ValueError("n_heads must be a positive multiple of n_kv_heads")
    if Q.shape[1] % n_heads:
        raise ValueError("Q width must be divisible by n_heads")
    if K.shape[1] % n_kv_heads or V.shape[1] % n_kv_heads:
        raise ValueError("K/V widths must be divisible by n_kv_heads")

    T = Q.shape[0]
    d_k = Q.shape[1] // n_heads
    if K.shape[1] // n_kv_heads != d_k:
        raise ValueError("Q and K head dimensions must match")
    d_v = V.shape[1] // n_kv_heads

    q = Q.reshape(T, n_heads, d_k).transpose(1, 0, 2)
    k = K.reshape(T, n_kv_heads, d_k).transpose(1, 0, 2)
    v = V.reshape(T, n_kv_heads, d_v).transpose(1, 0, 2)
    group_size = n_heads // n_kv_heads
    # repeat preserves groups [kv0, kv0, ..., kv1, kv1, ...].
    k = np.repeat(k, group_size, axis=0)
    v = np.repeat(v, group_size, axis=0)
    z = sdpa(q, k, v, mask)
    return z.transpose(1, 0, 2).reshape(T, n_heads * d_v)

H, H_kv, d_k, T = 8, 2, 16, 32
Q_flat = rng.standard_normal((T, H * d_k))
K_flat = rng.standard_normal((T, H_kv * d_k))
V_flat = rng.standard_normal((T, H_kv * d_k))
mask = make_mask(T)

out_gqa = gqa(Q_flat, K_flat, V_flat, H, H_kv, mask)
k_grouped = K_flat.reshape(T, H_kv, d_k)
v_grouped = V_flat.reshape(T, H_kv, d_k)
K_mha = np.repeat(k_grouped, H // H_kv, axis=1).reshape(T, H * d_k)
V_mha = np.repeat(v_grouped, H // H_kv, axis=1).reshape(T, H * d_k)
out_mha = gqa(Q_flat, K_mha, V_mha, H, H, mask)
assert np.allclose(out_gqa, out_mha, atol=1e-12)

# tile interleaves groups [kv0, kv1, kv0, kv1, ...] and changes the mapping.
K_bad = np.tile(k_grouped, (1, H // H_kv, 1)).reshape(T, H * d_k)
V_bad = np.tile(v_grouped, (1, H // H_kv, 1)).reshape(T, H * d_k)
out_bad = gqa(Q_flat, K_bad, V_bad, H, H, mask)
assert not np.allclose(out_gqa, out_bad, atol=1e-3)
print('The GQA invariant and wrong-order negative test passed', f"· max|Δ|={np.max(np.abs(out_gqa - out_mha)):.2e}")

def online_softmax_dot(
    logit_tiles: list[list[float]],
    value_tiles: list[list[float]],
) -> tuple[float, list[tuple[float, float, float]]]:
    """Stream scalar values through online softmax and return the update trace."""
    if len(logit_tiles) != len(value_tiles) or not logit_tiles:
        raise ValueError("logit_tiles and value_tiles must have the same non-zero length")
    m, normalizer, numerator = -np.inf, 0.0, 0.0
    trace: list[tuple[float, float, float]] = []
    for logits, values in zip(logit_tiles, value_tiles):
        if len(logits) != len(values) or not logits:
            raise ValueError("each logit tile must match a non-empty value tile")
        if not np.isfinite(logits).all() or not np.isfinite(values).all():
            raise ValueError("online-softmax inputs must be finite")
        m_new = max(m, max(logits))
        alpha = math.exp(m - m_new) if np.isfinite(m) else 0.0
        weights = [math.exp(x - m_new) for x in logits]
        normalizer = alpha * normalizer + sum(weights)
        numerator = alpha * numerator + sum(w * v for w, v in zip(weights, values))
        m = m_new
        trace.append((m, normalizer, numerator))
    return numerator / normalizer, trace

def tiled_attention(
    Q: np.ndarray,
    K: np.ndarray,
    V: np.ndarray,
    mask: np.ndarray | None = None,
    *,
    tile_size: int = 128,
    sink: float | None = None,
    empty_row: EmptyRowPolicy = "error",
    causal: bool = False,
    window: int | None = None,
    query_start: int = 0,
) -> np.ndarray:
    """Streaming 2D attention that never materializes a T_q x T_k score matrix."""
    Q, K, V = (np.asarray(x) for x in (Q, K, V))
    if any(x.ndim != 2 or 0 in x.shape for x in (Q, K, V)):
        raise ValueError("tiled_attention expects Q/K/V to be 2D")
    if Q.shape[1] != K.shape[1] or K.shape[0] != V.shape[0]:
        raise ValueError("incompatible Q/K/V shapes")
    if isinstance(tile_size, bool) or not isinstance(tile_size, int) or tile_size <= 0:
        raise ValueError("tile_size must be positive")
    if empty_row not in {"error", "zero"}:
        raise ValueError("empty_row must be 'error' or 'zero'")
    if not all(np.isfinite(x).all() for x in (Q, K, V)):
        raise ValueError("Q/K/V must be finite")
    if sink is not None and not np.isfinite(sink):
        raise ValueError("sink must be finite")

    # Accumulate at least in float32. This mirrors the direct reference and avoids
    # performing the dot product itself in float16 when learners try mixed precision.
    dtype = np.result_type(Q.dtype, K.dtype, V.dtype, np.float32)
    Q, K, V = (x.astype(dtype, copy=False) for x in (Q, K, V))
    T_q, d_k = Q.shape
    T_k, d_v = V.shape
    if window is not None and (not causal or isinstance(window, bool) or not isinstance(window, int) or window <= 0):
        raise ValueError("window requires causal=True and a positive integer")
    if isinstance(query_start, bool) or not isinstance(query_start, int) or query_start < 0:
        raise ValueError("query_start must be a nonnegative integer")
    if mask is not None:
        mask = np.asarray(mask, dtype=dtype)
        if mask.shape != (T_q, T_k):
            raise ValueError(f"mask must have shape {(T_q, T_k)}, got {mask.shape}")
        if np.isnan(mask).any() or np.isposinf(mask).any():
            raise ValueError("mask may not contain NaN/+inf")

    scale = np.asarray(1.0 / math.sqrt(d_k), dtype=dtype)
    out = np.zeros((T_q, d_v), dtype=dtype)
    for qi in range(T_q):
        # A sink starts with denominator mass 1 and zero value contribution.
        m = float(sink) if sink is not None else -np.inf
        normalizer = 1.0 if sink is not None else 0.0
        acc = np.zeros(d_v, dtype=out.dtype)

        for start in range(0, T_k, tile_size):
            stop = min(start + tile_size, T_k)
            with np.errstate(over="ignore", invalid="ignore"):
                logits = (Q[qi] @ K[start:stop].T) * scale
            if not np.isfinite(logits).all():
                raise FloatingPointError("Q @ K.T overflowed; rescale the inputs")
            mask_tile = 0.0 if mask is None else mask[qi, start:stop]
            if causal:
                distance = query_start + qi - np.arange(start, stop)
                allowed = distance >= 0
                if window is not None:
                    allowed &= distance < window
                mask_tile = mask_tile + np.where(allowed, 0.0, -np.inf)
            finite_bias = np.isfinite(mask_tile)
            with np.errstate(over="ignore", invalid="ignore"):
                logits = logits + mask_tile
            if np.any(finite_bias & ~np.isfinite(logits)):
                raise FloatingPointError(
                    "adding the finite mask bias overflowed; reduce the bias magnitude"
                )
            finite = np.isfinite(logits)
            if not finite.any():
                continue

            tile_max = float(logits[finite].max())
            m_new = max(m, tile_max)
            alpha = math.exp(m - m_new) if np.isfinite(m) else 0.0
            weights = np.where(finite, np.exp(logits - m_new), np.zeros_like(logits))
            normalizer = alpha * normalizer + float(weights.sum())
            acc = alpha * acc + weights @ V[start:stop]
            m = m_new

        if normalizer == 0.0:
            if empty_row == "error":
                raise ValueError('fully masked row in tiled_attention' + f": {qi}")
            continue
        out[qi] = acc / normalizer
    return out

logits = [2.0, 1.0, 3.0, 0.5]
values = [1.0, 2.0, 3.0, 4.0]
output, trace = online_softmax_dot([logits[:2], logits[2:]], [values[:2], values[2:]])
(_, l1, n1), (_, l2, n2) = trace
ref_trace = REFERENCE["online_softmax_trace"]
for got, key in [(l1, "l1"), (n1, "n1"), (l2, "l2"), (n2, "n2"), (output, "output")]:
    assert abs(got - ref_trace[key]) < 1e-12, (key, got, ref_trace[key])
direct_logits = np.asarray(logits)
direct_weights = np.exp(direct_logits - direct_logits.max())
direct_weights /= direct_weights.sum()
assert abs(output - direct_weights @ np.asarray(values)) < 1e-12

n, d = 40, 16
Q = rng.standard_normal((n, d))
K = rng.standard_normal((n, d))
V = rng.standard_normal((n, d))
for current_mask in (make_mask(n), make_mask(n, 8), make_mask(n, 8, 2)):
    reference = sdpa(Q, K, V, current_mask)
    streamed = tiled_attention(Q, K, V, current_mask, tile_size=7)
    assert np.allclose(reference, streamed, atol=1e-12)

empty_mask = np.full((n, n), -np.inf)
zero_ref = sdpa(Q, K, V, empty_mask, empty_row="zero")
zero_stream = tiled_attention(Q, K, V, empty_mask, tile_size=7, empty_row="zero")
assert np.all(zero_ref == 0) and np.all(zero_stream == 0)

# Float16 inputs are intentionally accumulated in float32 in both references.
Q16, K16, V16 = (x.astype(np.float16) for x in (Q, K, V))
ref16 = sdpa(Q16, K16, V16, make_mask(n))
stream16 = tiled_attention(Q16, K16, V16, make_mask(n), tile_size=7)
assert ref16.dtype == np.float32 and stream16.dtype == np.float32
assert np.allclose(ref16, stream16, atol=2e-6)

# A very large but finite additive bias may overflow only after it is added to
# a finite score. Both references reject that state instead of treating +inf
# as a forbidden position or allowing NaN to propagate.
f32_max = np.finfo(np.float32).max
q_over = np.asarray([[math.sqrt(0.75 * f32_max)]], dtype=np.float32)
v_over = np.ones((1, 1), dtype=np.float32)
for sign in (1.0, -1.0):
    k_over = sign * q_over
    bias_over = np.asarray([[sign * 0.75 * f32_max]], dtype=np.float32)
    for attention_fn in (sdpa, tiled_attention):
        try:
            attention_fn(q_over, k_over, v_over, bias_over)
            raise AssertionError("finite mask-bias overflow must be rejected")
        except FloatingPointError as exc:
            assert "mask bias" in str(exc)
print('Online softmax and tiled attention match direct computation', f"· O={output:.6f}")

KiB, MiB, GiB = 2**10, 2**20, 2**30

def kv_bytes_per_token(
    layers: int,
    kv_heads: int,
    head_dim: int,
    bytes_per_value: int = 2,
) -> int:
    """Return K+V cache bytes per generated token across all layers."""
    values = (layers, kv_heads, head_dim, bytes_per_value)
    if any(isinstance(x, bool) or not isinstance(x, int) or x <= 0 for x in values):
        raise ValueError("all KV-cache dimensions must be positive integers")
    return 2 * layers * kv_heads * head_dim * bytes_per_value

context = REFERENCE["context_tokens"]
rows = {
    "gpt_oss_gqa_all_full_equivalent": kv_bytes_per_token(36, 8, 64),
    "gpt_oss_mha_hypothetical": kv_bytes_per_token(36, 64, 64),
    "gpt_oss_mqa_hypothetical": kv_bytes_per_token(36, 1, 64),
}
assert rows == REFERENCE["kv_rows"]
labels = {
    "gpt_oss_gqa_all_full_equivalent": 'gpt-oss GQA all-full equivalent (36L, 8 KV, 64)',
    "gpt_oss_mha_hypothetical": 'gpt-oss hypothetical MHA (64 KV)',
    "gpt_oss_mqa_hypothetical": 'gpt-oss hypothetical MQA (1 KV)',
}
for key, byte_count in rows.items():
    print(f"{labels[key]:42s} {byte_count / KiB:7.1f} KiB/token · {byte_count * context / GiB:7.2f} GiB @131072")

mla = (512 + 64) * 61 * 2
mha_equivalent = 2 * 61 * 128 * 128 * 2
assert mla == REFERENCE["mla"]["bytes_per_token"]
assert mha_equivalent == REFERENCE["mla"]["mha_equivalent_bytes_per_token"]
assert abs(mha_equivalent / mla - REFERENCE["mla"]["raw_element_ratio"]) < 1e-12

per_layer_token = 2 * 8 * 64 * 2
full = 18 * context * per_layer_token
sliding = 18 * 128 * per_layer_token
hybrid = REFERENCE["gpt_oss_hybrid"]
assert full == hybrid["full_bytes"]
assert sliding == hybrid["sliding_bytes"]
assert full + sliding == hybrid["hybrid_total_bytes"]
all_full = rows["gpt_oss_gqa_all_full_equivalent"] * context
assert all_full == hybrid["all_full_bytes"]
print(f"KV payload with window-capped local caches: {(full + sliding) / GiB:.6f} GiB vs all-full {all_full / GiB:.3f} GiB · ×{all_full / (full + sliding):.3f}")

llama70 = kv_bytes_per_token(80, 8, 128)
qwen7 = kv_bytes_per_token(28, 4, 128)
assert llama70 == 320 * KiB and llama70 * context == 40 * GiB
assert qwen7 == 56 * KiB and qwen7 * context == 7 * GiB
print(f"E1: Llama-3.3-70B = 320 KiB/token → 40 GiB; Qwen2.5-7B = 56 KiB/token → 7 GiB")
print('KV calculator: every frozen Module-5 number reproduced')

for sink_logit, percent_ref in [(0.0, 0.78), (2.0, 5.46), (5.0, 53.69)]:
    share = math.exp(sink_logit) / (128 + math.exp(sink_logit)) * 100
    assert abs(share - percent_ref) < 0.01

n, d = 20, 8
Q = rng.standard_normal((n, d))
K = rng.standard_normal((n, d))
V = rng.standard_normal((n, d))
mask = make_mask(n)
sink_logit = 1.3
out_sink, token_weights = sdpa(Q, K, V, mask, sink=sink_logit, return_weights=True)

scores = Q @ K.T / np.sqrt(d) + mask
extended_scores = np.concatenate([scores, np.full((n, 1), sink_logit)], axis=1)
row_max = extended_scores.max(axis=-1, keepdims=True)
extended_weights = np.exp(extended_scores - row_max)
extended_weights /= extended_weights.sum(axis=-1, keepdims=True)
out_column = extended_weights[:, :-1] @ V  # the extra sink value is exactly zero
assert np.allclose(out_sink, out_column, atol=1e-12)
assert np.allclose(token_weights, extended_weights[:, :-1], atol=1e-12)
assert np.all(token_weights.sum(axis=-1) < 1.0)
# A common shift must include the sink. Shifting tokens alone changes the result.
shifted = sdpa(Q, K, V, mask + 2.0, sink=sink_logit + 2.0)
np.testing.assert_allclose(shifted, out_sink, atol=1e-12)
assert not np.allclose(sdpa(Q, K, V, mask + 2.0, sink=sink_logit), out_sink)

# With no visible token, all probability mass goes to the zero-valued sink.
all_masked = np.full((n, n), -np.inf)
sink_only, sink_only_weights = sdpa(
    Q, K, V, all_masked, sink=sink_logit, return_weights=True
)
assert np.all(sink_only == 0) and np.all(sink_only_weights == 0)
streamed_sink_only = tiled_attention(Q, K, V, all_masked, sink=sink_logit)
assert np.all(streamed_sink_only == 0)
print('Sink mass, column equivalence, and the empty-row case passed', f"· max|Δ|={np.max(np.abs(out_sink - out_column)):.2e}")

def naive_forward(Q: np.ndarray, K: np.ndarray, V: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """Materialize scores/probabilities for a 2D memory baseline.

    This path deliberately keeps the O(T_q*T_k) arrays, but still rejects a fully
    masked row instead of silently producing NaN.
    """
    Q, K, V, mask = map(np.asarray, (Q, K, V, mask))
    if any(x.ndim != 2 for x in (Q, K, V)):
        raise ValueError("naive_forward expects 2D Q/K/V")
    if Q.shape[1] != K.shape[1] or K.shape[0] != V.shape[0]:
        raise ValueError("incompatible Q/K/V shapes")
    if mask.shape != (Q.shape[0], K.shape[0]):
        raise ValueError("mask must have shape [T_q, T_k]")

    dtype = np.result_type(Q.dtype, K.dtype, V.dtype, np.float32)
    Q, K, V, mask = (x.astype(dtype, copy=False) for x in (Q, K, V, mask))
    d_k = Q.shape[-1]
    scale = np.asarray(1.0 / math.sqrt(d_k), dtype=dtype)
    scores = (Q @ K.T) * scale + mask
    if not np.isfinite(scores).any(axis=-1).all():
        raise ValueError("naive_forward received a fully masked row")
    row_max = scores.max(axis=-1, keepdims=True)
    probabilities = np.where(
        np.isfinite(scores), np.exp(scores - row_max), np.zeros_like(scores)
    )
    probabilities /= probabilities.sum(axis=-1, keepdims=True)
    return probabilities @ V


def peak_python_mib(function, *args, **kwargs) -> float:
    """Measure Python-visible allocations; native NumPy visibility may vary."""
    gc.collect()
    tracemalloc.start()
    try:
        function(*args, **kwargs)
        _, peak = tracemalloc.get_traced_memory()
    finally:
        tracemalloc.stop()
    return peak / MiB


# First establish correctness independently of any memory measurement.
n_check, d_check = 256, 32
Qc = rng.standard_normal((n_check, d_check))
Kc = rng.standard_normal((n_check, d_check))
Vc = rng.standard_normal((n_check, d_check))
Mc = make_mask(n_check)
assert np.allclose(
    naive_forward(Qc, Kc, Vc, Mc),
    tiled_attention(Qc, Kc, Vc, Mc, tile_size=64),
    atol=1e-12,
)

tile_size = 128
for n in (1024, 2048):
    itemsize = np.dtype(np.float64).itemsize
    naive_temporaries = 2 * n * n * itemsize  # scores + probabilities
    tiled_state = (2 * tile_size + 32) * itemsize  # logits + weights + d_v accumulator
    assert naive_temporaries / tiled_state > 100

    Qm = rng.standard_normal((n, 32))
    Km = rng.standard_normal((n, 32))
    Vm = rng.standard_normal((n, 32))
    observed_naive = peak_python_mib(lambda: naive_forward(Qm, Km, Vm, make_mask(n)))
    observed_tiled = peak_python_mib(tiled_attention, Qm, Km, Vm, tile_size=tile_size, causal=True)
    print(
        f"n={n} · selected buffers (not total peak): dense={naive_temporaries / MiB:.1f} MiB, "
        f"tiled≈{tiled_state / KiB:.1f} KiB · tracemalloc: "
        f"{observed_naive:.1f} vs {observed_tiled:.1f} MiB"
    )
print('Memory: selected dense O(n²) buffers versus O(tile + d_v) scratch state verified')

rng_cache = np.random.default_rng(505)
Qp = rng_cache.normal(size=(9, 6))
Kp = rng_cache.normal(size=(9, 6))
Vp = rng_cache.normal(size=(9, 4))
for window in (None, 3):
    complete = sdpa(Qp, Kp, Vp, make_mask(9, window))
    for prefix, count in ((0, 9), (5, 1), (4, 3)):
        end = prefix + count
        distance = np.arange(prefix, end)[:, None] - np.arange(end)[None, :]
        allowed = distance >= 0
        if window is not None:
            allowed &= distance < window
        mask = np.where(allowed, 0.0, -np.inf)
        direct = sdpa(Qp[prefix:end], Kp[:end], Vp[:end], mask)
        stream = tiled_attention(Qp[prefix:end], Kp[:end], Vp[:end],
            causal=True, query_start=prefix, window=window, tile_size=2)
        np.testing.assert_allclose(direct, complete[prefix:end], atol=1e-12)
        np.testing.assert_allclose(stream, direct, atol=1e-12)
wrong = sdpa(Qp[5:6], Kp[:6], Vp[:6], np.where(np.tri(1, 6), 0.0, -np.inf))
assert not np.allclose(wrong, sdpa(Qp[5:6], Kp[:6], Vp[:6]))
# Empty tiles must be skipped without forming -inf - (-inf).
late_mask = np.full((1, 9), -np.inf)
late_mask[0, -1] = 0
np.testing.assert_allclose(tiled_attention(Qp[:1], Kp, Vp, late_mask, tile_size=2), Vp[-1:])
print('Cached prefix, chunked decode, window, empty tiles: PASS')

rng_mla = np.random.default_rng(506)
H, Tq, Tk, dc, dk, dv, dr = 2, 3, 7, 4, 6, 5, 2
C = rng_mla.normal(size=(Tk, dc))
A = rng_mla.normal(size=(H, dk, dc))
B = rng_mla.normal(size=(H, dv, dc))
q = rng_mla.normal(size=(H, Tq, dk))
qr = rng_mla.normal(size=(H, Tq, dr))  # already rotated query branch
kr = rng_mla.normal(size=(Tk, dr))     # shared, already rotated key branch
Wout = rng_mla.normal(size=(H, dv, 8))
explicit, absorbed = [], []
for h in range(H):
    expanded_k, expanded_v = C @ A[h].T, C @ B[h].T
    scores = (q[h] @ expanded_k.T + qr[h] @ kr.T) / np.sqrt(dk + dr)
    latent_scores = ((q[h] @ A[h]) @ C.T + qr[h] @ kr.T) / np.sqrt(dk + dr)
    np.testing.assert_allclose(scores, latent_scores, atol=1e-12)
    p = np.exp(scores - scores.max(-1, keepdims=True))
    p /= p.sum(-1, keepdims=True)
    explicit.append((p @ expanded_v) @ Wout[h])
    absorbed.append((p @ C) @ (B[h].T @ Wout[h]))
np.testing.assert_allclose(sum(explicit), sum(absorbed), atol=1e-12)
# A rectangular A makes R @ A and A @ R different *shapes*, not a commutator.
# The correct obstruction: A.T @ R(n-m).T @ q changes with key position n.
def R(angle):
    return np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
A_small = np.array([[1., 2., 0.], [0., 1., 3.]])
q_small = np.array([1., .5])
c_small = np.array([.2, -.4, .7])
m, n = .3, 1.1
rotated_score = (R(m) @ q_small) @ (R(n) @ A_small @ c_small)
effective_q = A_small.T @ R(n-m).T @ q_small
np.testing.assert_allclose(rotated_score, effective_q @ c_small, atol=1e-12)
assert not np.allclose(effective_q, A_small.T @ R(.8-m).T @ q_small)
print('MLA key/value/output absorption and position dependence: PASS')

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_sdpa_explicit(
    q: np.ndarray,
    k: np.ndarray,
    v: np.ndarray,
    *,
    allowed_mask: np.ndarray | None = None,
) -> np.ndarray:
    """Minimal visible SDPA algorithm; ``True`` in the mask means allowed."""
    output_dtype = np.result_type(q.dtype, k.dtype, v.dtype)
    if output_dtype.kind not in "fc":
        output_dtype = np.dtype(np.float64)
    work_dtype = np.result_type(output_dtype, np.float32)
    q, k, v = (x.astype(work_dtype, copy=False) for x in (q, k, v))
    scale = 1.0 / math.sqrt(q.shape[-1])
    scores = np.matmul(q, np.swapaxes(k, -1, -2)) * scale

    if allowed_mask is None:
        row_has_key = np.ones_like(scores[..., :1], dtype=bool)
    else:
        mask = np.broadcast_to(np.asarray(allowed_mask, dtype=bool), scores.shape)
        row_has_key = mask.any(axis=-1, keepdims=True)
        scores = np.where(mask, scores, -np.inf)

    safe_scores = np.where(row_has_key, scores, 0.0)
    shifted = safe_scores - safe_scores.max(axis=-1, keepdims=True)
    exp = np.exp(shifted)
    weights = exp / exp.sum(axis=-1, keepdims=True)
    weights = np.where(row_has_key, weights, 0.0)
    return np.matmul(weights, v).astype(output_dtype, copy=False)

def torch_sdpa_explicit(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    allowed_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Line-for-line tensor counterpart of :func:`numpy_sdpa_explicit`."""
    output_dtype = q.dtype
    work_dtype = torch.float32 if q.dtype in (torch.float16, torch.bfloat16) else q.dtype
    q, k, v = (x.to(work_dtype) for x in (q, k, v))
    scale = 1.0 / math.sqrt(q.shape[-1])
    scores = q @ k.transpose(-2, -1) * scale

    if allowed_mask is None:
        row_has_key = torch.ones_like(scores[..., :1], dtype=torch.bool)
    else:
        mask = torch.broadcast_to(allowed_mask.to(torch.bool), scores.shape)
        row_has_key = mask.any(dim=-1, keepdim=True)
        scores = scores.masked_fill(~mask, -torch.inf)

    safe_scores = torch.where(row_has_key, scores, torch.zeros_like(scores))
    weights = torch.softmax(safe_scores, dim=-1)
    weights = torch.where(row_has_key, weights, torch.zeros_like(weights))
    return (weights @ v).to(output_dtype)

def torch_sdpa_optimized(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    allowed_mask: torch.Tensor | None = None,
    dropout_p: float = 0.0,
) -> torch.Tensor:
    """Public PyTorch SDPA API; the runtime may choose an optimized backend."""
    return F.scaled_dot_product_attention(
        q,
        k,
        v,
        attn_mask=allowed_mask,
        dropout_p=dropout_p,
    )

rng = np.random.default_rng(4)
q_np = rng.standard_normal((1, 2, 4, 8))
k_np = rng.standard_normal((1, 2, 4, 8))
v_np = rng.standard_normal((1, 2, 4, 6))
mask_np = np.tril(np.ones((4, 4), dtype=bool))

np_out = numpy_sdpa_explicit(q_np, k_np, v_np, allowed_mask=mask_np)
q = torch.tensor(q_np, dtype=torch.float64, requires_grad=True)
k = torch.tensor(k_np, dtype=torch.float64, requires_grad=True)
v = torch.tensor(v_np, dtype=torch.float64, requires_grad=True)
mask = torch.tensor(mask_np)
explicit_out = torch_sdpa_explicit(q, k, v, allowed_mask=mask)
api_out = torch_sdpa_optimized(q, k, v, allowed_mask=mask, dropout_p=0.0)

torch.testing.assert_close(explicit_out, torch.from_numpy(np_out), rtol=1e-10, atol=1e-10)
torch.testing.assert_close(api_out, explicit_out, rtol=1e-10, atol=1e-10)

mask_np[0] = False
np_empty = numpy_sdpa_explicit(q_np, k_np, v_np, allowed_mask=mask_np)
t_empty = torch_sdpa_explicit(q, k, v, allowed_mask=torch.tensor(mask_np))
api_empty = torch_sdpa_optimized(q, k, v, allowed_mask=torch.tensor(mask_np), dropout_p=0.0)
np.testing.assert_array_equal(np_empty[..., 0, :], 0.0)
torch.testing.assert_close(t_empty[..., 0, :], torch.zeros_like(t_empty[..., 0, :]))
torch.testing.assert_close(api_empty, t_empty, rtol=1e-10, atol=1e-10)

grads_explicit = torch.autograd.grad(explicit_out.square().sum(), (q, k, v), retain_graph=True)
grads_api = torch.autograd.grad(api_out.square().sum(), (q, k, v), retain_graph=True)
for ge, ga in zip(grads_explicit, grads_api):
    torch.testing.assert_close(ge, ga, rtol=1e-9, atol=1e-10)
explicit_out.sum().backward()
assert all(t.grad is not None and torch.isfinite(t.grad).all() for t in (q, k, v))
print("B04 explicit NumPy / explicit PyTorch / optimized API: PASS")

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

def numpy_sdpa_explicit(
    q: np.ndarray,
    k: np.ndarray,
    v: np.ndarray,
    *,
    allowed_mask: np.ndarray | None = None,
) -> np.ndarray:
    """Minimal visible SDPA algorithm; ``True`` in the mask means allowed."""
    output_dtype = np.result_type(q.dtype, k.dtype, v.dtype)
    if output_dtype.kind not in "fc":
        output_dtype = np.dtype(np.float64)
    work_dtype = np.result_type(output_dtype, np.float32)
    q, k, v = (x.astype(work_dtype, copy=False) for x in (q, k, v))
    scale = 1.0 / math.sqrt(q.shape[-1])
    scores = np.matmul(q, np.swapaxes(k, -1, -2)) * scale

    if allowed_mask is None:
        row_has_key = np.ones_like(scores[..., :1], dtype=bool)
    else:
        mask = np.broadcast_to(np.asarray(allowed_mask, dtype=bool), scores.shape)
        row_has_key = mask.any(axis=-1, keepdims=True)
        scores = np.where(mask, scores, -np.inf)

    safe_scores = np.where(row_has_key, scores, 0.0)
    shifted = safe_scores - safe_scores.max(axis=-1, keepdims=True)
    exp = np.exp(shifted)
    weights = exp / exp.sum(axis=-1, keepdims=True)
    weights = np.where(row_has_key, weights, 0.0)
    return np.matmul(weights, v).astype(output_dtype, copy=False)

def torch_sdpa_explicit(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    allowed_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Line-for-line tensor counterpart of :func:`numpy_sdpa_explicit`."""
    output_dtype = q.dtype
    work_dtype = torch.float32 if q.dtype in (torch.float16, torch.bfloat16) else q.dtype
    q, k, v = (x.to(work_dtype) for x in (q, k, v))
    scale = 1.0 / math.sqrt(q.shape[-1])
    scores = q @ k.transpose(-2, -1) * scale

    if allowed_mask is None:
        row_has_key = torch.ones_like(scores[..., :1], dtype=torch.bool)
    else:
        mask = torch.broadcast_to(allowed_mask.to(torch.bool), scores.shape)
        row_has_key = mask.any(dim=-1, keepdim=True)
        scores = scores.masked_fill(~mask, -torch.inf)

    safe_scores = torch.where(row_has_key, scores, torch.zeros_like(scores))
    weights = torch.softmax(safe_scores, dim=-1)
    weights = torch.where(row_has_key, weights, torch.zeros_like(weights))
    return (weights @ v).to(output_dtype)

def numpy_gqa_explicit(
    q: np.ndarray,
    k: np.ndarray,
    v: np.ndarray,
    *,
    allowed_mask: np.ndarray | None = None,
) -> np.ndarray:
    """Visible GQA: repeat K/V heads, then apply the explicit SDPA reference."""
    if q.shape[1] % k.shape[1]:
        raise ValueError("Hq must be divisible by Hkv")
    group = q.shape[1] // k.shape[1]
    k_repeated = np.repeat(k, group, axis=1)
    v_repeated = np.repeat(v, group, axis=1)
    return numpy_sdpa_explicit(q, k_repeated, v_repeated, allowed_mask=allowed_mask)

def torch_gqa_explicit(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    allowed_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Visible tensor GQA: repeat K/V heads, then use explicit tensor SDPA."""
    if q.shape[1] % k.shape[1]:
        raise ValueError("Hq must be divisible by Hkv")
    group = q.shape[1] // k.shape[1]
    k_repeated = k.repeat_interleave(group, dim=1)
    v_repeated = v.repeat_interleave(group, dim=1)
    return torch_sdpa_explicit(q, k_repeated, v_repeated, allowed_mask=allowed_mask)

def torch_gqa_optimized(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    *,
    allowed_mask: torch.Tensor | None = None,
    dropout_p: float = 0.0,
) -> torch.Tensor:
    """Public SDPA GQA path; availability depends on PyTorch backend support."""
    if q.shape[1] % k.shape[1]:
        raise ValueError("Hq must be divisible by Hkv")
    return F.scaled_dot_product_attention(
        q,
        k,
        v,
        attn_mask=allowed_mask,
        dropout_p=dropout_p,
        enable_gqa=True,
    )

rng = np.random.default_rng(5)
q_np = rng.standard_normal((1, 4, 3, 8))
k_np = rng.standard_normal((1, 2, 3, 8))
v_np = rng.standard_normal((1, 2, 3, 7))
np_out = numpy_gqa_explicit(q_np, k_np, v_np)
q, k, v = (torch.tensor(a, dtype=torch.float64, requires_grad=True) for a in (q_np, k_np, v_np))
explicit = torch_gqa_explicit(q, k, v)
api = torch_gqa_optimized(q, k, v, dropout_p=0.0)
torch.testing.assert_close(explicit, torch.from_numpy(np_out), rtol=1e-10, atol=1e-10)
torch.testing.assert_close(api, explicit, rtol=1e-10, atol=1e-10)
for ge, ga in zip(torch.autograd.grad(explicit.square().sum(), (q,k,v)),
                  torch.autograd.grad(api.square().sum(), (q,k,v))):
    torch.testing.assert_close(ge, ga, rtol=1e-9, atol=1e-10)
print("B05 explicit repetition / enable_gqa: PASS")

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'

lengths = np.array([1024, 2048, 4096, 8192, 16384])
heads, width, tile, bytes_per = 16, 64, 128, 4
output = heads * lengths * width * bytes_per
dense = output + 2 * heads * lengths**2 * bytes_per
stream = output + heads * (tile**2 + tile * width + 2 * tile) * bytes_per
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 12, "axes.spines.top": False, "axes.spines.right": False})
fig, ax = plt.subplots(figsize=(11, 5))
ax.loglog(lengths, dense / 2**20, "o-", color="#5E81AC", label="output + full scores/probabilities")
ax.loglog(lengths, stream / 2**20, "o-", color="#A3BE8C", label="output + selected tile buffers")
ax.set(xlabel="sequence length", ylabel="analytical MiB", title="Selected arrays, not measured GPU peak memory")
ax.legend(); ax.grid(alpha=.25); plt.show()
assert np.all(dense > stream)
print("Excludes Q/K/V, allocator overhead, backward, and external masks. Tile buffers depend on kernel scheduling.")

rng = np.random.default_rng(SEED)
trials, length, dim, index_rank, topk = 200, 2048, 32, 8, 64
recall = hits = full_hits = 0
relative_errors, retained_mass = [], []
def softmax_row(x):
    e = np.exp(x - x.max())
    return e / e.sum()
for _ in range(trials):
    keys = rng.normal(size=(length, dim))
    values = rng.normal(size=(length, 6))
    target = rng.integers(length)
    query = keys[target] + 0.05 * rng.normal(size=dim)
    index_scores = keys[:, :index_rank] @ query[:index_rank]
    shortlist = np.argpartition(index_scores, -topk)[-topk:]
    scores = keys @ query / np.sqrt(dim)
    recall += int(target in shortlist)
    hits += int(shortlist[np.argmax(scores[shortlist])] == target)
    full_hits += int(np.argmax(scores) == target)
    p_full, p_short = softmax_row(scores), softmax_row(scores[shortlist])
    out_full, out_short = p_full @ values, p_short @ values[shortlist]
    retained_mass.append(p_full[shortlist].sum())
    relative_errors.append(np.linalg.norm(out_short-out_full) / max(np.linalg.norm(out_full), 1e-12))
score_ratio = (length * index_rank + topk * dim) / (length * dim)
assert score_ratio == .28125
assert 0 <= hits <= recall <= trials
print(f"shortlist recall={recall/trials:.3f}; final retrieval={hits/trials:.3f}; full retrieval={full_hits/trials:.3f}")
print(f"mean retained probability={np.mean(retained_mass):.3f}; mean relative output error={np.mean(relative_errors):.3f}")
print(f"score multiplications ratio={score_ratio:.5f}; final scoring alone={topk/length:.5f}")
print("The ratio excludes top-k, learned projections, softmax, value aggregation and data movement; it is not a speedup.")