transformer_block.py

Transformer Block: from 2017 seq2seq to modern LLM blocks - minimal runnable implementation
在 GitHub 查看原始 .py python transformer_block.py
"""
Transformer Block: from 2017 seq2seq to modern LLM blocks - minimal runnable implementation
=============================================================================================

Educational PyTorch reference for how a Transformer *block* is assembled: residual
topology (Post-LN / serial Pre-LN / parallel Pre-LN), the encoder-decoder -> decoder-only
topological fork, RoPE + KV-cache assembly, MHA/GQA/MQA as one configurable axis, and
dense-vs-MoE FFN as an orthogonal, swappable interface.
Standalone script: runs sanity checks [D01]-[D15] on CPU in well under a minute
([D01]-[D09] are the core assembly/contract checks; [D10]-[D15] are targeted
regression tests, one per code-review bug fix, see §5 below).

Pairs with: docs/tutorials/transformer_block_tutorial.md (concept reference).

Scope note (assembly-level, not derivation-level): this file does NOT re-derive
softmax attention, FlashAttention, RoPE/YaRN/MLA algebra, MoE routing/capacity math,
or Xiong/DeepNorm gradient proofs -- those live in the sibling tutorials'
code/*.py files (normalization.py, linear_sparse_attention.py) and are only cited
here. This file only assembles known pieces into four canonical block classes and
verifies the *contracts* between them.

Entry classes (four clearly distinct block "eras"):
    VanillaEncoderLayer2017        - Post-LN, self-attn -> FFN (Vaswani et al. 2017)
    VanillaSeq2SeqDecoderLayer2017 - Post-LN, masked self-attn -> cross-attn -> FFN
    GPT2StyleDecoderLayer          - Pre-LN, LayerNorm, causal self-attn -> FFN, GELU, bias
    ModernDecoderBlock             - Pre-RMSNorm, RoPE, GQA-capable, SwiGLU/MoE-swappable

Sanity checks:
    [D01] Shape contract: all four blocks preserve residual width [B,T,d];
          cross-attn correctly decouples query length T from memory length S,
          AND changing memory actually changes the output (memory is read,
          not merely shape-compatible).
    [D02] Residual topology test: zeroing every sublayer's output projection makes
          a Pre-LN block return x EXACTLY; a real two-sublayer Post-LN block
          returns exactly ln2(ln1(x)) (not just "something != x").
    [D03] Causal non-leakage: perturbing tokens after position t must not change
          the output at positions <= t; perturbing position 0 instead must
          change positions > 0 (proves attention mixes forward, ruling out a
          dead/zeroed attention branch trivially "passing" the first check).
    [D04] Full forward vs cached decoding: eval(), no dropout, full-sequence
          forward vs token-by-token incremental KV-cache decoding must match
          numerically within a small tolerance (NOT bit-identical); the
          cache's actual K/V CONTENT is separately checked against an
          independently projected+rotated reference, not just cache shape.
    [D05] MHA is the n_kv_heads==n_q_heads degenerate case of GQA (compared
          against an independent plain-MHA reference function); a REAL grouped
          case (n_kv_heads=2, group=4) and MQA (n_kv_heads=1) are both checked
          against an independent repeat_interleave() reference, and MQA's
          per-query-head outputs are confirmed distinct (group dim isn't
          accidentally collapsed) while the cache never grows past 1 head.
    [D06] FFN parameter parity (d=96, dense_hidden=384, gated_hidden=256 both
          give exactly 73,728 matrix parameters) + attention parameter formula
          2*d^2*(1+n_kv/n_q), verified against real .numel() counts.
    [D07] Packed projection equivalence: a single packed QKV (or gate+up) Linear
          built from the same weights as three (or two) separate Linears must
          give numerically equivalent output (not a bit-identical guarantee
          across arbitrary GEMM shapes/backends).
    [D08] Parallel vs sequential dependency: forward hooks on BOTH the attn
          module (input+output) and the FFN module, captured during the SAME
          real forward call, prove that a serial Pre-LN block's FFN reads the
          attn-UPDATED residual stream while a GPT-J/PaLM-style parallel
          block's attn and FFN read the literal SAME normalized tensor.
    [D09] Backward smoke test: all block parameters (including the MoE
          router/gate and the rest of the block, not just expert weights) get
          finite AND non-zero gradients; for an FFN swapped to a minimal
          top-1 MoE, only experts that actually received a token this batch
          may have a non-None/non-zero gradient.

Additional regression tests [D10]-[D15] (§5): one per B1-class bug found
during code review -- config validation (indivisible d_model/heads, MoE
top_k), KV-cache/start_pos consistency, cross-attention + RoPE/cache
restrictions, masked-softmax NaN safety + mask-shape collision guards,
default-causal safety net, and RoPE cache device/dtype handling.

Run:
    python transformer_block.py
"""
import math
import sys

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)


# =============================================================================
# §1 Shared building blocks (used by more than one of the four block classes)
# =============================================================================

def build_causal_mask(q_positions: torch.Tensor, k_positions: torch.Tensor) -> torch.Tensor:
    """Boolean keep-mask [Tq, Tk]: True where key position <= query position.

    Works for both a full-sequence forward (q_positions == k_positions == arange(T))
    and a single incremental decode step (q_positions = [t], k_positions =
    arange(0, cache_len)) -- same formula, same function, see [D04].
    """
    return k_positions[None, :] <= q_positions[:, None]


def build_rope_cache(max_len: int, head_dim: int, base: float = 10000.0,
                      device=None, dtype=None):
    """cos/sin cache for RoPE (rotate-half convention), shape [max_len, head_dim].

    This is an assembly-level utility only -- for the *why* (complex-plane
    derivation, frequency choice, YaRN/NTK extensions), see
    long_context_rope_yarn_mla_tutorial.md §2/§5/§6.

    `device`/`dtype` are accepted explicitly (bug fix: a cache silently fixed
    to CPU/float32 breaks the moment the model is moved to GPU or run under
    mixed precision -- see [D15]). The trig itself is always computed in
    float32 for numerical stability, then cast to `dtype` at the end.
    """
    if head_dim % 2 != 0:
        raise ValueError(f"RoPE requires an even head_dim, got head_dim={head_dim}")
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))  # [dh/2]
    t = torch.arange(max_len, device=device).float()                              # [L]
    freqs = torch.outer(t, inv_freq)                                             # [L, dh/2]
    emb = torch.cat([freqs, freqs], dim=-1)                                      # [L, dh]
    cos, sin = emb.cos(), emb.sin()
    if dtype is not None:
        cos, sin = cos.to(dtype), sin.to(dtype)
    return cos, sin


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """x: [B, H, L, dh]; cos/sin: [L, dh]. Applied to Q/K ONLY -- V is never
    rotated (assembly rule, see markdown §3 / §4 [C007]).

    Defensively casts cos/sin to x's device/dtype (bug fix: a RoPE cache built
    once on CPU/float32 must not silently device-mismatch-crash on GPU, nor
    silently upcast a low-precision matmul -- see [D15])."""
    cos = cos[None, None, :, :].to(device=x.device, dtype=x.dtype)
    sin = sin[None, None, :, :].to(device=x.device, dtype=x.dtype)
    return x * cos + _rotate_half(x) * sin


class KVCache:
    """Minimal single-layer incremental KV cache.

    Stores K, V at shape [B, n_kv_heads, T, d_head] and grows along the T axis
    via .append(). The head dimension is ALWAYS n_kv_heads -- it must never be
    physically broadcast/copied up to n_q_heads before being written back
    (that broadcast happens only transiently inside attention, see
    MultiHeadAttention.forward and [D05]).

    Bug fix (see [D11]): a cache used to silently accept anything .append()-ed
    to it, with no record of "what RoPE position comes next" and no check that
    a new chunk actually belongs to the same (batch, n_kv_heads, d_head,
    device, dtype) sequence as what's already stored. That made two silent
    failure modes possible: (1) calling MultiHeadAttention.forward with a
    stale/mismatched `start_pos` against a non-empty cache re-uses the wrong
    RoPE angles for the new tokens; (2) reusing one KVCache instance across two
    different sequences/configs silently concatenates unrelated K/V. This
    class now tracks `next_pos` (the RoPE position the next appended chunk
    must start at) and validates shape/device/dtype identity on every append;
    call `.reset()` to reuse an instance for a new sequence.
    """

    def __init__(self):
        self.reset()

    def reset(self):
        self.k = None
        self.v = None
        self.next_pos = 0
        self._sig = None    # (batch, n_kv_heads, d_head, device, dtype) of the stored sequence

    def append(self, k_new: torch.Tensor, v_new: torch.Tensor):
        B, H, T, dh = k_new.shape
        sig = (B, H, dh, k_new.device, k_new.dtype)
        if self.k is None:
            self._sig = sig
            self.k, self.v = k_new, v_new
        else:
            if sig != self._sig:
                raise ValueError(
                    f"KVCache.append: shape/device/dtype {sig} does not match the cache's "
                    f"existing sequence {self._sig}. Reusing one KVCache across two different "
                    f"sequences or configs silently splices unrelated K/V together -- call "
                    f".reset() before starting a new sequence.")
            # Performance note (teaching limitation, not a correctness bug):
            # torch.cat here re-copies the ENTIRE history on every single
            # append, an O(T^2) total-copy cost over a T-token decode. A real
            # serving cache pre-allocates a fixed-capacity buffer and writes
            # new tokens in place (or manages paged/chunked blocks). Fine for
            # this file's 6-token demo; not representative of production KV
            # cache implementations -- see run-experiment/inference_serving
            # sibling material for the real thing.
            self.k = torch.cat([self.k, k_new], dim=2)
            self.v = torch.cat([self.v, v_new], dim=2)
        self.next_pos += T
        return self.k, self.v


