Decoder-Only Language Model

"""A decoder-only LM, transcribed the way the definition reads.

    log p(x) = sum_i log p(x_i | x_1..x_{i-1})

The chain rule factorizes the joint EXACTLY, so the likelihood is a sum of
terms one forward pass already computes - no bound, no invertibility
constraint, no extra computation. That is the cleanest exact likelihood in
this whole reference.

The cost is paid elsewhere: the factorization imposes an ORDER, so generation
is sequential. This version makes that cost visible in the worst possible way
- it recomputes the entire prefix for every new token, which is CUBIC in
sequence length. The fix is the key-value cache in the fast stage, and seeing
the naive version first is the point.

Plain loops, no libraries.
"""

import math
import random

SEED = 41


def softmax(scores):
    """Subtract the max before exponentiating, or a large logit overflows."""
    peak = max(scores)
    exponentials = [math.exp(score - peak) for score in scores]
    total = sum(exponentials)
    return [value / total for value in exponentials]


def causal_attention(queries, keys, values, head_dim):
    """Attention with a CAUSAL mask: position i sees 1..i and nothing later.

    The mask is what makes the chain-rule factorization exact. Without it
    position i would see the token it is supposed to predict, and the loss
    would go to zero while the model learned nothing - the classic and
    completely silent failure of this architecture.
    """
    tokens = len(queries)
    scale = 1.0 / math.sqrt(head_dim)
    output = []

    for i in range(tokens):
        # Only positions up to and including i. Everything later is masked
        # to negative infinity, which the softmax turns into exactly zero.
        scores = []
        for j in range(tokens):
            if j > i:
                scores.append(-math.inf)
                continue
            dot = sum(queries[i][d] * keys[j][d] for d in range(head_dim))
            scores.append(dot * scale)

        weights = softmax(scores)
        mixed = [0.0] * head_dim
        for j in range(i + 1):
            for d in range(head_dim):
                mixed[d] += weights[j] * values[j][d]
        output.append(mixed)

    return output


def linear(vectors, weight, bias):
    out_dim = len(weight[0])
    return [
        [bias[j] + sum(v * weight[k][j] for k, v in enumerate(vector)) for j in range(out_dim)]
        for vector in vectors
    ]


def layer_norm(vectors, gain, shift, epsilon=1e-5):
    output = []
    for vector in vectors:
        mean = sum(vector) / len(vector)
        variance = sum((value - mean) ** 2 for value in vector) / len(vector)
        denominator = math.sqrt(variance + epsilon)
        output.append([
            (value - mean) / denominator * gain[i] + shift[i]
            for i, value in enumerate(vector)
        ])
    return output


def gelu(value):
    inner = math.sqrt(2.0 / math.pi) * (value + 0.044715 * value ** 3)
    return 0.5 * value * (1.0 + math.tanh(inner))


def block(tokens, params, head_dim):
    """One pre-norm decoder block: causal attention, then an MLP."""
    normed = layer_norm(tokens, params['ln1_gain'], params['ln1_shift'])
    queries = linear(normed, params['wq'], params['bq'])
    keys = linear(normed, params['wk'], params['bk'])
    values = linear(normed, params['wv'], params['bv'])

    attended = causal_attention(queries, keys, values, head_dim)
    projected = linear(attended, params['wo'], params['bo'])
    tokens = [[a + b for a, b in zip(t, p)] for t, p in zip(tokens, projected)]

    normed = layer_norm(tokens, params['ln2_gain'], params['ln2_shift'])
    hidden = [[gelu(v) for v in row] for row in linear(normed, params['w1'], params['b1'])]
    out = linear(hidden, params['w2'], params['b2'])
    return [[a + b for a, b in zip(t, o)] for t, o in zip(tokens, out)]


def forward(token_ids, params, head_dim, depth):
    """Embed, add positions, run the blocks, project to the vocabulary.

    Note that this computes logits for EVERY position at once. Because the
    causal mask guarantees position i never sees beyond i, all n predictions
    are valid simultaneously - which is why training parallelizes completely
    across positions and generation does not.
    """
    tokens = [
        [e + p for e, p in zip(params['embedding'][token], params['position'][index])]
        for index, token in enumerate(token_ids)
    ]

    for layer in range(depth):
        tokens = block(tokens, params['blocks'][layer], head_dim)

    tokens = layer_norm(tokens, params['final_gain'], params['final_shift'])
    return linear(tokens, params['unembedding'], params['unembedding_bias'])


