modeling_microloop.py

"""Small, HF-compatible causal bootstrap model used by the M0 gate.

The diffusion objective and sampler are deliberately separate modules.  This model
provides the shared transformer backbone and a causal forward path so that the
project can validate shape correctness, parameter accounting, and reproducibility
before any expensive data work begins.
"""

from __future__ import annotations

from dataclasses import dataclass
from hashlib import sha256
from typing import Optional

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

from .configuration_microloop import MicroLoopConfig
from .prefix_lm import build_prefix_lm_mask

try:
    from transformers import PreTrainedModel
    from transformers.generation import GenerationMixin
    from transformers.utils import ModelOutput
except ImportError:  # pragma: no cover - only used in a minimal environment.

    class GenerationMixin:  # type: ignore[no-redef]
        pass

    class ModelOutput:  # type: ignore[no-redef]
        pass

    class PreTrainedModel(nn.Module):  # type: ignore[no-redef]
        config_class = MicroLoopConfig
        base_model_prefix = "microloop"

        def __init__(self, config: MicroLoopConfig) -> None:
            super().__init__()
            self.config = config


@dataclass
class MicroLoopCausalLMOutput(ModelOutput):
    """Minimal output object with both attribute and mapping-style access."""

    logits: Optional[Tensor] = None
    loss: Optional[Tensor] = None
    hidden_states: Optional[Tensor] = None
    loop_applications: Optional[int] = None

    def __getitem__(self, key: str):
        return getattr(self, key)


class RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, hidden_states: Tensor) -> Tensor:
        # Explicit computation rather than the fused ``F.rms_norm`` kernel:
        # the fused kernel selects implementations based on process-level state
        # and produces context-dependent numerics inside the training process
        # (the 2026-08-05 provenance incident).  This explicit path is
        # deterministic everywhere; the small speed cost is acceptable here.
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        return (hidden_states / torch.sqrt(variance + self.eps)) * self.weight


class UnweightedRMSNorm(nn.Module):
    """Parameter-free RMS normalization used by the mHC routing projections."""

    def __init__(self, eps: float) -> None:
        super().__init__()
        self.eps = eps

    def forward(self, hidden_states: Tensor) -> Tensor:
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        return hidden_states * torch.rsqrt(variance + self.eps).to(hidden_states.dtype)