def _infer_start_pos(kv_cache: "KVCache", start_pos):
    """Single source of truth for "what RoPE position does this call start at".

    Bug fix (see [D11]): the old default `start_pos: int = 0` meant that
    calling with a non-empty cache and no explicit start_pos silently reused
    RoPE angle 0 for new tokens instead of continuing from the cache. Passing
    `start_pos=None` (the new default) now means "ask the cache"; passing an
    explicit start_pos that disagrees with the cache's recorded position is a
    ValueError instead of a silent wrong-position bug.
    """
    if kv_cache is not None:
        inferred = kv_cache.next_pos
        if start_pos is None:
            return inferred
        if start_pos != inferred:
            raise ValueError(
                f"start_pos={start_pos} is inconsistent with kv_cache.next_pos={inferred}. "
                "Pass start_pos=None to let the cache supply the position, or call "
                "kv_cache.reset() if this is meant to start a new sequence.")
        return start_pos
    return 0 if start_pos is None else start_pos


class MultiHeadAttention(nn.Module):
    """Generic multi-head attention: self- or cross-attention, optional causal
    masking, optional GQA (n_kv_heads <= n_q_heads, MHA when equal, MQA when
    n_kv_heads==1), optional RoPE, optional incremental KV cache.

    Interface-level only: for *why* softmax(QK^T/sqrt(d)) and *why* GQA saves KV
    cache memory, see attention_tutorial.md §2 / §6. This class only assembles
    the pieces into the shape contract a Transformer block needs.

    Mask convention (bug fix, see [D13]): two SEPARATE, unambiguous mask
    arguments are accepted instead of one shape-guessed `mask` -- `causal_mask`
    is the [Tq, Tk] connectivity pattern, `key_padding_mask` is the [B, Tk]
    per-batch-element key validity. Both use "True = ATTEND / keep" (the
    OPPOSITE of the common PyTorch nn.MultiheadAttention `key_padding_mask`
    True-means-mask-OUT convention -- double check when porting real code).
    A single generic `mask` argument that could be either [Tq, Tk] or [B, Tk]
    is ambiguous and, worse, silently WRONG whenever B happens to equal Hkv or
    Tq (the tensor still broadcasts, just onto the wrong dimension).

    Combination limits that are NOT supported by this minimal reference (bug
    fix, see [D12]): cross-attention (`x_kv is not x_q`) together with RoPE or
    with a KV cache. Cross-attention's query and memory generally occupy
    different position spaces (a single `start_pos` cannot represent both),
    and repeatedly caching the same encoder memory would re-append it every
    call. Both combinations raise ValueError rather than silently computing
    something wrong.
    """

    def __init__(self, d_model: int, n_q_heads: int, n_kv_heads: int = None,
                 bias: bool = False, rope: bool = False):
        super().__init__()
        n_kv_heads = n_q_heads if n_kv_heads is None else n_kv_heads
        if n_q_heads <= 0 or n_kv_heads <= 0:
            raise ValueError(f"n_q_heads ({n_q_heads}) and n_kv_heads ({n_kv_heads}) must be positive")
        if n_q_heads % n_kv_heads != 0:
            raise ValueError(f"n_q_heads ({n_q_heads}) must be divisible by n_kv_heads ({n_kv_heads})")
        if d_model % n_q_heads != 0:
            raise ValueError(
                f"d_model ({d_model}) must be divisible by n_q_heads ({n_q_heads}); "
                f"floor-dividing here would silently narrow attention to "
                f"{n_q_heads * (d_model // n_q_heads)} of {d_model} dims while still "
                f"projecting back out to d_model, passing shape checks with a wrong answer.")
        self.n_q_heads, self.n_kv_heads = n_q_heads, n_kv_heads
        self.d_head = d_model // n_q_heads
        if rope and self.d_head % 2 != 0:
            raise ValueError(
                f"RoPE requires an even per-head dim, got d_head={self.d_head} "
                f"(d_model={d_model}, n_q_heads={n_q_heads})")
        self.group = n_q_heads // n_kv_heads          # Q heads sharing one KV head
        self.rope = rope
        self.q_proj = nn.Linear(d_model, n_q_heads * self.d_head, bias=bias)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=bias)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.d_head, bias=bias)
        self.o_proj = nn.Linear(n_q_heads * self.d_head, d_model, bias=bias)

    def _shape(self, x: torch.Tensor, n_heads: int) -> torch.Tensor:
        B, L, _ = x.shape
        return x.view(B, L, n_heads, self.d_head).transpose(1, 2)   # [B, H, L, dh]

    def forward(self, x_q, x_kv=None, causal_mask=None, key_padding_mask=None,
                cos=None, sin=None, kv_cache: KVCache = None, start_pos: int = None):
        is_cross_attn = x_kv is not None and x_kv is not x_q
        x_kv = x_q if x_kv is None else x_kv              # self-attn vs cross-attn source
        B, Tq, _ = x_q.shape

        if is_cross_attn:
            # Bug fix (see [D12]): cross-attention's Q and K/V generally live in
            # DIFFERENT position spaces. A single start_pos/cache cannot represent
            # both without silently mis-positioning RoPE or re-appending the same
            # memory into the cache on every call -- refuse instead of guessing.
            if self.rope:
                raise ValueError(
                    "RoPE is not supported for cross-attention in this minimal reference: "
                    "query and memory occupy different position spaces and a single "
                    "start_pos cannot represent both. Construct this module with rope=False "
                    "for cross-attention, or apply positions explicitly outside this class.")
            if kv_cache is not None:
                raise ValueError(
                    "KV caching is not supported for cross-attention in this minimal reference: "
                    "repeated decode-step calls would re-append the SAME encoder memory into the "
                    "cache every time. Encode memory once outside any KVCache, or extend KVCache "
                    "with a write-once static-memory mode.")

        resolved_start = _infer_start_pos(kv_cache, start_pos)

        q = self._shape(self.q_proj(x_q), self.n_q_heads)      # [B, Hq,  Tq, dh]
        k_new = self._shape(self.k_proj(x_kv), self.n_kv_heads)  # [B, Hkv, Tk_new, dh]
        v_new = self._shape(self.v_proj(x_kv), self.n_kv_heads)

        if self.rope:
            Tk_new = k_new.shape[2]
            q = apply_rope(q, cos[resolved_start:resolved_start + Tq], sin[resolved_start:resolved_start + Tq])
            # K carries RoPE BEFORE entering the cache; V is never rotated.
            k_new = apply_rope(k_new, cos[resolved_start:resolved_start + Tk_new],
                                sin[resolved_start:resolved_start + Tk_new])

        if kv_cache is not None:
            k, v = kv_cache.append(k_new, v_new)          # cache head dim stays n_kv_heads
        else:
            k, v = k_new, v_new

        Tk = k.shape[2]
        # GQA broadcast: view Q as (Hkv groups x group size), broadcast K/V across
        # the group dim via unsqueeze (a view, NOT a physical copy written to cache).
        q_g = q.view(B, self.n_kv_heads, self.group, Tq, self.d_head)
        k_g = k.unsqueeze(2)                               # [B, Hkv, 1, Tk, dh]
        v_g = v.unsqueeze(2)

        scores = (q_g @ k_g.transpose(-2, -1)) / math.sqrt(self.d_head)   # [B, Hkv, G, Tq, Tk]

        # Bug fix (see [D13]): causal_mask [Tq, Tk] and key_padding_mask [B, Tk]
        # are reshaped EXPLICITLY by their known rank/role, never shape-guessed,
        # so there is no B==Hkv / B==Tq collision with the 5-D `scores` tensor.
        keep = None
        if causal_mask is not None:
            keep = causal_mask[None, None, None, :, :]                 # -> [1, 1, 1, Tq, Tk]
        if key_padding_mask is not None:
            kpm = key_padding_mask[:, None, None, None, :]             # -> [B, 1, 1, 1, Tk]
            keep = kpm if keep is None else (keep & kpm)
        if keep is not None:
            scores = scores.masked_fill(~keep, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        # Bug fix (see [D13]): a query row with NO valid keys (all -inf) produces
        # softmax NaN, not a numerical accident but a real hazard for padded
        # batches/empty cross-memory. Safe masked softmax: treat a fully-masked
        # row as contributing zero attention weight rather than propagating NaN.
        attn = torch.nan_to_num(attn, nan=0.0)
        out = attn @ v_g                                    # [B, Hkv, G, Tq, dh]

        out = out.reshape(B, self.n_q_heads, Tq, self.d_head).transpose(1, 2)
        out = out.reshape(B, Tq, self.n_q_heads * self.d_head)
        return self.o_proj(out), (k, v)


class DenseFFN(nn.Module):
    """2-matrix FFN: down(act(up(x))). GELU by default (GPT-2-style)."""

    def __init__(self, d_model: int, d_ff: int, act=F.gelu, bias: bool = False):
        super().__init__()
        self.up = nn.Linear(d_model, d_ff, bias=bias)
        self.down = nn.Linear(d_ff, d_model, bias=bias)
        self.act = act

    def forward(self, x):
        return self.down(self.act(self.up(x)))


class GatedFFN(nn.Module):
    """3-matrix gated FFN: down(act(gate(x)) * up(x)). SiLU gate -> SwiGLU
    (Llama-style); GELU gate -> GeGLU. Full "why gated FFN helps" discussion:
    markdown §3 (only the parameter-matching arithmetic is re-derived here, [D06])."""

    def __init__(self, d_model: int, d_ff: int, act=F.silu, bias: bool = False):
        super().__init__()
        self.gate = nn.Linear(d_model, d_ff, bias=bias)
        self.up = nn.Linear(d_model, d_ff, bias=bias)
        self.down = nn.Linear(d_ff, d_model, bias=bias)
        self.act = act

    def forward(self, x):
        return self.down(self.act(self.gate(x)) * self.up(x))


class MoEFFN(nn.Module):
    """Minimal top-1 token-choice MoE FFN (Switch-style): same [B,T,d] -> [B,T,d]
    interface as DenseFFN/GatedFFN, so it drops into ModernDecoderBlock unchanged
    (markdown §5/§7 "FFN interface stays fixed" point). Routing / capacity /
    load-balancing-loss math is intentionally NOT re-derived here -- see
    moe_tutorial.md §2-§4 for the full treatment. `tokens_per_expert` is recorded
    each forward for the [D09] "only touched experts get gradients" check."""

    def __init__(self, d_model: int, d_ff: int, n_experts: int = 4, top_k: int = 1):
        super().__init__()
        if top_k != 1:
            # Bug fix (see [D10]): this used to be `assert top_k == 1`, which (a)
            # disappears entirely under `python -O`, and (b) even when it DOES
            # fire, `forward()` below never reads `top_k` at all -- it always does
            # a `.max(dim=-1)` top-1 route regardless. Any top_k != 1 must be
            # rejected explicitly rather than silently downgraded to top-1.
            raise NotImplementedError(
                f"MoEFFN in this minimal assembly demo only implements top-1 routing; "
                f"got top_k={top_k}. Extending to top_k>1 needs per-token weighted "
                f"multi-expert combination, intentionally out of scope here -- see "
                f"moe_tutorial.md §2-§4.")
        self.top_k = top_k
        self.experts = nn.ModuleList(
            [GatedFFN(d_model, d_ff, bias=False) for _ in range(n_experts)])
        self.gate = nn.Linear(d_model, n_experts, bias=False)
        self.tokens_per_expert = [0] * n_experts

    def forward(self, x):
        # Performance note (teaching limitation, not a correctness bug): this
        # is a plain Python for-loop over experts, and `sel.sum().item()` /
        # `sel.any()` each force a GPU->CPU sync every iteration. A real MoE
        # kernel dispatches tokens via a single vectorized gather/scatter (or
        # a fused grouped-GEMM) instead of per-expert boolean masking in a
        # Python loop -- see moe_tutorial.md for the real routing/dispatch
        # implementation. Fine for this file's few-token demo.
        B, T, d = x.shape
        flat = x.reshape(-1, d)                                # [B*T, d]
        gate_w, idx = self.gate(flat).softmax(-1).max(dim=-1)   # top-1 weight + expert id
        out = torch.zeros_like(flat)
        self.tokens_per_expert = [0] * len(self.experts)
        for e, expert in enumerate(self.experts):
            sel = idx == e
            self.tokens_per_expert[e] = int(sel.sum().item())
            if sel.any():
                out[sel] = expert(flat[sel]) * gate_w[sel].unsqueeze(-1)
        return out.view(B, T, d)


# =============================================================================
# §2 Four canonical block classes ("eras")
# =============================================================================

class VanillaEncoderLayer2017(nn.Module):
    """Post-LN encoder layer (Vaswani et al. 2017): self-attn -> FFN, TWO
    sublayers, each y = N(x + Sublayer(x)). Self-attn is bidirectional (mask, if
    any, is padding-only, never causal). Sinusoidal PE lives OUTSIDE this block
    (added to embeddings before the stack) -- see markdown §1."""

    def __init__(self, d_model: int = 96, n_heads: int = 8, d_ff: int = 384):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads, n_heads, bias=True)
        self.ffn = DenseFFN(d_model, d_ff, act=F.relu, bias=True)
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)

    def forward(self, x, pad_mask=None):
        a, _ = self.attn(x, key_padding_mask=pad_mask)
        x = self.ln1(x + a)          # Post-LN: y = N(x + F(x))
        f = self.ffn(x)
        x = self.ln2(x + f)
        return x


