keystone_mlx/dense_control.py

"""Matched dense causal-decoder controls for the Keystone Metal contest.

Keystone's claim is *stored-parameter* efficiency, so any comparison is only
meaningful against a conventional dense decoder holding the same number of
trainable scalars, reading the same tokenizer, and seeing the same batches.
This module supplies that control: a tied dense token table, full-rank
multi-head causal attention, and dense gated channel MLPs.  There is no weight
sharing, low-rank factorization, mixture routing, or quantization here — that is
the point of the control.

Two shapes are published rather than one.  A single hand-picked control shape
can silently flatter the candidate, so the contest trains every shape and the
*best* dense result is used as the control (strongest-control policy).
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
import math
from typing import Any

import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten
import numpy as np

try:
    from keystone_mlx.model import CenteredUnitNorm, DenseTiedTokens, deterministic_coordinates
except ModuleNotFoundError:
    from model import CenteredUnitNorm, DenseTiedTokens, deterministic_coordinates


@dataclass(frozen=True)
class DenseControlConfig:
    """Immutable specification for one matched dense decoder control.

    Args:
        vocab_size: Token IDs accepted by the tied dense interface.
        width: Residual stream width.
        blocks: Number of independent (non-shared) decoder blocks.
        head_dim: Channels per attention head; ``width`` must divide evenly.
        mlp_width: Hidden width of each block's dense gated MLP.
    """

    vocab_size: int = 2_048
    width: int = 128
    blocks: int = 6
    head_dim: int = 32
    mlp_width: int = 148

    def __post_init__(self) -> None:
        if self.vocab_size < 2:
            raise ValueError("vocab_size must be at least two")
        if self.width < 2 or self.blocks < 1 or self.mlp_width < 1:
            raise ValueError("width, blocks, and mlp_width must be positive")
        if self.head_dim < 1 or self.width % self.head_dim != 0:
            raise ValueError("head_dim must divide width exactly")


# The 1M shapes are matched to Keystone-1M-2K's exact 999,744 stored scalars.
# `deep` is 999,680 (0.0064% under) and `wide` is 999,360 (0.0384% under).
DENSE_CONTROL_DEEP = DenseControlConfig(vocab_size=2_048, width=128, blocks=6, head_dim=32, mlp_width=148)
DENSE_CONTROL_WIDE = DenseControlConfig(vocab_size=2_048, width=160, blocks=4, head_dim=32, mlp_width=135)
# 600K M-SimOW scale-down: preserve the 1M deep model's six blocks and 2K
# tied vocabulary while narrowing the residual stream to 96. This geometry
# lands on exactly 600,000 stored parameters.
DENSE_CONTROL_600K_DEEP = DenseControlConfig(vocab_size=2_048, width=96, blocks=6, head_dim=32, mlp_width=104)
# The 10M pair is matched to Trellis-10M-2K's 10,001,408 scalar budget:
# deep = 10,000,256 (-0.0115%); wide = 9,996,800 (-0.0461%).  Both retain the
# same 2K tokenizer interface and are inside the predeclared 0.05% tolerance.
DENSE_CONTROL_10M_DEEP = DenseControlConfig(vocab_size=2_048, width=448, blocks=8, head_dim=64, mlp_width=246)
DENSE_CONTROL_10M_WIDE = DenseControlConfig(vocab_size=2_048, width=512, blocks=5, head_dim=64, mlp_width=481)
# Under-3M deep geometry for the Ember-inspired XSA recurrent scale gate.
# It preserves the 10 physical / 14 effective-block layout with a healthy
# 4.22x gated MLP and totals exactly 2,996,480 stored parameters.
DENSE_CONTROL_3M_DEEP = DenseControlConfig(vocab_size=2_048, width=128, blocks=10, head_dim=32, mlp_width=540)
# Exact 2.9M scale-up: 2,896,640 stored parameters at the 2K vocabulary.
DENSE_CONTROL_2P9M_DEEP = DenseControlConfig(vocab_size=2_048, width=128, blocks=10, head_dim=32, mlp_width=514)
# Matched-budget wider recurrent-XSA candidate. At 2,893,440 stored parameters,
# it trades four physical blocks for a wider residual stream while retaining
# ten effective transformations under XSA's recurrent middle-block layout.
DENSE_CONTROL_2P9M_WIDE = DenseControlConfig(vocab_size=2_048, width=192, blocks=6, head_dim=32, mlp_width=466)
# The 5M scale-up keeps the same tied dense-deep family while increasing width
# and depth together: 5,001,728 stored parameters at the 2K vocabulary.
DENSE_CONTROL_5M_DEEP = DenseControlConfig(vocab_size=2_048, width=256, blocks=8, head_dim=64, mlp_width=386)
DENSE_CONTROL_SHAPES = {
    "1m-2k": {"deep": DENSE_CONTROL_DEEP, "wide": DENSE_CONTROL_WIDE},
    "600k-2k": {"deep": DENSE_CONTROL_600K_DEEP},
    "10m-2k": {"deep": DENSE_CONTROL_10M_DEEP, "wide": DENSE_CONTROL_10M_WIDE},
    "3m-2k": {"deep": DENSE_CONTROL_3M_DEEP},
    "2p9m-2k": {"deep": DENSE_CONTROL_2P9M_DEEP, "wide": DENSE_CONTROL_2P9M_WIDE},
    "5m-2k": {"deep": DENSE_CONTROL_5M_DEEP},
}


def control_for_shape(shape: str, model_size: str = "1m-2k") -> DenseControlConfig:
    """Select a published dense control shape at one fixed budget."""
    try:
        return DENSE_CONTROL_SHAPES[model_size][shape]
    except KeyError as error:
        choices = ", ".join(sorted(DENSE_CONTROL_SHAPES))
        raise ValueError(f"unknown dense control selection {shape!r} at {model_size!r}; budgets: {choices}") from error


class DenseCausalAttention(nn.Module):
    """Full-rank multi-head causal attention with an explicit additive mask."""

    def __init__(self, width: int, head_dim: int) -> None:
        super().__init__()
        if width % head_dim != 0:
            raise ValueError("head_dim must divide width exactly")
        self.heads = width // head_dim
        self.head_dim = head_dim
        self.query = nn.Linear(width, width, bias=False)
        self.key = nn.Linear(width, width, bias=False)
        self.value = nn.Linear(width, width, bias=False)
        self.output = nn.Linear(width, width, bias=False)

    def __call__(self, hidden: mx.array) -> mx.array:
        """Attend only to the current position and its strict prefix.

        Args:
            hidden: States shaped ``[batch, sequence, width]``.

        Returns:
            Context shaped like ``hidden``. Position ``i`` cannot read ``> i``.

        Side effects:
            Allocates per-head query/key/value and mask intermediates.
        """
        if hidden.ndim != 3:
            raise ValueError("hidden must have shape [batch, sequence, width]")
        batch, length, width = hidden.shape

        def split_heads(values: mx.array) -> mx.array:
            return values.reshape(batch, length, self.heads, self.head_dim).transpose(0, 2, 1, 3)

        queries = split_heads(self.query(hidden))
        keys = split_heads(self.key(hidden))
        values = split_heads(self.value(hidden))
        scores = (queries @ keys.transpose(0, 1, 3, 2)) * (self.head_dim**-0.5)
        positions = mx.arange(length)
        causal_block = mx.where(
            positions[:, None] >= positions[None, :], mx.array(0.0, dtype=scores.dtype), mx.array(-1e9, dtype=scores.dtype)
        )
        weights = mx.softmax(scores + causal_block[None, None, :, :], axis=-1)
        attended = (weights @ values).transpose(0, 2, 1, 3).reshape(batch, length, width)
        return self.output(attended)


class DenseGatedMLP(nn.Module):
    """Dense gated channel MLP with independent, unshared matrices."""

    def __init__(self, width: int, mlp_width: int) -> None:
        super().__init__()
        self.expand = nn.Linear(width, 2 * mlp_width, bias=False)
        self.contract = nn.Linear(mlp_width, width, bias=False)

    def __call__(self, state: mx.array) -> mx.array:
        """Apply one gated dense channel transform."""
        content, gate = mx.split(self.expand(state), 2, axis=-1)
        return self.contract(nn.silu(content) * mx.sigmoid(gate))


class DenseControlBlock(nn.Module):
    """One pre-normalized attention plus gated-MLP residual block."""

    def __init__(self, config: DenseControlConfig) -> None:
        super().__init__()
        self.attention_norm = CenteredUnitNorm(config.width)
        self.attention = DenseCausalAttention(config.width, config.head_dim)
        self.mlp_norm = CenteredUnitNorm(config.width)
        self.mlp = DenseGatedMLP(config.width, config.mlp_width)

    def __call__(self, hidden: mx.array) -> mx.array:
        """Run attention and channel mixing with residual connections."""
        hidden = hidden + self.attention(self.attention_norm(hidden))
        return hidden + self.mlp(self.mlp_norm(hidden))


class DenseControlLM(nn.Module):
    """Conventional dense decoder used as Keystone's parameter-matched control."""

    def __init__(self, config: DenseControlConfig = DENSE_CONTROL_DEEP) -> None:
        super().__init__()
        self.config = config
        self.tokens = DenseTiedTokens(config.vocab_size, config.width)
        self.blocks = [DenseControlBlock(config) for _ in range(config.blocks)]
        self.final_norm = CenteredUnitNorm(config.width)

    def __call__(self, token_ids: mx.array) -> mx.array:
        """Compute full-vocabulary causal next-token logits.

        Args:
            token_ids: Integer IDs shaped ``[batch, sequence]``.

        Returns:
            A ``[batch, sequence, vocab_size]`` logit tensor.

        Side effects:
            Allocates coordinate and per-block intermediate arrays.
        """
        if token_ids.ndim != 2:
            raise ValueError("token_ids must have shape [batch, sequence]")
        if token_ids.shape[1] < 1:
            raise ValueError("token_ids must include at least one position")
        hidden = self.tokens.encode(token_ids)
        coordinates = deterministic_coordinates(token_ids.shape[1], self.config.width)
        hidden = hidden + mx.expand_dims(coordinates, axis=0)
        for block in self.blocks:
            hidden = block(hidden)
        return self.tokens.decode(self.final_norm(hidden))


