alephlm.py

"""AlephLM — the full craft, config-driven.

Trigram byte (or BPE) embedding -> pre-norm stack (CausalSDPA majority,
CausalSplatHUB at the configured depths) -> LayerNorm -> DualHead.

Toggle surface (the causal contribution ledger, run at every eval):
    forward(idx, disable_bank=True)   dispatched experts off (exact C6 null)
    forward(idx, disable_hub=True)    hub attention residuals skipped
    forward(idx, disable_head_aleph=True)  gamma path off
"""
from __future__ import annotations

from typing import NamedTuple, Optional

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


class LMOutput(NamedTuple):
    """Still a tuple — `logits, loss = model(x)` keeps working — but also
    HF-duck-typed (`out.logits`, `out.loss`) so frozen-trunk tooling like
    amoe-lora drives the model natively."""
    logits: torch.Tensor
    loss: Optional[torch.Tensor]

from .presets import AlephLMConfig
from .attention import CausalSDPA, CausalSplatHUB
from .bank import AnchoredBank
from .embedding import TrigramByteEmbedding, TokenEmbedding
from .head import DualHead


class Block(nn.Module):
    def __init__(self, cfg: AlephLMConfig, layer_idx: int):
        super().__init__()
        d = cfg.d_model
        self.is_hub = layer_idx in cfg.hub_layers
        self.n1 = nn.LayerNorm(d)
        self.n2 = nn.LayerNorm(d)
        if self.is_hub:
            self.attn = CausalSplatHUB(d, cfg.hub_K, cfg.hub_D, cfg.tau,
                                       chunk=cfg.hub_chunk)
        else:
            self.attn = CausalSDPA(d, cfg.n_heads)
        self.bank = AnchoredBank(d, cfg.bank_experts, cfg.bank_ff, cfg.tau,
                                 cfg.gate_init)

    def forward(self, x, disable_bank=False, disable_hub=False):
        if not (disable_hub and self.is_hub):
            x = x + self.attn(self.n1(x))
        return x + self.bank(self.n2(x), disable_dispatch=disable_bank)

    def prefill(self, x):
        a, cache = self.attn.prefill(self.n1(x))
        x = x + a
        return x + self.bank(self.n2(x)), cache

    def step(self, x_t, cache):
        x_t = x_t + self.attn.step(self.n1(x_t), cache)
        return x_t + self.bank(self.n2(x_t))