class NGramMemory(nn.Module):
    """Small causal suffix-N-gram memory for controlled experiments.

    The lookup mode keeps raw token IDs unchanged, which is deliberately safe for
    digits and mathematical symbols.  Each order/head has its own deterministic
    hash table; the projected value is gated by the current hidden state and the
    value projection is zero-initialized by ``MicroLoopForDiffusionLM``.
    """

    def __init__(self, config: MicroLoopConfig) -> None:
        super().__init__()
        settings = dict(config.ngram_memory)
        self.mode = str(settings.get("mode", "lookup"))
        self.hash_version = str(settings.get("hash_version", "legacy_v1"))
        self.orders = tuple(int(order) for order in settings.get("orders", [2, 3]))
        self.num_hash_heads = int(settings.get("num_hash_heads", 2))
        self.num_buckets = int(settings.get("num_buckets", 16_384))
        self.embedding_dim = int(settings.get("embedding_dim", 16))
        self.composition_scale = float(settings.get("composition_scale", 0.1))
        self.eps = config.rms_norm_eps
        # Fixed, independent bases; no global RNG or checkpoint tensors needed.
        # Arithmetic modulo 2**31-1 keeps int64 products below 2**62.
        self.hash_bases = tuple(
            2 + int.from_bytes(sha256(f"microloop-ngram-v2:{i}".encode()).digest()[:4], "big")
            % (2**31 - 3)
            for i in range(len(self.orders) * self.num_hash_heads)
        )

        if self.mode == "lookup":
            table_count = len(self.orders) * self.num_hash_heads
            self.tables = nn.ModuleList(
                nn.Embedding(self.num_buckets, self.embedding_dim)
                for _ in range(table_count)
            )
            memory_width = table_count * self.embedding_dim
            self.key_proj = nn.Linear(memory_width, config.hidden_size, bias=False)
            self.value_proj = nn.Linear(memory_width, config.hidden_size, bias=False)
        else:
            self.tables = nn.ModuleList()
            self.key_proj = None
            self.value_proj = None

    def _hash_indices(
        self,
        input_ids: Tensor,
        order: int,
        table_index: int,
        document_ids: Tensor | None = None,
    ) -> tuple[Tensor, Tensor]:
        multiplier = 1_000_003 + table_index * 104_729
        coefficients = torch.arange(1, order + 1, device=input_ids.device, dtype=torch.long)
        coefficients = coefficients * multiplier
        if self.hash_version == "polynomial_v2":
            prime = 2**31 - 1
            base = self.hash_bases[table_index]
            padded = F.pad(input_ids, (order - 1, 0), value=0)
            windows = padded.unfold(1, order, 1)
            hashed = torch.zeros_like(input_ids)
            for position in range(order):
                hashed = (hashed * base + windows[..., position] + 1) % prime
            hashed = hashed % self.num_buckets
        elif order in (2, 3) and input_ids.size(1) >= order:
            hashed = input_ids * coefficients[-1]
            for lag in range(1, order):
                shifted = F.pad(input_ids[:, :-lag], (lag, 0), value=0)
                hashed = hashed + shifted * coefficients[-lag - 1]
        else:
            padded = F.pad(input_ids, (order - 1, 0), value=0)
            windows = padded.unfold(1, order, 1)
            hashed = (windows * coefficients).sum(dim=-1)
        if self.hash_version == "legacy_v1":
            hashed = (hashed + (table_index + 1) * 7_919) % self.num_buckets
        positions = torch.arange(input_ids.size(1), device=input_ids.device)
        valid = positions.unsqueeze(0) >= order - 1
        if document_ids is not None:
            if order in (2, 3) and input_ids.size(1) >= order:
                for lag in range(1, order):
                    shifted_documents = F.pad(document_ids[:, :-lag], (lag, 0), value=-1)
                    valid = valid & document_ids.eq(shifted_documents)
            else:
                padded_documents = F.pad(document_ids, (order - 1, 0), value=-1)
                document_windows = padded_documents.unfold(1, order, 1)
                valid = valid & document_windows.eq(document_windows[..., :1]).all(dim=-1)
        return hashed, valid

    def _rms_normalize(self, hidden_states: Tensor) -> Tensor:
        variance = hidden_states.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
        return hidden_states * torch.rsqrt(variance + self.eps).to(hidden_states.dtype)

    def _forward_flat(
        self,
        input_ids: Tensor,
        hidden_states: Tensor,
        token_embeddings: Tensor | None = None,
        document_ids: Tensor | None = None,
    ) -> Tensor:
        if self.mode == "parameter_free":
            if token_embeddings is None:
                raise ValueError("parameter-free ngram memory requires token embeddings")
            embeddings = token_embeddings
            outputs = []
            for order in self.orders:
                padded = F.pad(embeddings, (0, 0, order - 1, 0))
                windows = padded.unfold(1, order, 1)
                outputs.append(windows.mean(dim=-1))
            result = torch.stack(outputs, dim=0).mean(dim=0)
            positions = torch.arange(input_ids.size(1), device=input_ids.device)
            valid = positions.unsqueeze(0) >= max(self.orders) - 1
            if document_ids is not None:
                padded_documents = F.pad(
                    document_ids, (max(self.orders) - 1, 0), value=-1
                )
                document_windows = padded_documents.unfold(1, max(self.orders), 1)
                valid = valid & document_windows.eq(document_windows[..., :1]).all(dim=-1)
            return result * valid.unsqueeze(-1).to(result.dtype) * self.composition_scale

        retrieved = []
        table_index = 0
        for order in self.orders:
            for _ in range(self.num_hash_heads):
                indices, valid = self._hash_indices(
                    input_ids, order, table_index, document_ids=document_ids
                )
                values = self.tables[table_index](indices)
                retrieved.append(values * valid.unsqueeze(-1).to(values.dtype))
                table_index += 1
        memory = torch.cat(retrieved, dim=-1)
        assert self.key_proj is not None and self.value_proj is not None
        key = self.key_proj(memory)
        value = self.value_proj(memory)
        query_norm = self._rms_normalize(hidden_states)
        key_norm = self._rms_normalize(key)
        gate = torch.sigmoid((query_norm * key_norm).sum(dim=-1) / self.key_proj.out_features**0.5)
        return value * gate.unsqueeze(-1)

    def forward(
        self,
        input_ids: Tensor,
        hidden_states: Tensor,
        token_embeddings: Tensor | None = None,
        document_ids: Tensor | None = None,
    ) -> Tensor:
        """Return a residual update with the same shape as ``hidden_states``."""

        if input_ids.dim() != 2:
            raise ValueError("ngram memory input_ids must have shape [batch, sequence]")
        if hidden_states.dim() == 3:
            if hidden_states.shape[:2] != input_ids.shape:
                raise ValueError("ngram memory inputs must have matching batch and sequence")
            return self._forward_flat(
                input_ids, hidden_states, token_embeddings, document_ids
            )
        if hidden_states.dim() == 4:
            batch, sequence, streams, hidden = hidden_states.shape
            if (batch, sequence) != input_ids.shape:
                raise ValueError("ngram memory inputs must have matching batch and sequence")
            flat_ids = input_ids.unsqueeze(1).expand(-1, streams, -1).reshape(-1, sequence)
            flat_hidden = hidden_states.permute(0, 2, 1, 3).reshape(-1, sequence, hidden)
            flat_documents = None
            if document_ids is not None:
                if document_ids.shape != (batch, sequence):
                    raise ValueError("document_ids must match input_ids")
                flat_documents = document_ids.unsqueeze(1).expand(
                    -1, streams, -1
                ).reshape(-1, sequence)
            flat_embeddings = None
            if token_embeddings is not None:
                if token_embeddings.shape != (batch, sequence, hidden):
                    raise ValueError("token_embeddings must match the unstreamed hidden shape")
                flat_embeddings = token_embeddings.unsqueeze(1).expand(
                    -1, streams, -1, -1
                ).reshape(-1, sequence, hidden)
            update = self._forward_flat(
                flat_ids, flat_hidden, flat_embeddings, flat_documents
            )
            return update.view(batch, streams, sequence, hidden).permute(0, 2, 1, 3)
        raise ValueError("ngram memory hidden_states must have shape [batch, sequence, hidden]")


