Blame for source/src/deg/model.py

"""deg model: Qwen3.5-0.8B trunk + decision heads + span pointer + LM head (SPEC.md §1, §12).

Layout (token ids; plain text delimiters, no new tokens)::

    trunk : PREFIX  <context tokens>  "\\n"                      -- read once per request
    query : "FIELD /path\\n" ... "OPTIONS:\\n- a\\n- b\\n" "ANSWER:"   -- one per decide/extract
            field, array count, or jev question; continues a copy of the trunk state
    gen   : "KNOWN VALUES:\\n..." "WRITE /path\\n" ... "VALUE:\\n" <text> <|im_end|>

Field independence: the trunk holds the context only (never the schema). Every query is a
*fork* that continues its own copy of the trunk state (Gated DeltaNet recurrent state +
causal-conv tail, attention K/V); forks never see each other. At inference forks run in
fixed :data:`SLOTS`-row micro-batches padded to a :data:`BUCKET` multiple of their own
length, so a field's kernels and shapes depend only on that field (bitwise independence;
``tests/test_model.py``).

Padding is exact: a padded position has ``beta = 0`` and ``g = 0`` in the delta rule (no
write, no decay), its K/V are masked, and the conv tail is gathered at each row's length.

Heads (fp32, on the final-norm hidden state; ``a`` = hidden at the query's last token,
``o_i`` = hidden at the last token of option line i, ``c_t`` = trunk hidden at context
token t):

* choice: ``q(a) . k(o_i) / sqrt(d) + u . o_i`` over the options (jev v1 pointer head);
* score: a separate pointer head over the level lines;
* noul: ``w . a + b`` -> 2-way logits ``[l, 0]`` (``[p_true, p_false]``);
* extract: start/end pointers ``qs(a) . ks(c_t)`` over the context tokens only, plus a null
  logit ``ns(a)`` (``ne(a)``) as one extra outcome;
* count: a linear head over 0..:data:`MAX_COUNT` items, masked to the array bounds (index
  addressed arrays only: run-1 checkpoints, and arrays whose items have no extract leaf);
* anchor: per context token ``t`` an independent "an item starts here" logit
  ``qa(a) . ka(c_t) / sqrt(d) + ua . c_t + b`` (sigmoid, threshold 0.5), masked to the
  array's region (the whole context, or the parent item's region for a nested array). The
  item count is the number of anchors: no cap (SPEC.md §12.6);
* anchored extract: the span logits of a field inside an anchored item get a learned bias
  by bucket of ``t - anchor`` and by whether ``t`` lies inside the item's region;
* generate: the base LM head (tied embeddings), only for ``x-generate`` fields.

Calibration: one temperature per head kind (``choice, score, noul, extract, count``), fit
on validation data after training and stored with the checkpoint.
"""

from __future__ import annotations

import bisect
import glob
import json
import math
import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.utils.checkpoint import checkpoint

from deg import profile

# DEG_HF_HUB: another HF hub root (tests on a machine without /data)
HF_HUB = os.environ.get("DEG_HF_HUB", "/data/gwillen-cache/huggingface/hub")
BASE_GLOB = f"{HF_HUB}/models--Qwen--Qwen3.5-0.8B/snapshots/*"
PREFIX = ("<|im_start|>system\nFill JSON Schema fields from the context: decide, copy spans "
          "from the context, and write text only for fields marked x-generate.<|im_end|>\n"
          "<|im_start|>user\nCONTEXT:\n")
TRUNK_SUFFIX = "\n"
EOS_TEXT = "<|im_end|>"
CPU_BUCKET = 64  # CPU reference: per-row fork width grid
FORK_TOKENS = 16384  # padded fork tokens per batched forward (memory bound only)
FORK_BYTES_PER_TOKEN = 128 * 1024  # peak fork-pass work per padded token (CUDA; measured ~114-125 KiB)
FORK_MEM_FRACTION = 0.5  # share of free device memory one fork batch may plan to use
HEAD_LOGIT_BYTES = 64 * 2**20  # extract logits per head micro-batch; larger sets go to host
DENSE_PREFILL_MAX = 8192  # inputs up to this many tokens are read in one shot
PREFILL_CHUNK = 4096  # longer inputs stream in chunks of this many tokens (memory bound only)
HEAD_ROWS = 8192  # context rows per span-key projection block (memory bound only)
MAX_COUNT = 64  # count head capacity (index-addressed arrays only; anchors have no cap)
ARRAY_MODES = ("index", "anchor")  # model.meta["array_mode"]: how array items are addressed
ANCHOR_INIT_BIAS = -6.0  # anchor logits start near p = 0.0025 (few starts among many tokens)
# relative-offset buckets (t - anchor, in tokens) for the anchored span bias
REL_EDGES = (-256, -64, -16, -4, 0, 1, 2, 4, 8, 16, 32, 64, 128, 256)
HEAD_DIM = 256
KINDS = ("choice", "score", "noul", "extract", "count")




@dataclass(frozen=True)
class BackboneSpec:
    """One supported base model (``deg.train --backbone NAME``, ``deg.json["backbone"]``).

    The deg layout (queries, KNOWN VALUES) is the same text for every backbone; a spec
    says how to load the weights and tokenizer and the chat tokens that differ: ``bos``
    (prepended to every trunk), ``pad`` (padding token), ``gen_turn`` (closes the user turn
    and opens the assistant turn for base-model generation,
    :func:`deg.infer.base_generate_header`; thinking off), ``prefix`` (system turn +
    opened user turn before the context; ChatML :data:`PREFIX` by default) and ``eos``
    (the end-of-turn token every generated value ends with; ChatML :data:`EOS_TEXT`)."""

    name: str
    snapshot_glob: str
    lm_class: str  # transformers class that loads the text model
    bos: str
    pad: str
    gen_turn: str
    prefix: str = PREFIX
    eos: str = EOS_TEXT


QWEN_GEN_TURN = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"


def qwen35_spec(size: str) -> BackboneSpec:
    """Qwen3.5 of any size (``0.8B``, ``2B``, ``4B``, ...): one family, one tokenizer and
    chat format; layer counts, hidden/MLP widths and heads come from the snapshot's
    config (:class:`Backbone` reads them)."""
    return BackboneSpec(
        f"qwen3.5-{size.lower()}", f"{HF_HUB}/models--Qwen--Qwen3.5-{size.upper()}/snapshots/*",
        "Qwen3_5ForCausalLM", "", "<|endoftext|>", QWEN_GEN_TURN)


# Gemma 4 chat format (HF chat_template, thinking off): ``<bos><|turn>system\n...<turn|>\n
# <|turn>user\n...<turn|>\n<|turn>model\n``; a model turn ends with ``<turn|>``. Gemma's
# vocabulary has no ChatML tokens (``<|im_end|>`` would be <unk>), so the layout's chat
# text moves into the spec; the system/user wording is the ChatML layout's.
GEMMA_PREFIX = ("<|turn>system\nFill JSON Schema fields from the context: decide, copy spans "
                "from the context, and write text only for fields marked x-generate.<turn|>\n"
                "<|turn>user\nCONTEXT:\n")
GEMMA_GEN_TURN = "<turn|>\n<|turn>model\n"
# the revision of google/gemma-4-E2B-it every deg-gemma run uses (lab and the training
# box): pinned, so the frozen PLE table and embeddings a checkpoint reuses are its base's
GEMMA4_E2B_IT_REVISION = "3e22461f65e89153144f8adb70e3b8c2cc9845a7"


def gemma4_spec(repo: str) -> BackboneSpec:
    """google/gemma-4-E2B(-it): the text model only (Gemma4ForCausalLM; vision/audio
    towers are not loaded), Gemma's chat format."""
    return BackboneSpec(repo.split("/")[-1].lower(),
                        f"{HF_HUB}/models--{repo.replace('/', '--')}/snapshots/*",
                        "Gemma4ForCausalLM", "<bos>", "<pad>", GEMMA_GEN_TURN, GEMMA_PREFIX,
                        "<turn|>")


BACKBONES = {
    "qwen3.5-0.8b": qwen35_spec("0.8B"),
    "qwen3.5-2b": qwen35_spec("2B"),
    # LiquidAI LFM2.5-1.2B-Instruct: 10 gated short-conv + 6 GQA attention layers, hidden
    # 2048, tied embeddings, ChatML with a <|startoftext|> BOS, no thinking block
    "lfm2.5-1.2b": BackboneSpec(
        "lfm2.5-1.2b", f"{HF_HUB}/models--LiquidAI--LFM2.5-1.2B-Instruct/snapshots/*",
        "Lfm2ForCausalLM", "<|startoftext|>", "<|pad|>",
        "<|im_end|>\n<|im_start|>assistant\n"),
    # Google Gemma 4 E2B: 35 attention layers (28 sliding-window 512 + 7 global), the last
    # 20 reuse the K/V of layers 13/14, per-layer embeddings (PLE), hidden 1536, 262k vocab
    "gemma-4-e2b-it": gemma4_spec("google/gemma-4-E2B-it"),
    "gemma-4-e2b": gemma4_spec("google/gemma-4-E2B"),
}
DEFAULT_BACKBONE = "qwen3.5-0.8b"


def backbone_name(name: str) -> str:
    """Canonical backbone name: ``Qwen/Qwen3.5-2B`` (HF repo id, ``--base-model``) and
    ``qwen3.5-2b`` are the same backbone; ``LiquidAI/LFM2.5-1.2B-Instruct`` is
    ``lfm2.5-1.2b``."""
    low = name.strip().lower()
    if low == "liquidai/lfm2.5-1.2b-instruct":
        return "lfm2.5-1.2b"
    for org in ("qwen/", "google/"):
        if low.startswith(org):
            low = low[len(org):]
    return low


META_FILE = "deg.json"
# checkpoints saved before the rename to deg carry their metadata under this name
LEGACY_META_FILE = "deg.json"


def meta_path(ckpt: Path | str) -> Path:
    """The metadata file of checkpoint dir ``ckpt``: ``deg.json``, else the legacy name of
    checkpoints saved before the rename, else ``deg.json`` (absent)."""
    ckpt = Path(ckpt)
    for name in (META_FILE, LEGACY_META_FILE):
        if (ckpt / name).is_file():
            return ckpt / name
    return ckpt / META_FILE


def read_meta(ckpt: Path | str) -> dict[str, Any]:
    """The parsed metadata of checkpoint dir ``ckpt`` (see :func:`meta_path`)."""
    return json.loads(meta_path(ckpt).read_text())


def backbone_spec(name: str | None = None) -> BackboneSpec:
    name = backbone_name(name or DEFAULT_BACKBONE)
    if name in BACKBONES:
        return BACKBONES[name]
    m = re.fullmatch(r"qwen3\.5-(\d+(?:\.\d+)?b)", name)
    if m:  # any other Qwen3.5 size: the same family, sized by its config
        return qwen35_spec(m.group(1))
    raise ValueError(f"backbone must be one of {sorted(BACKBONES)} or qwen3.5-<size>: "
                     f"{name!r}")


def base_snapshot(backbone: str | None = None) -> Path:
    spec = backbone_spec(backbone)
    found = sorted(glob.glob(spec.snapshot_glob))
    if not found:
        raise FileNotFoundError(f"{spec.name} snapshot not found under {spec.snapshot_glob}")
    return Path(found[-1])


def load_base_lm(backbone: str | None = None, dtype: torch.dtype = torch.bfloat16
                 ) -> nn.Module:
    """The base causal LM of ``backbone`` (text model only) from the local HF cache."""
    import transformers

    spec = backbone_spec(backbone)
    snap = str(base_snapshot(spec.name))
    if spec.lm_class == "Gemma4ForCausalLM":
        # the hub checkpoint is the multimodal Gemma4ForConditionalGeneration: load its text
        # model only (model.language_model.* -> model.*; towers are skipped)
        cfg = transformers.AutoConfig.from_pretrained(snap)
        return transformers.Gemma4ForCausalLM.from_pretrained(
            snap, config=getattr(cfg, "text_config", cfg), dtype=dtype,
            key_mapping={r"^model\.language_model\.": "model."})
    return getattr(transformers, spec.lm_class).from_pretrained(snap, dtype=dtype)


def spec_of_lm(lm: nn.Module, backbone: str | None = None) -> BackboneSpec:
    """The backbone spec whose architecture ``lm`` is (tests pass tiny random models):
    ``backbone`` when it names that architecture's family (a tiny 2B-shaped Qwen3.5 may
    stand in for ``qwen3.5-2b``), else the family's default size."""
    kind = lm.config.model_type
    cls = {"lfm2": "Lfm2ForCausalLM", "gemma4_text": "Gemma4ForCausalLM"}.get(
        kind, "Qwen3_5ForCausalLM")
    if backbone is not None:
        spec = backbone_spec(backbone)
        if spec.lm_class != cls:
            raise ValueError(f"lm is a {kind} model, not {spec.name}")
        return spec
    return BACKBONES[{"Lfm2ForCausalLM": "lfm2.5-1.2b",
                      "Gemma4ForCausalLM": "gemma-4-e2b-it"}.get(cls, "qwen3.5-0.8b")]


# ------------------------------------------------------------------------------ encoding


@dataclass(frozen=True)
class Query:
    """One fork's input: what the field is, rendered from the field alone.

    ``anchor``/``region`` (context token indices; ``region`` = [lo, hi), hi = -1 means the
    end of the context) locate the item this query belongs to: an anchor query's logits
    are masked to ``region``; an extract field inside an anchored item gets the relative
    span bias from ``anchor`` and ``region[1]``. ``anchor = -1``: not inside an anchored
    item."""

    kind: str  # choice | score | noul | extract | count | anchor
    header: str  # FIELD/ITEM/TITLE/DESCRIPTION/KIND lines
    options: tuple[str, ...] = ()
    anchor: int = -1
    region: tuple[int, int] = (0, -1)


@dataclass(frozen=True)
class EncodedQuery:
    kind: str
    ids: tuple[int, ...]
    option_positions: tuple[int, ...]
    answer_position: int
    anchor: int = -1
    region: tuple[int, int] = (0, -1)


@dataclass(frozen=True)
class EncodedContext:
    ids: tuple[int, ...]
    ctx_start: int  # first context token index in ids
    offsets: tuple[tuple[int, int], ...]  # per context token: [start, end) code points
    # multimodal contexts (Gemma 4): per media item (index in ids of its first soft token,
    # soft-token embeddings [n, hidden]); its n placeholder tokens are replaced by them
    media: tuple[tuple[int, Tensor], ...] = ()


@dataclass(frozen=True)
class Media:
    """One image / audio item of a multimodal context: ``embeds`` [n, hidden] are its soft
    tokens in the language model's space (:meth:`MediaTowers.image_embeds` /
    :meth:`MediaTowers.audio_embeds`, or a precomputed cache)."""

    kind: str  # image | audio
    embeds: Tensor
    # grounding (deg.ground): the record's media id and the source geometry a pointer maps
    # back to (image (width, height) pixels, audio seconds); rendering ignores them
    id: str | None = None
    size: tuple[int, int] | None = None
    seconds: float | None = None


class Renderer:
    def __init__(self, tokenizer: Any, spec: BackboneSpec | None = None) -> None:
        spec = spec or backbone_spec()
        self.tokenizer = tokenizer
        self.eos_id = tokenizer.convert_tokens_to_ids(spec.eos)
        self.pad_id = tokenizer.convert_tokens_to_ids(spec.pad)
        unk = getattr(tokenizer, "unk_token_id", None)
        if unk is not None and unk in (self.eos_id, self.pad_id):
            raise ValueError(f"{spec.name}: eos {spec.eos!r} / pad {spec.pad!r} is not a "
                             "token of this tokenizer")
        self.gen_turn = spec.gen_turn  # base-model generation: assistant turn opener
        self._prefix = self.enc(spec.bos + spec.prefix)
        self._suffix = self.enc(TRUNK_SUFFIX)

    def enc(self, text: str) -> list[int]:
        return list(self.tokenizer(text, add_special_tokens=False)["input_ids"]) if text else []

    def context(self, text: str | list[str | Media]) -> EncodedContext:
        """The trunk of a context: text, or (multimodal, Gemma 4) a list of text parts and
        :class:`Media` items. Offsets are code points of the joined text parts; a media
        item renders as HF Gemma 4's placeholder span (``<boi> <image>*n <eoi>``, audio
        ``<boa> <audio>*n <eoa>``) whose tokens get zero-width offsets at their place in
        that text, so the span pointer never addresses them (text-only extraction)."""
        if isinstance(text, str):
            enc = self.tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
            ids = [*self._prefix, *enc["input_ids"], *self._suffix]
            return EncodedContext(tuple(ids), len(self._prefix),
                                  tuple((int(a), int(b)) for a, b in enc["offset_mapping"]))
        tok = self.tokenizer
        ids, offsets, media, pos = list(self._prefix), [], [], 0
        for part in text:
            if isinstance(part, str):
                enc = tok(part, add_special_tokens=False, return_offsets_mapping=True)
                ids += enc["input_ids"]
                offsets += [(pos + int(a), pos + int(b)) for a, b in enc["offset_mapping"]]
                pos += len(part)
                continue
            names = {"image": ("boi_token", "image_token", "eoi_token"),
                     "audio": ("boa_token", "audio_token", "eoa_token")}.get(part.kind)
            if names is None or not all(getattr(tok, n, None) for n in names):
                raise ValueError(f"no {part.kind!r} placeholder tokens in this tokenizer")
            bo, soft, eo = (tok.convert_tokens_to_ids(getattr(tok, n)) for n in names)
            n = int(part.embeds.shape[0])
            media.append((len(ids) + 1, part.embeds))
            ids += [bo, *[soft] * n, eo]
            offsets += [(pos, pos)] * (n + 2)
        ids += self._suffix
        return EncodedContext(tuple(ids), len(self._prefix), tuple(offsets), tuple(media))

    def query(self, q: Query) -> EncodedQuery:
        ids = self.enc(q.header)
        positions = []
        if q.options:
            ids += self.enc("OPTIONS:\n")
            for text in q.options:
                ids += self.enc(f"- {text}")
                positions.append(len(ids) - 1)
                ids += self.enc("\n")
        ids += self.enc("ANSWER:")
        return EncodedQuery(q.kind, tuple(ids), tuple(positions), len(ids) - 1, q.anchor,
                            q.region)

    def known_block(self, known: list[tuple[str, Any]]) -> list[int]:
        """The shared generate prefix: every decided/extracted value, one line each."""
        lines = ["KNOWN VALUES:"]
        lines += [f"{ptr} = {json.dumps(v, ensure_ascii=False)}" for ptr, v in known]
        return self.enc("\n".join(lines) + "\n")

    def generate_prompt(self, known: list[tuple[str, Any]], header: str) -> list[int]:
        """Training layout of one generate fork: the shared known block, then the field's
        own header (inference prefills the block once and forks per field)."""
        return self.known_block(known) + self.enc(header + "VALUE:\n")


def load_tokenizer(snapshot: Path | None = None, backbone: str | None = None):
    from transformers import AutoTokenizer

    return AutoTokenizer.from_pretrained(str(snapshot or base_snapshot(backbone)))


_SPAN_INDEX: list[Any] = [None, None, None]  # offsets object, starts, ends (last context)


def token_span(offsets: tuple[tuple[int, int], ...], start: int, end: int) -> tuple[int, int]:
    """Char span [start, end) -> inclusive context-token span covering it: the first token
    ending after ``start`` and the last token starting before ``end``. O(log T) by
    bisection (tokenizer offsets are non-decreasing); a non-monotone offset list falls
    back to the linear scan."""
    if _SPAN_INDEX[0] is not offsets:
        starts = [a for a, _ in offsets]
        ends = [b for _, b in offsets]
        mono = all(x <= y for x, y in zip(starts, starts[1:])) and all(
            x <= y for x, y in zip(ends, ends[1:]))
        _SPAN_INDEX[:] = [offsets, starts if mono else None, ends]
    starts, ends = _SPAN_INDEX[1], _SPAN_INDEX[2]
    if starts is None:
        i = next(k for k, (_, b) in enumerate(offsets) if b > start)
        j = max(k for k, (a, _) in enumerate(offsets) if a < end)
        return i, j
    i = bisect.bisect_right(ends, start)
    j = bisect.bisect_left(starts, end) - 1
    if i >= len(offsets):
        raise StopIteration
    if j < 0:
        raise ValueError("no token starts before end")
    return i, j


_OPEN = {"(": ")", "[": "]", "{": "}", "<": ">", '"': '"', "“": "”", "‘": "’", "«": "»"}
_CLOSE = {v: k for k, v in _OPEN.items()}
_TRAIL = set(",;:'\"”’»")
_LEAD = set(">'\"“‘«")
_QUOTES = set('"“‘«[<{')  # wrappers stripped as a pair


def _run(text: str, i: int, step: int, lo: int, hi: int) -> int:
    """Length of the '*' run starting at ``i`` walking by ``step`` within [lo, hi)."""
    n = 0
    while lo <= i < hi and text[i] == "*":
        n, i = n + 1, i + step
    return n


def _emphasis(text: str, start: int, end: int) -> tuple[int, int] | None:
    """Strip markdown ``*``/``**`` emphasis only as a symmetric wrapper: the same run of
    ``*`` opens and closes it, either inside the span or (when the token span stopped
    short) right outside it. A one-sided run (``****4821``, ``AB*``) is value text."""
    lead = _run(text, start, 1, start, end)
    trail = _run(text, end - 1, -1, start, end)
    if lead and lead == trail and lead < end - start - lead:
        return start + lead, end - lead
    if lead and not trail and _run(text, end, 1, end, len(text)) == lead \
            and end - start > lead:
        return start + lead, end
    if trail and not lead and _run(text, start - 1, -1, 0, start) == trail \
            and end - start > trail:
        return start, end - trail
    return None


def trim_span(text: str, start: int, end: int) -> tuple[int, int]:
    """Token span chars -> value chars: strip whitespace and boundary punctuation that a
    token carried in (unbalanced brackets/quotes, trailing ``,;:``, symmetric markdown
    ``*`` emphasis). ``*`` is never stripped one-sided: it can be value (``****4821``)."""
    while True:
        while start < end and text[start].isspace():
            start += 1
        while end > start and text[end - 1].isspace():
            end -= 1
        if end - start < 2:
            return start, end
        a, z = text[start], text[end - 1]
        inner = text[start:end]
        emph = _emphasis(text, start, end) if "*" in (a, z) else None
        if emph is not None:
            start, end = emph
        elif a in _QUOTES and z == _OPEN[a]:
            start, end = start + 1, end - 1
        elif a in _OPEN and _OPEN[a] not in inner[1:]:
            start += 1
        elif a in _LEAD and a not in _OPEN:
            start += 1
        elif z in _CLOSE and _CLOSE[z] not in inner[:-1]:
            end -= 1
        elif z in _TRAIL and not (z in _CLOSE and _CLOSE[z] in inner[:-1]):
            end -= 1
        else:
            return start, end