def training_loss(token_ids, params, head_dim, depth):
    """Cross-entropy against the NEXT token at every position.

    Teacher forcing: every position is predicted from the TRUE prefix. At
    generation time each position is predicted from the model's own output
    instead, and that gap is real - errors compound at inference in a way this
    loss never measures.
    """
    logits = forward(token_ids[:-1], params, head_dim, depth)

    total = 0.0
    for position, row in enumerate(logits):
        probabilities = softmax(row)
        target = token_ids[position + 1]
        total -= math.log(max(probabilities[target], 1e-12))

    return total / len(logits)


def exact_log_likelihood(token_ids, params, head_dim, depth):
    """The chain rule, summed. An EXACT log-density with no approximation.

    Worth contrasting: a VAE bounds this quantity, a normalizing flow buys it
    with an invertible architecture, a diffusion model recovers it from an ODE
    solve. Here it is a sum of terms the forward pass already produced.
    """
    logits = forward(token_ids[:-1], params, head_dim, depth)
    return sum(
        math.log(max(softmax(row)[token_ids[position + 1]], 1e-12))
        for position, row in enumerate(logits)
    )


def perplexity(token_ids, params, head_dim, depth):
    """exp of the mean negative log-likelihood. Nothing more than that."""
    log_likelihood = exact_log_likelihood(token_ids, params, head_dim, depth)
    return math.exp(-log_likelihood / (len(token_ids) - 1))


def generate(prompt, params, head_dim, depth, count, temperature, rng):
    """Sequential generation, recomputing EVERYTHING every step.

    This is the cost that matters. Each new token re-runs the full forward
    pass over the whole sequence so far, and attention inside it is quadratic,
    which makes generating n tokens CUBIC in n.

    The keys and values for positions already processed CANNOT CHANGE - the
    causal mask guarantees it - so recomputing them is pure waste. Caching
    them is the single largest algorithmic win available on this model, and it
    is what the fast stage does.
    """
    tokens = list(prompt)

    for _ in range(count):
        logits = forward(tokens, params, head_dim, depth)
        last = logits[-1]

        if temperature <= 0.0:
            tokens.append(max(range(len(last)), key=last.__getitem__))
            continue

        scaled = [value / temperature for value in last]
        probabilities = softmax(scaled)

        threshold = rng.random()
        cumulative = 0.0
        for token, probability in enumerate(probabilities):
            cumulative += probability
            if cumulative >= threshold:
                tokens.append(token)
                break

    return tokens