class ManifoldHyperConnection(nn.Module):
    """Dynamic manifold-constrained routing around one residual sublayer.

    The residual state has shape ``[batch, sequence, streams, hidden]``.  The
    routing matrix is projected onto the Birkhoff polytope with Sinkhorn-Knopp
    iterations, while sigmoid-constrained read/write weights collapse the
    streams for the inner sublayer and place its output back onto the streams.
    """

    def __init__(self, config: MicroLoopConfig) -> None:
        super().__init__()
        self.multiplier = config.mhc_multiplier
        self.sinkhorn_iterations = config.mhc_sinkhorn_iterations
        self.eps = config.mhc_eps
        self.input_norm = UnweightedRMSNorm(config.rms_norm_eps)
        mix = (2 + self.multiplier) * self.multiplier
        self.fn = nn.Parameter(torch.empty(mix, self.multiplier * config.hidden_size))
        self.base = nn.Parameter(torch.empty(mix))
        self.scale = nn.Parameter(torch.empty(3))

    def forward(self, hidden_streams: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        streams = self.multiplier
        flat = self.input_norm(hidden_streams.flatten(start_dim=2).float())
        raw = F.linear(flat, self.fn.float())
        pre_raw, post_raw, residual_raw = raw.split(
            [streams, streams, streams * streams], dim=-1
        )
        pre_base, post_base, residual_base = self.base.float().split(
            [streams, streams, streams * streams]
        )
        pre_scale, post_scale, residual_scale = self.scale.float().unbind(0)
        pre = torch.sigmoid(pre_raw * pre_scale + pre_base) + self.eps
        post = 2.0 * torch.sigmoid(post_raw * post_scale + post_base)
        residual_logits = (
            residual_raw.view(*residual_raw.shape[:-1], streams, streams) * residual_scale
            + residual_base.view(streams, streams)
        )
        residual = torch.softmax(residual_logits, dim=-1) + self.eps
        residual = residual / (residual.sum(dim=-2, keepdim=True) + self.eps)
        for _ in range(self.sinkhorn_iterations - 1):
            residual = residual / (residual.sum(dim=-1, keepdim=True) + self.eps)
            residual = residual / (residual.sum(dim=-2, keepdim=True) + self.eps)
        collapsed = (pre.unsqueeze(-1) * hidden_streams.float()).sum(dim=2)
        return (
            post.to(hidden_streams.dtype),
            residual.to(hidden_streams.dtype),
            collapsed.to(hidden_streams.dtype),
        )

    @staticmethod
    def merge(
        hidden_streams: Tensor, sublayer_output: Tensor, post: Tensor, residual: Tensor
    ) -> Tensor:
        mixed = torch.matmul(residual, hidden_streams)
        return mixed + post.unsqueeze(-1) * sublayer_output.unsqueeze(2)


class ManifoldHyperHead(nn.Module):
    """Dynamically collapse the final mHC residual streams back to model width."""

    def __init__(self, config: MicroLoopConfig) -> None:
        super().__init__()
        self.multiplier = config.mhc_multiplier
        self.eps = config.mhc_eps
        self.input_norm = UnweightedRMSNorm(config.rms_norm_eps)
        self.fn = nn.Parameter(
            torch.empty(self.multiplier, self.multiplier * config.hidden_size)
        )
        self.base = nn.Parameter(torch.empty(self.multiplier))
        self.scale = nn.Parameter(torch.empty(1))

    def forward(self, hidden_streams: Tensor) -> Tensor:
        flat = self.input_norm(hidden_streams.flatten(start_dim=2).float())
        raw = F.linear(flat, self.fn.float())
        pre = torch.sigmoid(raw * self.scale.float() + self.base.float()) + self.eps
        return (pre.unsqueeze(-1) * hidden_streams.float()).sum(dim=2).to(hidden_streams.dtype)


def _rotate_half(x: Tensor) -> Tensor:
    x_even = x[..., ::2]
    x_odd = x[..., 1::2]
    return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)


def _rope_tables(max_position: int, head_dim: int, theta: float) -> tuple[Tensor, Tensor]:
    """Build interleaved rotary tables once, in float32 for stable reuse."""

    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    positions = torch.arange(max_position, device=inv_freq.device, dtype=torch.float32)
    angles = positions.unsqueeze(-1) * inv_freq
    angles = torch.stack((angles, angles), dim=-1).flatten(-2)
    return angles.cos(), angles.sin()


def _apply_rope(q: Tensor, k: Tensor, cos: Tensor, sin: Tensor) -> tuple[Tensor, Tensor]:
    """Apply cached interleaved rotary embeddings to query and key tensors."""

    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin


class GroupedQueryAttention(nn.Module):
    """Grouped-query self-attention with explicit Q/K/V projections.

    Attention uses an explicit scaled-dot-product implementation (matmul +
    softmax) rather than ``F.scaled_dot_product_attention`` because the fused
    kernels select implementations based on process-level state and produce
    deterministic but context-dependent results: evaluations inside the
    training process then disagree with evaluations of the same saved
    checkpoint in a fresh process (the 2026-08-05 provenance incident).  The
    explicit math path is deterministic everywhere at the cost of a small
    amount of speed, which is acceptable at this model size.
    """

    def __init__(self, config: MicroLoopConfig) -> None:
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.head_dim = config.head_dimension
        self.num_groups = self.num_heads // self.num_key_value_heads
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(
            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
        self.output_gate = (
            nn.Linear(config.hidden_size, self.num_heads, bias=False)
            if config.attention_output_gate
            else None
        )
        self.output_gate_activation = config.attention_output_gate_activation
        self.rope_theta = config.rope_theta
        self.attention_implementation = config.attention_implementation
        self.qk_norm_position = config.qk_norm_position
        if config.qk_norm == "per_head":
            # One affine scale is shared by all Q heads and one by all K heads;
            # normalization itself is applied independently over each head.
            self.q_norm: RMSNorm | None = RMSNorm(self.head_dim, config.rms_norm_eps)
            self.k_norm: RMSNorm | None = RMSNorm(self.head_dim, config.rms_norm_eps)
        else:
            self.q_norm = None
            self.k_norm = None

    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        rope_embeddings: tuple[Tensor, Tensor] | None = None,
        first_value_states: Tensor | None = None,
        value_residual_scales: Tensor | None = None,
    ) -> Tensor:
        batch, sequence, _ = hidden_states.shape
        q = self.q_proj(hidden_states).view(batch, sequence, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(
            batch, sequence, self.num_key_value_heads, self.head_dim
        )
        v = self.v_proj(hidden_states).view(
            batch, sequence, self.num_key_value_heads, self.head_dim
        )
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        if first_value_states is not None:
            if value_residual_scales is None:
                raise ValueError("value residual states require value_residual_scales")
            if first_value_states.shape != v.shape:
                raise ValueError("first_value_states must match the current value-state shape")
            v = value_residual_scales[0] * v + value_residual_scales[1] * first_value_states
        if position_ids is None:
            position_ids = torch.arange(sequence, device=hidden_states.device).expand(batch, -1)
        if position_ids.shape != (batch, sequence):
            raise ValueError(
                f"position_ids must have shape [batch, sequence], got {tuple(position_ids.shape)}"
            )
        if rope_embeddings is None:
            table_cos, table_sin = _rope_tables(sequence, self.head_dim, self.rope_theta)
            cos = table_cos[position_ids].unsqueeze(1).to(q.dtype)
            sin = table_sin[position_ids].unsqueeze(1).to(q.dtype)
        else:
            cos, sin = rope_embeddings
        if self.q_norm is not None and self.qk_norm_position == "pre_rope":
            # Normalize in the unrotated head basis.  RMS normalization with
            # learned per-coordinate scales does not commute with RoPE; doing
            # it first preserves RoPE's relative-position shift equivariance.
            q = self.q_norm(q)
            assert self.k_norm is not None
            k = self.k_norm(k)
        q, k = _apply_rope(q, k, cos, sin)
        if self.q_norm is not None and self.qk_norm_position == "post_rope":
            q = self.q_norm(q)
            assert self.k_norm is not None
            k = self.k_norm(k)

        if self.num_heads % self.num_key_value_heads != 0:
            raise ValueError(
                f"num_heads {self.num_heads} must divide num_key_value_heads "
                f"{self.num_key_value_heads}"
            )
        if attention_mask is None:
            visible: Tensor | None = None
            is_causal = True
        else:
            # A 2-D mask uses the conventional HF meaning: one means visible.
            if attention_mask.shape == (batch, sequence):
                causal = torch.tril(
                    torch.ones(sequence, sequence, device=q.device, dtype=torch.bool)
                )
                visible = causal.unsqueeze(0).unsqueeze(0) & attention_mask.bool().unsqueeze(
                    1
                ).unsqueeze(2)
                is_causal = False
            elif attention_mask.shape == (batch, sequence, sequence):
                visible = attention_mask.bool().unsqueeze(1)
                is_causal = False
            else:
                raise ValueError(
                    "attention_mask must have shape [batch, sequence] or "
                    "[batch, sequence, sequence], "
                    f"got {tuple(attention_mask.shape)}"
                )
        # GQA: repeat the KV heads so every query head has its own K/V.
        repeat = self.num_heads // self.num_key_value_heads
        if repeat > 1:
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)
        if self.attention_implementation == "sdpa":
            attended = F.scaled_dot_product_attention(
                q, k, v, attn_mask=visible, dropout_p=0.0, is_causal=is_causal
            )
        else:
            scores = torch.matmul(q, k.transpose(-2, -1)) / float(self.head_dim) ** 0.5
            if is_causal:
                seq_ids = torch.arange(sequence, device=scores.device)
                visible = seq_ids.unsqueeze(0) <= seq_ids.unsqueeze(1)
                visible = visible.expand(batch, self.num_heads, sequence, sequence)
            if visible is not None:
                scores = scores.masked_fill(~visible, float("-inf"))
            probs = torch.softmax(scores, dim=-1)
            attended = torch.matmul(probs, v)
        if self.output_gate is not None:
            gate_logits = self.output_gate(hidden_states)
            gate = (
                torch.sigmoid(gate_logits)
                if self.output_gate_activation == "sigmoid"
                else F.silu(gate_logits)
            )
            gate = gate.transpose(1, 2).unsqueeze(-1)
            attended = attended * gate
        attended = attended.transpose(1, 2).contiguous().view(batch, sequence, -1)
        return self.o_proj(attended)


class LowRankFFNProjection(nn.Module):
    """CoLA-inspired B(activation(A(x))); outer SwiGLU remains in the block."""

    def __init__(self, in_features: int, out_features: int, rank: int, activation: str):
        super().__init__()
        self.reduce = nn.Linear(in_features, rank, bias=False)
        self.expand = nn.Linear(rank, out_features, bias=False)
        self.activation = nn.SiLU() if activation == "silu" else nn.Identity()

    def forward(self, hidden_states: Tensor) -> Tensor:
        return self.expand(self.activation(self.reduce(hidden_states)))