# ------------------------------------------------------------------------- trunk state


@dataclass
class State:
    """Prefix state for B rows. ``lin[i]`` = (recurrent state fp32, conv tail [B,K-1,C])
    per state layer (Qwen3.5 Gated DeltaNet: recurrent [B,H,K,V]; LFM2 short conv: no
    recurrent state, an empty [B,0] placeholder); ``kv[i]`` = (K, V) [B,Hkv,P,D] per attention layer;
    ``kv_mask`` [B,P] marks real prefix positions; ``lengths`` [B] = next position."""

    lin: list[tuple[Tensor, Tensor]]
    kv: list[tuple[Tensor, Tensor]]
    kv_mask: Tensor
    lengths: Tensor

    def select(self, index: Tensor) -> State:
        return State([(r.index_select(0, index), t.index_select(0, index)) for r, t in self.lin],
                     [(k.index_select(0, index), v.index_select(0, index)) for k, v in self.kv],
                     self.kv_mask.index_select(0, index), self.lengths.index_select(0, index))


def _l2norm(x: Tensor, eps: float = 1e-6) -> Tensor:
    return x * torch.rsqrt((x * x).sum(-1, keepdim=True) + eps)


def delta_rule_reference(q, k, v, g, beta, initial_state):
    """Exact per-token gated delta rule (CPU path; matches HF's recurrent reference)."""
    dtype = v.dtype
    q, k, v = _l2norm(q.float()), _l2norm(k.float()), v.float()
    q = q * (q.shape[-1] ** -0.5)
    B, L, H, Dk = k.shape
    state = (initial_state.float() if initial_state is not None
             else torch.zeros(B, H, Dk, v.shape[-1], device=v.device))
    outs = []
    for t in range(L):
        state = state * g[:, t].float().exp()[..., None, None]
        mem = (state * k[:, t, :, :, None]).sum(-2)
        delta = (v[:, t] - mem) * beta[:, t].float()[..., None]
        state = state + k[:, t, :, :, None] * delta[:, :, None, :]
        outs.append((state * q[:, t, :, :, None]).sum(-2))
    return torch.stack(outs, 1).to(dtype), state


def _delta_rule(q, k, v, g, beta, initial_state):
    if q.is_cuda:
        from fla.ops.gated_delta_rule import chunk_gated_delta_rule

        return chunk_gated_delta_rule(q, k, v, g=g, beta=beta, initial_state=initial_state,
                                      output_final_state=True, use_qk_l2norm_in_kernel=True)
    return delta_rule_reference(q, k, v, g, beta, initial_state)


# Elementwise pieces as pure tensor functions; on CUDA they run torch.compile'd (fused,
# dynamic shapes, one graph shared by all layers since weights are arguments).


def _rms(x: Tensor, w: Tensor, eps: float) -> Tensor:
    xf = x.float()
    return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float())
            ).type_as(x)


def _swiglu(gate: Tensor, up: Tensor) -> Tensor:
    return F.silu(gate) * up


def _dn_gates(b: Tensor, a: Tensor, A_log: Tensor, dt_bias: Tensor, valid: Tensor
              ) -> tuple[Tensor, Tensor]:
    """beta, g (log decay) with padded positions made exact no-ops (beta = 0, g = 0)."""
    v = valid[..., None]
    beta = b.sigmoid() * v
    g = -A_log.float().exp() * F.softplus(a.float() + dt_bias) * v
    return beta.to(b.dtype), g


def _conv_silu(ext: Tensor, w: Tensor) -> Tensor:
    """Depthwise causal conv over ext [B, K-1+L, C] with w [C, K] -> silu, [B, L, C]."""
    K = w.shape[1]
    L = ext.shape[1] - K + 1
    w = w.to(ext.dtype)
    out = ext[:, 0:L] * w[:, 0]
    for k in range(1, K):
        out = out + ext[:, k:k + L] * w[:, k]
    return F.silu(out)


def _gated_rms(x: Tensor, gate: Tensor, w: Tensor, eps: float) -> Tensor:
    xf = x.float()
    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
    return ((w * xf.to(x.dtype)) * F.silu(gate.float())).to(x.dtype)


def _attn_gate(o: Tensor, gate: Tensor) -> Tensor:
    return o * torch.sigmoid(gate)


def _rms_w(x: Tensor, w: Tensor, eps: float) -> Tensor:
    """LFM2 RMSNorm: ``w * norm(x).to(dtype)`` (weight not offset by 1, as in HF)."""
    xf = x.float()
    return w * (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)).to(x.dtype)


def _short_conv(ext: Tensor, w: Tensor) -> Tensor:
    """LFM2 depthwise causal conv over ext [B, K-1+L, C] with w [C, K], no activation,
    fp32 accumulation in a fixed tap order (elementwise: batch-invariant), [B, L, C]."""
    K = w.shape[1]
    L = ext.shape[1] - K + 1
    e, wf = ext.float(), w.float()
    out = e[:, 0:L] * wf[:, 0]
    for k in range(1, K):
        out = out + e[:, k:k + L] * wf[:, k]
    return out.to(ext.dtype)


def _rms_g(x: Tensor, w: Tensor, eps: float) -> Tensor:
    """Gemma 4 RMSNorm (HF ``Gemma4RMSNorm``): ``norm(x) * w`` in fp32, then cast."""
    xf = x.float()
    return (xf * torch.pow(xf.pow(2).mean(-1, keepdim=True) + eps, -0.5) * w.float()
            ).type_as(x)


def _rms_n(x: Tensor, eps: float) -> Tensor:
    """Gemma 4 RMSNorm without a scale (``with_scale=False``: the value norm)."""
    xf = x.float()
    return (xf * torch.pow(xf.pow(2).mean(-1, keepdim=True) + eps, -0.5)).type_as(x)


def _geglu(gate: Tensor, up: Tensor) -> Tensor:
    """``gelu_pytorch_tanh(gate) * up`` (Gemma's MLP and its per-layer-input gate)."""
    return F.gelu(gate, approximate="tanh") * up


_FUSED: dict[str, Any] = {}
_RAW = {"rms": _rms, "swiglu": _swiglu, "dn_gates": _dn_gates, "conv_silu": _conv_silu,
        "gated_rms": _gated_rms, "attn_gate": _attn_gate, "rms_w": _rms_w,
        "short_conv": _short_conv, "rms_g": _rms_g, "rms_n": _rms_n, "geglu": _geglu}


def _op(name: str, x: Tensor):
    if not x.is_cuda or not COMPILE:
        return _RAW[name]
    if name not in _FUSED:
        import torch._dynamo.config as dc

        dc.recompile_limit = max(dc.recompile_limit, 64)
        dc.accumulated_recompile_limit = max(dc.accumulated_recompile_limit, 512)
        _FUSED[name] = torch.compile(_RAW[name], dynamic=True)
    return _FUSED[name]


COMPILE = True


def _norm(module: nn.Module, x: Tensor) -> Tensor:
    return _op("rms", x)(x, module.weight, module.eps)


# ------------------------------------------------------------------------------ LoRA
# Low-rank adapters on the backbone projections (SPEC.md §12.7). Several NAMED adapters
# live side by side in one pair of matrices per projection: ``lora_A`` [n*r, in],
# ``lora_B`` [out, n*r]. A per-row selector ``sel`` [rows, n*r] (ones on the chosen
# adapter's r columns, zeros elsewhere; None = adapter off) picks each row's adapter, so
# a decide fork, an extract fork and an adapter-off (= base) row can share one batch: the
# delta of a row depends only on that row and its selector (batch-invariant on the fork
# path, whose GEMMs are the bi_linear kernels).
LORA_TARGETS = {"linear_attn": ("in_proj_qkv", "in_proj_z", "out_proj"),
                "self_attn": ("q_proj", "k_proj", "v_proj", "o_proj"),
                "mlp": ("gate_proj", "up_proj", "down_proj")}
# which adapter a fork kind uses when the adapters are split by capability
ADAPTER_OF_KIND = {"choice": "decide", "score": "decide", "noul": "decide",
                   "extract": "extract", "count": "extract", "anchor": "extract"}


# ------------------------------------------------------------------------------ deg-MoE
# Task-routed mixture of experts in ONE model (:func:`add_moe`, results/deg/moe-v1):
# every decoder layer has two FFN experts, ``mlp`` (R, the reader: deg's own fine-tuned
# MLP) and ``mlp_w`` (W, the writer: initialised from the base model's MLP), and the
# attention / linear-attention projections may carry a WRITER LoRA. Embeddings, attention,
# norms and the LM head are shared. Routing is deterministic by token role, never learned:
# a forward call whose ``sel`` is a :class:`WriteRoute` runs every row through W (+ the
# writer LoRA); any other ``sel`` (None or a LoRA selector tensor) runs R. The context
# trunk and decide/extract/count/anchor forks are R calls; the KNOWN VALUES block, the
# generate header and every generated token are W calls that read the R trunk's K/V and
# recurrent state (one prefill). Routing is per call, so it is decided on the host: no
# data-dependent branch, CUDA-graph capture sees a fixed path, rows stay batch-invariant.


@dataclass(frozen=True)
class WriteRoute:
    """``sel`` of a writer (W-expert) call: ``lora`` [B, r] ones (the writer LoRA on) or
    None (an MoE without a writer LoRA). Indexing slices the rows."""

    lora: Tensor | None = None

    def __getitem__(self, idx: Any) -> WriteRoute:
        return WriteRoute(None if self.lora is None else self.lora[idx])


# D/E/G task MoE (:func:`add_task_moe`): every decoder layer has an FFN expert per
# capability. E = the layer's own ``mlp`` (it also reads the context: the span / anchor
# keys are trunk hidden states, so the reader is trained by extract + decide), D =
# ``mlp_d`` (decide forks: choice / score / noul), G = ``mlp_g`` (KNOWN VALUES, the
# generate header and every decoded token). Attention, norms, embeddings and PLE are
# shared. A call's route is a :class:`TaskRoute` (one expert per call, decided on the
# host like :class:`WriteRoute`); fork batches are grouped by route.
TASK_EXPERT_OF_KIND = {"choice": "d", "score": "d", "noul": "d", "extract": "e",
                       "count": "e", "anchor": "e", "generate": "g"}


@dataclass(frozen=True)
class TaskRoute:
    """``sel`` of a D/E/G MoE call: every row runs FFN expert ``expert`` (``e`` = the
    layer's ``mlp``, else ``mlp_<expert>``). Indexing keeps the route."""

    expert: str

    def __getitem__(self, idx: Any) -> TaskRoute:
        return self


def _lsel(sel: Tensor | WriteRoute | TaskRoute | None) -> Tensor | None:
    """The LoRA selector inside a ``sel`` (a WriteRoute's writer LoRA columns)."""
    if isinstance(sel, TaskRoute):
        return None
    return sel.lora if isinstance(sel, WriteRoute) else sel


def _expert(layer, sel: Tensor | WriteRoute | TaskRoute | None, name: str = "mlp"
            ) -> nn.Module:
    """The FFN expert of this call: W (``<name>_w``) for a WriteRoute, ``<name>_<x>`` for a
    TaskRoute to expert x (``e``: ``<name>`` itself), else R / the dense FFN."""
    if isinstance(sel, WriteRoute):
        return getattr(layer, name + "_w")  # AttributeError: not an MoE model
    if isinstance(sel, TaskRoute) and sel.expert != "e":
        return getattr(layer, f"{name}_{sel.expert}")
    return getattr(layer, name)


def _lin(module: nn.Module, x: Tensor, sel: Tensor | WriteRoute | None) -> Tensor:
    """``module(x)`` plus the selected LoRA delta (training / trunk path, autograd)."""
    y = module(x)
    sel = _lsel(sel)
    if sel is None:
        return y
    A = getattr(module, "lora_A", None)
    if A is None:
        return y
    t = F.linear(x, A) * sel[:, None, :].to(x.dtype)
    return y + F.linear(t, module.lora_B) * module.lora_scale


def _flora(module: nn.Module, x: Tensor, y: Tensor, sel: Tensor | WriteRoute | None
           ) -> Tensor:
    """Fork path: ``y`` (= x @ W^T) plus the selected LoRA delta through the same
    batch-invariant GEMM config as the base projection."""
    sel = _lsel(sel)
    if sel is None:
        return y
    A = getattr(module, "lora_A", None)
    if A is None:
        return y
    from deg.kernels import bi_linear, bi_linear_small

    mm = bi_linear_small if _SMALL_GEMM else bi_linear
    t = mm(x, A) * sel[:, None, :].to(x.dtype)
    return y + mm(t, module.lora_B) * module.lora_scale


def _mlp(layer, h: Tensor, sel: Tensor | WriteRoute | None = None) -> Tensor:
    m = _expert(layer, sel)
    x = _norm(layer.post_attention_layernorm, h)
    a = _op("swiglu", x)(_lin(m.gate_proj, x, sel), _lin(m.up_proj, x, sel))
    return h + _lin(m.down_proj, a, sel)


def _linear_layer(layer, h, valid, lengths_new, rec, tail, rows=None, sel=None):
    """``rows``: ``rec``/``tail`` are the un-selected prefix states and row b continues
    prefix row ``rows[b]``; selecting here (inside the checkpoint) keeps the per-row copies
    out of the saved activations: they are rebuilt during the backward recompute."""
    attn = layer.linear_attn
    B, L, _ = h.shape
    if rows is not None:
        rec, tail = rec.index_select(0, rows), tail.index_select(0, rows)
    xn = _norm(layer.input_layernorm, h)
    mixed = _lin(attn.in_proj_qkv, xn, sel)
    z = _lin(attn.in_proj_z, xn, sel).reshape(B, L, -1, attn.head_v_dim)
    beta, g = _op("dn_gates", h)(attn.in_proj_b(xn), attn.in_proj_a(xn), attn.A_log,
                                 attn.dt_bias, valid)
    ext = torch.cat([tail.to(mixed.dtype), mixed], 1)  # [B, K-1+L, C]
    K1 = tail.shape[1]
    conv = _op("conv_silu", h)(ext, attn.conv1d.weight[:, 0])  # no conv bias in Qwen3.5
    gather = lengths_new[:, None] + torch.arange(K1, device=h.device)[None]  # last K-1 real
    new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1]))
    qq, kk, vv = torch.split(conv, [attn.key_dim, attn.key_dim, attn.value_dim], dim=-1)
    qq = qq.reshape(B, L, -1, attn.head_k_dim)
    kk = kk.reshape(B, L, -1, attn.head_k_dim)
    vv = vv.reshape(B, L, -1, attn.head_v_dim)
    rep = attn.num_v_heads // attn.num_k_heads
    if rep > 1:
        qq, kk = qq.repeat_interleave(rep, dim=2), kk.repeat_interleave(rep, dim=2)
    core, new_rec = _delta_rule(qq, kk, vv, g, beta.to(qq.dtype), rec)
    core = _op("gated_rms", h)(core, z, attn.norm.weight, attn.norm.variance_epsilon)
    h = h + _lin(attn.out_proj, core.reshape(B, L, -1), sel)
    return _mlp(layer, h, sel), new_rec, new_tail


def _attn_qkv(layer, h, cos, sin, sel=None):
    """Attention projections: q, k, v [B, H, L, D] (rotary applied) and the output gate."""
    from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb

    a = layer.self_attn
    B, L, _ = h.shape
    xn = _norm(layer.input_layernorm, h)
    q, gate = torch.chunk(_lin(a.q_proj, xn, sel).view(B, L, -1, a.head_dim * 2), 2, dim=-1)
    q = _norm(a.q_norm, q).transpose(1, 2)
    k = _norm(a.k_norm, _lin(a.k_proj, xn, sel).view(B, L, -1, a.head_dim)).transpose(1, 2)
    v = _lin(a.v_proj, xn, sel).view(B, L, -1, a.head_dim).transpose(1, 2)
    q, k = apply_rotary_pos_emb(q, k, cos, sin)
    return q, k, v, gate


def _attn_out(layer, h, o, gate, sel=None):
    """o [B, L, Hq*D] -> gated output projection, residual, MLP."""
    B, L, _ = h.shape
    h = h + _lin(layer.self_attn.o_proj, _op("attn_gate", h)(o, gate.reshape(B, L, -1)), sel)
    return _mlp(layer, h, sel)


def _attn_layer(fam, layer, h, cos, sin, k_prev, v_prev, mask, rows=None, sel=None):
    """``rows``: as in :func:`_linear_layer` (prefix K/V selected inside the checkpoint);
    ``fam``: the :class:`Family` whose projections run."""
    a = layer.self_attn
    B, L, _ = h.shape
    if rows is not None:
        k_prev, v_prev = k_prev.index_select(0, rows), v_prev.index_select(0, rows)
    q, k, v, gate = fam.attn_qkv(layer, h, cos, sin, sel)
    keys = torch.cat([k_prev.to(k.dtype), k], 2)
    values = torch.cat([v_prev.to(v.dtype), v], 2)
    o = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask[:, None],
                                       scale=a.scaling, enable_gqa=True)
    o = o.transpose(1, 2).reshape(B, L, -1)
    return fam.attn_out(layer, h, o, gate, sel), k, v


def _chunk_attn_layer(fam, layer, h, cos, sin, kparts, vparts, kv_mask, s, sel=None):
    """One attention layer over one chunk of a long right-padded input (training,
    :meth:`Backbone._forward_chunked_ckpt`): keys are the incoming prefix (``kv_mask``),
    the ``s`` earlier tokens of this call (all real for a real query: right padding) and
    the chunk's own tokens (causal). The [B, chunk, P+s+chunk] mask lives only inside the
    checkpoint. Returns (h, k, v) of the chunk."""
    a = layer.self_attn
    B, c, _ = h.shape
    q, k, v, gate = fam.attn_qkv(layer, h, cos, sin, sel)
    keys = torch.cat([*(t.to(k.dtype) for t in kparts), k], 2)
    values = torch.cat([*(t.to(v.dtype) for t in vparts), v], 2)
    causal = torch.ones(c, c, dtype=torch.bool, device=h.device).tril()
    mask = torch.cat([kv_mask[:, None, :].expand(-1, c, -1),
                      torch.ones(B, c, s, dtype=torch.bool, device=h.device),
                      causal[None].expand(B, -1, -1)], 2)
    o = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask[:, None],
                                       scale=a.scaling, enable_gqa=True)
    return fam.attn_out(layer, h, o.transpose(1, 2).reshape(B, c, -1), gate, sel), k, v


@dataclass
class ForkState:
    """B fork rows, each continuing one of N prefixes (SPEC.md §12.2).

    Prefixes (batch-1, unpadded trunk states, e.g. the contexts of several requests) are
    read by reference: attention reads row b's prefix K/V at ``pk``/``pv`` [Hkv, total, D]
    positions ``poff[b] .. poff[b]+plen[b]-1`` and the recurrent kernel starts row b from
    ``prefix_rec[prow[b]]``, so no per-row copy of a context exists. Per row: the
    recurrent state after the row's own tokens (``rec``; read once ``fresh`` is False),
    the conv tail, own K/V buffers [B, Hkv, cap, D] and ``own_len`` [B]. ``static``
    states (CUDA-graph buffers) are updated strictly in place."""

    prefix_rec: list[Tensor]
    pk: list[Tensor]
    pv: list[Tensor]
    poff: Tensor
    plen: Tensor
    prow: Tensor
    rec: list[Tensor | None]
    tail: list[Tensor]
    own_k: list[Tensor]
    own_v: list[Tensor]
    own_len: Tensor
    fresh: bool = True
    static: bool = False
    # Gemma 4 sliding layers: row b's prefix tail is ``pk[slot]`` positions
    # ``wpoff[b] .. wpoff[b]+wplen[b]-1`` (the last <= window-1 positions of its prefix)
    wpoff: Tensor | None = None
    wplen: Tensor | None = None
    ple_rows: Tensor | None = None  # Gemma 4 decode graph: this pass's PLE table rows

    @property
    def rows(self) -> int:
        return int(self.own_len.shape[0])


_SMALL_GEMM = False  # set by Backbone.fork_forward(decode=True)


def _fop(name: str, x: Tensor):
    """Fork-path ops, batch-invariant: row norms are fixed-config Triton kernels (one
    program per row) and elementwise ops run eagerly. torch.compile is NOT used here: it
    specializes size-1 dims, so one row alone and the same row in a batch would run
    different generated kernels."""
    if not x.is_cuda:
        return _RAW[name]
    from deg.kernels import gated_rms_rows, rms_rows

    if name == "rms":
        return rms_rows
    if name == "gated_rms":
        return gated_rms_rows
    if name == "rms_w":
        from deg.kernels import rms_rows_w

        return rms_rows_w
    if name in ("rms_g", "rms_n"):
        from deg.kernels import rms_rows_g

        return rms_rows_g if name == "rms_g" else (lambda x, eps: rms_rows_g(x, None, eps))
    return _RAW[name]


def _fnorm(module: nn.Module, x: Tensor) -> Tensor:
    return _fop("rms", x)(x, module.weight, module.eps)


def _bi(module: nn.Module, x: Tensor) -> Tensor:
    """Batch-invariant projection. Each call site always uses ONE kernel config: decode
    steps the small-M split kernel, every other fork pass the wide one."""
    from deg.kernels import bi_linear, bi_linear_small

    if _SMALL_GEMM:
        return bi_linear_small(x, module.weight, module.bias)
    return bi_linear(x, module.weight, module.bias)


def _bi_fused(owner: nn.Module, names: tuple[str, ...], x: Tensor) -> tuple[Tensor, ...]:
    """Several bias-free projections of the same input as ONE batch-invariant GEMM. The
    weights are concatenated once and the original parameters become views into the
    fused matrix (no extra memory; a later load_state_dict writes through). Each output
    column is still its own K-loop, so fusing does not change any row's independence."""
    from deg.kernels import bi_linear, bi_linear_small

    key = "_deg_fused_" + "_".join(names)
    cached = owner.__dict__.get(key)
    if cached is None:
        mods = [getattr(owner, n) for n in names]
        assert all(m.bias is None for m in mods)
        w = torch.cat([m.weight.data for m in mods]).contiguous()
        sizes, off = [], 0
        for m in mods:
            n = m.weight.shape[0]
            m.weight.data = w[off:off + n]
            sizes.append(n)
            off += n
        cached = owner.__dict__[key] = (w, sizes)
    w, sizes = cached
    y = (bi_linear_small if _SMALL_GEMM else bi_linear)(x, w)
    return tuple(torch.split(y, sizes, -1))


def _fork_mlp(layer, h: Tensor, sel: Tensor | WriteRoute | None = None) -> Tensor:
    m = _expert(layer, sel)
    x = _fnorm(layer.post_attention_layernorm, h)
    gate, up = _bi_fused(m, ("gate_proj", "up_proj"), x)
    gate, up = _flora(m.gate_proj, x, gate, sel), _flora(m.up_proj, x, up, sel)
    a = _fop("swiglu", x)(gate, up)
    return h + _flora(m.down_proj, a, _bi(m.down_proj, a), sel)