"""The same model, with the cache as a type and the context limit checked.

Two things change. The key-value cache becomes an explicit object with its own
memory accounting, because it is the model's real memory footprint at serving
time and it is what actually runs out rather than the weights. And the context
window becomes a checked limit, because truncation is otherwise silent and
removes the START of the prompt — which is where the instructions are.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass, field
from typing import NamedTuple, Sequence


class ContextOverflow(ValueError):
    """Raised when a sequence exceeds the context window.

    Its own type because the alternative is silent truncation from the FRONT,
    which removes the system prompt and the instructions while leaving the
    most recent tokens — so the model behaves as though it was never told what
    to do, and nothing is logged.
    """


class CacheExhausted(RuntimeError):
    """Raised when the key-value cache cannot hold another token.

    Its own type because the cache — not the weights — is what actually runs
    out of memory in production. It grows linearly in batch times context,
    which means capacity is a function of traffic shape rather than of the
    model.
    """


class CausalityViolation(AssertionError):
    """Raised when a position can attend to a later one.

    Its own type because the failure is catastrophic and completely silent:
    the loss goes to near zero because the model can see the answer, and
    nothing else indicates anything is wrong.
    """


class ShapeMismatch(ValueError):
    """Raised on a dimension violation instead of computing nonsense."""


@dataclass(frozen=True)
class ModelShape:
    """Derived once, so no call site recomputes head width or cache size."""

    depth: int
    model_dim: int
    num_heads: int
    context_length: int
    vocabulary: int

    def __post_init__(self) -> None:
        if self.model_dim % self.num_heads:
            raise ShapeMismatch(
                f'width {self.model_dim} does not divide into {self.num_heads} heads'
            )

    @property
    def head_dim(self) -> int:
        return self.model_dim // self.num_heads

    @property
    def scale(self) -> float:
        """1/sqrt(d_k). Omitting it saturates the softmax and stalls training."""
        return 1.0 / math.sqrt(self.head_dim)

    def cache_bytes(self, batch: int, bytes_per_value: int = 2) -> int:
        """The number that decides serving capacity.

        Two tensors — keys and values — per layer, per position, per batch
        element. Linear in batch times context, which is why capacity depends
        on traffic shape rather than on the model, and why this is what runs
        out rather than the weights.
        """
        return 2 * self.depth * batch * self.context_length * self.model_dim * bytes_per_value


class KeyValueCache:
    """Keys and values for positions already processed.

    The reason this is correct: the causal mask guarantees position i never
    attends to anything later, so the keys and values for earlier positions
    CANNOT CHANGE when a new token arrives. Recomputing them is pure waste,
    and caching turns generation from cubic in sequence length into quadratic.

    Preallocated to the context length rather than grown, because growing a
    buffer per token during generation is exactly the allocation pattern to
    avoid in the hot path.
    """

    __slots__ = ('_keys', '_values', '_length', '_capacity', '_head_dim')

    def __init__(self, shape: ModelShape) -> None:
        self._capacity = shape.context_length
        self._head_dim = shape.head_dim
        self._keys: list[list[float]] = [[] for _ in range(shape.context_length)]
        self._values: list[list[float]] = [[] for _ in range(shape.context_length)]
        self._length = 0

    def append(self, key: Sequence[float], value: Sequence[float]) -> None:
        if self._length >= self._capacity:
            raise CacheExhausted(
                f'the cache is full at {self._capacity} positions; generation cannot '
                'continue without evicting or restarting'
            )
        self._keys[self._length] = list(key)
        self._values[self._length] = list(value)
        self._length += 1

    @property
    def length(self) -> int:
        return self._length

    def keys(self) -> list[list[float]]:
        return self._keys[:self._length]

    def values(self) -> list[list[float]]:
        return self._values[:self._length]

    def reset(self) -> None:
        """Reuse the allocation across requests rather than reallocating."""
        self._length = 0


@dataclass(frozen=True)
class SamplingConfig:
    """Inference-time settings that change the output more than the model does.

    Frozen and explicit because that is what makes them dangerous: two
    deployments of one checkpoint at different temperatures produce materially
    different behaviour, and a result reported without them is not comparable
    to anything.
    """

    temperature: float = 1.0
    top_p: float = 1.0
    top_k: int = 0
    seed: int | None = None

    def __post_init__(self) -> None:
        if self.temperature < 0.0:
            raise ShapeMismatch('temperature cannot be negative')
        if not 0.0 < self.top_p <= 1.0:
            raise ShapeMismatch('top_p must lie in (0, 1]')
        if self.top_k < 0:
            raise ShapeMismatch('top_k cannot be negative')

    @property
    def deterministic(self) -> bool:
        """Temperature zero is greedy decoding, and needs no seed."""
        return self.temperature == 0.0


class TokenScore(NamedTuple):
    """A token AND its log-probability.

    The log-probability comes back because it is free — the forward pass
    computed it — and because per-token likelihood is what makes this model
    usable as an anomaly detector and as a scorer, not only as a generator.
    """

    token: int
    log_probability: float


def stable_log_softmax(logits: Sequence[float]) -> list[float]:
    """log-softmax without forming the softmax first.

    Taking the logarithm of a softmax underflows for a confident distribution,
    which is exactly the regime a trained language model lives in — so this is
    correctness rather than polish.
    """
    peak = max(logits)
    shifted = [value - peak for value in logits]
    total = math.log(math.fsum(math.exp(value) for value in shifted))
    return [value - total for value in shifted]


def assert_causal(attention_weights: Sequence[Sequence[float]]) -> None:
    """Verify no position attends to a later one.

    Worth having as a test: if the mask is wrong the loss collapses toward
    zero because the model can see the token it is predicting, and absolutely
    nothing else indicates a problem. This is the cheapest possible guard
    against the most catastrophic silent failure in the architecture.
    """
    for i, row in enumerate(attention_weights):
        for j, weight in enumerate(row):
            if j > i and abs(weight) > 1e-9:
                raise CausalityViolation(
                    f'position {i} attends to {j} with weight {weight}; the causal '
                    'mask is wrong and the loss will collapse to near zero'
                )


def fit_context(
    tokens: Sequence[int], shape: ModelShape, allow_truncation: bool = False
) -> list[int]:
    """Guard clause for the failure that is otherwise silent.

    Truncation removes the FRONT of the sequence, which is where the system
    prompt and instructions live. The model then behaves as though it was
    never told what to do, and nothing is logged — so the default here is to
    refuse rather than to truncate quietly.
    """
    if len(tokens) <= shape.context_length:
        return list(tokens)
    if not allow_truncation:
        raise ContextOverflow(
            f'{len(tokens)} tokens exceed the {shape.context_length}-token context; '
            'truncating would remove the instructions at the front'
        )
    return list(tokens[-shape.context_length:])


def sample_token(
    logits: Sequence[float], config: SamplingConfig, rng: random.Random
) -> TokenScore:
    """Temperature, then nucleus, then top-k — in that order, deliberately.

    The order matters and is easy to get wrong: temperature rescales the
    distribution, so applying it AFTER truncation changes which tokens were
    eligible. Scaling first and truncating second is the convention every
    implementation follows, and mixing the order gives subtly different output.
    """
    if config.deterministic:
        token = max(range(len(logits)), key=logits.__getitem__)
        return TokenScore(token=token, log_probability=stable_log_softmax(logits)[token])

    scaled = [value / config.temperature for value in logits]
    log_probabilities = stable_log_softmax(scaled)
    probabilities = [math.exp(value) for value in log_probabilities]

    order = sorted(range(len(probabilities)), key=lambda i: -probabilities[i])

    if config.top_k:
        order = order[:config.top_k]

    if config.top_p < 1.0:
        cumulative = 0.0
        kept: list[int] = []
        for token in order:
            kept.append(token)
            cumulative += probabilities[token]
            if cumulative >= config.top_p:
                break
        order = kept

    mass = math.fsum(probabilities[token] for token in order)
    threshold = rng.random() * mass
    cumulative = 0.0
    for token in order:
        cumulative += probabilities[token]
        if cumulative >= threshold:
            return TokenScore(token=token, log_probability=log_probabilities[token])

    return TokenScore(token=order[-1], log_probability=log_probabilities[order[-1]])


def exact_log_likelihood(
    logits_per_position: Sequence[Sequence[float]], targets: Sequence[int]
) -> float:
    """The chain rule, summed. An EXACT log-density with no approximation.

    Worth contrasting: a VAE bounds this quantity, a normalizing flow buys it
    with an invertible architecture, a diffusion model recovers it from an ODE
    solve. Here it is a sum of terms the forward pass already produced.
    """
    if len(logits_per_position) != len(targets):
        raise ShapeMismatch('one logit row per target is required')

    return math.fsum(
        stable_log_softmax(row)[target]
        for row, target in zip(logits_per_position, targets)
    )


@dataclass
class DecodeStats:
    """Prefill and decode counted separately.

    They are compute-bound and memory-bound respectively, so a single
    throughput number averages two entirely different problems into one
    meaningless figure — which is the standard mistake when profiling this
    model.
    """

    prefill_tokens: int = 0
    decode_tokens: int = 0
    cache_positions: int = 0
    weight_reads: int = 0

    def record_prefill(self, tokens: int) -> None:
        self.prefill_tokens += tokens
        self.cache_positions += tokens
        self.weight_reads += 1

    def record_decode(self, batch: int = 1) -> None:
        """Each decode step reads EVERY parameter to produce one token per
        batch element. That is why decode is memory-bound, and why batching
        helps so much: the weight read is shared across the batch."""
        self.decode_tokens += batch
        self.cache_positions += batch
        self.weight_reads += 1

    @property
    def tokens_per_weight_read(self) -> float:
        """Arithmetic intensity, roughly. Prefill is high and decode is near
        one, which is exactly why they need different optimizations."""
        total = self.prefill_tokens + self.decode_tokens
        return total / self.weight_reads if self.weight_reads else 0.0

"""The key-value cache, and why decode is memory-bound.

