Transformer

"""Scaled dot-product attention - the equation, transcribed.

Three loops: for each query position, score it against every key position,
softmax those scores, then take the weighted average of the values.
"""

import math


def softmax(scores):
    # Subtract the max before exponentiating - otherwise exp() overflows on
    # large logits and the whole row becomes NaN.
    largest = max(scores)
    exps = [math.exp(s - largest) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]


def attention(queries, keys, values, causal=False):
    """queries[n][d_k], keys[m][d_k], values[m][d_v] -> output[n][d_v]."""
    n = len(queries)
    m = len(keys)
    d_k = len(queries[0])
    scale = 1.0 / math.sqrt(d_k)

    output = []
    for i in range(n):
        # Score this query against every key: the all-pairs comparison.
        scores = []
        for j in range(m):
            dot = sum(queries[i][k] * keys[j][k] for k in range(d_k))
            # Causal masking: a position may not attend to the future.
            # Without this a sequence model trivially cheats.
            if causal and j > i:
                scores.append(float("-inf"))
            else:
                # Scaling by sqrt(d_k) keeps the softmax out of saturation,
                # where gradients vanish.
                scores.append(dot * scale)

        weights = softmax(scores)

        # Weighted average of the value vectors.
        mixed = [0.0] * len(values[0])
        for j in range(m):
            for v in range(len(mixed)):
                mixed[v] += weights[j] * values[j][v]
        output.append(mixed)

    return output

"""Transformer block - typed, pre-norm, with masking handled explicitly."""

from dataclasses import dataclass

import torch
from torch import Tensor, nn


@dataclass(frozen=True)
class TransformerConfig:
    d_model: int = 512
    n_heads: int = 8
    d_ff: int = 2048
    dropout: float = 0.1

    def __post_init__(self) -> None:
        if self.d_model % self.n_heads != 0:
            raise ValueError(
                f"d_model {self.d_model} must divide evenly by n_heads {self.n_heads}"
            )


class MultiHeadAttention(nn.Module):
    """Heads are a RESHAPE, not separate modules.

    One fused qkv projection then a view into (batch, heads, seq, head_dim) -
    running n_heads separate Linear layers would be the same arithmetic in
    n_heads times as many kernel launches.
    """

    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads

        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out = nn.Linear(config.d_model, config.d_model)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
        batch, seq, _ = x.shape

        qkv = self.qkv(x).reshape(batch, seq, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)     # (batch, heads, seq, head_dim)

        scores = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)

        if mask is not None:
            # -inf, not a large negative number: after softmax this is exactly
            # zero, whereas -1e9 leaves a small residue that accumulates over
            # long sequences.
            scores = scores.masked_fill(mask == 0, float("-inf"))

        weights = self.dropout(torch.softmax(scores, dim=-1))
        mixed = (weights @ v).transpose(1, 2).reshape(batch, seq, -1)
        return self.out(mixed)


class TransformerBlock(nn.Module):
    """PRE-norm: LayerNorm before each sublayer, not after.

    Post-norm (the original paper) needs careful warmup tuning to train deep.
    Pre-norm leaves the residual stream a clean identity path, which is what
    made deep transformers trainable in practice.
    """

    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(config.d_model)
        self.attention = MultiHeadAttention(config)
        self.norm2 = nn.LayerNorm(config.d_model)
        self.feedforward = nn.Sequential(
            nn.Linear(config.d_model, config.d_ff),
            nn.GELU(),
            nn.Linear(config.d_ff, config.d_model),
            nn.Dropout(config.dropout),
        )

    def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
        x = x + self.attention(self.norm1(x), mask)
        return x + self.feedforward(self.norm2(x))


def causal_mask(seq_len: int, device: str = "cpu") -> Tensor:
    """Lower-triangular mask. Omitting this in a sequence model lets every
    position read the answer, which validates spectacularly and means nothing."""
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))

"""Attention - FlashAttention, KV caching, length bucketing."""

import torch
from torch import Tensor, nn


class FastAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int) -> None:
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model)

    def forward(self, x: Tensor, is_causal: bool = True) -> Tensor:
        batch, seq, _ = x.shape
        qkv = self.qkv(x).reshape(batch, seq, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)

        # FlashAttention: never materializes the n x n score matrix at all.
        # It tiles the computation so scores stay in SRAM, which turns
        # attention from O(n^2) MEMORY to O(n) memory. That is the change that
        # made long context feasible - it is a memory-traffic win, not an
        # arithmetic one.
        mixed = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, is_causal=is_causal
        )

        return self.out(mixed.transpose(1, 2).reshape(batch, seq, -1))


class CachedAttention(nn.Module):
    """Incremental decoding with a KV cache.

    Autoregressive generation recomputes attention over the entire prefix for
    every new token, which is O(n^2) work across a sequence for something that
    only needs O(n). Caching past keys and values makes each step O(n) instead
    - the single most important inference optimization for generation.
    """

    def __init__(self, d_model: int, n_heads: int, max_seq: int, max_batch: int) -> None:
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model)

        # Pre-allocated to full size. Growing the cache per token would
        # reallocate and copy on every single step of generation.
        self.register_buffer(
            "k_cache",
            torch.zeros(max_batch, n_heads, max_seq, self.head_dim),
            persistent=False,
        )
        self.register_buffer(
            "v_cache",
            torch.zeros(max_batch, n_heads, max_seq, self.head_dim),
            persistent=False,
        )

    def forward(self, x: Tensor, position: int) -> Tensor:
        batch, seq, _ = x.shape
        qkv = self.qkv(x).reshape(batch, seq, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)

        # Write this step's keys/values into the cache, then attend over the
        # whole prefix. No recomputation of anything already seen.
        self.k_cache[:batch, :, position : position + seq] = k
        self.v_cache[:batch, :, position : position + seq] = v

        keys = self.k_cache[:batch, :, : position + seq]
        values = self.v_cache[:batch, :, : position + seq]

        mixed = torch.nn.functional.scaled_dot_product_attention(
            q, keys, values, is_causal=False   # prefix is all valid history
        )
        return self.out(mixed.transpose(1, 2).reshape(batch, seq, -1))


def bucket_by_length(sequences: list[Tensor], bucket_size: int = 64) -> list[list[Tensor]]:
    """Group sequences of similar length before batching.

    Attention is quadratic in the PADDED length, so one long sequence in a
    batch makes every short one pay its cost. Bucketing is usually a larger
    throughput win than any kernel-level optimization.
    """
    ordered = sorted(sequences, key=len)
    return [ordered[i : i + bucket_size] for i in range(0, len(ordered), bucket_size)]