def _fork_linear_layer(layer, h, valid, lengths, fs: ForkState, li: int, keep: bool,
                       sel: Tensor | None = None):
    from deg.kernels import fork_delta_rule

    attn = layer.linear_attn
    B, L, _ = h.shape
    xn = _fnorm(layer.input_layernorm, h)
    mixed, z, pb, pa = _bi_fused(attn, ("in_proj_qkv", "in_proj_z", "in_proj_b",
                                        "in_proj_a"), xn)
    mixed = _flora(attn.in_proj_qkv, xn, mixed, sel)
    z = _flora(attn.in_proj_z, xn, z, sel).reshape(B, L, -1, attn.head_v_dim)
    beta, g = _fop("dn_gates", h)(pb, pa, attn.A_log, attn.dt_bias, valid)
    tail = fs.tail[li]
    ext = torch.cat([tail.to(mixed.dtype), mixed], 1)
    K1 = tail.shape[1]
    conv = _fop("conv_silu", h)(ext, attn.conv1d.weight[:, 0])
    if keep:
        gather = lengths[:, None] + torch.arange(K1, device=h.device)[None]
        new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1]))
        if fs.static:
            fs.tail[li].copy_(new_tail)
        else:
            fs.tail[li] = new_tail
    qq, kk, vv = torch.split(conv, [attn.key_dim, attn.key_dim, attn.value_dim], dim=-1)
    qq = qq.reshape(B, L, -1, attn.head_k_dim)
    kk = kk.reshape(B, L, -1, attn.head_k_dim)
    vv = vv.reshape(B, L, -1, attn.head_v_dim)
    own = fs.rec[li]
    if fs.fresh:
        h0, h0_idx = fs.prefix_rec[li], fs.prow
    else:
        assert own is not None
        h0, h0_idx = own, torch.arange(B, device=h.device)
    store = None
    if keep:
        store = own if own is not None else torch.empty(
            B, *h0.shape[1:], device=h.device, dtype=torch.float32)
    core = fork_delta_rule(qq, kk, vv, g, beta.to(qq.dtype), h0, h0_idx, store)
    if keep:
        fs.rec[li] = store
    core = _fop("gated_rms", h)(core, z, attn.norm.weight, attn.norm.variance_epsilon)
    c2 = core.reshape(B, L, -1)
    h = h + _flora(attn.out_proj, c2, _bi(attn.out_proj, c2), sel)
    return _fork_mlp(layer, h, sel)


def _fork_attn_layer(layer, h, cos, sin, fs: ForkState, ai: int, sel: Tensor | None = None):
    from transformers.models.qwen3_5.modeling_qwen3_5 import apply_rotary_pos_emb

    from deg.kernels import fork_attention

    a = layer.self_attn
    B, L, _ = h.shape
    xn = _fnorm(layer.input_layernorm, h)
    pq, pk_, pv_ = _bi_fused(a, ("q_proj", "k_proj", "v_proj"), xn)
    pq, pk_ = _flora(a.q_proj, xn, pq, sel), _flora(a.k_proj, xn, pk_, sel)
    pv_ = _flora(a.v_proj, xn, pv_, sel)
    q, gate = torch.chunk(pq.reshape(B, L, -1, a.head_dim * 2), 2, dim=-1)
    q = _fnorm(a.q_norm, q).transpose(1, 2)
    k = _fnorm(a.k_norm, pk_.reshape(B, L, -1, a.head_dim)).transpose(1, 2)
    v = pv_.reshape(B, L, -1, a.head_dim).transpose(1, 2)
    q, k = apply_rotary_pos_emb(q, k, cos, sin)
    o = _fork_attend(fs, ai, q, k, v, a.scaling)
    o = _fop("attn_gate", h)(o.reshape(B, L, -1), gate.reshape(B, L, -1))
    h = h + _flora(a.o_proj, o, _bi(a.o_proj, o), sel)
    return _fork_mlp(layer, h, sel)


def _fork_attend(fs: ForkState, ai: int, q: Tensor, k: Tensor, v: Tensor, scaling: float
                 ) -> Tensor:
    """Fork attention of attention layer ``ai``: q/k/v [B, H, L, D] (rotary applied); the
    rows' new K/V are written into their own buffers at ``own_len``, then every query
    attends its prefix (by reference) and its own causal keys -> [B, L, Hq, D]."""
    from deg.kernels import fork_attention

    L = q.shape[2]
    ok, ov = fs.own_k[ai], fs.own_v[ai]
    idx = (fs.own_len[:, None] + torch.arange(L, device=q.device)[None])  # [B, L]
    idx = idx[:, None, :, None].expand(-1, k.shape[1], -1, k.shape[-1])
    ok.scatter_(2, idx, k.to(ok.dtype))
    ov.scatter_(2, idx, v.to(ov.dtype))
    return fork_attention(q.transpose(1, 2).contiguous(), fs.pk[ai], fs.pv[ai], ok, ov,
                          fs.own_len, fs.poff, fs.plen, scaling)


# --------------------------------------------------------------------------- LFM2 layers
# LFM2 / LFM2.5 (HF ``modeling_lfm2``): every layer is ``h + op(operator_norm(h))`` then
# ``h + w2(silu(w1 x) * w3 x)`` on ``ffn_norm``; ``op`` is a gated short conv
# (``in_proj`` -> B, C, x; ``C * conv(B * x)``, depthwise kernel 3, no activation;
# state = the last 2 tokens of ``B * x``) or GQA attention with per-head q/k RMSNorm and
# plain 1-D RoPE. Norms are ``w * norm(x)``. There is no recurrent state: a conv layer's
# ``State.lin`` entry is (an empty [B, 0] tensor, conv tail [B, 2, C]).


def _lnorm(module: nn.Module, x: Tensor) -> Tensor:
    return _op("rms_w", x)(x, module.weight, module.variance_epsilon)


def _lfnorm(module: nn.Module, x: Tensor) -> Tensor:
    return _fop("rms_w", x)(x, module.weight, module.variance_epsilon)


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


def _rope_1d(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> tuple[Tensor, Tensor]:
    """HF ``apply_rotary_pos_emb`` (full rotary) for q/k [B, H, L, D], cos/sin [B, L, D]."""
    cos, sin = cos[:, None], sin[:, None]
    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin


def _lfm_mlp(layer, h: Tensor, sel: Tensor | None = None) -> Tensor:
    m = layer.feed_forward
    x = _lnorm(layer.ffn_norm, h)
    a = _op("swiglu", x)(_lin(m.w1, x, sel), _lin(m.w3, x, sel))
    return h + _lin(m.w2, a, sel)


def _lfm_conv_layer(layer, h, valid, lengths_new, rec, tail, rows=None, sel=None):
    """Trunk/training path of a short-conv layer; same contract as :func:`_linear_layer`
    (``rec`` is the empty placeholder and passes through)."""
    cv = layer.conv
    if rows is not None:
        rec, tail = rec.index_select(0, rows), tail.index_select(0, rows)
    xn = _lnorm(layer.operator_norm, h)
    b, c, x = _lin(cv.in_proj, xn, sel).chunk(3, -1)
    bx = b * x
    ext = torch.cat([tail.to(bx.dtype), bx], 1)  # [B, K-1+L, C]
    K1 = tail.shape[1]
    conv = _op("short_conv", h)(ext, cv.conv.weight[:, 0])
    gather = lengths_new[:, None] + torch.arange(K1, device=h.device)[None]  # last K-1 real
    new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1]))
    h = h + _lin(cv.out_proj, c * conv, sel)
    return _lfm_mlp(layer, h, sel), rec, new_tail


def _lfm_attn_qkv(layer, h, cos, sin, sel=None):
    a = layer.self_attn
    B, L, _ = h.shape
    xn = _lnorm(layer.operator_norm, h)
    q = _lnorm(a.q_layernorm, _lin(a.q_proj, xn, sel).view(B, L, -1, a.head_dim))
    k = _lnorm(a.k_layernorm, _lin(a.k_proj, xn, sel).view(B, L, -1, a.head_dim))
    v = _lin(a.v_proj, xn, sel).view(B, L, -1, a.head_dim).transpose(1, 2)
    q, k = _rope_1d(q.transpose(1, 2), k.transpose(1, 2), cos, sin)
    return q, k, v, None


def _lfm_attn_out(layer, h, o, gate, sel=None):
    return _lfm_mlp(layer, h + _lin(layer.self_attn.out_proj, o, sel), sel)


def _lfm_fork_mlp(layer, h: Tensor, sel: Tensor | None = None) -> Tensor:
    m = layer.feed_forward
    x = _lfnorm(layer.ffn_norm, h)
    g, u = _bi_fused(m, ("w1", "w3"), x)
    g, u = _flora(m.w1, x, g, sel), _flora(m.w3, x, u, sel)
    a = _fop("swiglu", x)(g, u)
    return h + _flora(m.w2, a, _bi(m.w2, a), sel)


def _lfm_fork_conv_layer(layer, h, valid, lengths, fs: ForkState, li: int, keep: bool,
                         sel: Tensor | None = None):
    """Fork path of a short-conv layer: row b continues its own conv tail (copied per row
    by :meth:`Backbone.fork`, 2 tokens); everything is per row (bi_linear GEMMs,
    elementwise conv), so batch-invariant."""
    cv = layer.conv
    B = h.shape[0]
    xn = _lfnorm(layer.operator_norm, h)
    b, c, x = _flora(cv.in_proj, xn, _bi(cv.in_proj, xn), sel).chunk(3, -1)
    bx = b * x
    tail = fs.tail[li]
    ext = torch.cat([tail.to(bx.dtype), bx], 1)
    K1 = tail.shape[1]
    conv = _fop("short_conv", h)(ext, cv.conv.weight[:, 0])
    if keep:
        gather = lengths[:, None] + torch.arange(K1, device=h.device)[None]
        new_tail = ext.gather(1, gather[..., None].expand(-1, -1, ext.shape[-1]))
        if fs.static:
            fs.tail[li].copy_(new_tail)
        else:
            fs.tail[li] = new_tail
            fs.rec[li] = h.new_zeros(B, 0, dtype=torch.float32)
    y = c * conv
    h = h + _flora(cv.out_proj, y, _bi(cv.out_proj, y), sel)
    return _lfm_fork_mlp(layer, h, sel)


def _lfm_fork_attn_layer(layer, h, cos, sin, fs: ForkState, ai: int,
                         sel: Tensor | None = None):
    a = layer.self_attn
    B, L, _ = h.shape
    xn = _lfnorm(layer.operator_norm, h)
    pq, pk_, pv_ = _bi_fused(a, ("q_proj", "k_proj", "v_proj"), xn)
    pq, pk_ = _flora(a.q_proj, xn, pq, sel), _flora(a.k_proj, xn, pk_, sel)
    pv_ = _flora(a.v_proj, xn, pv_, sel)
    q = _lfnorm(a.q_layernorm, pq.reshape(B, L, -1, a.head_dim)).transpose(1, 2)
    k = _lfnorm(a.k_layernorm, pk_.reshape(B, L, -1, a.head_dim)).transpose(1, 2)
    v = pv_.reshape(B, L, -1, a.head_dim).transpose(1, 2)
    q, k = _rope_1d(q, k, cos, sin)
    o = _fork_attend(fs, ai, q, k, v, a.scaling).reshape(B, L, -1)
    h = h + _flora(a.out_proj, o, _bi(a.out_proj, o), sel)
    return _lfm_fork_mlp(layer, h, sel)


# ------------------------------------------------------------------------ Gemma 4 layers
# Gemma 4 text model (HF ``modeling_gemma4``; google/gemma-4-E2B): every layer is attention,
# sliding-window (a query sees itself and the ``sliding_window - 1`` positions before it) or
# global, with q/k RMSNorm, an unscaled v RMSNorm, attention scale 1, one rotary table per
# layer type and Hkv K/V heads. A layer is ``h + post_attn_norm(o_proj(attn(in_norm(h))))``,
# then ``h + post_ff_norm(mlp(pre_ff_norm(h)))`` (GeGLU), then the per-layer input (PLE):
# ``h + post_ple_norm(proj(gelu(gate(h)) * ple_l))``, then ``h *= layer_scalar``. The last
# ``num_kv_shared_layers`` layers compute no K/V: they attend the K/V of the last earlier
# layer of their type (the KV-sharing map, :class:`GemmaBackbone`). PLE ``ple_l`` is a
# per-token side input: ``(norm(proj(embed)) + table[token]) / sqrt(2)``, table frozen.


def _gnorm(module: nn.Module, x: Tensor) -> Tensor:
    """Trunk/training-path Gemma norm, eager (not torch.compile'd: the compiled backward of
    this norm returned NaN weight gradients on long checkpointed reads on A100, torch
    2.14 + dynamic shapes; eager is exact and the norms are a small share of the step)."""
    if getattr(module, "with_scale", True):
        return _rms_g(x, module.weight, module.eps)
    return _rms_n(x, module.eps)


def _gfnorm(module: nn.Module, x: Tensor) -> Tensor:
    if getattr(module, "with_scale", True):
        return _fop("rms_g", x)(x, module.weight, module.eps)
    return _fop("rms_n", x)(x, module.eps)