This is the most consequential optimization in the reference, and it is worth
stating precisely why it is correct. The causal mask guarantees position i
never attends to anything after i. So when a new token arrives, the keys and
values for every earlier position CANNOT CHANGE. Recomputing them is not an
approximation being avoided - it is arithmetic whose result is already known.

    naive generation:   O(n^3)   full quadratic forward pass per token
    with the cache:     O(n^2)   one new row of attention per token

Then the second observation, which is the one people miss. With the cache in
place, producing one token requires reading EVERY parameter and doing one
multiply-accumulate each. Arithmetic intensity is about one - so decode is
bound by MEMORY BANDWIDTH, not compute, and the arithmetic units sit idle.

Every remaining serving technique follows from that: batching amortizes the
weight read, quantization shrinks the bytes read, speculative decoding gets
more than one token per read. None of them is about doing less arithmetic.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FLOAT = np.float32


class KvCache:
    """Preallocated keys and values, one slab per layer.

    Preallocated to the full context rather than grown: growing a buffer per
    token is the allocation pattern to avoid in a generation loop, and the
    peak footprint is the same either way since the context is a hard limit.

    Layout is (layer, batch, head, position, head_dim). Position is the
    second-to-last axis on purpose - it is the axis that grows, and putting it
    there keeps each head's history contiguous for the attention product.
    """

    def __init__(self, depth: int, batch: int, heads: int, context: int, head_dim: int) -> None:
        shape = (depth, batch, heads, context, head_dim)
        self._keys = np.zeros(shape, dtype=FLOAT)
        self._values = np.zeros(shape, dtype=FLOAT)
        self._length = 0
        self._context = context

    @property
    def length(self) -> int:
        return self._length

    def bytes(self) -> int:
        """What actually runs out of memory in production.

        Linear in batch times context, so serving capacity is a function of
        traffic shape rather than of the model — and it is usually this rather
        than the weights that exhausts the device.
        """
        return self._keys.nbytes + self._values.nbytes

    def append(
        self, layer: int, keys: NDArray[np.float32], values: NDArray[np.float32]
    ) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
        """Write the new position and return views over the whole history.

        Views, not copies: the attention product below reads straight out of
        the cache slab, so nothing is duplicated per step.
        """
        width = keys.shape[2]
        end = self._length + width
        if end > self._context:
            raise RuntimeError(
                f'cache full at {self._context} positions; generation cannot continue'
            )

        self._keys[layer, :, :, self._length:end, :] = keys
        self._values[layer, :, :, self._length:end, :] = values

        # Note `advance` is separate: the length must move once per TOKEN, not
        # once per layer, and every layer writes at the same position.
        return (
            self._keys[layer, :, :, :end, :],
            self._values[layer, :, :, :end, :],
        )

    def advance(self, width: int) -> None:
        self._length += width

    def reset(self) -> None:
        """Reuse the allocation across requests rather than reallocating."""
        self._length = 0