def parameter_accounting(config: DenseControlConfig = DENSE_CONTROL_DEEP) -> dict[str, int]:
    """Return the exact stored-parameter budget, counting the tied table once."""
    vocabulary = config.vocab_size * config.width
    attention = config.blocks * 4 * config.width * config.width
    mlp = config.blocks * 3 * config.width * config.mlp_width
    normalization = (config.blocks * 2 + 1) * 2 * config.width
    total = vocabulary + attention + mlp + normalization
    return {
        "dense_tied_vocabulary_interface": vocabulary,
        "independent_attention": attention,
        "independent_gated_mlps": mlp,
        "normalization": normalization,
        "total": total,
    }


def parameter_count(model: nn.Module) -> int:
    """Count trainable MLX scalars, including each tied array exactly once."""
    return sum(int(parameter.size) for _, parameter in tree_flatten(model.trainable_parameters()))


def future_token_causality_error(
    model: DenseControlLM,
    sequence_length: int = 11,
    batch_size: int = 2,
    seed: int = 23,
) -> float:
    """Return the largest changed prefix logit after deterministic suffix edits."""
    if sequence_length < 3 or batch_size < 1:
        raise ValueError("sequence_length must be at least three and batch_size must be positive")
    generator = np.random.default_rng(seed)
    values = generator.integers(0, model.config.vocab_size, size=(batch_size, sequence_length), dtype=np.int32)
    boundary = sequence_length // 2
    mutated = values.copy()
    mutated[:, boundary:] = (mutated[:, boundary:] + 1) % model.config.vocab_size
    original = model(mx.array(values))
    changed = model(mx.array(mutated))
    mx.eval(original, changed)
    return float(mx.max(mx.abs(original[:, :boundary] - changed[:, :boundary])))