def _rope_g(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    """HF ``apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim=2)``: x [B, L, H, D]."""
    cos, sin = cos[:, :, None], sin[:, :, None]
    return x * cos + _rotate_half(x) * sin


def _cat_kv(parts: Tensor | tuple[Tensor, ...], new: Tensor) -> Tensor:
    if isinstance(parts, Tensor):
        parts = (parts,)
    return torch.cat([*(t.to(new.dtype) for t in parts), new], 2)


def _gemma_mask(mask: Tensor | tuple[Tensor, int], c: int) -> Tensor:
    """A global layer's attention mask: given, or built from ``(kv_mask, s)`` inside the
    checkpoint (chunked training read: incoming prefix mask, ``s`` earlier tokens of this
    call, causal chunk) so no [B, c, P+s+c] mask is saved."""
    if isinstance(mask, Tensor):
        return mask
    kv_mask, s = mask
    B = kv_mask.shape[0]
    causal = torch.ones(c, c, dtype=torch.bool, device=kv_mask.device).tril()
    return torch.cat([kv_mask[:, None, :].expand(-1, c, -1),
                      torch.ones(B, c, s, dtype=torch.bool, device=kv_mask.device),
                      causal[None].expand(B, -1, -1)], 2)


class _RowSelect(torch.autograd.Function):
    """``x.index_select(0, rows)`` whose backward sums the rows' gradients with one GEMM
    (``onehot(rows)^T @ grad``, fp32 accumulation) instead of an atomic scatter-add: fork
    rows re-select their trunk's prefix K/V in every layer, and the scatter-add backward of
    [rows, Hkv, P, D] was ~17% of a long-context training step on A100."""

    @staticmethod
    def forward(ctx, x: Tensor, rows: Tensor) -> Tensor:  # type: ignore[override]
        ctx.save_for_backward(rows)
        ctx.n = x.shape[0]
        return x.index_select(0, rows)

    @staticmethod
    def backward(ctx, grad: Tensor):  # type: ignore[override]
        (rows,) = ctx.saved_tensors
        onehot = F.one_hot(rows, ctx.n).to(grad.dtype).T  # [n, R]
        out = (onehot @ grad.reshape(grad.shape[0], -1)).reshape(ctx.n, *grad.shape[1:])
        return out, None


def _select_rows(x: Tensor, rows: Tensor) -> Tensor:
    if torch.is_grad_enabled() and x.requires_grad:
        return _RowSelect.apply(x, rows)
    return x.index_select(0, rows)


def _gemma_sdpa(q: Tensor, k: Tensor, v: Tensor, mask: Tensor, scale: float,
                window: int = 0) -> Tensor:
    """SDPA of q [B, H, L, D] over keys [B, Hkv, K, D] with a bool mask [B, L, K].

    ``window`` > 0 (sliding layer; keys = the T = window-1 tail then the L new keys, as
    :meth:`GemmaBackbone._window_mask` lays them out): queries run in blocks of
    ``window``, each against only the keys its window can reach (T + window), so a long
    read costs O(L * 2 window) instead of O(L * (T + L)). Same math: every dropped key is
    masked for every query of the block."""
    B, H, L, D = q.shape
    T = k.shape[2] - L
    if window <= 0 or L <= window or T != window - 1:
        return F.scaled_dot_product_attention(q, k, v, attn_mask=mask[:, None], scale=scale,
                                              enable_gqa=True)
    W = window
    nb = -(-L // W)
    pad = nb * W - L
    if pad:
        q = F.pad(q, (0, 0, 0, pad))
        k, v = F.pad(k, (0, 0, 0, pad)), F.pad(v, (0, 0, 0, pad))
        mask = F.pad(mask, (0, pad, 0, pad))
        mask[:, L:, 0] = True  # padded queries see one key (output dropped): no NaN
    hkv = k.shape[1]
    # key block n = keys [n*W, n*W + W + T) (tail-relative), query block n = [n*W, n*W + W)
    kb = k.unfold(2, W + T, W).permute(0, 2, 1, 4, 3).reshape(B * nb, hkv, W + T, D)
    vb = v.unfold(2, W + T, W).permute(0, 2, 1, 4, 3).reshape(B * nb, hkv, W + T, D)
    qb = q.reshape(B, H, nb, W, D).transpose(1, 2).reshape(B * nb, H, W, D)
    mc = mask.contiguous()
    s0, s1, s2 = mc.stride()
    mb = mc.as_strided((B, nb, W, W + T), (s0, W * (s1 + s2), s1, s2)).reshape(B * nb, W, W + T)
    o = F.scaled_dot_product_attention(qb, kb, vb, attn_mask=mb[:, None], scale=scale,
                                       enable_gqa=True)
    o = o.reshape(B, nb, H, W, D).transpose(1, 2).reshape(B, H, nb * W, D)
    return o[:, :, :L]


def _gemma_rest(layer, h: Tensor, attn: Tensor, ple: Tensor | None, sel=None) -> Tensor:
    """Everything after the attention output projection: post-attention norm + residual,
    GeGLU MLP, the per-layer input, the layer scalar (trunk/training path)."""
    h = h + _gnorm(layer.post_attention_layernorm, attn)
    m = _expert(layer, sel)
    x = _gnorm(layer.pre_feedforward_layernorm, h)
    y = _lin(m.down_proj, _op("geglu", x)(_lin(m.gate_proj, x, sel), _lin(m.up_proj, x, sel)),
             sel)
    h = h + _gnorm(layer.post_feedforward_layernorm, y)
    if ple is not None:
        g = _op("geglu", h)(layer.per_layer_input_gate(h), ple)
        h = h + _gnorm(layer.post_per_layer_input_norm, layer.per_layer_projection(g))
    return h * layer.layer_scalar


def _gemma_layer(layer, h, ple, cos, sin, kp, vp, mask, rows=None, sel=None, kn=None,
                 vn=None, window=0):
    """One Gemma 4 layer (trunk / training path). ``kp``/``vp`` [B', Hkv, P, D] (or a tuple
    of parts concatenated here): the incoming keys of this layer's K/V slot (a global
    layer's whole prefix, a sliding layer's tail); ``rows``: as in :func:`_linear_layer`
    (prefix rows selected inside the checkpoint). A K/V-sharing layer gets the source
    layer's new K/V of this call as ``kn``/``vn`` and computes only q. ``mask`` [B, L, P+L]
    (or ``(kv_mask, s)``, :func:`_gemma_mask`); ``window``: a sliding layer's window (its
    keys = tail + new keys, :func:`_gemma_sdpa` blocks long reads). Returns (h, k, v of
    this call's tokens)."""
    a = layer.self_attn
    B, L, _ = h.shape
    if rows is not None:
        kp = tuple(_select_rows(t, rows) for t in (kp if isinstance(kp, tuple) else (kp,)))
        vp = tuple(_select_rows(t, rows) for t in (vp if isinstance(vp, tuple) else (vp,)))
    xn = _gnorm(layer.input_layernorm, h)
    q = _gnorm(a.q_norm, _lin(a.q_proj, xn, sel).view(B, L, -1, a.head_dim))
    q = _rope_g(q, cos, sin).transpose(1, 2)
    if kn is None:
        k = _gnorm(a.k_norm, _lin(a.k_proj, xn, sel).view(B, L, -1, a.head_dim))
        k = _rope_g(k, cos, sin).transpose(1, 2)
        v = _gnorm(a.v_norm, _lin(a.v_proj, xn, sel).view(B, L, -1, a.head_dim)).transpose(1, 2)
    else:
        k, v = kn, vn
    keys, values = _cat_kv(kp, k), _cat_kv(vp, v)
    o = _gemma_sdpa(q, keys, values, _gemma_mask(mask, L), a.scaling, window)
    o = o.transpose(1, 2).reshape(B, L, -1)
    return _gemma_rest(layer, h, _lin(a.o_proj, o, sel), ple, sel), k, v


def _gemma_fork_layer(layer, h, ple, cos, sin, fs: ForkState, slot: int, owner: bool,
                      window: int, sel=None):
    """Fork path of one Gemma 4 layer (batch-invariant, like :func:`_fork_attn_layer`). An
    owner layer writes the rows' new K/V into slot ``slot``'s own buffers; a K/V-sharing
    layer reads them (written by its source layer earlier in the same pass). A sliding
    layer (``window`` > 0) attends its prefix tail (``fs.wpoff``/``fs.wplen``) and its own
    keys within the window; a global layer the whole prefix by reference."""
    from deg.kernels import fork_attention

    a = layer.self_attn
    B, L, _ = h.shape
    xn = _gfnorm(layer.input_layernorm, h)
    if owner:
        pq, pk_, pv_ = _bi_fused(a, ("q_proj", "k_proj", "v_proj"), xn)
        pk_, pv_ = _flora(a.k_proj, xn, pk_, sel), _flora(a.v_proj, xn, pv_, sel)
    else:
        pq = _bi(a.q_proj, xn)
    pq = _flora(a.q_proj, xn, pq, sel)
    q = _rope_g(_gfnorm(a.q_norm, pq.reshape(B, L, -1, a.head_dim)), cos, sin)
    ok, ov = fs.own_k[slot], fs.own_v[slot]
    if owner:
        k = _rope_g(_gfnorm(a.k_norm, pk_.reshape(B, L, -1, a.head_dim)), cos, sin)
        v = _gfnorm(a.v_norm, pv_.reshape(B, L, -1, a.head_dim))
        idx = (fs.own_len[:, None] + torch.arange(L, device=h.device)[None])  # [B, L]
        idx = idx[:, None, :, None].expand(-1, k.shape[2], -1, k.shape[-1])
        ok.scatter_(2, idx, k.transpose(1, 2).to(ok.dtype))
        ov.scatter_(2, idx, v.transpose(1, 2).to(ov.dtype))
    if window > 0:
        assert fs.wpoff is not None and fs.wplen is not None
        poff, plen = fs.wpoff, fs.wplen
    else:
        poff, plen = fs.poff, fs.plen
    o = fork_attention(q.contiguous(), fs.pk[slot], fs.pv[slot], ok, ov, fs.own_len, poff,
                       plen, a.scaling, window).reshape(B, L, -1)
    h = h + _gfnorm(layer.post_attention_layernorm, _flora(a.o_proj, o, _bi(a.o_proj, o), sel))
    m = _expert(layer, sel)
    x = _gfnorm(layer.pre_feedforward_layernorm, h)
    gate, up = _bi_fused(m, ("gate_proj", "up_proj"), x)
    gate, up = _flora(m.gate_proj, x, gate, sel), _flora(m.up_proj, x, up, sel)
    y = _fop("geglu", x)(gate, up)
    h = h + _gfnorm(layer.post_feedforward_layernorm, _flora(m.down_proj, y, _bi(m.down_proj, y),
                                                               sel))
    if ple is not None:
        g = _fop("geglu", h)(_bi(layer.per_layer_input_gate, h), ple)
        h = h + _gfnorm(layer.post_per_layer_input_norm, _bi(layer.per_layer_projection, g))
    return h * layer.layer_scalar


# ------------------------------------------------------------------- backbone families


@dataclass(frozen=True)
class Family:
    """The layer functions of one architecture behind :class:`Backbone` (SPEC.md §12.8).
    Every family is a hybrid of state-carrying layers (``rec_type``) and attention
    layers; :class:`State` / :class:`ForkState` hold, per state layer, (recurrent state,
    conv tail) and, per attention layer, K/V, so the trunk/fork/chunk/decode machinery is
    shared and only these functions differ.

    * ``rec_layer(layer, h, valid, lengths, rec, tail, rows, sel) -> (h, rec, tail)``
    * ``attn_qkv(layer, h, cos, sin, sel) -> (q, k, v [B, H, L, D] rotary applied, gate)``
    * ``attn_out(layer, h, o [B, L, Hq*D], gate, sel) -> h`` (residual + MLP)
    * ``fork_rec_layer(layer, h, valid, lengths, fs, li, keep, sel) -> h`` and
      ``fork_attn_layer(layer, h, cos, sin, fs, ai, sel) -> h``: batch-invariant fork path
    * ``rec_shapes(layer) -> (recurrent state shape, conv tail shape)`` per row
    * ``final_norm``: attribute of the text model; ``fork_norm(module, h)`` its
      batch-invariant form; ``mrope``: rotary takes [3, B, L] positions."""

    name: str
    rec_type: str
    rec_layer: Any
    attn_qkv: Any
    attn_out: Any
    fork_rec_layer: Any
    fork_attn_layer: Any
    rec_shapes: Any
    final_norm: str
    fork_norm: Any
    mrope: bool
    lora_targets: dict[str, tuple[str, ...]]


def _qwen_rec_shapes(layer) -> tuple[tuple[int, ...], tuple[int, ...]]:
    at = layer.linear_attn
    return ((at.num_v_heads, at.head_k_dim, at.head_v_dim),
            (at.conv_kernel_size - 1, at.conv_dim))


def _lfm_rec_shapes(layer) -> tuple[tuple[int, ...], tuple[int, ...]]:
    cv = layer.conv
    return (0,), (cv.conv_kernel_size - 1, cv.conv.weight.shape[0])


QWEN35 = Family("qwen3_5", "linear_attention", _linear_layer, _attn_qkv, _attn_out,
                _fork_linear_layer, _fork_attn_layer, _qwen_rec_shapes, "norm", _fnorm, True,
                LORA_TARGETS)
LFM2 = Family("lfm2", "conv", _lfm_conv_layer, _lfm_attn_qkv, _lfm_attn_out,
              _lfm_fork_conv_layer, _lfm_fork_attn_layer, _lfm_rec_shapes, "embedding_norm",
              _lfnorm, False,
              {"conv": ("in_proj", "out_proj"),
               "self_attn": ("q_proj", "k_proj", "v_proj", "out_proj"),
               "feed_forward": ("w1", "w3", "w2")})


# Gemma 4: no state-carrying layer (rec_type ""), its own Backbone (:class:`GemmaBackbone`);
# LoRA on q/k/v/o (k/v only where a layer computes K/V) and the MLP
GEMMA4 = Family("gemma4", "", None, None, None, None, None, None, "norm", _gfnorm, False,
                {"self_attn": ("q_proj", "k_proj", "v_proj", "o_proj"),
                 "mlp": ("gate_proj", "up_proj", "down_proj")})


def family_of(lm: nn.Module) -> Family:
    kind = lm.config.model_type
    if kind == "lfm2":
        return LFM2
    if kind == "gemma4_text":
        return GEMMA4
    if kind.startswith("qwen3_5"):
        return QWEN35
    raise ValueError(f"no deg backbone family for model_type {kind!r}")


class Backbone(nn.Module):
    """Functional hybrid text stack (Qwen3.5 or LFM2.5, :class:`Family`) over explicit
    prefix state (no HF cache objects, so gradients flow from forks into the shared
    trunk). ``Backbone(lm)`` of a Gemma 4 model is a :class:`GemmaBackbone`."""

    def __new__(cls, lm: nn.Module | None = None, *args: Any, **kwargs: Any) -> Backbone:
        if cls is Backbone and lm is not None and family_of(lm) is GEMMA4:
            return super().__new__(GemmaBackbone)
        return super().__new__(cls)

    def __init__(self, lm: nn.Module) -> None:
        super().__init__()
        self.lm = lm
        text = lm.model
        self.text = getattr(text, "language_model", text)
        self.config = self.text.config
        self.fam = family_of(lm)
        # "rec" (state-carrying: Gated DeltaNet / short conv) or "attn" per layer
        self.types = ["rec" if t == self.fam.rec_type else "attn"
                      for t in self.config.layer_types]
        self.grad_ckpt = False

    @property
    def hidden(self) -> int:
        return self.config.hidden_size

    def rope(self, h: Tensor, pos: Tensor) -> tuple[Tensor, Tensor]:
        """cos/sin for positions ``pos`` [B, L]."""
        return self.text.rotary_emb(h, pos[None].expand(3, -1, -1) if self.fam.mrope else pos)

    def final_norm(self, h: Tensor) -> Tensor:
        return getattr(self.text, self.fam.final_norm)(h)

    def fork_final_norm(self, h: Tensor) -> Tensor:
        return self.fam.fork_norm(getattr(self.text, self.fam.final_norm), h)

    def rec_buffers(self, layer, batch: int, device, dtype) -> tuple[Tensor, Tensor]:
        """Zero (recurrent state fp32, conv tail) of one state layer for ``batch`` rows."""
        rec, tail = self.fam.rec_shapes(layer)
        return (torch.zeros(batch, *rec, device=device, dtype=torch.float32),
                torch.zeros(batch, *tail, device=device, dtype=dtype))

    def empty_state(self, batch: int, device, dtype) -> State:
        lin, kv = [], []
        for layer, kind in zip(self.text.layers, self.types):
            if kind == "rec":
                lin.append(self.rec_buffers(layer, batch, device, dtype))
            else:
                a = layer.self_attn
                shape = (batch, self.config.num_key_value_heads, 0, a.head_dim)
                kv.append((torch.zeros(shape, device=device, dtype=dtype),
                           torch.zeros(shape, device=device, dtype=dtype)))
        return State(lin, kv, torch.zeros(batch, 0, dtype=torch.bool, device=device),
                     torch.zeros(batch, dtype=torch.long, device=device))

    def forward(self, ids: Tensor, lengths: Tensor, state: State | None = None,
                consume: bool = False, rows: Tensor | None = None, sel: Tensor | None = None
                ) -> tuple[Tensor, State | None]:
        """ids [B,L] right-padded, lengths [B] real tokens -> (final-norm hidden [B,L,C],
        state after each row's real tokens). Inputs longer than :data:`DENSE_PREFILL_MAX`
        stream through :meth:`_forward_chunked` (no O(L^2) tensor, no length cap).

        ``consume``: the caller gives up ``state``; each attention layer's prefix K/V is
        released (``state.kv[i] = None``) as soon as its extended copy exists, so
        extending a long context peaks at its K/V plus one layer, not twice its K/V.

        ``rows`` (training forks): row b continues row ``rows[b]`` of ``state`` (the
        un-selected trunk state). Each layer selects its prefix rows inside the gradient
        checkpoint, so the saved activations hold no per-row copy of the prefix (K/V of
        every attention layer, recurrent state of every linear layer): memory per fork
        row is independent of the context length except during one layer's recompute.
        Returns ``(hidden, None)``: no continued state is built.

        ``sel`` [B, n*r]: per-row LoRA adapter selector (None = adapters off)."""
        if rows is not None and ids.shape[1] > DENSE_PREFILL_MAX:
            assert state is not None
            state, rows = state.select(rows), None  # long fork input: plain select path
        if ids.shape[1] > DENSE_PREFILL_MAX:
            if self.grad_ckpt and torch.is_grad_enabled():
                return self._forward_chunked_ckpt(ids, lengths, state, sel)
            return self._forward_chunked(ids, lengths, state, consume, sel)
        B, L = ids.shape
        device = ids.device
        if profile._ACTIVE.get() is not None:
            profile.count(forwards=1, forward_rows=B, forward_tokens=B * L,
                          forward_real_tokens=int(lengths.sum()))
        h = self.text.embed_tokens(ids)
        if state is None:
            state = self.empty_state(B, device, h.dtype)
        kv_mask, plens = state.kv_mask, state.lengths
        if rows is not None:
            kv_mask, plens = kv_mask.index_select(0, rows), plens.index_select(0, rows)
        valid = torch.arange(L, device=device)[None] < lengths[:, None]  # [B, L]
        pos = plens[:, None] + torch.arange(L, device=device)[None]
        cos, sin = self.rope(h, pos)
        causal = torch.ones(L, L, dtype=torch.bool, device=device).tril()
        mask = torch.cat([kv_mask[:, None, :].expand(-1, L, -1),
                          causal[None].expand(B, -1, -1)], 2)
        validf = valid.to(torch.float32)
        lin_out, kv_out = [], []
        li = ai = 0
        for layer, kind in zip(self.text.layers, self.types):
            if kind == "rec":
                rec, tail = state.lin[li]
                args = (layer, h, validf, lengths, rec, tail, rows, sel)
                if self.grad_ckpt and torch.is_grad_enabled():
                    h, rec, tail = checkpoint(self.fam.rec_layer, *args, use_reentrant=False)
                else:
                    h, rec, tail = self.fam.rec_layer(*args)
                if rows is None:
                    lin_out.append((rec, tail))
                del rec, tail
                li += 1
            else:
                kp, vp = state.kv[ai]
                args = (self.fam, layer, h, cos, sin, kp, vp, mask, rows, sel)
                if self.grad_ckpt and torch.is_grad_enabled():
                    h, k, v = checkpoint(_attn_layer, *args, use_reentrant=False)
                else:
                    h, k, v = _attn_layer(*args)
                if rows is None:
                    kv_out.append((torch.cat([kp.to(k.dtype), k], 2),
                                   torch.cat([vp.to(v.dtype), v], 2)))
                del k, v
                if consume:
                    del kp, vp, args
                    state.kv[ai] = None  # type: ignore[call-overload]
                ai += 1
        h = self.final_norm(h)
        if rows is not None:
            return h, None
        new = State(lin_out, kv_out, torch.cat([state.kv_mask, valid], 1),
                    state.lengths + lengths)
        return h, new

    def _forward_chunked(self, ids: Tensor, lengths: Tensor, state: State | None,
                         consume: bool = False, sel: Tensor | None = None,
                         soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State]:
        """Chunked prefill, same contract as :meth:`forward`: the input is read
        :data:`PREFILL_CHUNK` tokens at a time, carrying each linear layer's recurrent
        state and conv tail and appending each attention layer's K/V. Rows are
        right-padded, so a real query's earlier positions in this call are all real:
        attention needs only the incoming prefix mask plus causality inside the chunk.

        Peak memory is O(L) (K/V, hidden) plus O(chunk) work. On CUDA without autograd,
        attention runs the flash-style fork kernel (prefix K/V read in place from one
        preallocated buffer per layer, causal own keys) whenever the incoming prefix is
        unpadded; otherwise (CPU, autograd, padded incoming prefix) SDPA with an explicit
        [B, chunk, P+chunk] mask: O(chunk * (P+chunk)), never L x L."""
        from deg.kernels import fork_attention

        B, L = ids.shape
        device = ids.device
        C = PREFILL_CHUNK
        if profile._ACTIVE.get() is not None:
            profile.count(forwards=1, forward_rows=B, forward_tokens=B * L,
                          forward_real_tokens=int(lengths.sum()))
        if state is None:
            state = self.empty_state(B, device, self.text.embed_tokens.weight.dtype)
        P0 = state.kv_mask.shape[1]
        fast = (ids.is_cuda and not torch.is_grad_enabled()
                and bool(state.kv_mask.all()))
        lin = list(state.lin)
        if fast:  # [Hkv, B, P0+L, D] so row b's prefix is one [Hkv, P, D] strided slab
            bufs = []
            for i in range(len(state.kv)):
                kp, vp = state.kv[i]
                shape = (kp.shape[1], B, P0 + L, kp.shape[3])
                kb = torch.empty(shape, device=device, dtype=self.text.embed_tokens.weight.dtype)
                vb = torch.empty_like(kb)
                kb[:, :, :P0] = kp.transpose(0, 1)
                vb[:, :, :P0] = vp.transpose(0, 1)
                bufs.append((kb, vb))
                if consume:
                    del kp, vp
                    state.kv[i] = None  # type: ignore[call-overload]
            poff = torch.arange(B, device=device) * (P0 + L)
            qstart = torch.zeros(B, dtype=torch.long, device=device)
        else:
            kv = list(state.kv)
        outs = []
        for s in range(0, L, C):
            e = min(L, s + C)
            c = e - s
            lens = (lengths - s).clamp(0, c)
            h = self.text.embed_tokens(ids[:, s:e])
            valid = torch.arange(c, device=device)[None] < lens[:, None]
            pos = state.lengths[:, None] + s + torch.arange(c, device=device)[None]
            cos, sin = self.rope(h, pos)
            validf = valid.to(torch.float32)
            if not fast:
                causal = torch.ones(c, c, dtype=torch.bool, device=device).tril()
                mask = torch.cat([state.kv_mask[:, None, :].expand(-1, c, -1),
                                  torch.ones(B, c, s, dtype=torch.bool, device=device),
                                  causal[None].expand(B, -1, -1)], 2)
            li = ai = 0
            for layer, kind in zip(self.text.layers, self.types):
                if kind == "rec":
                    rec, tail = lin[li]
                    h, rec, tail = self.fam.rec_layer(layer, h, validf, lens, rec, tail, None,
                                                      sel)
                    lin[li] = (rec, tail)
                    li += 1
                    continue
                a = layer.self_attn
                q, k, v, gate = self.fam.attn_qkv(layer, h, cos, sin, sel)
                if fast:
                    kb, vb = bufs[ai]
                    o = fork_attention(q.transpose(1, 2).contiguous(),
                                       kb.view(kb.shape[0], -1, kb.shape[3]),
                                       vb.view(vb.shape[0], -1, vb.shape[3]),
                                       k.contiguous(), v.contiguous(), qstart, poff,
                                       torch.full((B,), P0 + s, device=device), a.scaling)
                    kb[:, :, P0 + s:P0 + e] = k.transpose(0, 1)
                    vb[:, :, P0 + s:P0 + e] = v.transpose(0, 1)
                else:
                    kp, vp = kv[ai]
                    keys = torch.cat([kp.to(k.dtype), k], 2)
                    values = torch.cat([vp.to(v.dtype), v], 2)
                    o = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask[:, None],
                                                       scale=a.scaling, enable_gqa=True)
                    o = o.transpose(1, 2)
                    kv[ai] = (keys, values)
                    if consume and s == 0:
                        state.kv[ai] = None  # type: ignore[call-overload]
                h = self.fam.attn_out(layer, h, o.reshape(B, c, -1), gate, sel)
                ai += 1
            outs.append(self.final_norm(h))
        if fast:
            kv = [(kb.transpose(0, 1), vb.transpose(0, 1)) for kb, vb in bufs]
        valid = torch.arange(L, device=device)[None] < lengths[:, None]
        new = State(lin, kv, torch.cat([state.kv_mask, valid], 1), state.lengths + lengths)
        return torch.cat(outs, 1), new

    def _forward_chunked_ckpt(self, ids: Tensor, lengths: Tensor, state: State | None,
                              sel: Tensor | None = None) -> tuple[Tensor, State]:
        """Training read of an input longer than :data:`DENSE_PREFILL_MAX` (autograd +
        ``grad_ckpt``): :data:`PREFILL_CHUNK` tokens at a time like :meth:`_forward_chunked`,
        every layer of every chunk gradient-checkpointed. A checkpoint saves only its
        inputs (the chunk's hidden, the carried recurrent state / conv tail, and the
        earlier chunks' K/V as a tuple of the tensors that exist anyway); the attention
        mask is built inside the checkpoint, so nothing O(L^2) is saved and one layer's
        recompute is O(chunk * L). The returned state's K/V concatenates the chunks."""
        B, L = ids.shape
        device = ids.device
        C = PREFILL_CHUNK
        if state is None:
            state = self.empty_state(B, device, self.text.embed_tokens.weight.dtype)
        lin = list(state.lin)
        kparts: list[list[Tensor]] = [[kp] for kp, _ in state.kv]
        vparts: list[list[Tensor]] = [[vp] for _, vp in state.kv]
        outs = []
        for s in range(0, L, C):
            e = min(L, s + C)
            c = e - s
            lens = (lengths - s).clamp(0, c)
            h = self.text.embed_tokens(ids[:, s:e])
            valid = torch.arange(c, device=device)[None] < lens[:, None]
            pos = state.lengths[:, None] + s + torch.arange(c, device=device)[None]
            cos, sin = self.rope(h, pos)
            validf = valid.to(torch.float32)
            li = ai = 0
            for layer, kind in zip(self.text.layers, self.types):
                if kind == "rec":
                    rec, tail = lin[li]
                    h, rec, tail = checkpoint(self.fam.rec_layer, layer, h, validf, lens, rec,
                                              tail, None, sel, use_reentrant=False)
                    lin[li] = (rec, tail)
                    li += 1
                    continue
                h, k, v = checkpoint(_chunk_attn_layer, self.fam, layer, h, cos, sin,
                                     tuple(kparts[ai]), tuple(vparts[ai]), state.kv_mask, s, sel,
                                     use_reentrant=False)
                kparts[ai].append(k)
                vparts[ai].append(v)
                ai += 1
            outs.append(self.final_norm(h))
        kv = [(torch.cat(kp, 2), torch.cat(vp, 2)) for kp, vp in zip(kparts, vparts)]
        valid = torch.arange(L, device=device)[None] < lengths[:, None]
        new = State(lin, kv, torch.cat([state.kv_mask, valid], 1), state.lengths + lengths)
        return torch.cat(outs, 1), new

    def fork(self, prefixes: State | list[State], prow: list[int] | int, cap: int
             ) -> ForkState:
        """Fresh fork rows; row b continues ``prefixes[prow[b]]`` (batch-1, unpadded
        states). ``prow`` as an int = that many rows of a single prefix. Own K/V capacity
        ``cap``. A single prefix is read in place; several are concatenated once."""
        if isinstance(prefixes, State):
            prefixes = [prefixes]
        if isinstance(prow, int):
            prow = [0] * prow
        for p in prefixes:
            assert p.lengths.shape[0] == 1 and bool(p.kv_mask.all())
        rows = len(prow)
        device = prefixes[0].lengths.device
        if device.type == "cpu":
            cap += CPU_BUCKET  # rows run padded to CPU_BUCKET (see _fork_forward_rows)
        own_k, own_v = [], []
        for k, _ in prefixes[0].kv:
            shape = (rows, k.shape[1], cap, k.shape[3])
            own_k.append(torch.zeros(shape, device=device, dtype=k.dtype))
            own_v.append(torch.zeros(shape, device=device, dtype=k.dtype))
        lens = [int(p.kv[0][0].shape[2]) if p.kv else int(p.lengths[0]) for p in prefixes]
        offs = [sum(lens[:i]) for i in range(len(lens))]
        if len(prefixes) == 1:
            pk = [k[0] for k, _ in prefixes[0].kv]
            pv = [v[0] for _, v in prefixes[0].kv]
            prec = [r.contiguous() for r, _ in prefixes[0].lin]
        else:
            pk = [torch.cat([p.kv[i][0][0] for p in prefixes], 1)
                  for i in range(len(prefixes[0].kv))]
            pv = [torch.cat([p.kv[i][1][0] for p in prefixes], 1)
                  for i in range(len(prefixes[0].kv))]
            prec = [torch.cat([p.lin[i][0] for p in prefixes]).contiguous()
                    for i in range(len(prefixes[0].lin))]
        idx = torch.tensor(prow, dtype=torch.long, device=device)
        tails = [torch.cat([p.lin[i][1] for p in prefixes]).index_select(0, idx)
                 for i in range(len(prefixes[0].lin))]
        return ForkState(prec, pk, pv, torch.tensor(offs, device=device)[idx],
                         torch.tensor(lens, device=device)[idx], idx,
                         [None] * len(prefixes[0].lin), tails, own_k, own_v,
                         torch.zeros(rows, dtype=torch.long, device=device))

    def fork_forward(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool = False,
                     decode: bool = False, sel: Tensor | None = None) -> Tensor:
        """See :meth:`_fork_forward`; ``decode`` (one token per row) selects the small-M
        GEMM config for every projection of this pass; ``sel`` [B, n*r] selects each
        row's LoRA adapter (None = off)."""
        global _SMALL_GEMM
        prev, _SMALL_GEMM = _SMALL_GEMM, decode
        try:
            return self._fork_forward(ids, lengths, fs, keep, sel)
        finally:
            _SMALL_GEMM = prev

    def _fork_forward(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool = False,
                      sel: Tensor | None = None) -> Tensor:
        """Run fork rows ids [B,L] (right-padded, ``lengths`` real tokens) on top of the
        shared prefix and their own earlier tokens -> final-norm hidden [B,L,C].

        Batch-invariant: every op is per row (bi_linear GEMMs, fork kernels, elementwise),
        so a row's hidden states are bitwise the same in any batch. ``keep`` stores the
        rows' new recurrent state, conv tail and own length (for continued decoding)."""
        B, L = ids.shape
        if not ids.is_cuda and (B > 1 or (L > 1 and L % CPU_BUCKET)):
            return self._fork_forward_rows(ids, lengths, fs, keep, sel)  # CPU reference
        device = ids.device
        if profile._ACTIVE.get() is not None and not torch.cuda.is_current_stream_capturing():
            profile.count(forwards=1, forward_rows=B, forward_tokens=B * L)
        h = self.text.embed_tokens(ids)
        valid = (torch.arange(L, device=device)[None] < lengths[:, None]).to(torch.float32)
        pos = fs.plen[:, None] + fs.own_len[:, None] + torch.arange(L, device=device)[None]
        cos, sin = self.rope(h, pos)
        li = ai = 0
        for layer, kind in zip(self.text.layers, self.types):
            if kind == "rec":
                h = self.fam.fork_rec_layer(layer, h, valid, lengths, fs, li, keep, sel)
                li += 1
            else:
                h = self.fam.fork_attn_layer(layer, h, cos, sin, fs, ai, sel)
                ai += 1
        if keep:
            if fs.static:
                fs.own_len.add_(lengths)
            else:
                fs.own_len = fs.own_len + lengths
            fs.fresh = False
        return self.fork_final_norm(h)

    def _fork_forward_rows(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool,
                           sel: Tensor | None = None) -> Tensor:
        """CPU reference: torch's CPU elementwise kernels switch between SIMD body and
        scalar tail by tensor size, so a row's numerics would depend on the batch. Run
        each row alone at a width padded to :data:`CPU_BUCKET` (decode steps: width 1)."""
        B, L = ids.shape
        W = 1 if L == 1 else -(-L // CPU_BUCKET) * CPU_BUCKET  # decode steps stay 1 wide
        outs, recs, tails, lens = [], [], [], []
        for b in range(B):
            sub = ForkState(fs.prefix_rec, fs.pk, fs.pv, fs.poff[b:b + 1], fs.plen[b:b + 1],
                            fs.prow[b:b + 1],
                            [r[b:b + 1] if r is not None else None for r in fs.rec],
                            [t[b:b + 1] for t in fs.tail], [k[b:b + 1] for k in fs.own_k],
                            [v[b:b + 1] for v in fs.own_v], fs.own_len[b:b + 1], fs.fresh,
                            wpoff=None if fs.wpoff is None else fs.wpoff[b:b + 1],
                            wplen=None if fs.wplen is None else fs.wplen[b:b + 1])
            row = torch.full((1, W), self.lm.config.pad_token_id or 0, dtype=ids.dtype)
            row[0, :L] = ids[b]
            outs.append(self._fork_forward(row, lengths[b:b + 1], sub, keep,
                                           None if sel is None else sel[b:b + 1])[:, :L])
            recs.append(sub.rec)
            tails.append(sub.tail)
            lens.append(sub.own_len)
        if keep:
            fs.rec = [torch.cat([r[i] for r in recs]) for i in range(len(fs.rec))]
            fs.tail = [torch.cat([t[i] for t in tails]) for i in range(len(fs.tail))]
            fs.own_len = torch.cat(lens)
            fs.fresh = False
        return torch.cat(outs)

    def lm_logits(self, hidden: Tensor) -> Tensor:
        """LM head through the batch-invariant GEMM (a row's logits never depend on the
        other rows decoding with it). The Triton GEMM has no backward: when autograd needs
        a gradient (training) the head is a plain differentiable ``F.linear``."""
        weight = self.lm.lm_head.weight
        if torch.is_grad_enabled() and (hidden.requires_grad or weight.requires_grad):
            return F.linear(hidden, weight)
        from deg.kernels import bi_linear_small

        return bi_linear_small(hidden, weight)


class GemmaBackbone(Backbone):
    """Gemma 4 text stack (:data:`GEMMA4`) on the same State / ForkState contract.

    K/V slots: one per layer that computes K/V (layers ``0 .. n - num_kv_shared_layers - 1``;
    E2B: 15 = 12 sliding + 3 global); ``kv_slot[i]`` maps every layer to its slot (a
    K/V-sharing layer to the last earlier owner of its type). ``State.kv[slot]`` holds, per
    row, a global slot's whole prefix K/V [B, Hkv, P, D] (``kv_mask`` marks real positions)
    and a sliding slot's **tail**: the K/V of the row's last ``T = sliding_window - 1`` real
    positions [B, Hkv, T, D], right-aligned at the row's length (slot j = position
    ``lengths - T + j``; slots before position 0 are zeros and masked). A query never sees
    more than T earlier positions in a sliding layer, so the tail is exact and a fork's
    sliding state is bounded (T per layer, whatever the context length). ``State.lin`` is
    empty (no recurrent layer).

    PLE: :meth:`ple_inputs` gathers the frozen per-layer token table (on the device, or in
    host memory after :meth:`ple_to_host`) and adds the projected embedding."""

    def __init__(self, lm: nn.Module) -> None:
        super().__init__(lm)
        cfg = self.config
        if getattr(cfg, "enable_moe_block", False):
            raise ValueError("Gemma 4 MoE blocks are not supported")
        if getattr(cfg, "attention_k_eq_v", False):
            raise ValueError("Gemma 4 attention_k_eq_v (v = k) is not supported")
        if getattr(cfg, "use_bidirectional_attention", None):
            raise ValueError("Gemma 4 bidirectional attention is not supported")
        n = cfg.num_hidden_layers
        first_shared = n - int(getattr(cfg, "num_kv_shared_layers", 0) or 0)
        self.kinds = ["sliding" if t == "sliding_attention" else "full"
                      for t in cfg.layer_types]
        self.window = int(cfg.sliding_window)
        self.tail = self.window - 1
        self.owner = [i < first_shared for i in range(n)]
        slot_of_layer: dict[int, int] = {}
        self.slot_kinds: list[str] = []
        self.slot_layers: list[int] = []
        last_of_kind: dict[str, int] = {}
        for i in range(n):
            if self.owner[i]:
                slot_of_layer[i] = len(self.slot_kinds)
                self.slot_kinds.append(self.kinds[i])
                self.slot_layers.append(i)
                last_of_kind[self.kinds[i]] = slot_of_layer[i]
            else:  # HF: the last non-shared layer of the same type stores the K/V it shares
                slot_of_layer[i] = last_of_kind[self.kinds[i]]
        self.kv_slot = [slot_of_layer[i] for i in range(n)]
        # last layer reading each slot (consume: a slot's prefix is released after it)
        self.last_use = [max(i for i in range(n) if self.kv_slot[i] == k)
                         for k in range(len(self.slot_kinds))]
        self.ple_dim = int(getattr(cfg, "hidden_size_per_layer_input", 0) or 0)
        if self.ple_dim:  # the PLE table is token identity: never trained
            self.text.embed_tokens_per_layer.weight.requires_grad_(False)
        self.softcap = getattr(cfg, "final_logit_softcapping", None)
        self.types = ["attn"] * n

    @property
    def ple_on_host(self) -> bool:
        t = getattr(self.text, "embed_tokens_per_layer", None)
        return t is not None and t.weight.device != self.text.embed_tokens.weight.device

    def ple_to_host(self) -> None:
        """Keep the frozen PLE table (262144 x 35*256 bf16 = 4.4 GiB on E2B) in pinned host
        memory: each pass gathers its tokens' rows on the host and copies them in. Call
        after the model is on its device; never trained."""
        t = self.text.embed_tokens_per_layer
        t.weight.requires_grad_(False)
        if t.weight.device.type == "cpu":
            return
        w = t.weight.detach().cpu()
        t.weight = nn.Parameter(w.pin_memory() if torch.cuda.is_available() else w,
                                requires_grad=False)

    def ple_rows(self, ids: Tensor) -> Tensor:
        """Unscaled PLE table rows of ``ids`` [..] -> [.., n_layers * ple_dim] on ids' device
        (gathered on the host when the table lives there)."""
        w = self.text.embed_tokens_per_layer.weight
        if w.device == ids.device:
            return F.embedding(ids, w)
        rows = F.embedding(ids.cpu(), w)
        if ids.is_cuda:
            rows = rows.pin_memory()
        return rows.to(ids.device, non_blocking=True)

    def ple_inputs(self, ids: Tensor, h: Tensor, fork: bool = False,
                   rows: Tensor | None = None) -> Tensor | None:
        """Per-layer inputs [B, L, n_layers, ple_dim] (HF ``get_per_layer_inputs`` +
        ``project_per_layer_inputs``) for token ids [B, L] whose scaled embeddings are
        ``h``; ``fork``: batch-invariant GEMM/norm (fork path); ``rows``: precomputed
        :meth:`ple_rows` (a decode graph's static buffer)."""
        if not self.ple_dim:
            return None
        text = self.text
        B, L = ids.shape
        nl = self.config.num_hidden_layers
        emb = text.embed_tokens_per_layer
        tok = (self.ple_rows(ids) if rows is None else rows) * emb.embed_scale.to(
            device=h.device, dtype=emb.weight.dtype)
        proj_mod = text.per_layer_model_projection
        proj = _bi(proj_mod, h) if fork else proj_mod(h)
        proj = (proj * text.per_layer_model_projection_scale).reshape(B, L, nl, self.ple_dim)
        proj = (_gfnorm if fork else _gnorm)(text.per_layer_projection_norm, proj)
        return (proj + tok.reshape(B, L, nl, self.ple_dim).to(proj.dtype)) \
            * text.per_layer_input_scale

    def ropes(self, h: Tensor, pos: Tensor) -> dict[str, tuple[Tensor, Tensor]]:
        """cos/sin [B, L, D] per layer kind for positions ``pos`` [B, L]."""
        out = {}
        for kind, name in (("sliding", "sliding_attention"), ("full", "full_attention")):
            if kind in self.kinds:
                out[kind] = self.text.rotary_emb(h, pos, name)
        return out

    def rope(self, h: Tensor, pos: Tensor) -> tuple[Tensor, Tensor]:
        raise TypeError("Gemma 4 has one rotary table per layer kind: use ropes()")

    def _slot_shape(self, slot: int) -> tuple[int, int]:
        a = self.text.layers[self.slot_layers[slot]].self_attn
        return a.k_proj.out_features // a.head_dim, a.head_dim

    def empty_state(self, batch: int, device, dtype) -> State:
        kv = []
        for slot, kind in enumerate(self.slot_kinds):
            hkv, d = self._slot_shape(slot)
            n = self.tail if kind == "sliding" else 0
            kv.append((torch.zeros(batch, hkv, n, d, device=device, dtype=dtype),
                       torch.zeros(batch, hkv, n, d, device=device, dtype=dtype)))
        return State([], kv, torch.zeros(batch, 0, dtype=torch.bool, device=device),
                     torch.zeros(batch, dtype=torch.long, device=device))

    def _window_mask(self, plens: Tensor, L: int) -> Tensor:
        """Sliding layers' mask [B, L, T+L] over (tail, this call's tokens): tail slot j
        (position ``plens - T + j``) is visible to query i iff it exists and ``j >= i``;
        own key j iff ``i - window < j <= i``."""
        device = plens.device
        T = self.tail
        i = torch.arange(L, device=device)[:, None]
        j = torch.arange(T, device=device)[None]
        tail_ok = (j >= i)[None] & (j[None] >= (T - plens)[:, None, None])
        jj = torch.arange(L, device=device)[None]
        own = (jj <= i) & (i - jj < self.window)
        return torch.cat([tail_ok, own[None].expand(plens.shape[0], -1, -1)], 2)

    def _tail_of(self, ext: Tensor, lengths: Tensor) -> Tensor:
        """The last T real positions of ext [B, Hkv, T+L, D] (tail + this call's tokens,
        right-padded): gathered at each row's length, like a conv tail."""
        T = self.tail
        idx = lengths[:, None] + torch.arange(T, device=ext.device)[None]
        return ext.gather(2, idx[:, None, :, None].expand(-1, ext.shape[1], -1, ext.shape[3]))

    def forward(self, ids: Tensor, lengths: Tensor, state: State | None = None,
                consume: bool = False, rows: Tensor | None = None, sel: Tensor | None = None,
                soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State | None]:
        """Same contract as :meth:`Backbone.forward`. Inputs longer than
        :data:`DENSE_PREFILL_MAX` are read :data:`PREFILL_CHUNK` tokens at a time (the
        state, sliding tails included, is carried across chunks).

        ``soft`` = (mask [B, L] bool, embeds [N, C]): multimodal soft tokens (HF Gemma 4
        placeholders): the N masked positions (row-major) take ``embeds`` in place of their
        token embeddings, and their PLE token part is the pad token's (as HF does)."""
        if rows is not None and ids.shape[1] > DENSE_PREFILL_MAX:
            assert state is not None
            state, rows = state.select(rows), None
        if soft is not None and (rows is not None or soft[0].shape != ids.shape
                                 or int(soft[0].sum()) != soft[1].shape[0]):
            raise ValueError("soft tokens: mask [B, L] over ids with one embedding per "
                             "masked position (not on fork rows)")
        if ids.shape[1] > DENSE_PREFILL_MAX:
            return self._forward_chunked(ids, lengths, state, consume, sel, soft)
        if profile._ACTIVE.get() is not None:
            profile.count(forwards=1, forward_rows=ids.shape[0],
                          forward_tokens=ids.numel(), forward_real_tokens=int(lengths.sum()))
        return self._read(ids, lengths, state, consume, rows, sel, soft=soft)

    def embed(self, ids: Tensor, soft: tuple[Tensor, Tensor] | None = None,
              fork: bool = False) -> tuple[Tensor, Tensor | None]:
        """(scaled token embeddings [B, L, C] with soft tokens merged in, PLE inputs)."""
        if soft is None:
            h = self.text.embed_tokens(ids)
            return h, self.ple_inputs(ids, h, fork=fork)
        mask, feats = soft
        tids = ids.masked_fill(mask, int(self.config.pad_token_id or 0))
        h = self.text.embed_tokens(tids)
        h = h.masked_scatter(mask[..., None], feats.to(h.dtype))
        return h, self.ple_inputs(tids, h, fork=fork)

    def _read(self, ids: Tensor, lengths: Tensor, state: State | None, consume: bool,
              rows: Tensor | None, sel: Tensor | None,
              parts: list[tuple[list[Tensor], list[Tensor]]] | None = None, s: int = 0,
              soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State | None]:
        """One dense read of ids [B, L] on top of ``state``. ``parts`` (chunked training
        read): per global slot, the incoming prefix K/V and the earlier chunks' K/V of this
        call (concatenated inside each checkpoint, so no O(L^2/C) copies are saved; the new
        chunk's K/V is appended); ``s`` = tokens of this call before ids. ``state`` then
        carries the base ``kv_mask``, the carried sliding tails and each row's real tokens
        so far (``lengths``); its global entries are returned unchanged."""
        B, L = ids.shape
        device = ids.device
        h, ple = self.embed(ids, soft)
        if state is None:
            state = self.empty_state(B, device, h.dtype)
        kv_mask, plens = state.kv_mask, state.lengths
        if rows is not None:
            kv_mask, plens = kv_mask.index_select(0, rows), plens.index_select(0, rows)
        pos = plens[:, None] + torch.arange(L, device=device)[None]
        ropes = self.ropes(h, pos)
        ckpt = self.grad_ckpt and torch.is_grad_enabled()
        masks: dict[str, Any] = {"sliding": self._window_mask(plens, L)}
        if parts is not None:
            masks["full"] = (kv_mask, s)  # built inside the checkpoint
        else:
            causal = torch.ones(L, L, dtype=torch.bool, device=device).tril()
            masks["full"] = torch.cat([kv_mask[:, None, :].expand(-1, L, -1),
                                       causal[None].expand(B, -1, -1)], 2)
        new_kv: list[tuple[Tensor, Tensor] | None] = [None] * len(self.slot_kinds)
        cur: dict[int, tuple[Tensor, Tensor]] = {}
        for i, layer in enumerate(self.text.layers):
            slot, kind = self.kv_slot[i], self.kinds[i]
            kp, vp = state.kv[slot]
            if parts is not None and kind == "full":
                kp, vp = tuple(parts[slot][0]), tuple(parts[slot][1])
            kn, vn = (None, None) if self.owner[i] else cur[slot]
            args = (layer, h, None if ple is None else ple[:, :, i], *ropes[kind], kp, vp,
                    masks[kind], rows, sel, kn, vn, self.window if kind == "sliding" else 0)
            if ckpt:
                h, k, v = checkpoint(_gemma_layer, *args, use_reentrant=False)
            else:
                h, k, v = _gemma_layer(*args)
            del args
            if self.owner[i]:
                cur[slot] = (k, v)
                if rows is None:
                    if kind == "sliding":
                        new_kv[slot] = (self._tail_of(_cat_kv(kp, k), lengths),
                                        self._tail_of(_cat_kv(vp, v), lengths))
                    elif parts is None:
                        new_kv[slot] = (_cat_kv(kp, k), _cat_kv(vp, v))
            del k, v
            if consume and i == self.last_use[slot]:
                del kp, vp
                state.kv[slot] = None  # type: ignore[call-overload]
        h = self.final_norm(h)
        if rows is not None:
            return h, None
        if parts is not None:  # chunked training read: the caller keeps the global parts
            for slot, kind in enumerate(self.slot_kinds):
                if kind == "full":
                    parts[slot][0].append(cur[slot][0])
                    parts[slot][1].append(cur[slot][1])
                    new_kv[slot] = state.kv[slot]
        valid = torch.arange(L, device=device)[None] < lengths[:, None]
        new = State([], new_kv, torch.cat([state.kv_mask, valid], 1),  # type: ignore[arg-type]
                    state.lengths + lengths)
        return h, new

    def _forward_chunked(self, ids: Tensor, lengths: Tensor, state: State | None,
                         consume: bool = False, sel: Tensor | None = None,
                         soft: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, State]:
        """Chunked read: :data:`PREFILL_CHUNK` tokens at a time, each chunk a dense read on
        top of the state left by the previous one (sliding tails carried, global K/V
        appended; rows are right-padded, so a row's real tokens stay contiguous). With
        autograd + ``grad_ckpt``, the global K/V of the earlier chunks are passed to the
        checkpoints as parts (never re-concatenated outside them); the returned state
        concatenates them once."""
        B, L = ids.shape
        device = ids.device
        C = PREFILL_CHUNK
        if profile._ACTIVE.get() is not None:
            profile.count(forwards=1, forward_rows=B, forward_tokens=B * L,
                          forward_real_tokens=int(lengths.sum()))
        if state is None:
            state = self.empty_state(B, device, self.text.embed_tokens.weight.dtype)
        train = self.grad_ckpt and torch.is_grad_enabled()
        outs = []
        index = None
        if soft is not None:  # each soft position's row in embeds (row-major over [B, L])
            index = soft[0].flatten().cumsum(0).reshape(soft[0].shape) - 1

        def soft_of(s0: int, e: int) -> tuple[Tensor, Tensor] | None:
            if soft is None or index is None:
                return None
            m = soft[0][:, s0:e]
            return m, soft[1][index[:, s0:e][m]]

        if not train:
            for s0 in range(0, L, C):
                e = min(L, s0 + C)
                h, nxt = self._read(ids[:, s0:e], (lengths - s0).clamp(0, e - s0), state,
                                    consume or s0 > 0, None, sel, soft=soft_of(s0, e))
                assert nxt is not None
                state = nxt
                outs.append(h)
            return torch.cat(outs, 1), state
        parts = [([kp], [vp]) if kind == "full" else ([], [])
                 for (kp, vp), kind in zip(state.kv, self.slot_kinds)]
        cur = State([], list(state.kv), state.kv_mask, state.lengths)
        for s0 in range(0, L, C):
            e = min(L, s0 + C)
            lens = (lengths - s0).clamp(0, e - s0)
            # sliding slots read the carried tails, global slots the parts (prefix + chunks)
            h, nxt = self._read(ids[:, s0:e], lens, cur, False, None, sel, parts, s0,
                                soft_of(s0, e))
            assert nxt is not None
            # each row's real tokens so far (right-padded: a row's tokens are contiguous)
            cur = State([], nxt.kv, state.kv_mask, state.lengths + lengths.clamp(max=e))
            outs.append(h)
        kv = [(torch.cat(parts[k][0], 2), torch.cat(parts[k][1], 2)) if kind == "full"
              else cur.kv[k] for k, kind in enumerate(self.slot_kinds)]
        valid = torch.arange(L, device=device)[None] < lengths[:, None]
        return torch.cat(outs, 1), State([], kv, torch.cat([state.kv_mask, valid], 1),
                                         state.lengths + lengths)

    def _forward_chunked_ckpt(self, ids: Tensor, lengths: Tensor, state: State | None,
                              sel: Tensor | None = None) -> tuple[Tensor, State]:
        return self._forward_chunked(ids, lengths, state, False, sel)


    def fork(self, prefixes: State | list[State], prow: list[int] | int, cap: int
             ) -> ForkState:
        """As :meth:`Backbone.fork`: global slots read the prefixes' K/V by reference
        (several concatenated once); sliding slots read their tails (T per prefix,
        concatenated: row b's live tail is ``wplen[b] = min(len, T)`` positions at
        ``wpoff[b]``)."""
        if isinstance(prefixes, State):
            prefixes = [prefixes]
        if isinstance(prow, int):
            prow = [0] * prow
        for p in prefixes:
            assert p.lengths.shape[0] == 1 and bool(p.kv_mask.all())
        rows = len(prow)
        device = prefixes[0].lengths.device
        if device.type == "cpu":
            cap += CPU_BUCKET
        T = self.tail
        own_k, own_v, pk, pv = [], [], [], []
        for slot in range(len(self.slot_kinds)):
            k0 = prefixes[0].kv[slot][0]
            shape = (rows, k0.shape[1], cap, k0.shape[3])
            own_k.append(torch.zeros(shape, device=device, dtype=k0.dtype))
            own_v.append(torch.zeros(shape, device=device, dtype=k0.dtype))
            if len(prefixes) == 1:
                pk.append(k0[0])
                pv.append(prefixes[0].kv[slot][1][0])
            else:
                pk.append(torch.cat([p.kv[slot][0][0] for p in prefixes], 1))
                pv.append(torch.cat([p.kv[slot][1][0] for p in prefixes], 1))
        lens = [int(p.lengths[0]) for p in prefixes]
        offs = [sum(lens[:i]) for i in range(len(lens))]
        live = [min(n, T) for n in lens]
        woff = [i * T + T - n for i, n in enumerate(live)]
        idx = torch.tensor(prow, dtype=torch.long, device=device)

        def at(xs: list[int]) -> Tensor:
            return torch.tensor(xs, dtype=torch.long, device=device)[idx]

        return ForkState([], pk, pv, at(offs), at(lens), idx, [], [], own_k, own_v,
                         torch.zeros(rows, dtype=torch.long, device=device),
                         wpoff=at(woff), wplen=at(live))

    def _fork_forward(self, ids: Tensor, lengths: Tensor, fs: ForkState, keep: bool = False,
                      sel: Tensor | None = None) -> Tensor:
        """:meth:`Backbone._fork_forward` over the K/V slots (sliding layers on their
        prefix tails, K/V-sharing layers on their source's buffers). ``fs.ple_rows``: the
        PLE table rows of ``ids`` gathered beforehand (a decode graph's static buffer when
        the table lives in host memory)."""
        B, L = ids.shape
        if not ids.is_cuda and (B > 1 or (L > 1 and L % CPU_BUCKET)):
            return self._fork_forward_rows(ids, lengths, fs, keep, sel)  # CPU reference
        device = ids.device
        if profile._ACTIVE.get() is not None and not torch.cuda.is_current_stream_capturing():
            profile.count(forwards=1, forward_rows=B, forward_tokens=B * L)
        h = self.text.embed_tokens(ids)
        ple = self.ple_inputs(ids, h, fork=True, rows=fs.ple_rows)
        pos = fs.plen[:, None] + fs.own_len[:, None] + torch.arange(L, device=device)[None]
        ropes = self.ropes(h, pos)
        for i, layer in enumerate(self.text.layers):
            kind = self.kinds[i]
            h = _gemma_fork_layer(layer, h, None if ple is None else ple[:, :, i], *ropes[kind],
                                  fs, self.kv_slot[i], self.owner[i],
                                  self.window if kind == "sliding" else 0, sel)
        if keep:
            if fs.static:
                fs.own_len.add_(lengths)
            else:
                fs.own_len = fs.own_len + lengths
            fs.fresh = False
        return self.fork_final_norm(h)

    def lm_logits(self, hidden: Tensor) -> Tensor:
        """:meth:`Backbone.lm_logits` with Gemma's final logit soft-cap
        (``cap * tanh(logits / cap)``, fp32; argmax unchanged)."""
        logits = super().lm_logits(hidden)
        if not self.softcap:
            return logits
        return torch.tanh(logits.float() / self.softcap) * self.softcap

    def static_fork_state(self, rows: int, pcap: int, cap: int) -> ForkState:
        """CUDA-graph buffers (:class:`DecodeGraph`): global slots' prefixes up to ``pcap``
        tokens in total, sliding slots' tails for up to ``rows`` prefixes, own K/V of
        ``cap`` tokens per row."""
        dev = self.text.embed_tokens.weight.device
        dtype = self.text.embed_tokens.weight.dtype
        pk, pv, ok, ov = [], [], [], []
        for slot, kind in enumerate(self.slot_kinds):
            hkv, d = self._slot_shape(slot)
            n = rows * self.tail if kind == "sliding" else pcap
            pk.append(torch.zeros(hkv, n, d, device=dev, dtype=dtype))
            pv.append(torch.zeros(hkv, n, d, device=dev, dtype=dtype))
            ok.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype))
            ov.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype))
        z = torch.zeros(rows, dtype=torch.long, device=dev)
        return ForkState([], pk, pv, z.clone(), z.clone(), z.clone(), [], [], ok, ov,
                         z.clone(), fresh=False, static=True, wpoff=z.clone(), wplen=z.clone())

    def load_static(self, fs: ForkState, prefixes: list[State], prow: list[int]) -> None:
        """Copy batch-1 prefixes into :meth:`static_fork_state` buffers; row b (<
        len(prow)) continues ``prefixes[prow[b]]``, spare rows idle on prefix 0."""
        rows = fs.own_len.shape[0]
        T = self.tail
        lens = [int(p.lengths[0]) for p in prefixes]
        assert len(prefixes) <= rows and len(prow) <= rows
        offs = [sum(lens[:i]) for i in range(len(lens))]
        for pi, (p, o, n) in enumerate(zip(prefixes, offs, lens)):
            for slot, kind in enumerate(self.slot_kinds):
                k, v = p.kv[slot]
                if kind == "sliding":
                    fs.pk[slot][:, pi * T:(pi + 1) * T].copy_(k[0])
                    fs.pv[slot][:, pi * T:(pi + 1) * T].copy_(v[0])
                else:
                    assert o + n <= fs.pk[slot].shape[1]
                    fs.pk[slot][:, o:o + n].copy_(k[0])
                    fs.pv[slot][:, o:o + n].copy_(v[0])
        dev = fs.own_len.device
        idx = torch.tensor(list(prow) + [0] * (rows - len(prow)), dtype=torch.long, device=dev)
        live = [min(n, T) for n in lens]
        assert fs.wpoff is not None and fs.wplen is not None
        fs.poff.copy_(torch.tensor(offs, device=dev)[idx])
        fs.plen.copy_(torch.tensor(lens, device=dev)[idx])
        fs.prow.copy_(idx)
        fs.wpoff.copy_(torch.tensor([i * T + T - n for i, n in enumerate(live)], device=dev)[idx])
        fs.wplen.copy_(torch.tensor(live, device=dev)[idx])
        fs.own_len.zero_()
        fs.fresh = True