class MicroLoopBlock(nn.Module):
    """Pre-norm transformer block with SwiGLU feed-forward network."""

    def __init__(self, config: MicroLoopConfig, layer_index: int) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = GroupedQueryAttention(config)
        self.ffn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        if config.ffn_rank is None:
            self.ffn_gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
            self.ffn_up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
            self.ffn_down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        else:
            self.ffn_gate = LowRankFFNProjection(
                config.hidden_size, config.intermediate_size,
                config.ffn_rank, config.ffn_factor_activation,
            )
            self.ffn_up = LowRankFFNProjection(
                config.hidden_size, config.intermediate_size,
                config.ffn_rank, config.ffn_factor_activation,
            )
            self.ffn_down = LowRankFFNProjection(
                config.intermediate_size, config.hidden_size,
                config.ffn_rank, config.ffn_factor_activation,
            )
        self.use_mhc = config.mhc_multiplier > 1
        if self.use_mhc:
            self.attn_mhc = ManifoldHyperConnection(config)
            self.ffn_mhc = ManifoldHyperConnection(config)
        self.use_attn_residuals = config.attn_res_block_size is not None
        self.value_residual_scales = (
            nn.Parameter(torch.tensor([1.0, 0.0]))
            if config.value_residual.get("enabled", False) and layer_index > 0
            else None
        )
        if self.use_attn_residuals:
            self.attn_res_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
            self.ffn_res_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
            self.attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False)
            self.ffn_res_proj = nn.Linear(config.hidden_size, 1, bias=False)
        clamp = dict(config.swiglu_clamp)
        self.swiglu_clamp_enabled = bool(clamp.get("enabled", False))
        self.swiglu_linear_min = float(clamp.get("linear_min", -10.0))
        self.swiglu_linear_max = float(clamp.get("linear_max", 10.0))
        self.swiglu_gate_max = float(clamp.get("gate_max", 10.0))

    def forward(
        self,
        hidden_states: Tensor,
        attention_mask: Tensor | None = None,
        position_ids: Tensor | None = None,
        rope_embeddings: tuple[Tensor, Tensor] | None = None,
        block_residuals: list[Tensor] | None = None,
        first_value_states: Tensor | None = None,
    ) -> Tensor:
        if self.use_mhc:
            attn_post, attn_residual, attn_input = self.attn_mhc(hidden_states)
            attn_output = self.attn(
                self.attn_norm(attn_input),
                attention_mask,
                position_ids,
                rope_embeddings,
                first_value_states=first_value_states,
                value_residual_scales=self.value_residual_scales,
            )
            hidden_states = self.attn_mhc.merge(
                hidden_states, attn_output, attn_post, attn_residual
            )
            ffn_post, ffn_residual, ffn_input = self.ffn_mhc(hidden_states)
            normalized = self.ffn_norm(ffn_input)
            gate_linear = self.ffn_gate(normalized)
            up_linear = self.ffn_up(normalized)
            if self.swiglu_clamp_enabled:
                gate_linear = gate_linear.clamp(self.swiglu_linear_min, self.swiglu_linear_max)
                up_linear = up_linear.clamp(self.swiglu_linear_min, self.swiglu_linear_max)
            gate = F.silu(gate_linear)
            if self.swiglu_clamp_enabled:
                gate = gate.clamp(max=self.swiglu_gate_max)
            ffn_output = self.ffn_down(gate * up_linear)
            return self.ffn_mhc.merge(hidden_states, ffn_output, ffn_post, ffn_residual)
        if self.use_attn_residuals and block_residuals:
            residual_stack = torch.stack(block_residuals, dim=-2)
            scores = torch.cat(
                [self.attn_res_proj(self.attn_res_norm(state)) for state in block_residuals], dim=-1
            )
            hidden_states = hidden_states + (
                torch.softmax(scores, dim=-1).unsqueeze(-1) * residual_stack
            ).sum(dim=-2)
        hidden_states = hidden_states + self.attn(
            self.attn_norm(hidden_states),
            attention_mask,
            position_ids,
            rope_embeddings,
            first_value_states=first_value_states,
            value_residual_scales=self.value_residual_scales,
        )
        if self.use_attn_residuals and block_residuals:
            residual_stack = torch.stack(block_residuals, dim=-2)
            scores = torch.cat(
                [self.ffn_res_proj(self.ffn_res_norm(state)) for state in block_residuals], dim=-1
            )
            hidden_states = hidden_states + (
                torch.softmax(scores, dim=-1).unsqueeze(-1) * residual_stack
            ).sum(dim=-2)
        ffn_input = self.ffn_norm(hidden_states)
        gate_linear = self.ffn_gate(ffn_input)
        up_linear = self.ffn_up(ffn_input)
        if self.swiglu_clamp_enabled:
            gate_linear = gate_linear.clamp(self.swiglu_linear_min, self.swiglu_linear_max)
            up_linear = up_linear.clamp(self.swiglu_linear_min, self.swiglu_linear_max)
        gate = F.silu(gate_linear)
        if self.swiglu_clamp_enabled:
            gate = gate.clamp(max=self.swiglu_gate_max)
        ffn_output = self.ffn_down(gate * up_linear)
        return hidden_states + ffn_output


class MicroLoopPreTrainedModel(PreTrainedModel):
    config_class = MicroLoopConfig
    base_model_prefix = "microloop"