class AlephLM(nn.Module):
    def __init__(self, cfg: AlephLMConfig):
        super().__init__()
        self.cfg = cfg
        if cfg.tokenizer == "byte-trigram":
            assert cfg.vocab_size == 256, "byte crafts use vocab 256"
            self.embed = TrigramByteEmbedding(cfg.d_model, cfg.context)
            tied = None
        else:
            self.embed = TokenEmbedding(cfg.vocab_size, cfg.d_model, cfg.context)
            tied = self.embed.emb.weight if cfg.tie_embeddings else None
        self.blocks = nn.ModuleList(
            Block(cfg, i) for i in range(cfg.n_layers))
        self.nf = nn.LayerNorm(cfg.d_model)
        self.head = DualHead(cfg.d_model, cfg.vocab_size, cfg.head_K,
                             cfg.head_D, cfg.tau, tied_weight=tied)

    def forward(self, idx=None, targets=None, disable_bank=False,
                disable_hub=False, disable_head_aleph=False,
                input_ids=None, labels=None, attention_mask=None):
        """HF-style aliases are accepted so frozen-trunk tooling drives the
        model unchanged, WITH HF semantics: `labels` are same-position and
        shifted internally (logits[:-1] vs labels[1:]); `targets` are the
        package's own pre-shifted convention and used as-is. attention_mask
        is deliberately ignored: under causal attention with right-padding
        and -100 label masking, pads can never influence a scored position."""
        if idx is None:
            idx = input_ids
        x = self.embed(idx)
        for b in self.blocks:
            x = b(x, disable_bank=disable_bank, disable_hub=disable_hub)
        h = self.nf(x)
        logits = self.head(h, disable_aleph=disable_head_aleph)
        if targets is not None:                     # pre-shifted (ours)
            loss = F.cross_entropy(
                logits.reshape(-1, logits.shape[-1]).float(),
                targets.reshape(-1), ignore_index=-100)
        elif labels is not None:                    # HF: shift internally
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.shape[-1]).float(),
                labels[:, 1:].reshape(-1), ignore_index=-100)
        else:
            return LMOutput(logits, None)
        return LMOutput(logits, loss)

    # ---------------------------------------------------- incremental decode
    @torch.no_grad()
    def prefill(self, idx):
        """Run the prompt once, return (last-position logits, decode cache).
        The cache carries per-layer attention state, the trigram history
        bytes, and the absolute position cursor."""
        self.eval()
        from .embedding import PAD_ROW
        caches = []
        x = self.embed(idx)
        for b in self.blocks:
            x, c = b.prefill(x)
            caches.append(c)
        h = self.nf(x)
        logits = self.head(h[:, -1:])
        n = idx.shape[1]
        prev2 = idx[:, -2] if n >= 2 else torch.full_like(idx[:, -1], PAD_ROW)
        return logits, {"layers": caches, "t": n,
                        "prev1": idx[:, -1], "prev2": prev2}

    @torch.no_grad()
    def decode_step(self, next_id, cache):
        """One token through the cached path. next_id: (B,) or (B,1)."""
        next_id = next_id.reshape(-1)
        t = cache["t"]
        assert t < self.cfg.context, "decode exceeded the position table"
        if isinstance(self.embed, TrigramByteEmbedding):
            e = (self.embed.emb0(next_id) + self.embed.emb1(cache["prev1"])
                 + self.embed.emb2(cache["prev2"])).unsqueeze(1) \
                + self.embed.pos[:, t:t + 1]
            cache["prev2"] = cache["prev1"]
            cache["prev1"] = next_id
        else:
            e = self.embed.emb(next_id).unsqueeze(1) + self.embed.pos[:, t:t + 1]
        x = e
        for b, c in zip(self.blocks, cache["layers"]):
            x = b.step(x, c)
        cache["t"] = t + 1
        return self.head(self.nf(x))

    @staticmethod
    def _sample(logits, temperature, top_p):
        logits = logits[:, -1].float()
        if temperature <= 0.02:
            return logits.argmax(-1, keepdim=True)
        probs = F.softmax(logits / temperature, dim=-1)
        sp, si = probs.sort(dim=-1, descending=True)
        keep = (sp.cumsum(-1) - sp) < top_p
        keep[..., :1] = True   # top-1 always survives: top_p<=0 must never
        sp = sp * keep         # yield an all-zero row (CUDA multinomial on
        return si.gather(-1, torch.multinomial(  # zeros poisons the context)
            sp / sp.sum(-1, keepdim=True), 1))

    @torch.no_grad()
    def generate(self, idx, max_new: int = 128, temperature: float = 1.0,
                 top_p: float = 0.95, use_cache: bool = True):
        """Cached decode while the sequence fits the position table; any
        remainder (long prompts, fills past the context) continues through
        the sliding-window parallel path — the hub's constant-size state
        cannot evict, so sliding continuation must recompute."""
        self.eval()
        ctx = self.cfg.context
        if use_cache and idx.shape[1] < ctx and max_new > 0:
            n_cached = min(max_new, ctx - idx.shape[1])
            logits, cache = self.prefill(idx)
            for i in range(n_cached):
                nxt = self._sample(logits, temperature, top_p)
                idx = torch.cat([idx, nxt], dim=1)
                if i + 1 < n_cached:
                    logits = self.decode_step(nxt, cache)
            max_new -= n_cached
        for _ in range(max_new):
            logits, _ = self(idx[:, -ctx:])
            nxt = self._sample(logits, temperature, top_p)
            idx = torch.cat([idx, nxt], dim=1)
        return idx

    def param_count(self) -> int:
        seen, total = set(), 0
        for p in self.parameters():
            if id(p) not in seen:
                seen.add(id(p))
                total += p.numel()
        return total