class MediaTowers(nn.Module):
    """Gemma 4's frozen vision and audio encoders with their embedders (HF
    ``Gemma4Model.vision_tower`` / ``audio_tower`` / ``embed_vision`` / ``embed_audio``),
    loaded from the base snapshot on demand (:func:`load_media_towers`; optional: a text
    deployment never loads them). Never trained or saved by deg: a checkpoint's media
    path is its base snapshot's (``meta["base"]``). Outputs are soft-token embeddings in
    the language model's space for :class:`Media` / :meth:`Renderer.context`."""

    def __init__(self, config: Any) -> None:
        super().__init__()
        from transformers import AutoModel
        from transformers.models.gemma4.modeling_gemma4 import Gemma4MultimodalEmbedder

        self.config = config
        text = config.text_config
        self.vision_tower = (AutoModel.from_config(config.vision_config)
                             if config.vision_config is not None else None)
        self.audio_tower = (AutoModel.from_config(config.audio_config)
                            if config.audio_config is not None else None)
        self.embed_vision = (Gemma4MultimodalEmbedder(config.vision_config, text)
                             if config.vision_config is not None else None)
        self.embed_audio = (Gemma4MultimodalEmbedder(config.audio_config, text)
                            if config.audio_config is not None else None)

    @torch.no_grad()
    def image_embeds(self, pixel_values: Tensor, position_ids: Tensor) -> list[Tensor]:
        """HF ``get_image_features``: processor outputs (``pixel_values`` [I, patches,
        3*p*p], ``image_position_ids`` [I, patches, 2]) -> per image [n_i, hidden]."""
        assert self.vision_tower is not None and self.embed_vision is not None
        out = self.vision_tower(pixel_values=pixel_values, pixel_position_ids=position_ids)
        pooled = self.embed_vision(inputs_embeds=out.last_hidden_state)
        k = self.config.vision_config.pooling_kernel_size
        output_length = pixel_values.shape[-2] // (k * k)
        k_squared = int((position_ids.shape[1] // output_length) ** 0.5) ** 2
        sizes = ((position_ids != -1).all(-1).sum(-1) // k_squared).tolist()
        return list(torch.split(pooled.reshape(-1, pooled.shape[-1]), sizes))

    @torch.no_grad()
    def audio_embeds(self, features: Tensor, mask: Tensor) -> list[Tensor]:
        """HF ``get_audio_features``: feature-extractor outputs (``input_features`` [A, T,
        mels], ``input_features_mask`` [A, T]) -> per clip its valid soft tokens."""
        assert self.audio_tower is not None and self.embed_audio is not None
        out = self.audio_tower(features, mask, return_dict=True)
        pooled = self.embed_audio(inputs_embeds=out.last_hidden_state)
        valid = out.attention_mask
        return [pooled[i][valid[i]] for i in range(pooled.shape[0])]


MEDIA_PREFIXES = ("model.vision_tower.", "model.audio_tower.", "model.embed_vision.",
                  "model.embed_audio.")


def load_media_towers(snapshot: Path | str | None = None, *, backbone: str | None = None,
                      device: str | torch.device = "cpu",
                      dtype: torch.dtype = torch.bfloat16) -> MediaTowers:
    """The frozen vision/audio towers of a Gemma 4 snapshot (default: the backbone's base
    snapshot), read from its safetensors without loading the text model."""
    import transformers
    from safetensors import safe_open

    snap = Path(snapshot) if snapshot is not None else base_snapshot(backbone or
                                                                     "gemma-4-e2b-it")
    cfg = transformers.AutoConfig.from_pretrained(str(snap))
    towers = MediaTowers(cfg)
    state: dict[str, Tensor] = {}
    for f in sorted(snap.glob("*.safetensors")):
        with safe_open(str(f), "pt") as fh:
            for k in fh.keys():
                if k.startswith(MEDIA_PREFIXES):
                    state[k[len("model."):]] = fh.get_tensor(k)
    missing, unexpected = towers.load_state_dict(state, strict=False)
    if unexpected or [k for k in missing if not k.endswith(("inv_freq",))]:
        raise RuntimeError(f"media towers mismatch: {(list(unexpected) + missing)[:5]}")
    towers.to(device=device, dtype=dtype)
    towers.requires_grad_(False)
    towers.eval()
    return towers


class DecodeGraph:
    """One generate decode step (``rows`` fork rows x 1 token, LM head, argmax) captured
    as a CUDA graph over static buffers: prefix K/V for up to ``pcap`` tokens in total
    (row offsets/lengths are device tensors), per-row recurrent state / conv tail / own
    K/V (``cap`` own tokens). :meth:`load` copies the prefixes in; the header pass runs
    eagerly on the same buffers; every decode step is then one ``replay``. The step is
    the same batch-invariant fork forward as the eager path, so replay changes speed,
    not semantics. ``sel``: the fixed route of every step (a deg-MoE's
    :class:`WriteRoute` for ``rows`` rows; None = the backbone as is)."""

    def __init__(self, backbone: Backbone, rows: int, pcap: int, cap: int,
                 sel: WriteRoute | None = None) -> None:
        self.backbone, self.rows, self.pcap, self.cap = backbone, rows, pcap, cap
        self.sel = sel
        dev = backbone.text.embed_tokens.weight.device
        dtype = backbone.text.embed_tokens.weight.dtype
        self.ids = torch.zeros(rows, 1, dtype=torch.long, device=dev)
        self.lens = torch.zeros(rows, dtype=torch.long, device=dev)
        self.graph: torch.cuda.CUDAGraph | None = None
        self.logits: Tensor | None = None
        self.top1: Tensor | None = None
        # Gemma 4 with the PLE table in host memory: each step's table rows are gathered on
        # the host and copied into this static buffer before the replay
        self.ple_rows: Tensor | None = None
        if isinstance(backbone, GemmaBackbone):
            self.fs = backbone.static_fork_state(rows, pcap, cap)
            if backbone.ple_on_host:
                w = backbone.text.embed_tokens_per_layer.weight
                self.ple_rows = torch.zeros(rows, 1, w.shape[1], device=dev, dtype=w.dtype)
            return
        rec, tail, pk, pv, ok, ov = [], [], [], [], [], []
        for layer, kind in zip(backbone.text.layers, backbone.types):
            if kind == "rec":
                r, t = backbone.rec_buffers(layer, rows, dev, dtype)
                rec.append(r)
                tail.append(t)
            else:
                hkv, d = backbone.config.num_key_value_heads, layer.self_attn.head_dim
                pk.append(torch.zeros(hkv, pcap, d, device=dev, dtype=dtype))
                pv.append(torch.zeros(hkv, pcap, d, device=dev, dtype=dtype))
                ok.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype))
                ov.append(torch.zeros(rows, hkv, cap, d, device=dev, dtype=dtype))
        zeros = torch.zeros(rows, dtype=torch.long, device=dev)
        self.fs = ForkState([r[:1] for r in rec], pk, pv, zeros.clone(), zeros.clone(),
                            zeros.clone(), rec, tail, ok, ov, zeros.clone(), fresh=False,
                            static=True)

    def _step(self) -> None:
        self.fs.ple_rows = self.ple_rows  # read by the captured kernels on every replay
        try:
            hid = self.backbone.fork_forward(self.ids, self.lens, self.fs, keep=True,
                                             decode=True, sel=self.sel)
        finally:
            self.fs.ple_rows = None  # the eager header pass gathers its own rows
        self.logits = self.backbone.lm_logits(hid[:, 0]).float()
        self.top1 = self.logits.argmax(-1)

    def capture(self) -> None:
        self.fs.fresh = False
        side = torch.cuda.Stream()
        side.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(side):
            for _ in range(3):  # compile / autotune outside the capture
                self._step()
        torch.cuda.current_stream().wait_stream(side)
        self.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self.graph):
            self._step()

    def load(self, prefixes: list[State], prow: list[int]) -> ForkState:
        """Point the static rows at new batch-1 prefixes: row b (< len(prow)) continues
        ``prefixes[prow[b]]``; spare rows idle on prefix 0. Fresh rows, own_len 0."""
        lens = [int(p.lengths[0]) for p in prefixes]
        assert sum(lens) <= self.pcap and len(prow) <= self.rows
        if isinstance(self.backbone, GemmaBackbone):
            self.backbone.load_static(self.fs, prefixes, prow)
            return self.fs
        offs = [sum(lens[:i]) for i in range(len(lens))]
        for p, o, n in zip(prefixes, offs, lens):
            for i, (k, v) in enumerate(p.kv):
                self.fs.pk[i][:, o:o + n].copy_(k[0])
                self.fs.pv[i][:, o:o + n].copy_(v[0])
        rows = list(prow) + [0] * (self.rows - len(prow))
        dev = self.fs.own_len.device
        idx = torch.tensor(rows, dtype=torch.long, device=dev)
        self.fs.poff.copy_(torch.tensor(offs, device=dev)[idx])
        self.fs.plen.copy_(torch.tensor(lens, device=dev)[idx])
        self.fs.prow.copy_(idx)
        self.fs.prefix_rec = [torch.cat([p.lin[i][0] for p in prefixes]).contiguous()
                              for i in range(len(prefixes[0].lin))]
        for i, t in enumerate(self.fs.tail):
            t.copy_(torch.cat([p.lin[i][1] for p in prefixes]).index_select(0, idx))
        self.fs.own_len.zero_()
        self.fs.fresh = True
        return self.fs

    def step(self, ids: Tensor, lens: Tensor) -> tuple[Tensor, Tensor]:
        """One decode step for every row (``lens`` 0 = idle row) -> (logits, argmax)."""
        if self.ple_rows is not None:
            assert isinstance(self.backbone, GemmaBackbone)
            self.ple_rows.copy_(self.backbone.ple_rows(ids.cpu()), non_blocking=True)
        self.ids.copy_(ids, non_blocking=True)
        self.lens.copy_(lens, non_blocking=True)
        assert self.graph is not None and self.logits is not None and self.top1 is not None
        self.fs.fresh = False
        self.graph.replay()
        return self.logits, self.top1