class VanillaSeq2SeqDecoderLayer2017(nn.Module):
    """Post-LN decoder layer (Vaswani et al. 2017): masked self-attn -> cross-attn
    -> FFN, THREE sublayers -- the sublayer that decoder-only LLM blocks later
    drop entirely (not merely reconfigure, see markdown §2). Cross-attn: Q comes
    from THIS stream (length T), K/V come from encoder memory (length S); T and
    S need not match."""

    def __init__(self, d_model: int = 96, n_heads: int = 8, d_ff: int = 384):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, n_heads, bias=True)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, n_heads, bias=True)
        self.ffn = DenseFFN(d_model, d_ff, act=F.relu, bias=True)
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.ln3 = nn.LayerNorm(d_model)

    def forward(self, x, memory, causal_mask=None, cross_mask=None):
        # Bug fix (see [D14]): self_attn here is documented as MASKED
        # (causal) self-attention. The old code passed `mask=None` straight
        # through whenever the caller omitted it, which silently ran
        # BIDIRECTIONAL self-attention instead of causal -- no error, just a
        # quietly wrong contract. Auto-build the causal mask when omitted.
        if causal_mask is None:
            T = x.shape[1]
            causal_mask = build_causal_mask(torch.arange(T), torch.arange(T))
        a, _ = self.self_attn(x, causal_mask=causal_mask)
        x = self.ln1(x + a)
        c, _ = self.cross_attn(x, x_kv=memory, causal_mask=cross_mask)   # Q len T, K/V len S
        x = self.ln2(x + c)
        f = self.ffn(x)
        x = self.ln3(x + f)
        return x


class GPT2StyleDecoderLayer(nn.Module):
    """Pre-LN decoder-only block (GPT-2, Radford et al. 2019): causal self-attn
    -> FFN, TWO sublayers -- cross-attn is structurally ABSENT, not merely
    unused (markdown §2 explains why this is a topological fork). GPT-2 is
    ALREADY Pre-LN with LayerNorm, learned absolute PE added outside the stack,
    a GELU dense FFN, and bias everywhere. A final LayerNorm after the LAST
    block in the stack is required (not modeled per-block here) -- see
    normalization_init_tutorial.md §5.3 and markdown §2.2."""

    def __init__(self, d_model: int = 96, n_heads: int = 8, d_ff: int = 384):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads, n_heads, bias=True)
        self.ffn = DenseFFN(d_model, d_ff, act=F.gelu, bias=True)
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)

    def forward(self, x, causal_mask=None):
        # Bug fix (see [D14]): same silent-bidirectional hazard as above --
        # this block's docstring/class name promise CAUSAL self-attention, so
        # omitting the mask must not silently fall back to full attention.
        if causal_mask is None:
            T = x.shape[1]
            causal_mask = build_causal_mask(torch.arange(T), torch.arange(T))
        a, _ = self.attn(self.ln1(x), causal_mask=causal_mask)   # Pre-LN: y = x + F(N(x))
        x = x + a
        x = x + self.ffn(self.ln2(x))
        return x