def cached_attention(
    query: NDArray[np.float32],
    keys: NDArray[np.float32],
    values: NDArray[np.float32],
) -> NDArray[np.float32]:
    """One new query against the whole cached history.

    No causal mask is needed here, and that is the point: the cache contains
    only positions up to and including the current one, so causality is
    enforced by what is IN the cache rather than by masking a full matrix.
    That removes both the mask and the wasted computation behind it.

    query is (batch, heads, 1, head_dim); keys and values are
    (batch, heads, length, head_dim). One new row of attention per token
    rather than a full quadratic matrix — this is the O(n^3) to O(n^2) change.
    """
    scale = FLOAT(1.0 / np.sqrt(query.shape[-1]))
    scores = np.einsum('bhqd,bhkd->bhqk', query, keys, optimize=True)
    scores *= scale

    # Stable softmax in place over the single query row.
    scores -= scores.max(axis=-1, keepdims=True)
    np.exp(scores, out=scores)
    scores /= scores.sum(axis=-1, keepdims=True)

    return np.einsum('bhqk,bhkd->bhqd', scores, values, optimize=True)


def prefill(
    tokens: NDArray[np.int32],
    embedding: NDArray[np.float32],
    position: NDArray[np.float32],
) -> NDArray[np.float32]:
    """Process the whole prompt in ONE parallel pass.

    This half is compute-bound: every position is computed simultaneously
    because the causal mask makes all n predictions valid at once, so the
    arithmetic units are saturated and this is the phase where FLOP
    optimization actually helps.

    Confusing this phase with decode is the standard profiling mistake — they
    sit on opposite sides of the memory-versus-compute boundary and need
    opposite optimizations.
    """
    hidden = embedding[tokens]
    hidden += position[: tokens.shape[1]]
    return hidden