class MicroLoopForDiffusionLM(MicroLoopPreTrainedModel, GenerationMixin):
    """Backbone plus tied output head for causal bootstrap and diffusion training."""

    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}

    def __init__(self, config: MicroLoopConfig) -> None:
        config.validate()
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [MicroLoopBlock(config, index) for index in range(config.num_hidden_layers)]
        )
        looping = config.looping
        self.loop_gates = (
            nn.Parameter(
                torch.zeros(
                    max(
                        0,
                        int(
                            looping.get(
                                "max_loop_count", looping.get("maximum_serving_loops", 3)
                            )
                        )
                        - 1,
                    ),
                    config.hidden_size,
                )
            )
            if looping.get("mode", "layer") == "block" and looping.get("gated", False)
            else None
        )
        self.ngram_memory = (
            NGramMemory(config) if config.ngram_memory.get("enabled", False) else None
        )
        self.ngram_memory_layer = (
            int(config.ngram_memory["insertion_layer"])
            if self.ngram_memory is not None
            else None
        )
        digit_settings = config.digit_position_embedding
        self.digit_position_embedding = (
            nn.Embedding(
                int(digit_settings.get("max_positions", 128)), config.hidden_size
            )
            if digit_settings.get("enabled", False)
            else None
        )
        self.digit_token_ids = frozenset(
            int(token_id) for token_id in digit_settings.get("digit_token_ids", [])
        )
        self.register_buffer(
            "digit_token_ids_tensor",
            torch.tensor(sorted(self.digit_token_ids), dtype=torch.long),
            persistent=False,
        )
        self.mhc_head = ManifoldHyperHead(config) if config.mhc_multiplier > 1 else None
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.mtp_proj = (
            nn.Linear(config.hidden_size, config.hidden_size, bias=False)
            if config.mtp_enabled
            else None
        )
        rope_cos, rope_sin = _rope_tables(
            config.max_position_embeddings, config.head_dimension, config.rope_theta
        )
        # PERSISTENT buffers on purpose: transformers 5.x ``from_pretrained``
        # re-initializes non-persistent buffers that are missing from the
        # checkpoint (``_initialize_missing_keys`` -> ``initialize_weights``),
        # overwriting the rotary tables with garbage.  That made every fresh-
        # process evaluation of saved checkpoints compute with corrupted rope
        # tables (the historical "external collapse" at ~0.03 accuracy was this
        # artifact).  Persisting the tables makes the saved checkpoint carry
        # exactly the tables used in training.
        self.register_buffer("rope_cos", rope_cos, persistent=True)
        self.register_buffer("rope_sin", rope_sin, persistent=True)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight
        # Canonical HF pattern: post_init() installs all_tied_weights_keys and
        # dispatches _initialize_weights per module.
        self.post_init()
        if self.digit_position_embedding is not None:
            nn.init.zeros_(self.digit_position_embedding.weight)
        if self.ngram_memory is not None and self.ngram_memory.value_proj is not None:
            nn.init.zeros_(self.ngram_memory.value_proj.weight)

    def _reset_rope_buffers(self) -> None:
        """Recompute the rotary tables from the config (they are deterministic)."""

        rope_cos, rope_sin = _rope_tables(
            self.config.max_position_embeddings,
            self.config.head_dimension,
            self.config.rope_theta,
        )
        self.rope_cos.copy_(rope_cos)
        self.rope_sin.copy_(rope_sin)

    def _initialize_weights(self, module: nn.Module, is_custom_code: bool = False) -> None:
        # The caller controls the RNG through seed_everything; this method performs
        # no hidden reseeding and is therefore reproducible by construction.
        if getattr(module, "_is_hf_initialized", False):
            return
        if module is self:
            # The main module owns the rotary tables.  transformers 5.x
            # re-initializes buffers that are missing from a loaded checkpoint
            # ("_initialize_missing_keys"), which zeroes/garbles the tables;
            # restore the deterministic canonical tables here instead.
            self._reset_rope_buffers()
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if getattr(module, "bias", None) is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, ManifoldHyperConnection):
            nn.init.normal_(module.fn, mean=0.0, std=0.02)
            nn.init.zeros_(module.base)
            nn.init.constant_(module.scale, self.config.mhc_init_scale)
        elif isinstance(module, ManifoldHyperHead):
            nn.init.normal_(module.fn, mean=0.0, std=0.02)
            nn.init.zeros_(module.base)
            nn.init.constant_(module.scale, self.config.mhc_init_scale)
        elif isinstance(module, RMSNorm):
            nn.init.ones_(module.weight)

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embed_tokens

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.embed_tokens = value
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight

    def digit_position_ids(
        self,
        input_ids: Tensor,
        document_ids: Tensor | None = None,
        prior_run_length: Tensor | None = None,
        prior_active: Tensor | None = None,
    ) -> tuple[Tensor, Tensor, Tensor]:
        """Return zero-based contiguous digit positions and streaming state.

        Positions reset after non-digit tokens and document boundaries.  The two
        state tensors contain the run length and whether the final token is a
        digit, allowing the cache to continue a number across blocks.
        """

        batch, sequence = input_ids.shape
        device = input_ids.device
        enabled = self.digit_position_embedding is not None
        if not enabled:
            zeros = torch.zeros((batch, sequence), dtype=torch.long, device=device)
            return zeros, zeros[:, -1], zeros[:, -1].bool()
        is_digit = torch.isin(input_ids, self.digit_token_ids_tensor)
        positions = torch.zeros_like(input_ids, dtype=torch.long)
        run = (
            prior_run_length.to(device=device, dtype=torch.long)
            if prior_run_length is not None
            else torch.zeros(batch, dtype=torch.long, device=device)
        )
        active = (
            prior_active.to(device=device, dtype=torch.bool)
            if prior_active is not None
            else torch.zeros(batch, dtype=torch.bool, device=device)
        )
        previous_digit = F.pad(is_digit[:, :-1], (1, 0), value=False)
        same_document = torch.ones_like(is_digit)
        if document_ids is not None:
            previous_document = F.pad(document_ids[:, :-1], (1, 0), value=-1)
            same_document = document_ids.eq(previous_document)
            same_document[:, 0] = False
        continuation = is_digit & previous_digit & same_document
        if prior_active is not None:
            continuation[:, 0] = is_digit[:, 0] & active
        starts = is_digit & ~continuation
        max_positions = self.digit_position_embedding.num_embeddings
        base = sequence + max_positions + 1
        markers = torch.where(starts, torch.arange(sequence, device=device) + base, 0)
        if prior_active is not None:
            prior_start = base - run
            markers[:, 0] = torch.where(
                is_digit[:, 0] & active, prior_start, markers[:, 0]
            )
        last_start = torch.cummax(markers, dim=1).values - base
        positions = (torch.arange(sequence, device=device) - last_start).clamp_min(0)
        positions = positions.expand(batch, -1).clamp_max(max_positions - 1)
        positions = torch.where(is_digit, positions, torch.zeros_like(positions))
        run = torch.where(is_digit, positions + 1, 0)[:, -1]
        active = is_digit[:, -1]
        return positions, run, active

    def add_digit_position_embeddings(
        self,
        input_ids: Tensor,
        token_embeddings: Tensor,
        document_ids: Tensor | None = None,
        prior_run_length: Tensor | None = None,
        prior_active: Tensor | None = None,
    ) -> tuple[Tensor, Tensor, Tensor]:
        positions, run, active = self.digit_position_ids(
            input_ids, document_ids=document_ids, prior_run_length=prior_run_length
            , prior_active=prior_active
        )
        if self.digit_position_embedding is not None:
            is_digit = torch.isin(input_ids, self.digit_token_ids_tensor)
            token_embeddings = token_embeddings + self.digit_position_embedding(
                positions
            ) * is_digit.unsqueeze(-1)
        return token_embeddings, run, active

    def forward(
        self,
        input_ids: Tensor | None = None,
        inputs_embeds: Tensor | None = None,
        attention_mask: Tensor | None = None,
        document_ids: Tensor | None = None,
        prefix_lengths: Tensor | None = None,
        position_ids: Tensor | None = None,
        labels: Tensor | None = None,
        logit_mask: Tensor | None = None,
        mtp_loss_weight: float = 0.0,
        use_cut_cross_entropy: bool = False,
        loop_count: int = 1,
        output_hidden_states: bool = False,
        **_: object,
    ) -> MicroLoopCausalLMOutput:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError("exactly one of input_ids or inputs_embeds must be provided")
        if loop_count < 1:
            raise ValueError("loop_count must be at least one")
        if use_cut_cross_entropy and labels is None:
            raise ValueError("use_cut_cross_entropy requires labels")
        if input_ids is not None:
            if input_ids.dim() != 2:
                raise ValueError(
                    f"input_ids must have shape [batch, sequence], got {input_ids.dim()}-D"
                )
            batch, sequence = input_ids.shape
        else:
            if inputs_embeds is None or inputs_embeds.dim() != 3:
                raise ValueError(
                    "inputs_embeds must have shape [batch, sequence, hidden], got "
                    f"{None if inputs_embeds is None else inputs_embeds.dim()}-D"
                )
            batch, sequence, _ = inputs_embeds.shape
        if self.ngram_memory is not None and input_ids is None:
            raise ValueError("ngram memory requires input_ids for deterministic suffix lookup")
        device = (input_ids if input_ids is not None else inputs_embeds).device
        if document_ids is not None:
            if document_ids.shape != (batch, sequence):
                raise ValueError("document_ids must have the same shape as the input")
            if attention_mask is not None and attention_mask.dim() != 2:
                raise ValueError(
                    "document_ids cannot be combined with a precomputed attention mask"
                )
            valid = (
                attention_mask.bool()
                if attention_mask is not None
                else torch.ones((batch, sequence), dtype=torch.bool, device=device)
            )
            if prefix_lengths is None:
                causal = torch.tril(torch.ones(sequence, sequence, dtype=torch.bool, device=device))
                attention_mask = (
                    causal.unsqueeze(0)
                    & document_ids.unsqueeze(2).eq(document_ids.unsqueeze(1))
                    & valid.unsqueeze(1)
                    & valid.unsqueeze(2)
                )
            else:
                attention_mask = build_prefix_lm_mask(
                    prefix_lengths.to(device),
                    sequence,
                    valid=valid,
                    document_ids=document_ids,
                )
        elif prefix_lengths is not None:
            if attention_mask is not None and attention_mask.dim() != 2:
                raise ValueError("prefix_lengths requires a 2-D padding mask")
            valid = (
                attention_mask.bool()
                if attention_mask is not None
                else torch.ones((batch, sequence), dtype=torch.bool, device=device)
            )
            attention_mask = build_prefix_lm_mask(
                prefix_lengths.to(device), sequence, valid=valid
            )
        if position_ids is None:
            position_ids = torch.arange(sequence, device=device).expand(batch, -1)
        rope_dtype = (
            torch.get_autocast_dtype("cuda")
            if device.type == "cuda" and torch.is_autocast_enabled("cuda")
            else self.embed_tokens.weight.dtype
        )
        rope_embeddings = (
            self.rope_cos[position_ids].unsqueeze(1).to(rope_dtype),
            self.rope_sin[position_ids].unsqueeze(1).to(rope_dtype),
        )
        token_embeddings = self.embed_tokens(input_ids) if input_ids is not None else inputs_embeds
        if input_ids is not None:
            token_embeddings, _, _ = self.add_digit_position_embeddings(
                input_ids, token_embeddings, document_ids=document_ids
            )
        hidden_states = token_embeddings
        if self.config.mhc_multiplier > 1:
            hidden_states = hidden_states.unsqueeze(2).expand(
                -1, -1, self.config.mhc_multiplier, -1
            ).contiguous()
        first_value_states = None
        if self.config.value_residual.get("enabled", False):
            first_attention = self.layers[0].attn
            first_values = first_attention.v_proj(self.layers[0].attn_norm(hidden_states))
            first_value_states = first_values.view(
                batch,
                sequence,
                first_attention.num_key_value_heads,
                first_attention.head_dim,
            ).transpose(1, 2)
        loop_layers = list(self.config.looping.get("layers", [4, 5, 6]))
        loop_mode = str(self.config.looping.get("mode", "layer"))
        if (
            loop_mode == "block"
            and self.config.looping.get("gated", False)
            and loop_count
            > int(
                self.config.looping.get(
                    "max_loop_count", self.config.looping.get("maximum_serving_loops", 3)
                )
            )
        ):
            raise ValueError("loop_count exceeds configured max_loop_count")
        block_residuals: list[Tensor] = []
        block_size = self.config.attn_res_block_size
        total_applications = 0
        if loop_mode == "block":
            valid_layers = sorted(
                {layer for layer in loop_layers if 1 <= int(layer) <= len(self.layers)}
            )
            if valid_layers:
                loop_set = set(valid_layers)
                layer_order = [
                    *range(1, valid_layers[0]),
                    *valid_layers * loop_count,
                    *range(valid_layers[-1] + 1, len(self.layers) + 1),
                ]
            else:
                loop_set = set()
                layer_order = list(range(1, len(self.layers) + 1))
        else:
            loop_set = {int(layer) for layer in loop_layers}
            layer_order = [
                layer_number
                for layer_number in range(1, len(self.layers) + 1)
                for _ in range(loop_count if layer_number in loop_set else 1)
            ]
        block_start = min(loop_set) if loop_set else None
        block_end = max(loop_set) if loop_set else None
        block_size_layers = len(loop_set)
        block_pass = 0
        block_input: Tensor | None = None
        for application_index, layer_number in enumerate(layer_order):
            layer = self.layers[layer_number - 1]
            if (
                loop_mode == "block"
                and block_start is not None
                and application_index >= block_start - 1
                and (application_index - (block_start - 1)) % block_size_layers == 0
            ):
                block_input = hidden_states
            if block_size is not None and (layer_number - 1) % block_size == 0:
                block_residuals.append(hidden_states)
            prior_block_residuals = block_residuals[:-1] if block_size is not None else None
            hidden_states = layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                rope_embeddings=rope_embeddings,
                block_residuals=prior_block_residuals,
                first_value_states=first_value_states if layer_number > 1 else None,
            )
            if self.ngram_memory is not None and layer_number == self.ngram_memory_layer:
                assert input_ids is not None
                hidden_states = hidden_states + self.ngram_memory(
                    input_ids,
                    hidden_states,
                    token_embeddings=token_embeddings,
                    document_ids=document_ids,
                )
            total_applications += 1
            if (
                loop_mode == "block"
                and block_end is not None
                and layer_number == block_end
                and block_start is not None
                and (application_index - (block_start - 1) + 1) % block_size_layers == 0
            ):
                block_pass += 1
                if block_pass < loop_count:
                    assert block_input is not None
                    hidden_states = self._apply_loop_gate(
                        block_input, hidden_states, block_pass - 1
                    )
        if self.mhc_head is not None:
            hidden_states = self.mhc_head(hidden_states)
        hidden_states = self.norm(hidden_states)
        if use_cut_cross_entropy:
            if logit_mask is not None:
                raise ValueError("use_cut_cross_entropy cannot be combined with logit_mask")
            logits = None
        elif logit_mask is not None:
            # Diffusion training: only the masked positions carry loss, so the
            # output head runs on those rows instead of the full sequence.
            if labels is not None:
                raise ValueError("logit_mask cannot be combined with labels")
            if logit_mask.shape != (batch, sequence):
                raise ValueError("logit_mask must have the same shape as the input")
            flat = hidden_states.reshape(-1, self.config.hidden_size)
            logits = self.lm_head(flat[logit_mask.reshape(-1)])
        else:
            logits = self.lm_head(hidden_states)
        loss = None
        if labels is not None:
            if labels.shape != (batch, sequence):
                raise ValueError("labels must have the same shape as the input")
            if labels.size(1) < 2:
                raise ValueError("causal training requires sequences with at least two tokens")
            # Causal next-token prediction: position t predicts the label at t + 1.
            if use_cut_cross_entropy:
                try:
                    from cut_cross_entropy import linear_cross_entropy
                except ImportError as exc:  # pragma: no cover - optional dependency
                    raise RuntimeError(
                        "Cut Cross Entropy is optional; install the fused-loss project extra"
                    ) from exc
                cce_embeddings = hidden_states[:, :-1, :]
                if torch.is_autocast_enabled("cuda"):
                    cce_embeddings = cce_embeddings.to(torch.get_autocast_dtype("cuda"))
                if cce_embeddings.dtype not in (torch.float16, torch.bfloat16):
                    raise ValueError("Cut Cross Entropy requires CUDA fp16 or bf16 autocast")
                loss = linear_cross_entropy(
                    cce_embeddings,
                    self.lm_head.weight,
                    labels[:, 1:],
                    ignore_index=-100,
                )
            else:
                assert logits is not None
                loss = F.cross_entropy(
                    logits[:, :-1, :].reshape(-1, logits.size(-1)),
                    labels[:, 1:].reshape(-1),
                    ignore_index=-100,
                )
            if mtp_loss_weight:
                if self.mtp_proj is None:
                    raise ValueError("mtp_loss_weight requires mtp_enabled=true")
                if mtp_loss_weight < 0:
                    raise ValueError("mtp_loss_weight must be non-negative")
                mtp_logits = self.lm_head(self.mtp_proj(hidden_states[:, :-2, :]))
                mtp_loss = F.cross_entropy(
                    mtp_logits.reshape(-1, mtp_logits.size(-1)),
                    labels[:, 2:].reshape(-1),
                    ignore_index=-100,
                )
                loss = loss + float(mtp_loss_weight) * mtp_loss
        return MicroLoopCausalLMOutput(
            logits=logits,
            loss=loss,
            hidden_states=hidden_states if output_hidden_states else None,
            loop_applications=total_applications,
        )

    def _apply_loop_gate(
        self, previous: Tensor, candidate: Tensor, pass_index: int
    ) -> Tensor:
        """Blend a repeated shared block into the previous pass."""

        if self.loop_gates is None:
            return candidate
        if pass_index < 0 or pass_index >= self.loop_gates.size(0):
            raise ValueError("loop_count exceeds configured max_loop_count")
        gate = torch.tanh(self.loop_gates[pass_index]).to(candidate.dtype)
        return previous + gate.view(1, 1, -1) * (candidate - previous)

    @torch.no_grad()
    def generate_greedy(
        self, input_ids: Tensor, max_new_tokens: int, eos_token_id: int | None = None
    ) -> Tensor:
        """Small causal smoke decoder; diffusion sampling belongs in ``sampler.py``."""

        return self.generate_causal(
            input_ids, max_new_tokens=max_new_tokens, eos_token_id=eos_token_id, do_sample=False
        )

    @torch.no_grad()
    def generate_causal(
        self,
        input_ids: Tensor,
        *,
        max_new_tokens: int,
        eos_token_id: int | None = None,
        do_sample: bool = False,
        temperature: float = 1.0,
        top_k: int | None = None,
    ) -> Tensor:
        """Generate a batched causal continuation for M2 validation and serving.

        This deliberately recomputes the context on every step. KV caching is a
        later optimization; keeping this reference path simple makes M2 output
        semantics straightforward to test.
        """

        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative")
        if temperature <= 0:
            raise ValueError("temperature must be positive")
        if top_k is not None and top_k <= 0:
            raise ValueError("top_k must be positive when supplied")
        generated = input_ids
        finished = torch.zeros(input_ids.size(0), dtype=torch.bool, device=input_ids.device)
        for _ in range(max_new_tokens):
            next_logits = self(generated).logits[:, -1, :]
            if do_sample:
                next_logits = next_logits / temperature
                if top_k is not None and top_k < next_logits.size(-1):
                    threshold = torch.topk(next_logits, top_k, dim=-1).values[:, -1:]
                    next_logits = next_logits.masked_fill(next_logits < threshold, float("-inf"))
                next_token = torch.multinomial(torch.softmax(next_logits, dim=-1), 1)
            else:
                next_token = next_logits.argmax(dim=-1, keepdim=True)
            if eos_token_id is not None:
                next_token = torch.where(
                    finished.unsqueeze(1),
                    torch.full_like(next_token, eos_token_id),
                    next_token,
                )
                finished |= next_token.squeeze(1).eq(eos_token_id)
            generated = torch.cat((generated, next_token), dim=1)
            if eos_token_id is not None and bool(finished.all()):
                break
        return generated