class ModernDecoderBlock(nn.Module):
    """Pre-RMSNorm + RoPE + (GQA-capable) causal self-attn + gated FFN
    ("Llama-style core recipe", markdown §3/§7). No bias. RoPE applied inside
    attention to Q/K only. Supports incremental KV-cache decoding. `ffn` is
    swappable (DenseFFN / GatedFFN / MoEFFN) -- the block's own interface never
    changes, only what sits behind it (markdown §5)."""

    def __init__(self, d_model: int = 96, n_q_heads: int = 8, n_kv_heads: int = 2,
                 d_ff: int = 256, ffn: nn.Module = None):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_q_heads, n_kv_heads, bias=False, rope=True)
        self.ffn = ffn if ffn is not None else GatedFFN(d_model, d_ff, act=F.silu, bias=False)
        self.norm1 = nn.RMSNorm(d_model)
        self.norm2 = nn.RMSNorm(d_model)

    def forward(self, x, cos, sin, causal_mask=None, kv_cache: KVCache = None, start_pos: int = None):
        # Bug fix (see [D14]): resolve the position BEFORE deciding the mask, so
        # an omitted causal_mask still gets built over the correct absolute
        # positions when a non-empty KV cache is present (query positions
        # continue from where the cache left off, not from 0).
        resolved_start = _infer_start_pos(kv_cache, start_pos)
        Tq = x.shape[1]
        if causal_mask is None:
            q_pos = torch.arange(resolved_start, resolved_start + Tq)
            k_pos = torch.arange(0, resolved_start + Tq)
            causal_mask = build_causal_mask(q_pos, k_pos)
        a, _ = self.attn(self.norm1(x), causal_mask=causal_mask, cos=cos, sin=sin,
                          kv_cache=kv_cache, start_pos=resolved_start)
        x = x + a
        x = x + self.ffn(self.norm2(x))
        return x


# =============================================================================
# §3 Two extra minimal blocks, ONLY for the [D08] parallel-vs-sequential test
# =============================================================================

class SequentialPreLNBlock(nn.Module):
    """Serial Pre-LN, independent norms: attn updates the residual stream BEFORE
    FFN reads it. x1 = x + Attn(N1(x)); y = x1 + FFN(N2(x1))."""

    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads, n_heads, bias=False)
        self.ffn = GatedFFN(d_model, d_ff, bias=False)
        self.norm1 = nn.RMSNorm(d_model)
        self.norm2 = nn.RMSNorm(d_model)

    def forward(self, x, causal_mask=None):
        a, _ = self.attn(self.norm1(x), causal_mask=causal_mask)
        x1 = x + a
        return x1 + self.ffn(self.norm2(x1))


class ParallelPreLNBlock(nn.Module):
    """GPT-J/PaLM-style parallel residual: y = x + Attn(N(x)) + FFN(N(x)) -- FFN
    reads the SAME normalized ORIGINAL x that attention reads, never the
    attn-updated stream. A single shared norm is used here to match the
    GPT-J/PaLM core formula; sharing norm PARAMETERS is not required for the
    "no dependency" property this test checks (markdown §6)."""

    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads, n_heads, bias=False)
        self.ffn = GatedFFN(d_model, d_ff, bias=False)
        self.norm = nn.RMSNorm(d_model)

    def forward(self, x, causal_mask=None):
        n = self.norm(x)
        a, _ = self.attn(n, causal_mask=causal_mask)
        f = self.ffn(n)
        return x + a + f