def self_check(config: DenseControlConfig = DENSE_CONTROL_DEEP, reference_total: int = 999_744) -> dict[str, Any]:
    """Validate accounting, shape, causality, and the match tolerance.

    Args:
        config: Dense control shape to validate.
        reference_total: Stored-parameter budget the control must match.

    Returns:
        A JSON-serializable diagnostic record including the signed mismatch.

    Side effects:
        Builds the model and evaluates two Metal inference passes.
    """
    model = DenseControlLM(config)
    formula = parameter_accounting(config)
    observed = parameter_count(model)
    if observed != formula["total"]:
        raise RuntimeError(f"formula counted {formula['total']}, model stores {observed}")
    relative = abs(observed - reference_total) / reference_total
    if relative > 5e-4:
        raise RuntimeError(f"control stores {observed} scalars, {relative:.4%} from reference {reference_total}")
    probe = model(mx.zeros((2, 11), dtype=mx.int32))
    mx.eval(probe)
    if tuple(probe.shape) != (2, 11, config.vocab_size):
        raise RuntimeError(f"unexpected logit shape {tuple(probe.shape)}")
    causal_error = future_token_causality_error(model)
    if causal_error != 0.0:
        raise RuntimeError(f"future-token leakage: {causal_error}")
    return {
        "architecture": "DenseControl",
        "config": asdict(config),
        **formula,
        "reference_total": reference_total,
        "relative_mismatch": relative,
        "causal_forward_shape": list(probe.shape),
        "future_token_causality_error": causal_error,
    }