def _pow2(n: int, lo: int) -> int:
    out = lo
    while out < n:
        out *= 2
    return out


# ---------------------------------------------------------------------------------- heads


class PointerHead(nn.Module):
    def __init__(self, hidden: int, dim: int = HEAD_DIM) -> None:
        super().__init__()
        self.query = nn.Linear(hidden, dim)
        self.key = nn.Linear(hidden, dim)
        self.prior = nn.Linear(hidden, 1)
        self.scale = 1.0 / math.sqrt(dim)

    def forward(self, answer: Tensor, options: Tensor) -> Tensor:
        """answer [..., C], options [..., O, C] -> logits [..., O]."""
        q = self.query(answer).unsqueeze(-2)
        return ((self.key(options) * q).sum(-1) * self.scale
                + self.prior(options).squeeze(-1))


class SpanHead(nn.Module):
    def __init__(self, hidden: int, dim: int = HEAD_DIM) -> None:
        super().__init__()
        self.q_start, self.k_start = nn.Linear(hidden, dim), nn.Linear(hidden, dim)
        self.q_end, self.k_end = nn.Linear(hidden, dim), nn.Linear(hidden, dim)
        self.null_start, self.null_end = nn.Linear(hidden, 1), nn.Linear(hidden, 1)
        self.scale = 1.0 / math.sqrt(dim)

    def forward(self, answer: Tensor, context: Tensor) -> tuple[Tensor, Tensor]:
        """answer [n, C], context [T, C] -> start/end logits [n, T+1] (last = null)."""
        s = (self.q_start(answer) @ self.k_start(context).T) * self.scale
        e = (self.q_end(answer) @ self.k_end(context).T) * self.scale
        return (torch.cat([s, self.null_start(answer)], -1),
                torch.cat([e, self.null_end(answer)], -1))


class AnchorHead(nn.Module):
    """Item anchors: one independent "an item of this array starts here" logit per context
    token, ``q(a) . k(c_t) / sqrt(d) + u . c_t + b`` (the prior ``u . c_t + b`` starts at
    :data:`ANCHOR_INIT_BIAS`)."""

    def __init__(self, hidden: int, dim: int = HEAD_DIM) -> None:
        super().__init__()
        self.query = nn.Linear(hidden, dim)
        self.key = nn.Linear(hidden, dim)
        self.prior = nn.Linear(hidden, 1)
        with torch.no_grad():
            self.prior.weight.mul_(0.1)
            self.prior.bias.fill_(ANCHOR_INIT_BIAS)
        self.scale = 1.0 / math.sqrt(dim)

    def forward(self, answer: Tensor, context: Tensor) -> Tensor:
        """answer [n, C], context [T, C] -> logits [n, T]."""
        return ((self.query(answer) @ self.key(context).T) * self.scale
                + self.prior(context)[:, 0][None])


def rel_bucket(d: Tensor) -> Tensor:
    """Bucket of a token offset ``t - anchor`` (:data:`REL_EDGES`; 0 .. len(REL_EDGES))."""
    return torch.bucketize(d, torch.tensor(REL_EDGES, device=d.device), right=True)


def span_rel_bias(table: Tensor, anchor: int, hi: int, r0: int, r1: int
                  ) -> tuple[Tensor, Tensor]:
    """Anchored span bias for context tokens [r0, r1): ``table[which, inside, bucket]``
    with ``inside = anchor <= t < hi`` (the item's region) -> (start bias, end bias)."""
    t = torch.arange(r0, r1, device=table.device)
    b = rel_bucket(t - anchor)
    inside = ((t >= anchor) & (t < hi)).long()
    return table[0, inside, b], table[1, inside, b]


def region_bounds(region: tuple[int, int], T: int) -> tuple[int, int]:
    lo, hi = region
    return max(0, lo), (T if hi < 0 else min(hi, T))


class Heads(nn.Module):
    def __init__(self, hidden: int) -> None:
        super().__init__()
        self.choice = PointerHead(hidden)
        self.score = PointerHead(hidden)
        self.noul = nn.Linear(hidden, 1)
        self.span = SpanHead(hidden)
        self.count = nn.Linear(hidden, MAX_COUNT + 1)
        self.anchor = AnchorHead(hidden)
        # [start/end, outside/inside the item region, offset bucket]; zero = no bias
        self.span_bias = nn.Parameter(torch.zeros(2, 2, len(REL_EDGES) + 1))
        self.register_buffer("temperature", torch.ones(len(KINDS)))

    def temp(self, kind: str) -> Tensor:
        return self.temperature[KINDS.index(kind)]


class _Picked:
    """The rows of one fork's hidden states [Lq, C] that the heads read (answer and
    option positions), indexable by original position; the rest is dropped."""

    __slots__ = ("rows", "at")

    def __init__(self, hidden: Tensor, positions: tuple[int, ...]) -> None:
        keep = sorted(set(positions))
        self.at = {p: k for k, p in enumerate(keep)}
        self.rows = hidden[keep]

    def __getitem__(self, p: int) -> Tensor:
        return self.rows[self.at[p]]