def fused_qkv_decode(
    hidden: NDArray[np.float32],
    qkv_weight: NDArray[np.float32],
    heads: int,
) -> tuple[NDArray[np.float32], ...]:
    """One GEMV for all three projections of a single new token.

    Note what this is: a matrix-VECTOR product, not matrix-matrix. During
    decode there is one token, so every weight matrix is read to multiply
    against a single vector — arithmetic intensity of one, which is the whole
    memory-bound story in one line.
    """
    batch, _, model_dim = hidden.shape
    head_dim = model_dim // heads

    projected = hidden.reshape(batch, model_dim) @ qkv_weight
    projected = projected.reshape(batch, 3, heads, head_dim)
    # (3, batch, heads, 1, head_dim): the singleton query axis is what makes
    # the cached attention above a one-row product.
    projected = projected.transpose(1, 0, 2, 3)[:, :, :, None, :]
    return projected[0], projected[1], projected[2]


def batched_sample(
    logits: NDArray[np.float32],
    temperature: float,
    top_p: float,
    rng: np.random.Generator,
) -> NDArray[np.int32]:
    """Nucleus sampling for a whole batch at once.

    Temperature is applied BEFORE truncation, deliberately: it rescales the
    distribution, so truncating first would change which tokens were
    eligible. Every implementation follows this order and mixing it gives
    subtly different output.
    """
    if temperature <= 0.0:
        return logits.argmax(axis=-1).astype(np.int32)

    scaled = logits / FLOAT(temperature)
    scaled -= scaled.max(axis=-1, keepdims=True)
    np.exp(scaled, out=scaled)
    scaled /= scaled.sum(axis=-1, keepdims=True)

    if top_p >= 1.0:
        return _categorical(scaled, rng)

    # Sort descending, find the nucleus boundary per row, zero the tail.
    order = np.argsort(-scaled, axis=-1)
    ordered = np.take_along_axis(scaled, order, axis=-1)
    cumulative = np.cumsum(ordered, axis=-1)

    # Keep everything up to and INCLUDING the token that crosses the
    # threshold, or a very peaked distribution keeps nothing at all.
    keep = cumulative - ordered < top_p
    ordered *= keep
    ordered /= ordered.sum(axis=-1, keepdims=True)

    picked = _categorical(ordered, rng)
    return np.take_along_axis(order, picked[:, None], axis=-1).ravel().astype(np.int32)


def _categorical(probabilities: NDArray[np.float32], rng: np.random.Generator) -> NDArray[np.int32]:
    """One uniform draw per row against the cumulative distribution.

    searchsorted per row rather than a Python loop: the inverse-CDF sample is
    a binary search, and doing it as an array operation keeps the whole batch
    in one call.
    """
    cumulative = np.cumsum(probabilities, axis=-1)
    draws = rng.random((len(probabilities), 1), dtype=FLOAT) * cumulative[:, -1:]
    return (cumulative < draws).sum(axis=-1).astype(np.int32)


def arithmetic_intensity(parameters: int, batch: int, bytes_per_parameter: int = 2) -> float:
    """FLOPs per byte read during decode. The diagnostic that matters.

    One decode step reads every parameter and does roughly two FLOPs per
    parameter per batch element, so intensity is about batch divided by
    bytes-per-parameter. At batch one with 16-bit weights that is around one
    FLOP per byte — against hardware that wants hundreds.

    That gap IS the memory-bound problem, and it is why batching helps so much
    here: the weight read is shared, so intensity scales linearly with batch
    until the cache runs out of memory.
    """
    flops = 2 * parameters * batch
    bytes_read = parameters * bytes_per_parameter
    return flops / bytes_read