def reference_mha(x, q_w, k_w, v_w, o_w, n_heads, causal_mask, cos, sin):
    """Plain, non-grouped reference MHA (every head has its own independent K/V,
    no group broadcasting at all) -- written independently of MultiHeadAttention,
    used only to validate that MultiHeadAttention(n_kv_heads=n_q_heads) is
    mathematically the same computation ([D05])."""
    B, T, d = x.shape
    dh = d // n_heads

    def shape(w):
        return (x @ w.t()).view(B, T, n_heads, dh).transpose(1, 2)

    q, k, v = shape(q_w), shape(k_w), shape(v_w)
    q = apply_rope(q, cos[:T], sin[:T])
    k = apply_rope(k, cos[:T], sin[:T])
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
    scores = scores.masked_fill(~causal_mask, float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    out = (attn @ v).transpose(1, 2).reshape(B, T, d)
    return out @ o_w.t()


def reference_gqa(x, q_w, k_w, v_w, o_w, n_q_heads, n_kv_heads, causal_mask, cos, sin,
                   return_per_head=False):
    """Independent GQA/MQA reference that PHYSICALLY repeat_interleave()s K/V
    heads up to n_q_heads (group = n_q_heads // n_kv_heads) and runs ordinary
    per-head attention -- no view+unsqueeze broadcasting trick at all. Used to
    validate that MultiHeadAttention's view-based GQA broadcast is
    mathematically equivalent to literally copying K/V ([D05]). Unlike
    reference_mha, this covers group sizes > 1 (n_kv_heads < n_q_heads), which
    reference_mha (group == 1 only) cannot exercise."""
    B, T, d = x.shape
    dh = d // n_q_heads
    group = n_q_heads // n_kv_heads

    def shape(w, n_heads):
        return (x @ w.t()).view(B, T, n_heads, dh).transpose(1, 2)

    q = shape(q_w, n_q_heads)
    k = shape(k_w, n_kv_heads)
    v = shape(v_w, n_kv_heads)
    q = apply_rope(q, cos[:T], sin[:T])
    k = apply_rope(k, cos[:T], sin[:T])
    k_rep = k.repeat_interleave(group, dim=1)     # [B, Hq, T, dh] -- physical copy, not a view
    v_rep = v.repeat_interleave(group, dim=1)
    scores = (q @ k_rep.transpose(-2, -1)) / math.sqrt(dh)
    scores = scores.masked_fill(~causal_mask, float("-inf"))
    attn = torch.softmax(scores, dim=-1)
    per_head_out = attn @ v_rep                    # [B, Hq, T, dh], BEFORE o_proj
    out = per_head_out.transpose(1, 2).reshape(B, T, d)
    result = out @ o_w.t()
    if return_per_head:
        return result, per_head_out
    return result


# =============================================================================
# §4 Sanity checks [D01]-[D09]
# =============================================================================

def check_d01_shape_contract():
    B, T, S, d, n_q, n_kv, d_ff = 2, 7, 5, 96, 8, 2, 384
    x = torch.randn(B, T, d)
    memory = torch.randn(B, S, d)          # deliberately S != T
    causal = build_causal_mask(torch.arange(T), torch.arange(T))

    enc = VanillaEncoderLayer2017(d, n_q, d_ff)
    dec2017 = VanillaSeq2SeqDecoderLayer2017(d, n_q, d_ff)
    gpt2 = GPT2StyleDecoderLayer(d, n_q, d_ff)
    cos, sin = build_rope_cache(T, d // n_q)
    modern = ModernDecoderBlock(d, n_q, n_kv, d_ff=256)

    out_enc = enc(x)
    out_dec = dec2017(x, memory, causal_mask=causal)                  # Q len T, K/V len S
    out_gpt2 = gpt2(x, causal_mask=causal)
    out_modern = modern(x, cos, sin, causal_mask=causal)

    shapes = [out_enc.shape, out_dec.shape, out_gpt2.shape, out_modern.shape]
    shapes_ok = all(s == (B, T, d) for s in shapes)
    # cross-attn ran with S=5 != T=7 memory and still returned length-T output:
    # this only works if cross-attn correctly decouples query length from K/V length.
    cross_len_ok = out_dec.shape[1] == T and memory.shape[1] == S and S != T

    # Bug-strength fix (see [D01] in the review ledger): matching shapes alone
    # passes even if cross-attn completely IGNORES memory (e.g. a bug that
    # accidentally reads x_q instead of x_kv). Re-run with different memory and
    # require the output to actually change -- this proves memory is read, not
    # just shape-compatible.
    memory2 = torch.randn(B, S, d)
    out_dec2 = dec2017(x, memory2, causal_mask=causal)
    memory_matters = not torch.allclose(out_dec, out_dec2, atol=1e-6)

    ok = shapes_ok and cross_len_ok and memory_matters
    print(f"[D01] shapes (enc/dec2017/gpt2/modern) = {shapes} (expect 4x ({B},{T},{d})); "
          f"cross-attn ran with Tq={T} != S={S}; changing memory changes cross-attn output "
          f"(memory is actually read, not ignored) = {memory_matters}  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d02_residual_topology():
    d = 96
    x = torch.randn(4, 5, d)
    ln = nn.LayerNorm(d)

    def zero_sublayer(_x):
        return torch.zeros_like(_x)

    # Formula-level proof (§0's three canonical equations, F(.) == 0):
    post_ln_out = ln(x + zero_sublayer(x))       # y = N(x + F(x)) -> y = N(x)
    pre_ln_out = x + zero_sublayer(ln(x))        # y = x + F(N(x)) -> y = x

    # rtol=0 below: these are declared-EXACT claims (F=0 should give EXACTLY
    # N(x) / EXACTLY x), so the default rtol=1e-5 fudge-factor must not be
    # allowed to paper over a real discrepancy.
    post_matches_Nx = torch.allclose(post_ln_out, ln(x), atol=1e-6, rtol=0)
    post_matches_x = torch.allclose(post_ln_out, x, atol=1e-4)
    pre_matches_x = torch.allclose(pre_ln_out, x, atol=1e-6, rtol=0)

    # Class-level corroboration: zero out every sublayer's OUTPUT projection
    # (weight AND bias) in a real Pre-LN block and a real Post-LN block.
    modern = ModernDecoderBlock(d, 8, 2, d_ff=256)
    with torch.no_grad():
        modern.attn.o_proj.weight.zero_()
        modern.ffn.down.weight.zero_()
    cos, sin = build_rope_cache(5, d // 8)
    causal = build_causal_mask(torch.arange(5), torch.arange(5))
    modern_out = modern(x, cos, sin, causal_mask=causal)
    modern_is_identity = torch.allclose(modern_out, x, atol=1e-6, rtol=0)

    enc = VanillaEncoderLayer2017(d, 8, 384)
    with torch.no_grad():
        enc.attn.o_proj.weight.zero_()
        enc.attn.o_proj.bias.zero_()
        enc.ffn.down.weight.zero_()
        enc.ffn.down.bias.zero_()
    enc_out = enc(x)
    enc_not_identity = not torch.allclose(enc_out, x, atol=1e-3)
    # Bug-strength fix (see [D02] in the review ledger): the old check only
    # verified enc_out != x, which any arbitrary wrong transform would also
    # satisfy. A REAL two-sublayer Post-LN block with both sublayer outputs
    # zeroed must equal exactly ln2(ln1(x) + 0) = ln2(ln1(x)) -- compare
    # against that closed form directly, not a generic "!= x" sanity check.
    enc_expected = enc.ln2(enc.ln1(x))
    enc_matches_expected = torch.allclose(enc_out, enc_expected, atol=1e-6, rtol=0)

    ok = (pre_matches_x and post_matches_Nx and (not post_matches_x)
          and modern_is_identity and enc_not_identity and enc_matches_expected)
    print(f"[D02] formula-level: Pre-LN(F=0)==x: {pre_matches_x}; Post-LN(F=0)==N(x): {post_matches_Nx} "
          f"(and != x: {not post_matches_x}); class-level: zeroed ModernDecoderBlock==x: {modern_is_identity}, "
          f"zeroed Post-LN encoder layer != x: {enc_not_identity}, == ln2(ln1(x)) exactly: "
          f"{enc_matches_expected}  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d03_causal_non_leakage():
    d, n_heads, d_ff, T, t = 96, 8, 384, 6, 3
    x = torch.randn(2, T, d)

    # Perturb ONLY a future token (after t) -- tests that changes after
    # position t do not leak BACKWARD into positions <= t.
    x_future_perturbed = x.clone()
    x_future_perturbed[:, t + 1:, :] = torch.randn_like(x_future_perturbed[:, t + 1:, :])

    # Bug-strength fix (see [D03] in the review ledger): "future positions
    # changed" is true even with a completely dead/zeroed attention branch,
    # because each future token's OWN value changed and still feeds its own
    # FFN through the residual stream -- that alone says nothing about causal
    # MIXING. Perturb an EARLIER token instead and require LATER positions to
    # change: that can only happen if attention is actually carrying
    # information FORWARD from position 0 to positions > 0.
    x_past_perturbed = x.clone()
    x_past_perturbed[:, 0, :] = torch.randn_like(x_past_perturbed[:, 0, :])

    causal = build_causal_mask(torch.arange(T), torch.arange(T))
    blk = GPT2StyleDecoderLayer(d, n_heads, d_ff)
    blk.eval()
    with torch.no_grad():
        o1 = blk(x, causal_mask=causal)
        o2 = blk(x_future_perturbed, causal_mask=causal)
        o3 = blk(x_past_perturbed, causal_mask=causal)

    # rtol=0: "unchanged" is a declared-EXACT claim (no backward leakage at all).
    past_unchanged = torch.allclose(o1[:, :t + 1], o2[:, :t + 1], atol=1e-6, rtol=0)
    future_did_change = not torch.allclose(o1[:, t + 1:], o2[:, t + 1:], atol=1e-6)
    forward_mixing_present = not torch.allclose(o1[:, 1:], o3[:, 1:], atol=1e-6)

    ok = past_unchanged and future_did_change and forward_mixing_present
    print(f"[D03] perturbing tokens after t={t}: positions 0..{t} unchanged = {past_unchanged}, "
          f"positions {t + 1}..{T - 1} changed = {future_did_change}; perturbing position 0 changes "
          f"positions 1..{T - 1} (attention actually mixes forward, not a dead branch) = "
          f"{forward_mixing_present}  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d04_full_vs_cached():
    d, n_q, n_kv, T = 96, 8, 2, 6
    x = torch.randn(2, T, d)
    cos, sin = build_rope_cache(T, d // n_q)
    blk = ModernDecoderBlock(d, n_q, n_kv, d_ff=256)
    blk.eval()

    causal = build_causal_mask(torch.arange(T), torch.arange(T))
    with torch.no_grad():
        out_full = blk(x, cos, sin, causal_mask=causal)

    cache = KVCache()
    outs = []
    with torch.no_grad():
        for t in range(T):
            xt = x[:, t:t + 1, :]
            q_pos = torch.tensor([t])
            k_pos = torch.arange(t + 1)               # cache so far, INCLUDING this new token
            mask_t = build_causal_mask(q_pos, k_pos)   # a single new query attends to all of 0..t
            ot = blk(xt, cos, sin, causal_mask=mask_t, kv_cache=cache, start_pos=t)
            outs.append(ot)
    out_cached = torch.cat(outs, dim=1)

    TOL = 1e-4   # bug-fix wording (see [D04] ledger): this is a numerical-equivalence
                 # tolerance, NOT bit-exactness -- do not call it "exact" in the report.
    max_diff = (out_full - out_cached).abs().max().item()
    cache_len_ok = cache.k.shape[2] == T and cache.k.shape[1] == n_kv

    # Bug-strength fix (see [D04] in the review ledger): full-vs-cached agreement
    # alone can't catch a bug shared by BOTH code paths (they route through the
    # same apply_rope/k_proj/v_proj), and only checks cache SHAPE, never cache
    # CONTENT. Independently project+rotate K/V for the WHOLE sequence in one
    # shot (bypassing the incremental per-step append path entirely) and
    # compare against what actually ended up sitting in the cache.
    with torch.no_grad():
        x_norm = blk.norm1(x)
        k_expected = blk.attn._shape(blk.attn.k_proj(x_norm), n_kv)
        k_expected = apply_rope(k_expected, cos[:T], sin[:T])
        v_expected = blk.attn._shape(blk.attn.v_proj(x_norm), n_kv)
    cache_k_diff = (cache.k - k_expected).abs().max().item()
    cache_v_diff = (cache.v - v_expected).abs().max().item()
    cache_content_ok = cache_k_diff < TOL and cache_v_diff < TOL

    ok = max_diff < TOL and cache_len_ok and cache_content_ok
    print(f"[D04] full-forward vs token-by-token cached decode: max|Δ| = {max_diff:.2e} "
          f"(tol {TOL:.0e}, numerically equivalent, NOT bit-identical); cache CONTENT vs "
          f"independently projected+rotated K/V: max|ΔK|={cache_k_diff:.2e}, "
          f"max|ΔV|={cache_v_diff:.2e}; final cache shape head_dim={cache.k.shape[1]} "
          f"(expect {n_kv}) len={cache.k.shape[2]} (expect {T})  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d05_mha_gqa_mqa():
    d, n_heads, T = 96, 8, 5
    x = torch.randn(2, T, d)
    cos, sin = build_rope_cache(T, d // n_heads)
    causal = build_causal_mask(torch.arange(T), torch.arange(T))

    # n_kv_heads == n_q_heads must match an INDEPENDENT plain-MHA reference.
    mha = MultiHeadAttention(d, n_heads, n_heads, bias=False, rope=True)
    out_mha, _ = mha(x, causal_mask=causal, cos=cos, sin=sin)
    out_ref = reference_mha(x, mha.q_proj.weight, mha.k_proj.weight, mha.v_proj.weight,
                             mha.o_proj.weight, n_heads, causal, cos, sin)
    degenerate_to_mha_ok = torch.allclose(out_mha, out_ref, atol=1e-5)

    # Bug-strength fix (see [D05] in the review ledger): n_kv==n_q above only
    # exercises group==1 and can't catch a wrong repeat ORDER (repeat vs
    # repeat_interleave) or a wrong group-to-head assignment. Test a REAL
    # grouping (n_kv_heads=2, group size 4) against an independent
    # repeat_interleave reference.
    n_kv2 = 2
    gqa2 = MultiHeadAttention(d, n_heads, n_kv2, bias=False, rope=True)
    out_gqa2, _ = gqa2(x, causal_mask=causal, cos=cos, sin=sin)
    out_gqa2_ref = reference_gqa(x, gqa2.q_proj.weight, gqa2.k_proj.weight, gqa2.v_proj.weight,
                                  gqa2.o_proj.weight, n_heads, n_kv2, causal, cos, sin)
    gqa2_ok = torch.allclose(out_gqa2, out_gqa2_ref, atol=1e-5)

    # n_kv_heads == 1 (MQA): cache head dim must stay 1 (never physically copied
    # up to n_q_heads); output must match the SAME repeat_interleave reference
    # specialized to n_kv_heads=1 (not just "some shape"); and different query
    # heads must produce genuinely DIFFERENT per-head outputs (proving the
    # group dimension is broadcast independently per head, not collapsed to a
    # single shared result before o_proj mixes everything back together).
    mqa = MultiHeadAttention(d, n_heads, 1, bias=False, rope=True)
    cache = KVCache()
    out_mqa, (k_cached, _) = mqa(x, causal_mask=causal, cos=cos, sin=sin, kv_cache=cache)
    out_mqa_ref, per_head_out = reference_gqa(x, mqa.q_proj.weight, mqa.k_proj.weight, mqa.v_proj.weight,
                                               mqa.o_proj.weight, n_heads, 1, causal, cos, sin,
                                               return_per_head=True)
    mqa_matches_ref = torch.allclose(out_mqa, out_mqa_ref, atol=1e-5)
    mqa_cache_ok = k_cached.shape[1] == 1
    mqa_shape_ok = out_mqa.shape == (2, T, d)
    heads_are_distinct = not torch.allclose(per_head_out[:, 0], per_head_out[:, 1], atol=1e-6)

    ok = (degenerate_to_mha_ok and gqa2_ok and mqa_matches_ref and mqa_cache_ok
          and mqa_shape_ok and heads_are_distinct)
    print(f"[D05] GQA(n_kv=n_q) == independent reference MHA: max|Δ|="
          f"{(out_mha - out_ref).abs().max():.2e}; GQA(n_kv=2, group=4) == independent "
          f"repeat_interleave reference: max|Δ|={(out_gqa2 - out_gqa2_ref).abs().max():.2e}; "
          f"MQA == same reference (n_kv=1): max|Δ|={(out_mqa - out_mqa_ref).abs().max():.2e}, "
          f"cache head_dim={k_cached.shape[1]} (expect 1), output shape={tuple(out_mqa.shape)}, "
          f"per-query-head outputs distinct (not collapsed) = {heads_are_distinct}  "
          f"{'PASS' if ok else 'FAIL'}")
    assert ok


def check_d06_param_formulas():
    # Note: this formula only holds for a config where d_model divides evenly
    # by n_q_heads (as validated at construction time -- see [D10]); an
    # indivisible config is now rejected with ValueError instead of silently
    # producing a DIFFERENT, wrong parameter count that would fail this check
    # for the wrong reason.
    d = 96
    dense_hidden, gated_hidden = 384, 256                 # 4d and 8d/3 respectively
    dense = DenseFFN(d, dense_hidden, act=F.relu, bias=False)
    gated = GatedFFN(d, gated_hidden, act=F.silu, bias=False)
    dense_params = sum(p.numel() for p in dense.parameters())
    gated_params = sum(p.numel() for p in gated.parameters())
    ffn_ok = dense_params == 73728 and gated_params == 73728

    n_q, n_kv = 8, 2
    gqa = MultiHeadAttention(d, n_q, n_kv, bias=False)
    gqa_params = sum(p.numel() for p in gqa.parameters())
    gqa_formula = int(2 * d ** 2 * (1 + n_kv / n_q))
    gqa_ok = gqa_params == gqa_formula == 23040

    mha = MultiHeadAttention(d, n_q, n_q, bias=False)
    mha_params = sum(p.numel() for p in mha.parameters())
    mha_ok = mha_params == 4 * d ** 2 == 36864

    ok = ffn_ok and gqa_ok and mha_ok
    print(f"[D06] FFN params: dense(4d)={dense_params}, gated(8d/3)={gated_params} (both expect 73728); "
          f"attn params: GQA={gqa_params} (formula 2d^2(1+n_kv/n_q)={gqa_formula}), MHA={mha_params} "
          f"(formula 4d^2={4 * d ** 2})  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d07_packed_projection():
    d, T = 96, 5
    x = torch.randn(2, T, d)

    n_q, n_kv, dh = 8, 2, 12
    q_lin = nn.Linear(d, n_q * dh, bias=False)
    k_lin = nn.Linear(d, n_kv * dh, bias=False)
    v_lin = nn.Linear(d, n_kv * dh, bias=False)
    packed_qkv = nn.Linear(d, n_q * dh + 2 * n_kv * dh, bias=False)
    with torch.no_grad():
        packed_qkv.weight.copy_(torch.cat([q_lin.weight, k_lin.weight, v_lin.weight], dim=0))
    sep_qkv = torch.cat([q_lin(x), k_lin(x), v_lin(x)], dim=-1)
    max_diff_qkv = (sep_qkv - packed_qkv(x)).abs().max().item()

    d_ff = 256
    gate_lin = nn.Linear(d, d_ff, bias=False)
    up_lin = nn.Linear(d, d_ff, bias=False)
    packed_gu = nn.Linear(d, 2 * d_ff, bias=False)
    with torch.no_grad():
        packed_gu.weight.copy_(torch.cat([gate_lin.weight, up_lin.weight], dim=0))
    sep_gu = torch.cat([gate_lin(x), up_lin(x)], dim=-1)
    max_diff_gu = (sep_gu - packed_gu(x)).abs().max().item()

    ok = max_diff_qkv < 1e-6 and max_diff_gu < 1e-6
    print(f"[D07] packed QKV vs separate: max|Δ|={max_diff_qkv:.2e}; packed gate+up vs separate: "
          f"max|Δ|={max_diff_gu:.2e}  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d08_parallel_vs_sequential():
    d, n_heads, d_ff, T = 96, 8, 256, 5
    x = torch.randn(2, T, d)
    causal = build_causal_mask(torch.arange(T), torch.arange(T))

    seq_blk = SequentialPreLNBlock(d, n_heads, d_ff)
    par_blk = ParallelPreLNBlock(d, n_heads, d_ff)

    # Bug-strength fix (see [D08] in the review ledger): the old test hooked
    # only the FFN's input and derived "what attn saw" by a SEPARATE, redundant
    # call to seq_blk.attn(...) outside the real forward pass -- for the
    # parallel block it never captured attn's input at all, so "both attn and
    # FFN read the SAME normalized stream" was never actually verified end to
    # end. Hook BOTH attn (input+output) and FFN (input) DURING the one real
    # forward call, for both blocks.
    captured_seq, captured_par = {}, {}

    def make_attn_io_hook(store):
        def hook(_module, inputs, output):
            store["attn_in"] = inputs[0].detach().clone()
            store["attn_out"] = output[0].detach().clone()   # (out, (k, v))
        return hook

    def make_ffn_in_hook(store):
        def hook(_module, inputs, _output):
            store["ffn_in"] = inputs[0].detach().clone()
        return hook

    h1 = seq_blk.attn.register_forward_hook(make_attn_io_hook(captured_seq))
    h2 = seq_blk.ffn.register_forward_hook(make_ffn_in_hook(captured_seq))
    seq_blk(x, causal)
    h1.remove(); h2.remove()

    # rtol=0 throughout: these are declared "exact reconstruction" claims.
    seq_attn_sees_norm1 = torch.allclose(captured_seq["attn_in"], seq_blk.norm1(x), atol=1e-6, rtol=0)
    expected_seq_ffn_in = seq_blk.norm2(x + captured_seq["attn_out"])   # FFN should see the UPDATED stream
    seq_sees_updated_stream = torch.allclose(captured_seq["ffn_in"], expected_seq_ffn_in, atol=1e-6, rtol=0)
    seq_not_original_norm = not torch.allclose(captured_seq["ffn_in"], seq_blk.norm2(x), atol=1e-3)

    h1 = par_blk.attn.register_forward_hook(make_attn_io_hook(captured_par))
    h2 = par_blk.ffn.register_forward_hook(make_ffn_in_hook(captured_par))
    par_blk(x, causal)
    h1.remove(); h2.remove()

    par_attn_sees_norm = torch.allclose(captured_par["attn_in"], par_blk.norm(x), atol=1e-6, rtol=0)
    par_sees_original_stream = torch.allclose(captured_par["ffn_in"], par_blk.norm(x), atol=1e-6, rtol=0)
    # The strongest form of "SAME stream": compare the two captured tensors to
    # EACH OTHER directly, not just each individually to norm(x).
    par_both_read_same_tensor = torch.allclose(captured_par["attn_in"], captured_par["ffn_in"],
                                                atol=1e-6, rtol=0)

    ok = (seq_attn_sees_norm1 and seq_sees_updated_stream and seq_not_original_norm
          and par_attn_sees_norm and par_sees_original_stream and par_both_read_same_tensor)
    print(f"[D08] sequential: attn input == norm1(x): {seq_attn_sees_norm1}, FFN input == "
          f"norm2(x+attn_out) (updated stream): {seq_sees_updated_stream} (and != norm2(x) alone: "
          f"{seq_not_original_norm}); parallel: attn input == norm(x): {par_attn_sees_norm}, FFN "
          f"input == norm(x): {par_sees_original_stream}, attn and FFN inputs are the literal SAME "
          f"tensor: {par_both_read_same_tensor}  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d09_backward_smoke():
    d, n_q, n_kv, T = 96, 8, 2, 6
    x = torch.randn(2, T, d)
    cos, sin = build_rope_cache(T, d // n_q)
    causal = build_causal_mask(torch.arange(T), torch.arange(T))

    dense_blk = ModernDecoderBlock(d, n_q, n_kv, d_ff=256)
    out = dense_blk(x, cos, sin, causal_mask=causal)
    out.pow(2).mean().backward()
    dense_finite = all(p.grad is not None and torch.isfinite(p.grad).all()
                        for p in dense_blk.parameters())
    # Bug-strength fix (see [D09] in the review ledger): "finite" alone passes
    # for an accidentally-detached branch whose grad is a finite ZERO tensor.
    # Require non-zero too.
    dense_nonzero = all(p.grad is not None and p.grad.abs().sum().item() > 0
                         for p in dense_blk.parameters())

    # MoE case: force n_tokens < n_experts so at least one expert is GUARANTEED
    # to receive zero tokens this batch, regardless of routing outcome (pigeonhole).
    n_experts, Tm = 4, 3
    xm = torch.randn(1, Tm, d)
    cos_m, sin_m = build_rope_cache(Tm, d // n_q)
    causal_m = build_causal_mask(torch.arange(Tm), torch.arange(Tm))
    moe_ffn = MoEFFN(d, d_ff=256, n_experts=n_experts, top_k=1)
    moe_blk = ModernDecoderBlock(d, n_q, n_kv, ffn=moe_ffn)
    out_m = moe_blk(xm, cos_m, sin_m, causal_mask=causal_m)
    out_m.pow(2).mean().backward()

    tokens_per_expert = moe_ffn.tokens_per_expert
    unused_exists = any(c == 0 for c in tokens_per_expert)
    moe_experts_ok = True
    for e_idx, expert in enumerate(moe_ffn.experts):
        hit = tokens_per_expert[e_idx] > 0
        for p in expert.parameters():
            if hit:
                # touched experts must get a NON-ZERO gradient, not merely finite
                grad_ok = (p.grad is not None and torch.isfinite(p.grad).all().item()
                           and p.grad.abs().sum().item() > 0)
                moe_experts_ok = moe_experts_ok and grad_ok
            else:
                moe_experts_ok = moe_experts_ok and (
                    p.grad is None or torch.allclose(p.grad, torch.zeros_like(p.grad)))

    # Bug-strength fix (see [D09] in the review ledger): the old check only
    # inspected expert parameters. A detached/STE-broken router, or a bug that
    # accidentally cuts the FFN branch out of the residual stream, would leave
    # the router and/or the rest of the block gradient-free while this
    # experts-only check still passes. Check the router AND the rest of the
    # block (attention + both norms) too.
    router_ok = all(p.grad is not None and torch.isfinite(p.grad).all().item()
                     and p.grad.abs().sum().item() > 0
                     for p in moe_ffn.gate.parameters())
    other_params = (list(moe_blk.attn.parameters()) + list(moe_blk.norm1.parameters())
                     + list(moe_blk.norm2.parameters()))
    rest_of_block_ok = all(p.grad is not None and torch.isfinite(p.grad).all().item()
                            and p.grad.abs().sum().item() > 0
                            for p in other_params)

    moe_ok = moe_experts_ok and router_ok and rest_of_block_ok
    ok = dense_finite and dense_nonzero and unused_exists and moe_ok
    print(f"[D09] dense block: all grads finite = {dense_finite}, all non-zero = {dense_nonzero}; "
          f"MoE tokens_per_expert={tokens_per_expert} (<{n_experts} tokens, so >=1 expert unused by "
          f"pigeonhole); touched experts non-zero+finite grad, untouched experts no/zero grad: "
          f"{moe_experts_ok}; router gate non-zero+finite grad: {router_ok}; rest-of-block "
          f"(attn+norms) non-zero+finite grad: {rest_of_block_ok}  {'PASS' if ok else 'FAIL'}")
    assert ok


# =============================================================================
# §5 Additional regression tests [D10]-[D15], one per B1-class bug found in
# code review (each targets a real silent-failure mode, not just formula
# correctness -- see the review ledger).
# =============================================================================

def check_d10_config_validation():
    """[D10] targets bug #2 (indivisible d_model/n_q_heads silently narrows
    attention) and the MoEFFN top_k silent-downgrade bug."""
    raised_divisibility = False
    try:
        MultiHeadAttention(100, 8)     # 100 % 8 != 0
    except ValueError:
        raised_divisibility = True

    raised_group = False
    try:
        MultiHeadAttention(96, 8, n_kv_heads=3)   # 8 % 3 != 0
    except ValueError:
        raised_group = True

    raised_odd_head = False
    try:
        MultiHeadAttention(48, 16, rope=True)     # d_head = 48//16 = 3 (odd)
    except ValueError:
        raised_odd_head = True

    raised_topk = False
    try:
        MoEFFN(96, d_ff=256, n_experts=4, top_k=2)
    except NotImplementedError:
        raised_topk = True

    ok = raised_divisibility and raised_group and raised_odd_head and raised_topk
    print(f"[D10] d_model % n_q_heads != 0 raises ValueError: {raised_divisibility}; "
          f"n_q_heads % n_kv_heads != 0 raises ValueError: {raised_group}; "
          f"RoPE + odd head_dim raises ValueError: {raised_odd_head}; "
          f"MoEFFN(top_k=2) raises NotImplementedError (no silent top-1 downgrade): {raised_topk}  "
          f"{'PASS' if ok else 'FAIL'}")
    assert ok


def check_d11_kv_cache_consistency():
    """[D11] targets bug #4 (KV cache + start_pos has no consistency
    contract, and the cache never validates shape/device/dtype identity)."""
    d, n_heads, T = 32, 4, 3
    x = torch.randn(1, T, d)
    cos, sin = build_rope_cache(10, d // n_heads)
    mha = MultiHeadAttention(d, n_heads, n_heads, bias=False, rope=True)

    cache = KVCache()
    causal1 = build_causal_mask(torch.arange(T), torch.arange(T))
    mha(x, causal_mask=causal1, cos=cos, sin=sin, kv_cache=cache)   # cache.next_pos becomes T

    # A stale/wrong explicit start_pos against a non-empty cache must raise --
    # NOT silently reuse the wrong RoPE angles for the new chunk.
    x2 = torch.randn(1, 1, d)
    k_pos = torch.arange(T + 1)
    mask_wrong = build_causal_mask(torch.tensor([0]), k_pos)   # deliberately wrong query position
    raised_inconsistent = False
    try:
        mha(x2, causal_mask=mask_wrong, cos=cos, sin=sin, kv_cache=cache, start_pos=0)
    except ValueError:
        raised_inconsistent = True

    # The CORRECT explicit start_pos (== cache.next_pos) must still work fine.
    mask_ok = build_causal_mask(torch.tensor([T]), k_pos)
    no_error_when_consistent = True
    try:
        mha(x2, causal_mask=mask_ok, cos=cos, sin=sin, kv_cache=cache, start_pos=T)
    except ValueError:
        no_error_when_consistent = False

    cache.reset()
    reset_ok = cache.k is None and cache.next_pos == 0

    cache2 = KVCache()
    dh = d // n_heads
    cache2.append(torch.randn(1, n_heads, 2, dh), torch.randn(1, n_heads, 2, dh))
    raised_shape_mismatch = False
    try:
        k_wrong_batch = torch.randn(2, n_heads, 1, dh)   # different batch size than what's cached
        cache2.append(k_wrong_batch, k_wrong_batch)
    except ValueError:
        raised_shape_mismatch = True

    ok = raised_inconsistent and no_error_when_consistent and reset_ok and raised_shape_mismatch
    print(f"[D11] inconsistent start_pos vs non-empty cache raises ValueError: {raised_inconsistent}; "
          f"consistent start_pos does not raise: {no_error_when_consistent}; cache.reset() clears "
          f"state: {reset_ok}; append() with mismatched batch/shape raises ValueError: "
          f"{raised_shape_mismatch}  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d12_cross_attention_guards():
    """[D12] targets bug #5 (self/cross + RoPE + cache "any combination" isn't
    actually a coherent interface -- cross-attention + RoPE/cache must be
    refused, plain cross-attention must keep working)."""
    d, n_heads, T, S = 32, 4, 4, 3
    x_q = torch.randn(1, T, d)
    memory = torch.randn(1, S, d)
    cos, sin = build_rope_cache(10, d // n_heads)

    cross_rope = MultiHeadAttention(d, n_heads, n_heads, bias=False, rope=True)
    raised_rope_cross = False
    try:
        cross_rope(x_q, x_kv=memory, cos=cos, sin=sin)
    except ValueError:
        raised_rope_cross = True

    cross_plain = MultiHeadAttention(d, n_heads, n_heads, bias=False, rope=False)
    cache = KVCache()
    raised_cache_cross = False
    try:
        cross_plain(x_q, x_kv=memory, kv_cache=cache)
    except ValueError:
        raised_cache_cross = True

    # Plain cross-attention (no RoPE, no cache) must still work -- this is
    # exactly what VanillaSeq2SeqDecoderLayer2017 relies on.
    out, _ = cross_plain(x_q, x_kv=memory)
    plain_cross_ok = out.shape == (1, T, d)

    ok = raised_rope_cross and raised_cache_cross and plain_cross_ok
    print(f"[D12] cross-attention + RoPE raises ValueError: {raised_rope_cross}; cross-attention + "
          f"KV cache raises ValueError: {raised_cache_cross}; plain cross-attention (no RoPE/cache) "
          f"still works, shape={tuple(out.shape)}: {plain_cross_ok}  {'PASS' if ok else 'FAIL'}")
    assert ok


def check_d13_masking_safety_and_broadcast():
    """[D13] targets bug #6 (fully-masked row -> NaN) and bug #3 (mask
    broadcasting collides when B == Hkv or B == Tq)."""
    d, n_heads, T = 32, 4, 5

    x1 = torch.randn(1, T, d)
    mha_nan = MultiHeadAttention(d, n_heads, n_heads, bias=False)
    bad_causal = build_causal_mask(torch.arange(T), torch.arange(T)).clone()
    bad_causal[0, :] = False   # query 0 has NO valid key at all (not even itself)
    out_nan, _ = mha_nan(x1, causal_mask=bad_causal)
    no_nan = torch.isfinite(out_nan).all().item()

    # B == n_kv_heads: a shape-guessed [B, Tk] mask could be silently
    # misread as a [1, B, 1, Tq, Tk] KV-head-dim mask instead of the intended
    # per-BATCH mask. Verify per-sample isolation against an independent
    # per-sample reference computed one batch element at a time.
    B = n_kv_heads_eq_B = 2
    mha_b = MultiHeadAttention(d, n_heads, n_kv_heads_eq_B, bias=False)
    x2 = torch.randn(B, T, d)
    causal = build_causal_mask(torch.arange(T), torch.arange(T))
    key_padding_mask = torch.tensor([
        [True, True, True, False, False],
        [True, True, True, True, True],
    ])
    out_batched, _ = mha_b(x2, causal_mask=causal, key_padding_mask=key_padding_mask)
    per_sample_outs = [mha_b(x2[b:b + 1], causal_mask=causal,
                              key_padding_mask=key_padding_mask[b:b + 1])[0] for b in range(B)]
    out_per_sample = torch.cat(per_sample_outs, dim=0)
    batch_hkv_isolation_ok = torch.allclose(out_batched, out_per_sample, atol=1e-5)

    # B == Tq: another special shape where a shape-guessed mask could collide.
    Tq_eq_B = 3
    mha_t = MultiHeadAttention(d, n_heads, n_heads, bias=False)
    x3 = torch.randn(Tq_eq_B, Tq_eq_B, d)
    causal3 = build_causal_mask(torch.arange(Tq_eq_B), torch.arange(Tq_eq_B))
    kpm3 = torch.tensor([[True, True, False], [True, True, True], [False, True, True]])
    out_batched3, _ = mha_t(x3, causal_mask=causal3, key_padding_mask=kpm3)
    per_sample_outs3 = [mha_t(x3[b:b + 1], causal_mask=causal3,
                               key_padding_mask=kpm3[b:b + 1])[0] for b in range(Tq_eq_B)]
    out_per_sample3 = torch.cat(per_sample_outs3, dim=0)
    batch_tq_isolation_ok = torch.allclose(out_batched3, out_per_sample3, atol=1e-5)

    ok = no_nan and batch_hkv_isolation_ok and batch_tq_isolation_ok
    print(f"[D13] fully-masked query row produces finite output (no NaN): {no_nan}; "
          f"key_padding_mask correctly isolated per-batch-sample when B==n_kv_heads ({B}): "
          f"max|Δ| vs per-sample reference={(out_batched - out_per_sample).abs().max():.2e} -> "
          f"{batch_hkv_isolation_ok}; same check when B==Tq ({Tq_eq_B}): "
          f"max|Δ|={(out_batched3 - out_per_sample3).abs().max():.2e} -> {batch_tq_isolation_ok}  "
          f"{'PASS' if ok else 'FAIL'}")
    assert ok


def check_d14_default_causal_behavior():
    """[D14] targets bug #1 (decoder blocks defaulted to mask=None, silently
    running bidirectional attention instead of causal)."""
    d, n_heads, d_ff, T = 32, 4, 128, 5
    x = torch.randn(2, T, d)

    gpt2 = GPT2StyleDecoderLayer(d, n_heads, d_ff)
    gpt2.eval()
    causal = build_causal_mask(torch.arange(T), torch.arange(T))
    with torch.no_grad():
        out_explicit = gpt2(x, causal_mask=causal)
        out_omitted = gpt2(x)                              # no mask argument at all
    gpt2_matches = torch.allclose(out_explicit, out_omitted, atol=1e-6, rtol=0)

    t = 2
    x_future_perturbed = x.clone()
    x_future_perturbed[:, t + 1:, :] = torch.randn_like(x_future_perturbed[:, t + 1:, :])
    with torch.no_grad():
        o1 = gpt2(x)
        o2 = gpt2(x_future_perturbed)
    gpt2_default_is_causal = torch.allclose(o1[:, :t + 1], o2[:, :t + 1], atol=1e-6, rtol=0)

    n_q, n_kv, d_modern = 8, 2, 32
    modern = ModernDecoderBlock(d_modern, n_q, n_kv, d_ff=64)
    modern.eval()
    cos, sin = build_rope_cache(T, d_modern // n_q)
    x_m = torch.randn(2, T, d_modern)
    with torch.no_grad():
        out_m_explicit = modern(x_m, cos, sin, causal_mask=causal)
        out_m_omitted = modern(x_m, cos, sin)               # no mask argument at all
    modern_matches = torch.allclose(out_m_explicit, out_m_omitted, atol=1e-6, rtol=0)

    ok = gpt2_matches and gpt2_default_is_causal and modern_matches
    print(f"[D14] GPT2StyleDecoderLayer: omitted causal_mask == explicit causal_mask: {gpt2_matches}, "
          f"default is genuinely causal (no future->past leakage): {gpt2_default_is_causal}; "
          f"ModernDecoderBlock: omitted causal_mask == explicit causal_mask: {modern_matches}  "
          f"{'PASS' if ok else 'FAIL'}")
    assert ok


def check_d15_rope_device_dtype():
    """[D15] targets bug #7 (RoPE cache hard-fixed to CPU/float32 with no
    device/dtype parameters)."""
    d, n_heads, T = 32, 4, 4
    head_dim = d // n_heads

    cos64, sin64 = build_rope_cache(T, head_dim, dtype=torch.float64)
    dtype_honored = cos64.dtype == torch.float64 and sin64.dtype == torch.float64

    cos32, sin32 = build_rope_cache(T, head_dim)              # default float32 cache
    x64 = torch.randn(1, n_heads, T, head_dim, dtype=torch.float64)
    out64 = apply_rope(x64, cos32, sin32)                     # deliberately mismatched dtype
    cast_ok = out64.dtype == torch.float64

    out64_native = apply_rope(x64, cos64, sin64)
    matches_native_cache = torch.allclose(out64, out64_native, atol=1e-6)

    ok = dtype_honored and cast_ok and matches_native_cache
    print(f"[D15] build_rope_cache(dtype=float64) produces float64 cos/sin: {dtype_honored}; "
          f"apply_rope auto-casts a mismatched (float32) cache to a float64 input's dtype "
          f"(no crash, correct output dtype): {cast_ok}; matches a natively-float64-built cache: "
          f"{matches_native_cache} (max|Δ|={(out64 - out64_native).abs().max():.2e})  "
          f"{'PASS' if ok else 'FAIL'}")
    assert ok


def main():
    if not __debug__:
        print("WARNING: running under `python -O` -- Python's `assert` statement is a no-op in "
              "this mode, which silently disables EVERY correctness check in this script "
              "(it would print 'all ... passed' unconditionally regardless of actual results). "
              "Re-run without -O for the checks below to mean anything.", file=sys.stderr)

    check_d01_shape_contract()
    check_d02_residual_topology()
    check_d03_causal_non_leakage()
    check_d04_full_vs_cached()
    check_d05_mha_gqa_mqa()
    check_d06_param_formulas()
    check_d07_packed_projection()
    check_d08_parallel_vs_sequential()
    check_d09_backward_smoke()
    check_d10_config_validation()
    check_d11_kv_cache_consistency()
    check_d12_cross_attention_guards()
    check_d13_masking_safety_and_broadcast()
    check_d14_default_causal_behavior()
    check_d15_rope_device_dtype()
    print("\nall transformer block sanity checks passed ✓")


if __name__ == "__main__":
    main()