def fork_token_budget(device: torch.device) -> int:
    """Padded fork tokens per batched fork pass: :data:`FORK_TOKENS`, lowered on CUDA to
    what fits in :data:`FORK_MEM_FRACTION` of the memory free right now (device free +
    cached-but-unused), so a fork pass next to a big context (or another process on the
    GPU) is planned to fit instead of OOM-and-split. Batch composition never changes a
    row's outputs (batch-invariant kernels), only speed."""
    if device.type != "cuda":
        return FORK_TOKENS
    free, _ = torch.cuda.mem_get_info(device)
    free += torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device)
    return max(256, min(FORK_TOKENS, int(free * FORK_MEM_FRACTION) // FORK_BYTES_PER_TOKEN))


@dataclass
class QueryOutput:
    """Raw (uncalibrated) logits for one query."""

    kind: str
    logits: Tensor | None = None  # decide outcomes / count / anchor [T] (-inf off-region)
    start: Tensor | None = None  # extract: [T+1]
    end: Tensor | None = None


def soft_inputs(contexts: list[EncodedContext], shape: tuple[int, int], device: Any
                ) -> tuple[Tensor, Tensor] | None:
    """The backbone's ``soft`` = (mask [B, L], embeds [N, C]) of the contexts' media (row i
    = contexts[i], read from position 0), or None when none has media."""
    if not any(c.media for c in contexts):
        return None
    mask = torch.zeros(shape, dtype=torch.bool, device=device)
    feats = []
    for i, c in enumerate(contexts):  # row-major, as the soft-token mask enumerates
        for start, emb in sorted(c.media, key=lambda m: m[0]):
            mask[i, start:start + emb.shape[0]] = True
            feats.append(emb.to(device))
    return mask, torch.cat(feats)


@dataclass
class DegModel:
    backbone: Backbone
    heads: Heads
    renderer: Renderer
    meta: dict[str, Any] = field(default_factory=dict)

    graphs: dict[tuple[int, int, int], DecodeGraph] = field(default_factory=dict)
    # frozen base Qwen3.5 body (:func:`attach_base`): the x-generate decoder when
    # ``generate_with == "base"`` (the hybrid), and the --lm-kl-base anchor in training
    base: Backbone | None = None
    generate_with: str = "self"  # self | base (:data:`GENERATE_WITH`)
    _gen_view: DegModel | None = None
    # LoRA (:func:`add_lora`): {"rank", "alpha", "adapters": [names], "context": "base" or
    # an adapter name}; None = full fine-tune layout (no adapters)
    lora: dict[str, Any] | None = None
    # deg-MoE (:func:`add_moe`): {"writer_rank", "writer_alpha", "writer_targets"}; None =
    # one FFN per layer. An MoE model generates with its W experts (generate_with "moe")
    moe: dict[str, Any] | None = None

    def gen_route(self, rows: int) -> WriteRoute | None:
        """``sel`` of x-generate calls (KNOWN VALUES prefill, generate header, decode):
        an MoE's W route for ``rows`` rows; None for every other layout."""
        if self.moe is None:
            return None
        if "experts" in self.moe:  # D/E/G task MoE
            return TaskRoute("g")
        r = int(self.moe["writer_rank"])
        return WriteRoute(torch.ones(rows, r, dtype=torch.float32, device=self.device)
                          if r else None)

    def adapter_of(self, kind: str) -> str | None:
        """The adapter a fork kind runs with (None = adapter off, the base model)."""
        if self.lora is None:
            return None
        names = self.lora["adapters"]
        if names == ["shared"]:
            return "shared" if kind in ADAPTER_OF_KIND else None
        name = ADAPTER_OF_KIND.get(kind)
        return name if name in names else None

    def _sel(self, names: list[str | None]) -> Tensor | None:
        if self.lora is None or all(n is None for n in names):
            return None
        r, order = int(self.lora["rank"]), self.lora["adapters"]
        sel = torch.zeros(len(names), r * len(order), dtype=torch.float32)
        for i, n in enumerate(names):
            if n is not None:
                k = order.index(n)
                sel[i, k * r:(k + 1) * r] = 1.0
        return sel.to(self.device)

    def lora_sel(self, kinds: list[str]) -> Tensor | TaskRoute | None:
        """Per-row adapter selector for fork rows of these kinds (None: all off); for a
        D/E/G task MoE the route of the call (every kind must share it)."""
        if self.moe is not None and "experts" in self.moe:
            routes = {self.fork_route(k) for k in kinds}
            assert len(routes) == 1, f"one expert per fork call, got {routes}"
            return TaskRoute(routes.pop())
        return self._sel([self.adapter_of(k) for k in kinds])

    def fork_route(self, kind: str) -> str | None:
        """The expert a fork kind runs in a D/E/G task MoE (None: not one); fork batches
        are grouped by it (:meth:`answer_many`, ``deg.train.forward_losses``)."""
        if self.moe is None or "experts" not in self.moe:
            return None
        return TASK_EXPERT_OF_KIND[kind]

    def context_sel(self, rows: int) -> Tensor | None:
        """Selector of the context (trunk) read: None when the context is read by the base
        weights (``lora["context"] == "base"``: one shared prefill for every fork,
        generate included)."""
        if self.lora is None or self.lora["context"] == "base":
            return None
        return self._sel([self.lora["context"]] * rows)

    def generator(self) -> tuple[DegModel, bool, bool]:
        """How x-generate fields decode: ``(model, trunk_shared, base_layout)``.
        ``model`` decodes (adapter-off rows); ``trunk_shared``: the request's trunk state
        can be extended by the KNOWN VALUES block (else the context is read again by
        ``model``); ``base_layout``: the base model writes (base header + assistant turn).

        * full fine-tune, generate_with self: (self, True, False) -- deg's own LM
          (``gen_layout == "chat"``: (self, True, True), its own weights in the chat layout);
        * full fine-tune, generate_with base: (base copy, False, True) -- two-copy hybrid;
        * LoRA, context base: (self, True, True) -- one prefill, generation = base;
        * LoRA, context adapter: (self, False, True) -- adapter-off re-read of the context;
        * deg-MoE: (self, True, True) -- one prefill (R), the W experts write in the base
          layout (calls routed by :meth:`gen_route`).
        """
        if self.moe is not None:
            return self, True, True
        if self.lora is not None:
            return self, self.lora["context"] == "base", True
        gm = self.generate_model()
        if gm is self:
            return self, True, self.gen_layout == "chat"
        return gm, False, True

    @property
    def gen_layout(self) -> str:
        """``meta["gen_layout"]`` (:data:`GEN_LAYOUTS`; default ``deg``)."""
        layout = self.meta.get("gen_layout", "deg")
        assert layout in GEN_LAYOUTS, layout
        return layout

    @property
    def device(self) -> torch.device:
        return next(self.heads.parameters()).device

    def generate_model(self) -> DegModel:
        """The model that decodes x-generate fields: ``self``, or (hybrid) a view with the
        frozen base backbone, the same tokenizer and its own decode-graph cache."""
        if self.generate_with in ("self", "moe"):
            return self
        assert self.generate_with == "base", self.generate_with
        return self.base_view()

    def base_view(self) -> DegModel:
        """This model with the frozen base backbone (:func:`attach_base`) in place of its
        own: same heads, tokenizer and meta, its own decode-graph cache."""
        if self.base is None:
            raise RuntimeError("no base backbone: attach_base(model) first")
        if self._gen_view is None or self._gen_view.backbone is not self.base:
            self._gen_view = DegModel(self.base, self.heads, self.renderer, self.meta)
        return self._gen_view

    def decode_graph(self, rows: int, prefix_len: int, cap: int) -> DecodeGraph:
        """Cached decode graph for at least ``rows`` rows (1, 2, 4, ...), a prefix of
        ``prefix_len`` tokens (pow2 >= 2048) and ``cap`` own tokens (pow2 >= 512)."""
        key = (_pow2(rows, 1), _pow2(prefix_len, 2048), _pow2(cap, 512))
        if key not in self.graphs:
            graph = DecodeGraph(self.backbone, *key, sel=self.gen_route(key[0]))
            graph.capture()  # idle rows (lens 0): the warm-up steps are no-ops on state
            self.graphs[key] = graph
        return self.graphs[key]

    def modules(self) -> list[nn.Module]:
        return [self.backbone, self.heads]

    @property
    def array_mode(self) -> str:
        """How array items are addressed (:data:`ARRAY_MODES`): ``anchor`` (position
        anchors, no count cap) or ``index`` ("ITEM i of" + the capped count head; run-1)."""
        mode = self.meta.get("array_mode", "index")
        assert mode in ARRAY_MODES, mode
        return mode

    def array_mode_for(self, doc_tokens: int) -> str:
        """The array mode of one document of ``doc_tokens`` context tokens: an ``anchor``
        model routes documents shorter than ``meta["anchor_min_tokens"]`` (default 0 =
        never) to ``index`` (SPEC.md §12.6 Routing); inference and training share it."""
        mode = self.array_mode
        if mode == "anchor" and doc_tokens < int(self.meta.get("anchor_min_tokens", 0)):
            return "index"
        return mode

    # ------------------------------------------------------------------ query heads

    def apply_heads(self, eq: EncodedQuery, hidden: Tensor, context_hidden: Tensor | None
                    ) -> QueryOutput:
        """hidden [Lq, C] of the fork (fp32 heads)."""
        a = hidden[eq.answer_position].float()
        if eq.kind in ("choice", "score"):
            opts = hidden[list(eq.option_positions)].float()
            head = self.heads.choice if eq.kind == "choice" else self.heads.score
            return QueryOutput(eq.kind, logits=head(a, opts))
        if eq.kind == "noul":
            lg = self.heads.noul(a)[0]
            return QueryOutput("noul", logits=torch.stack([lg, torch.zeros_like(lg)]))
        if eq.kind == "count":
            return QueryOutput("count", logits=self.heads.count(a))
        assert context_hidden is not None
        T = context_hidden.shape[0]
        if eq.kind == "anchor":
            lg = self.heads.anchor(a[None], context_hidden.float())[0]
            lo, hi = region_bounds(eq.region, T)
            off = torch.ones(T, dtype=torch.bool, device=lg.device)
            off[lo:hi] = False
            return QueryOutput("anchor", logits=lg.masked_fill(off, float("-inf")))
        assert eq.kind == "extract"
        start, end = self.heads.span(a[None], context_hidden.float())
        start, end = start[0], end[0]
        if eq.anchor >= 0:
            bs, be = span_rel_bias(self.heads.span_bias, eq.anchor,
                                   region_bounds(eq.region, T)[1], 0, T)
            start = torch.cat([start[:T] + bs, start[T:]])
            end = torch.cat([end[:T] + be, end[T:]])
        return QueryOutput("extract", start=start, end=end)

    def run_forks(self, state: State, row_of: Tensor, seqs: list[tuple[int, ...]],
                  pad_to: int | None = None, kinds: list[str] | None = None) -> list[Tensor]:
        """Run forks (row_of[i] = trunk row of seqs[i]) as one right-padded batch -> per
        fork hidden [len_i, C]. ``kinds`` (per fork) select the LoRA adapters; a kind with
        no adapter (``generate``) runs adapter-off (= the base model)."""
        device = self.device
        width = pad_to or max(len(s) for s in seqs)
        ids = torch.full((len(seqs), width), self.renderer.pad_id, dtype=torch.long,
                         device=device)
        for i, s in enumerate(seqs):
            ids[i, :len(s)] = torch.tensor(s, dtype=torch.long, device=device)
        lengths = torch.tensor([len(s) for s in seqs], dtype=torch.long, device=device)
        sel = self.lora_sel(kinds) if kinds is not None else None
        hidden, _ = self.backbone(ids, lengths, state, rows=row_of, sel=sel)
        return [hidden[i, :len(s)] for i, s in enumerate(seqs)]

    def trunk(self, contexts: list[EncodedContext]) -> tuple[Tensor, State]:
        device = self.device
        width = max(len(c.ids) for c in contexts)
        ids = torch.full((len(contexts), width), self.renderer.pad_id, dtype=torch.long,
                         device=device)
        for i, c in enumerate(contexts):
            ids[i, :len(c.ids)] = torch.tensor(c.ids, dtype=torch.long, device=device)
        lengths = torch.tensor([len(c.ids) for c in contexts], dtype=torch.long, device=device)
        if not any(c.media for c in contexts):
            return self.backbone(ids, lengths, sel=self.context_sel(len(contexts)))
        if not isinstance(self.backbone, GemmaBackbone):
            raise ValueError(f"{self.meta.get('backbone')} reads text contexts only")
        return self.backbone(ids, lengths, sel=self.context_sel(len(contexts)),
                             soft=soft_inputs(contexts, tuple(ids.shape), device))

    def answer_batched(self, state: State, trunk_hidden: Tensor, ctx: EncodedContext,
                       queries: list[EncodedQuery]) -> list[QueryOutput]:
        """Inference path for one context: :meth:`answer_many` with a single trunk."""
        return self.answer_many([(state, trunk_hidden, ctx)], [(0, q) for q in queries])

    def answer_many(self, trunks: list[tuple[State, Tensor, EncodedContext]],
                    queries: list[tuple[int, EncodedQuery]]) -> list[QueryOutput]:
        """Every query (``(trunk index, query)``; trunks may belong to different requests)
        is a fork row of ONE batch (chunked only by :data:`FORK_TOKENS` padded tokens for
        memory; a batch that hits a CUDA OOM is split in half and retried) on top of its
        own trunk, then the heads run batched per trunk. All ops are batch-invariant
        (:mod:`deg.kernels`), so a query's raw outputs are bitwise identical in any batch
        composition (alone, reordered, inside a bigger schema, next to other requests)."""
        if not queries:
            return []
        device = self.device
        budget = fork_token_budget(device)
        route = [self.fork_route(q.kind) or "" for _, q in queries]
        order = sorted(range(len(queries)), key=lambda i: (route[i], len(queries[i][1].ids)))
        chunks: list[list[int]] = [[]]
        for i in order:
            width = len(queries[i][1].ids)
            if chunks[-1] and (width * (len(chunks[-1]) + 1) > budget
                               or route[chunks[-1][0]] != route[i]):  # one expert per call
                chunks.append([])
            chunks[-1].append(i)
        hidden_of: dict[int, _Picked] = {}
        chunks.reverse()
        while chunks:
            chunk = chunks.pop()
            try:
                self._fork_chunk(trunks, queries, chunk, hidden_of)
            except torch.OutOfMemoryError:
                if len(chunk) == 1:
                    raise
                split = True
            else:
                split = False
            if split:  # same rows in two smaller batches: bitwise the same outputs
                if device.type == "cuda":
                    torch.cuda.empty_cache()
                chunks += [chunk[len(chunk) // 2:], chunk[:len(chunk) // 2]]
        return self._heads_per_trunk(trunks, queries, hidden_of)

    def _fork_chunk(self, trunks: list[tuple[State, Tensor, EncodedContext]],
                    queries: list[tuple[int, EncodedQuery]], chunk: list[int],
                    hidden_of: dict[int, _Picked]) -> None:
        """Fork rows ``chunk`` as one batch -> ``hidden_of[i]``: only the positions the
        heads read (answer + option tokens), so a big schema's fork hidden states are not
        all held until the heads run."""
        device = self.device
        seqs = [queries[i][1].ids for i in chunk]
        width = max(len(x) for x in seqs)
        ids = torch.full((len(seqs), width), self.renderer.pad_id, dtype=torch.long)
        for r, x in enumerate(seqs):
            ids[r, :len(x)] = torch.tensor(x, dtype=torch.long)
        lengths = torch.tensor([len(x) for x in seqs], dtype=torch.long)
        used = sorted({queries[i][0] for i in chunk})
        fs = self.backbone.fork([trunks[t][0] for t in used],
                                [used.index(queries[i][0]) for i in chunk], width)
        sel = self.lora_sel([queries[i][1].kind for i in chunk])
        hid = self.backbone.fork_forward(ids.to(device), lengths.to(device), fs, sel=sel)
        for r, i in enumerate(chunk):
            q = queries[i][1]
            hidden_of[i] = _Picked(hid[r], (q.answer_position, *q.option_positions))
        del hid, fs

    def _heads_per_trunk(self, trunks: list[tuple[State, Tensor, EncodedContext]],
                         queries: list[tuple[int, EncodedQuery]],
                         hidden_of: dict[int, _Picked]) -> list[QueryOutput]:
        out: list[QueryOutput | None] = [None] * len(queries)
        for t, (_, trunk_hidden, ctx) in enumerate(trunks):
            idx = [i for i, (tt, _) in enumerate(queries) if tt == t]
            if not idx:
                continue
            ctx_hidden = trunk_hidden[0, ctx.ctx_start: ctx.ctx_start + len(ctx.offsets)]
            for i, o in zip(idx, self.batched_heads([queries[i][1] for i in idx],
                                                    [hidden_of[i] for i in idx], ctx_hidden)):
                out[i] = o
        return [o for o in out if o is not None]

    def batched_heads(self, queries: list[EncodedQuery], hidden: list[Any],
                      ctx_hidden: Tensor) -> list[QueryOutput]:
        """All heads over all forks at once, fp32, through the batch-invariant GEMM: each
        fork's logits depend only on its own hidden states (and the context)."""
        from deg.kernels import bi_linear

        hd = self.heads

        def lin(m: nn.Linear, x: Tensor) -> Tensor:
            return bi_linear(x, m.weight, m.bias)

        out: list[QueryOutput | None] = [None] * len(queries)
        ans = torch.stack([h[q.answer_position] for q, h in zip(queries, hidden)]).float()
        by_kind: dict[str, list[int]] = {}
        for i, q in enumerate(queries):
            by_kind.setdefault(q.kind, []).append(i)
        for kind, idx in by_kind.items():
            a = ans[idx]
            if kind in ("choice", "score"):
                head = hd.choice if kind == "choice" else hd.score
                owner = [j for j, i in enumerate(idx) for _ in queries[i].option_positions]
                opts = torch.stack([hidden[i][p] for i in idx
                                    for p in queries[i].option_positions]).float()
                qa = lin(head.query, a)
                ko = lin(head.key, opts)
                dots = bi_linear(ko, qa)  # [n_opt, n_fork]
                own = torch.tensor(owner, device=dots.device)
                logits = (dots.gather(1, own[:, None])[:, 0] * head.scale
                          + lin(head.prior, opts)[:, 0])
                n0 = 0
                for i in idx:
                    n = len(queries[i].option_positions)
                    out[i] = QueryOutput(kind, logits=logits[n0:n0 + n])
                    n0 += n
            elif kind == "noul":
                lg = lin(hd.noul, a)[:, 0]
                for j, i in enumerate(idx):
                    out[i] = QueryOutput("noul", logits=torch.stack(
                        [lg[j], torch.zeros_like(lg[j])]))
            elif kind == "count":
                lg = lin(hd.count, a)
                for j, i in enumerate(idx):
                    out[i] = QueryOutput("count", logits=lg[j])
            elif kind == "anchor":
                for i, lg in zip(idx, self._anchor_logits_blocked(
                        a, ctx_hidden, [queries[i].region for i in idx])):
                    out[i] = QueryOutput("anchor", logits=lg)
            else:
                assert kind == "extract"
                sp = hd.span
                T = ctx_hidden.shape[0]
                locs = [(queries[i].anchor, region_bounds(queries[i].region, T)[1])
                        for i in idx]
                if T <= HEAD_ROWS and len(idx) * (T + 1) * 8 <= HEAD_LOGIT_BYTES:
                    ctx = ctx_hidden.float()
                    start = torch.cat([bi_linear(lin(sp.q_start, a), lin(sp.k_start, ctx))
                                       * sp.scale, lin(sp.null_start, a)], -1)
                    end = torch.cat([bi_linear(lin(sp.q_end, a), lin(sp.k_end, ctx))
                                     * sp.scale, lin(sp.null_end, a)], -1)
                    for j, i in enumerate(idx):
                        st, en = start[j], end[j]
                        if locs[j][0] >= 0:
                            bs, be = span_rel_bias(hd.span_bias, *locs[j], 0, T)
                            st = torch.cat([st[:T] + bs, st[T:]])
                            en = torch.cat([en[:T] + be, en[T:]])
                        out[i] = QueryOutput("extract", start=st, end=en)
                    continue
                for i, (st, en) in zip(idx, self._span_logits_blocked(a, ctx_hidden, locs)):
                    out[i] = QueryOutput("extract", start=st, end=en)
        return [o for o in out if o is not None]

    def _anchor_logits_blocked(self, a: Tensor, ctx_hidden: Tensor,
                               regions: list[tuple[int, int]]) -> list[Tensor]:
        """Anchor logits [T] per answer row of ``a`` (off-region tokens = -inf), in answer
        micro-batches of at most :data:`HEAD_LOGIT_BYTES` and :data:`HEAD_ROWS`-row context
        blocks (keys re-projected per block); finished rows move to host memory on CUDA.
        Every element is a fixed-config dot product of its own row and token, so the values
        do not depend on the batch, the block size, or the other rows."""
        from deg.kernels import bi_linear

        an = self.heads.anchor
        T = ctx_hidden.shape[0]

        def lin(m: nn.Linear, x: Tensor) -> Tensor:
            return bi_linear(x, m.weight, m.bias)

        qa = lin(an.query, a)
        rows = max(1, HEAD_LOGIT_BYTES // (4 * max(1, T)))
        res: list[Tensor] = []
        for q0 in range(0, a.shape[0], rows):
            q1 = min(a.shape[0], q0 + rows)
            lg = torch.empty(q1 - q0, T, device=a.device, dtype=torch.float32)
            for r0 in range(0, T, HEAD_ROWS):
                r1 = min(T, r0 + HEAD_ROWS)
                ctx = ctx_hidden[r0:r1].float()
                lg[:, r0:r1] = (bi_linear(qa[q0:q1], lin(an.key, ctx)) * an.scale
                                + lin(an.prior, ctx)[:, 0][None])
                del ctx
            for j in range(q1 - q0):
                lo, hi = region_bounds(regions[q0 + j], T)
                lg[j, :lo] = float("-inf")
                lg[j, hi:] = float("-inf")
            if a.is_cuda:
                lg = lg.cpu()
            res += [lg[j] for j in range(q1 - q0)]
        return res

    def _span_logits_blocked(self, a: Tensor, ctx_hidden: Tensor,
                             locs: list[tuple[int, int]] | None = None
                             ) -> list[tuple[Tensor, Tensor]]:
        """Span start/end logits [T+1] per answer row of ``a`` for a long context or a big
        schema, never holding more than :data:`HEAD_LOGIT_BYTES` of logits or a
        :data:`HEAD_ROWS`-row block of fp32 context on the device: answer rows run in
        micro-batches, context keys are projected per row block (recomputed per
        micro-batch: no [T, HEAD_DIM] key copy), and finished logits move to host memory.
        Every element is the same fixed-config dot product as the one-shot path
        (:func:`deg.kernels.bi_linear` is row- and column-invariant), so the values are
        bitwise those of the one-shot path."""
        from deg.kernels import bi_linear

        sp = self.heads.span
        T = ctx_hidden.shape[0]

        def lin(m: nn.Linear, x: Tensor) -> Tensor:
            return bi_linear(x, m.weight, m.bias)

        qs, qe = lin(sp.q_start, a), lin(sp.q_end, a)
        ns, ne = lin(sp.null_start, a), lin(sp.null_end, a)
        rows = max(1, HEAD_LOGIT_BYTES // (8 * (T + 1)))
        host = a.is_cuda
        res: list[tuple[Tensor, Tensor]] = []
        for q0 in range(0, a.shape[0], rows):
            q1 = min(a.shape[0], q0 + rows)
            start = torch.empty(q1 - q0, T + 1, device=a.device, dtype=torch.float32)
            end = torch.empty_like(start)
            for r0 in range(0, T, HEAD_ROWS):
                r1 = min(T, r0 + HEAD_ROWS)
                ctx = ctx_hidden[r0:r1].float()
                start[:, r0:r1] = bi_linear(qs[q0:q1], lin(sp.k_start, ctx)) * sp.scale
                end[:, r0:r1] = bi_linear(qe[q0:q1], lin(sp.k_end, ctx)) * sp.scale
                del ctx
                for j in range(q1 - q0):  # anchored items: same bias as the one-shot path
                    if locs is not None and locs[q0 + j][0] >= 0:
                        bs, be = span_rel_bias(self.heads.span_bias, *locs[q0 + j], r0, r1)
                        start[j, r0:r1] += bs
                        end[j, r0:r1] += be
            start[:, T:] = ns[q0:q1]
            end[:, T:] = ne[q0:q1]
            if host:
                start, end = start.cpu(), end.cpu()
            res += [(start[j], end[j]) for j in range(q1 - q0)]
        return res

    # ------------------------------------------------------------------ persistence

    def save(self, path: Path, extra: dict[str, Any] | None = None) -> None:
        from safetensors.torch import save_file

        path.mkdir(parents=True, exist_ok=True)
        lm_state = {k: v.detach().to(torch.bfloat16).contiguous()
                    for k, v in self.backbone.lm.state_dict().items()
                    if k != "lm_head.weight" and ".lora_" not in k  # tied to embed_tokens
                    and not _is_moe_key(k)
                    and not k.endswith(PLE_TABLE_KEY)}  # frozen = the base's (Gemma 4)
        save_file(lm_state, str(path / "backbone.safetensors"))
        if self.moe is not None:  # W experts (bf16) + writer LoRA (fp32) beside R
            save_file({k: (v.detach().float() if ".lora_" in k else
                           v.detach().to(torch.bfloat16)).contiguous()
                       for k, v in moe_state(self).items()}, str(path / "moe.safetensors"))
        if self.lora is not None:  # adapters beside the (frozen = base) backbone
            save_file({k: v.detach().float().contiguous() for k, v in lora_state(self).items()},
                      str(path / "adapters.safetensors"))
        save_file({k: v.detach().float().contiguous() for k, v in self.heads.state_dict().items()},
                  str(path / "heads.safetensors"))
        backbone = self.meta.get("backbone", DEFAULT_BACKBONE)
        meta = {**self.meta, **(extra or {}), "lora": self.lora, "moe": self.moe,
                "kinds": list(KINDS), "max_count": MAX_COUNT,
                "array_mode": self.array_mode, "backbone": backbone,
                "base": str(self.meta.get("base", base_snapshot(backbone)))}
        (path / META_FILE).write_text(json.dumps(meta, indent=2))


# Gemma 4's per-layer token table: frozen in deg training, so checkpoints do not store
# it (2.35B params on E2B); it is always the base snapshot's
PLE_TABLE_KEY = "embed_tokens_per_layer.weight"
GENERATE_WITH = ("self", "base", "moe")
# meta["gen_layout"]: how a full fine-tune writes x-generate text with its OWN weights
# (generate_with "self"): "deg" (WRITE/VALUE layout, deg's own LM) or "chat" (the base
# model's chat layout, :func:`deg.infer.base_generate_header`: an instruction-tuned base
# writes well there from step 0, and training keeps it there with LM + KL to the base)
GEN_LAYOUTS = ("deg", "chat")
# ``generate="base+lora:<dir>"``: the hybrid's frozen base copy with a GENERATE LoRA
# (:mod:`deg.gentrain`) merged into it at load; deg's own backbone is untouched
GENERATOR_PREFIX = "base+lora:"
# frozen identifier (pre-rename name): stored in generator checkpoints and checked on load
GENERATOR_FORMAT = "deg-generator-lora-v1"


def parse_generate(spec: str | None) -> tuple[str | None, Path | None]:
    """``(mode, generator dir)`` of a ``generate`` option: ``"self"``, ``"base"`` or
    ``"base+lora:<dir>"`` (mode ``"base"`` with that generator)."""
    if spec is not None and spec.startswith(GENERATOR_PREFIX):
        path = Path(spec[len(GENERATOR_PREFIX):])
        if not (path / "generator.json").is_file():
            raise ValueError(f"no generator.json in {path}")
        return "base", path
    if spec is not None and spec not in GENERATE_WITH:
        raise ValueError(f"generate must be one of {GENERATE_WITH} or "
                         f"{GENERATOR_PREFIX}<dir>: {spec!r}")
    return spec, None


def merge_generator(lm: nn.Module, path: Path | str, backbone: str | None = None
                    ) -> dict[str, Any]:
    """Merge a GENERATE LoRA (``generator.json`` + ``lora.safetensors``, keys
    ``<module>.lora_A`` [r, in] / ``<module>.lora_B`` [out, r] named as in ``lm``) into
    ``lm``'s weights in place: ``W += alpha / rank * B @ A`` in fp32 on the CPU, then cast
    to ``W``'s dtype (deterministic). ``backbone`` given: the generator must be for that
    base (``generator.json`` ``backbone``; older generators are qwen3.5-0.8b), else
    ValueError. Returns the generator config."""
    from safetensors.torch import load_file

    path = Path(path)
    cfg = json.loads((path / "generator.json").read_text())
    if cfg.get("format") != GENERATOR_FORMAT:
        raise ValueError(f"{path}: not a {GENERATOR_FORMAT} generator")
    gen_backbone = cfg.get("backbone", DEFAULT_BACKBONE)
    want = backbone_spec(backbone).name if backbone is not None else gen_backbone
    if want != gen_backbone:
        raise ValueError(f"{path}: a generator for {gen_backbone}, not {want}")
    ad = load_file(str(path / "lora.safetensors"))
    scale = float(cfg["alpha"]) / int(cfg["rank"])
    names = sorted({k.rsplit(".lora_", 1)[0] for k in ad})
    if not names or len(ad) != 2 * len(names):
        raise ValueError(f"{path}: lora.safetensors needs lora_A and lora_B per module")
    for name in names:
        w = lm.get_submodule(name).weight
        delta = ad[f"{name}.lora_B"].float() @ ad[f"{name}.lora_A"].float()
        if delta.shape != w.shape:
            raise ValueError(f"{name}: delta {tuple(delta.shape)} vs {tuple(w.shape)}")
        with torch.no_grad():
            w.copy_((w.detach().float().cpu() + scale * delta).to(w.dtype).to(w.device))
    return cfg


def add_lora(model: DegModel, rank: int, alpha: float | None = None,
             adapters: list[str] | tuple[str, ...] = ("shared",), context: str = "base",
             seed: int = 0) -> list[nn.Parameter]:
    """Add named LoRA adapters (rank ``rank``, scale ``alpha / rank``; default alpha =
    2 rank) to every :data:`LORA_TARGETS` projection of ``model.backbone``; B starts at 0,
    so the adapted model starts as the base. ``adapters``: ``["shared"]`` (one adapter for
    every decide/extract fork) or a subset of ``decide``/``extract`` (by fork kind,
    :data:`ADAPTER_OF_KIND`). ``context``: ``"base"`` (the context is read once by the
    base weights; adapters only on fork tokens) or the adapter that also reads the
    context. Returns the new parameters (the trainable backbone set)."""
    names = list(adapters)
    if names != ["shared"] and not (names and set(names) <= {"decide", "extract"}
                                    and len(set(names)) == len(names)):
        raise ValueError(f"adapters must be ['shared'] or from decide,extract: {names}")
    if context != "base" and context not in names:
        raise ValueError(f"lora context must be 'base' or one of {names}: {context!r}")
    if rank < 1:
        raise ValueError(f"lora rank must be >= 1: {rank}")
    alpha = float(alpha if alpha is not None else 2 * rank)
    gen = torch.Generator().manual_seed(seed)
    n = rank * len(names)
    params = []
    for layer in model.backbone.text.layers:
        for owner_name, targets in model.backbone.fam.lora_targets.items():
            owner = getattr(layer, owner_name, None)
            if owner is None:
                continue
            for t in targets:
                lin = getattr(owner, t, None)
                if lin is None:  # Gemma 4 K/V-sharing layers have no k/v projections
                    continue
                w = lin.weight
                a = torch.randn(n, w.shape[1], generator=gen) / math.sqrt(w.shape[1])
                lin.lora_A = nn.Parameter(a.to(w.device, w.dtype))
                lin.lora_B = nn.Parameter(torch.zeros(w.shape[0], n, device=w.device,
                                                      dtype=w.dtype))
                lin.lora_scale = alpha / rank
                params += [lin.lora_A, lin.lora_B]
    model.lora = {"rank": int(rank), "alpha": alpha, "adapters": names, "context": context}
    model.meta["lora"] = dict(model.lora)
    return params


MOE_FFN = "mlp"  # Qwen3.5's FFN attribute; the W expert is ``mlp_w``
WRITER_TARGETS = {"linear_attn": ("in_proj_qkv", "in_proj_z", "out_proj"),
                  "self_attn": ("q_proj", "k_proj", "v_proj", "o_proj")}


def _is_moe_key(key: str) -> bool:
    """A key of an MoE expert beside the dense FFN: W (``mlp_w``) or a D/E/G expert
    (``mlp_d``, ``mlp_g``)."""
    return re.search(rf"\.{MOE_FFN}_[a-z]\.", key) is not None


def add_moe(model: DegModel, base_lm: nn.Module | None = None, writer_rank: int = 32,
            writer_alpha: float | None = None, seed: int = 0) -> list[nn.Parameter]:
    """Make ``model`` a deg-MoE (see :class:`WriteRoute`): give every decoder layer a W
    FFN expert ``mlp_w`` = a copy of ``base_lm``'s MLP of that layer (None: a copy of the
    layer's own R MLP, a placeholder that :func:`load_model` overwrites from
    ``moe.safetensors``) and, with ``writer_rank`` > 0, a writer LoRA (rank r, scale
    ``alpha / r``, default alpha 2r, B = 0) on the attention / linear-attention projections
    (:data:`WRITER_TARGETS`) that only W calls switch on. R (everything deg had) is not
    touched, so decide/extract outputs are unchanged by construction. Returns the W
    parameters (W experts + writer LoRA): the trainable set of :mod:`deg.moetrain`.
    Qwen3.5 only; refuses a LoRA-layout model (its adapters use the same slots)."""
    import copy

    if model.moe is not None:
        raise ValueError("already a deg-MoE")
    if model.lora is not None:
        raise ValueError("deg-MoE needs a full fine-tune checkpoint, not a LoRA one")
    if model.backbone.fam is not QWEN35:
        raise ValueError(f"deg-MoE supports Qwen3.5 backbones, not {model.backbone.fam.name}")
    if writer_rank < 0:
        raise ValueError(f"writer rank must be >= 0: {writer_rank}")
    layers = model.backbone.text.layers
    if base_lm is not None:
        btext = getattr(base_lm.model, "language_model", base_lm.model)
        if family_of(base_lm) is not QWEN35 or len(btext.layers) != len(layers):
            raise ValueError("base model does not match the deg backbone")
        src = [blayer.mlp for blayer in btext.layers]
    else:
        src = [layer.mlp for layer in layers]
    w = layers[0].mlp.gate_proj.weight
    params: list[nn.Parameter] = []
    for layer, m in zip(layers, src):
        exp = copy.deepcopy(m)
        exp.__dict__.pop("_deg_fused_gate_proj_up_proj", None)  # own fused cache
        exp.to(device=w.device, dtype=w.dtype)
        exp.requires_grad_(False)
        layer.mlp_w = exp
        params += list(exp.parameters())
    alpha = float(writer_alpha if writer_alpha is not None else 2 * max(writer_rank, 1))
    if writer_rank:
        gen = torch.Generator().manual_seed(seed)
        for layer in layers:
            for owner_name, targets in WRITER_TARGETS.items():
                owner = getattr(layer, owner_name, None)
                if owner is None:
                    continue
                for t in targets:
                    lin = getattr(owner, t)
                    pw = lin.weight
                    a = (torch.randn(writer_rank, pw.shape[1], generator=gen)
                         / math.sqrt(pw.shape[1]))
                    lin.lora_A = nn.Parameter(a.to(pw.device, pw.dtype), requires_grad=False)
                    lin.lora_B = nn.Parameter(torch.zeros(pw.shape[0], writer_rank,
                                                          device=pw.device, dtype=pw.dtype),
                                              requires_grad=False)
                    lin.lora_scale = alpha / writer_rank
                    params += [lin.lora_A, lin.lora_B]
    model.moe = {"writer_rank": int(writer_rank), "writer_alpha": alpha,
                 "writer_targets": {k: list(v) for k, v in WRITER_TARGETS.items()}}
    model.meta["moe"] = dict(model.moe)
    model.generate_with = "moe"
    model.graphs.clear()
    return params


TASK_EXPERTS = ("d", "g")  # beside E = the layer's own mlp


def add_task_moe(model: DegModel, experts: tuple[str, ...] | list[str] = TASK_EXPERTS
                 ) -> list[nn.Parameter]:
    """Make ``model`` a D/E/G task MoE (:class:`TaskRoute`): every decoder layer gets FFN
    experts ``mlp_<x>`` for x in ``experts``, copies of the layer's ``mlp`` (on a fresh
    model: the base's FFN, so every expert starts as the base and the model's outputs are
    unchanged), trainable. E (``mlp``) also reads the context. Routing: decide forks -> D,
    extract / count / anchor forks and the context -> E, generate (KNOWN VALUES, header,
    decoded tokens) -> G. Returns the new parameters. Refuses an MoE or LoRA layout."""
    import copy

    if model.moe is not None or model.lora is not None:
        raise ValueError("task MoE needs a dense full fine-tune layout")
    if sorted(experts) != sorted(TASK_EXPERTS):
        raise ValueError(f"task experts must be {TASK_EXPERTS}: {experts}")
    params: list[nn.Parameter] = []
    for layer in model.backbone.text.layers:
        m = getattr(layer, MOE_FFN)
        for x in experts:
            exp = copy.deepcopy(m)
            for key in [k for k in exp.__dict__ if k.startswith("_deg_fused_")]:
                exp.__dict__.pop(key)  # its own fused-GEMM cache
            setattr(layer, f"{MOE_FFN}_{x}", exp)
            params += list(exp.parameters())
    model.moe = {"experts": list(experts), "reader": "e",
                 "routes": dict(TASK_EXPERT_OF_KIND)}
    model.meta["moe"] = dict(model.moe)
    model.generate_with = "moe"
    model.graphs.clear()
    return params


def moe_state(model: DegModel) -> dict[str, Tensor]:
    """The W side of a deg-MoE: W experts + writer LoRA (keys as in the backbone lm)."""
    return {k: v for k, v in model.backbone.lm.state_dict().items()
            if _is_moe_key(k) or ".lora_" in k}


def moe_params(model: DegModel) -> dict[str, int]:
    """Parameter counts: ``total`` in memory (tied embeddings once), ``r_ffn``/``w_ffn``
    (the experts), ``writer_lora``, and ``active`` per token (R or W path; a W token also
    runs the writer LoRA)."""
    seen: set[int] = set()
    total = r = w = lora = 0
    for name, p in model.backbone.lm.named_parameters():
        if id(p) in seen:
            continue
        seen.add(id(p))
        n = p.numel()
        total += n
        if _is_moe_key(name):  # W (R/W MoE) or D + G (task MoE)
            w += n
        elif ".lora_" in name:
            lora += n
        elif f".{MOE_FFN}." in name:
            r += n
    return {"total": total, "r_ffn": r, "w_ffn": w, "writer_lora": lora,
            "active_reader": total - w - lora, "active_writer": total - r}


def task_moe_params(model: DegModel) -> dict[str, int]:
    """D/E/G task MoE parameter counts (tied embeddings once): ``total`` in memory with
    and without the PLE table (``ple``, frozen, host-able), one FFN copy per expert
    (``ffn_e`` = ``ffn_d`` = ``ffn_g``), and ``active`` per token (one expert)."""
    seen: set[int] = set()
    out = {"total": 0, "ple": 0, "ffn_e": 0, "ffn_d": 0, "ffn_g": 0}
    for name, p in model.backbone.lm.named_parameters():
        if id(p) in seen:
            continue
        seen.add(id(p))
        n = p.numel()
        out["total"] += n
        if name.endswith(PLE_TABLE_KEY):
            out["ple"] += n
        for x in ("d", "g"):
            if f".{MOE_FFN}_{x}." in name:
                out[f"ffn_{x}"] += n
        if f".{MOE_FFN}." in name:
            out["ffn_e"] += n
    out["total_without_ple"] = out["total"] - out["ple"]
    out["active_without_ple"] = out["total_without_ple"] - out["ffn_d"] - out["ffn_g"]
    return out


def lora_state(model: DegModel) -> dict[str, Tensor]:
    return {k: v for k, v in model.backbone.lm.state_dict().items() if ".lora_" in k}


def attach_base(model: DegModel, lm: nn.Module | None = None) -> Backbone:
    """Load the frozen base model of ``model``'s backbone (``meta["backbone"]``: Qwen3.5
    or LFM2.5; ``lm`` overrides it; tests) beside ``model`` as ``model.base``. Its token
    embeddings (= the tied LM head) are shared with ``model`` when they are equal (deg
    trains with them frozen), so the copy costs only the body (~1 GB bf16 for the 0.8B,
    ~2.1 GB for LFM2.5-1.2B); otherwise it keeps its own. Never trained: no grad, eval
    mode, not in ``model.modules()`` (not saved, not optimised)."""
    emb = model.backbone.text.embed_tokens.weight
    if lm is None:
        lm = load_base_lm(model.meta.get("backbone"), emb.dtype)
    if family_of(lm) is not model.backbone.fam:
        raise ValueError(f"base {lm.config.model_type} does not match the deg backbone "
                         f"{model.backbone.fam.name}")
    text = getattr(lm.model, "language_model", lm.model)
    own = text.embed_tokens.weight
    if own.shape == emb.shape and torch.equal(own.detach().to("cpu", emb.dtype),
                                              emb.detach().to("cpu")):
        text.embed_tokens.weight = emb  # same Parameter: one copy on the device
        lm.lm_head.weight = emb
        shared = True
    else:
        shared = False
    ple = getattr(model.backbone.text, "embed_tokens_per_layer", None)
    if ple is not None:  # Gemma 4: the frozen PLE table is token identity: one copy
        bple = text.embed_tokens_per_layer
        if bple.weight.shape != ple.weight.shape or not torch.equal(
                bple.weight.detach().to("cpu", ple.weight.dtype), ple.weight.detach().cpu()):
            raise ValueError("base PLE table differs from the deg model's (frozen in "
                             "deg training): not this model's base")
        text.embed_tokens_per_layer = None  # moved below, kept off the device
    lm.to(device=emb.device, dtype=emb.dtype)
    if ple is not None:
        text.embed_tokens_per_layer = ple
    lm.requires_grad_(False)
    lm.eval()
    model.base = Backbone(lm)
    model.base.shared_embeddings = shared  # type: ignore[attr-defined]
    model._gen_view = None
    return model.base


def load_model(path: Path | str | None = None, *, device: str | torch.device = "cuda",
               dtype: torch.dtype = torch.bfloat16, lm: nn.Module | None = None,
               tokenizer: Any = None, generate: str | None = None,
               base_lm: nn.Module | None = None, backbone: str | None = None,
               ple_host: bool | None = None) -> DegModel:
    """:func:`_load_model`; ``ple_host`` (default: env ``DEG_PLE_HOST=1``): a Gemma 4
    model keeps its frozen PLE table in host memory (:meth:`GemmaBackbone.ple_to_host`;
    ~2.3B params off the GPU, ~4.4 GiB on E2B)."""
    model = _load_model(path, device=device, dtype=dtype, lm=lm, tokenizer=tokenizer,
                        generate=generate, base_lm=base_lm, backbone=backbone)
    if ple_host is None:
        ple_host = os.environ.get("DEG_PLE_HOST", "") == "1"
    if ple_host and isinstance(model.backbone, GemmaBackbone):
        model.backbone.ple_to_host()
        if model.base is not None:  # shares the same table module
            model.base.text.embed_tokens_per_layer = model.backbone.text.embed_tokens_per_layer
    return model


def _load_model(path: Path | str | None = None, *, device: str | torch.device = "cuda",
                dtype: torch.dtype = torch.bfloat16, lm: nn.Module | None = None,
                tokenizer: Any = None, generate: str | None = None,
                base_lm: nn.Module | None = None, backbone: str | None = None) -> DegModel:
    """``path=None`` -> fresh heads on the base model; else a saved deg checkpoint.
    ``lm``/``tokenizer`` override the base (tests use a tiny random Qwen3.5 / LFM2).

    ``backbone`` (:data:`BACKBONES`; default: the checkpoint's ``meta["backbone"]``, else
    ``qwen3.5-0.8b``) picks the base model; a checkpoint refuses a different one. With
    ``lm`` given, the backbone is ``lm``'s architecture.

    ``generate`` (default: the checkpoint's ``meta["generate_backbone"]``, else
    ``"self"``): ``"base"`` decodes x-generate fields with the frozen base weights
    (:func:`attach_base`; ``base_lm`` overrides them) while decide/extract use deg's;
    ``"base+lora:<dir>"`` is that hybrid with a GENERATE LoRA merged into the base copy
    (:func:`merge_generator`)."""
    from safetensors.torch import load_file

    generate, generator = parse_generate(generate)
    saved: dict[str, Any] | None = None
    if path is not None:
        path = Path(path)
        saved = read_meta(path)
        # checkpoints from before pluggable backbones are Qwen3.5-0.8B
        ckpt_backbone = saved.get("backbone", DEFAULT_BACKBONE)
        if backbone is not None and backbone_spec(backbone).name != ckpt_backbone:
            raise ValueError(f"checkpoint {path} is a {ckpt_backbone} model, not {backbone}")
        backbone = ckpt_backbone
    spec = spec_of_lm(lm, backbone) if lm is not None else backbone_spec(backbone)
    base = base_snapshot(spec.name) if lm is None or tokenizer is None else None
    if lm is None:
        lm = load_base_lm(spec.name, dtype)
    tokenizer = tokenizer or load_tokenizer(base)
    meta: dict[str, Any] = {"base": str(base) if base else "custom", "backbone": spec.name}
    backbone_mod = Backbone(lm)
    heads = Heads(backbone_mod.hidden)
    if path is not None:
        assert saved is not None
        state = load_file(str(path / "backbone.safetensors"))
        missing, unexpected = lm.load_state_dict(state, strict=False)
        bad = [k for k in missing if k != "lm_head.weight"
               and not k.endswith(PLE_TABLE_KEY)] + list(unexpected)
        if bad:
            raise RuntimeError(f"checkpoint mismatch: {bad[:5]}")
        missing, unexpected = heads.load_state_dict(load_file(str(path / "heads.safetensors")),
                                                    strict=False)
        new_heads = [k for k in missing if k.startswith(("anchor.", "span_bias"))]
        if unexpected or len(new_heads) != len(missing):
            raise RuntimeError(f"heads mismatch: {(list(unexpected) + missing)[:5]}")
        # a checkpoint without anchor heads (run-1) addresses items by index
        meta["array_mode"] = "index" if new_heads else "anchor"
        meta.update(saved)
        meta["backbone"] = spec.name
    else:
        meta["array_mode"] = "anchor"
    lm.to(device=device, dtype=dtype)
    heads.to(device=device, dtype=torch.float32)
    lm.eval()
    heads.eval()
    model = DegModel(backbone_mod, heads, Renderer(tokenizer, spec), meta)
    if meta.get("moe"):
        assert path is not None
        cfg = meta["moe"]
        if "experts" in cfg:
            add_task_moe(model, cfg["experts"])
        else:
            add_moe(model, None, cfg["writer_rank"], cfg["writer_alpha"])
        ws = load_file(str(path / "moe.safetensors"))
        missing, unexpected = lm.load_state_dict(ws, strict=False)
        lost = [k for k in missing if _is_moe_key(k) or ".lora_" in k]
        if unexpected or lost:
            raise RuntimeError(f"moe mismatch: {(list(unexpected) + lost)[:5]}")
        lm.to(device=device, dtype=dtype)
        lm.eval()
        if generate not in (None, "moe") or generator is not None:
            raise ValueError("a deg-MoE checkpoint generates with its W experts (moe)")
        return model
    if meta.get("lora"):
        cfg = meta["lora"]
        add_lora(model, cfg["rank"], cfg["alpha"], cfg["adapters"], cfg["context"])
        assert path is not None
        ad = load_file(str(path / "adapters.safetensors"))
        missing, unexpected = lm.load_state_dict(ad, strict=False)
        if unexpected or any(".lora_" in k for k in missing):
            raise RuntimeError(f"adapters mismatch: {(list(unexpected) + missing)[:5]}")
        lm.to(device=device, dtype=dtype)
        lm.eval()
        if generate == "self" or generator is not None:
            raise ValueError("a LoRA checkpoint always generates with the base (adapter off)")
        return model
    mode = generate or meta.get("generate_backbone", "self")
    if mode == "moe":
        raise ValueError("generate moe needs a deg-MoE checkpoint (deg.moetrain)")
    if mode not in GENERATE_WITH:
        raise ValueError(f"generate must be one of {GENERATE_WITH}: {mode!r}")
    if mode == "base":
        if generator is not None:
            if base_lm is None:
                base_lm = load_base_lm(spec.name, dtype)
            cfg = merge_generator(base_lm, generator, spec.name)
            model.meta["generator"] = {"path": str(generator), "rank": cfg["rank"],
                                       "alpha": cfg["alpha"]}
        attach_base(model, base_lm)
        model.generate_with = "base"
    return model