model.py

import math
import importlib.util
import os
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    PreTrainedModel,
    PretrainedConfig,
)
from transformers.modeling_outputs import CausalLMOutput, CausalLMOutputWithPast


# argonne4.5: the SDPA OOM guard below must work on CUDA *and* Intel XPU, where
# `torch.cuda.OutOfMemoryError` is not defined at all. torch.OutOfMemoryError is
# the backend-agnostic form (torch >= 2.5); fall back for older CUDA builds.
_OOM_ERRORS = tuple(e for e in (
    getattr(torch, "OutOfMemoryError", None),
    getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None),
) if isinstance(e, type)) or (RuntimeError,)

flash_attn_func = None
_flash_attn_available = False
if importlib.util.find_spec("flash_attn") is not None:
    try:
        from flash_attn.flash_attn_interface import flash_attn_func
        _flash_attn_available = True
    except ImportError:
        _flash_attn_available = False

# One-time startup log of which attention kernel is in effect, because the
# sliding-window local attention is only implemented on the flash-attn path
# and silently degrades to full attention on the SDPA/math fallbacks.
_attention_path_logged = False


class ArgonneConfig(PretrainedConfig):
    """Configuration for the Argonne v3.0 family of models."""

    model_type = "argonne2"

    def __init__(
        self,
        vocab_size: int = 151936,
        hidden_size: int = 3072,
        num_hidden_layers: int = 24,
        num_attention_heads: int = 12,
        num_key_value_heads: Optional[int] = 4,
        intermediate_size: Optional[int] = 8192,
        max_position_embeddings: int = 1024,
        attention_dropout: float = 0.0,
        hidden_dropout: float = 0.0,
        rms_norm_eps: float = 1e-6,
        rope_theta: float = 1000000.0,
        sliding_window: Optional[int] = None,
        use_flash_attention: bool = True,
        use_gradient_checkpointing: bool = False,
        qk_norm: bool = True,
        v_norm: bool = True,
        sandwich_norm: bool = True,
        z_loss_weight: float = 0.0,
        mtp_horizon: int = 1,
        mtp_loss_weight: float = 0.0,
        interleaved_local_attention: bool = True,
        local_attention_window: Optional[int] = 256,
        # --- argonne4.5 additions (all default-OFF: an argonne4.0 config reproduces 4.0 exactly) ---
        attn_pattern: Optional[str] = None,
        sliding_window_size: int = 2048,
        nope_global: bool = False,
        attn_gate: bool = False,
        mlp_type: str = "swiglu",
        mtp_module_layers: int = 0,
        doc_mask: bool = False,
        logit_softcap: float = 15.0,
        loss_chunk_size: int = 0,
        tie_word_embeddings: bool = True,
        attention_bias: bool = False,
        mlp_bias: bool = False,
        pad_token_id: Optional[int] = None,
        bos_token_id: Optional[int] = None,
        eos_token_id: Optional[int] = None,
        **kwargs,
    ) -> None:
        pad_token_id = pad_token_id if pad_token_id is not None else kwargs.pop("pad_token_id", None)
        bos_token_id = bos_token_id if bos_token_id is not None else kwargs.pop("bos_token_id", None)
        eos_token_id = eos_token_id if eos_token_id is not None else kwargs.pop("eos_token_id", None)

        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            **kwargs,
        )
        # Backwards compatibility with Argonne 1.x naming.
        if "n_layer" in kwargs:
            num_hidden_layers = kwargs["n_layer"]
        if "n_head" in kwargs:
            num_attention_heads = kwargs["n_head"]
        if "n_embd" in kwargs:
            hidden_size = kwargs["n_embd"]
        if "block_size" in kwargs:
            max_position_embeddings = kwargs["block_size"]

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = (
            num_key_value_heads if num_key_value_heads is not None else num_attention_heads // 2
        )
        if self.num_key_value_heads < 1:
            self.num_key_value_heads = 1
        if num_attention_heads % self.num_key_value_heads != 0:
            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")

        if intermediate_size is None:
            width = int(8 * hidden_size / 3)
            self.intermediate_size = ((width + 255) // 256) * 256
        else:
            self.intermediate_size = intermediate_size

        self.max_position_embeddings = max_position_embeddings
        self.attention_dropout = attention_dropout
        self.hidden_dropout = hidden_dropout
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.sliding_window = sliding_window
        self.use_flash_attention = use_flash_attention
        self.use_gradient_checkpointing = use_gradient_checkpointing
        self.qk_norm = qk_norm
        self.v_norm = v_norm
        self.sandwich_norm = sandwich_norm
        self.z_loss_weight = z_loss_weight
        self.mtp_horizon = max(1, int(mtp_horizon))
        self.mtp_loss_weight = float(mtp_loss_weight)
        self.interleaved_local_attention = bool(interleaved_local_attention)
        self.local_attention_window = (
            int(local_attention_window) if local_attention_window is not None and int(local_attention_window) > 0 else None
        )
        # argonne4.5 attention layout. `attn_pattern` is a cycled string of 'L' (local/sliding)
        # and 'G' (global/full) -- Muse Glimmer's layout is "LLLG" (3 sliding@2048: 1 full).
        # When set it OVERRIDES interleaved_local_attention/local_attention_window. `nope_global`
        # strips RoPE from the global layers (Muse sets layer_rope_theta=0 on exactly those),
        # so the layers that see the whole sequence carry no frequency basis to extrapolate wrong.
        self.attn_pattern = attn_pattern.upper() if attn_pattern else None
        if self.attn_pattern is not None and set(self.attn_pattern) - {"L", "G"}:
            raise ValueError(f"attn_pattern must contain only 'L'/'G', got {attn_pattern!r}")
        self.sliding_window_size = int(sliding_window_size)
        self.nope_global = bool(nope_global)
        self.attn_gate = bool(attn_gate)
        if mlp_type not in ("swiglu", "relu2"):
            raise ValueError(f"mlp_type must be 'swiglu' or 'relu2', got {mlp_type!r}")
        self.mlp_type = mlp_type
        self.mtp_module_layers = max(0, int(mtp_module_layers))
        self.doc_mask = bool(doc_mask)
        self.logit_softcap = float(logit_softcap)
        # >0 enables chunked cross-entropy (frees the full (batch*seq, vocab) fp32 logit transient
        # that caps batch at long context). 0 = off (default; unchanged full-logit path).
        self.loss_chunk_size = int(loss_chunk_size)
        self.tie_word_embeddings = tie_word_embeddings
        self.attention_bias = attention_bias
        self.mlp_bias = mlp_bias

        if self.pad_token_id is None and self.eos_token_id is not None:
            self.pad_token_id = self.eos_token_id

        # Backwards compatibility aliases
        self.n_embd = self.hidden_size
        self.n_layer = self.num_hidden_layers
        self.n_head = self.num_attention_heads
        self.block_size = self.max_position_embeddings


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        orig_dtype = x.dtype
        x = x.to(torch.float32)
        # Clamp values to prevent overflow in pow(2)
        x = torch.clamp(x, min=-65504.0, max=65504.0)
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return (self.weight * x.to(orig_dtype))


class RotaryEmbedding(nn.Module):
    def __init__(
        self,
        dim: int,
        max_position_embeddings: int = 2048,
        base: float = 10000.0,
        device: Optional[torch.device] = None,
    ) -> None:
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base

        # Always compute inv_freq on CPU to avoid meta-device issues during from_pretrained
        # (low_cpu_mem_usage=True initializes model on meta device, which would give uninitialized values)
        inv_freq = 1.0 / (
            self.base
            ** (torch.arange(0, dim, 2, dtype=torch.float32, device="cpu") / dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._set_cos_sin_cache(max_position_embeddings, torch.device("cpu"), torch.float32)

    def _set_cos_sin_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
        self.max_seq_len_cached = seq_len
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        # Always store cache in float32 to avoid precision loss when model dtype is bfloat16
        # (large freqs lose precision in bfloat16, corrupting cos/sin values)
        self.register_buffer("cos_cached", emb.cos().to(torch.float32), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(torch.float32), persistent=False)

    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len, x.device, x.dtype)
        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype, device=x.device),
            self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device),
        )

    def init_buffers(self) -> None:
        """Re-initialize inv_freq, cos_cached, sin_cached. Call after from_pretrained
        to fix meta-device corruption (low_cpu_mem_usage=True destroys buffer values)."""
        device = self.inv_freq.device
        if device.type == "meta":
            device = torch.device("cpu")
        inv_freq = 1.0 / (
            self.base
            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device="cpu") / self.dim)
        )
        self.inv_freq = inv_freq.to(device)
        self._set_cos_sin_cache(
            self.max_position_embeddings,
            device,
            torch.float32,
        )


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


def apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    position_ids: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    if position_ids is None:
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
    else:
        cos = cos[position_ids].unsqueeze(1)
        sin = sin[position_ids].unsqueeze(1)

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


class GroupedQueryAttention(nn.Module):
    def __init__(self, config: ArgonneConfig) -> None:
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_groups = self.num_heads // self.num_kv_heads
        self.sliding_window = config.sliding_window
        self.qk_norm = config.qk_norm
        self.v_norm = config.v_norm

        self.q_proj = nn.Linear(
            self.hidden_size,
            self.num_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.k_proj = nn.Linear(
            self.hidden_size,
            self.num_kv_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.v_proj = nn.Linear(
            self.hidden_size,
            self.num_kv_heads * self.head_dim,
            bias=config.attention_bias,
        )
        self.o_proj = nn.Linear(
            self.num_heads * self.head_dim,
            self.hidden_size,
            bias=config.attention_bias,
        )
        self.o_proj._is_residual = True

        # argonne4.5: gated attention (Muse Glimmer's `self_attn.gate_proj`). A per-(head, position)
        # sigmoid gate on the attention output before o_proj. Cheap, and reported to suppress
        # attention sinks / stabilize high-LR training -- the regime [[argonne-next-arch-search]]
        # already put us in ("qk_norm is ESSENTIAL at this LR").
        self.attn_gate = getattr(config, "attn_gate", False)
        if self.attn_gate:
            self.gate_proj = nn.Linear(
                self.hidden_size,
                self.num_heads * self.head_dim,
                bias=config.attention_bias,
            )
        # Set per-layer by ArgonneModel when attn_pattern/nope_global are in play. False => this
        # layer is NoPE (no rotary applied at all), which is what Muse does on its global layers.
        self.use_rope = True

        self.attention_dropout = config.attention_dropout
        self.use_flash_attention = config.use_flash_attention
        if self.qk_norm:
            self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        if self.v_norm:
            self.v_norm_layer = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

    def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
        if self.num_key_value_groups == 1:
            return x
        bsz, num_kv, seqlen, head_dim = x.shape
        x = x[:,:, None,:,:].expand(bsz, num_kv, self.num_key_value_groups, seqlen, head_dim)
        return x.reshape(bsz, num_kv * self.num_key_value_groups, seqlen, head_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
    ):
        bsz, seqlen, _ = hidden_states.shape

        query = self.q_proj(hidden_states)
        key = self.k_proj(hidden_states)
        value = self.v_proj(hidden_states)

        query = query.view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
        key = key.view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)
        value = value.view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)
        if self.qk_norm:
            query = self.q_norm(query)
            key = self.k_norm(key)
        if self.v_norm:
            value = self.v_norm_layer(value)

        if self.use_rope and position_embeddings is not None:
            cos, sin = position_embeddings
            query, key = apply_rotary_pos_emb(query, key, cos, sin)

        # ---- KV cache (inference only; training keeps past_kv=None/use_cache=False) ----
        # Cache pre-repeat (num_kv_heads), post-rotary K/V. The SDPA/math path
        # ignores the sliding window (full causal), so storing all keys and having
        # the new query attend to everything reproduces the no-cache result.
        is_decode = past_kv is not None
        if is_decode:
            past_k, past_v = past_kv
            key = torch.cat([past_k, key], dim=2)
            value = torch.cat([past_v, value], dim=2)
        new_kv = (key, value) if use_cache else None
        kv_len = key.shape[2]

        # Additive mask used only when decoding with a cache: query at absolute
        # position (kv_len - seqlen + i) attends to keys j <= that position.
        decode_mask = None
        if is_decode:
            q_pos = torch.arange(kv_len - seqlen, kv_len, device=hidden_states.device)
            k_pos = torch.arange(kv_len, device=hidden_states.device)
            allowed = k_pos[None,:] <= q_pos[:, None]
            decode_mask = torch.zeros(seqlen, kv_len, dtype=query.dtype, device=hidden_states.device)
            decode_mask = decode_mask.masked_fill(~allowed, -65504.0)[None, None]

        key = self._repeat_kv(key)
        value = self._repeat_kv(value)

        # Argonne-3.5 speedup (hardware accel, training-neutral): qk/v-norm return fp32
        # (fp32 norm-weight * bf16 activation type-promotes), which trips the dtype gate below
        # and forces the hand-rolled O(seq^2) math attention. Under bf16 autocast the matmuls
        # already run in bf16, so casting q/k/v to the autocast dtype here is numerically
        # equivalent (measured rel diff 3.7e-3 < bf16 roundoff 7.8e-3) and lets SDPA pick a fused
        # FLASH/cuDNN kernel: ~2x lower attention memory + faster. Training-under-autocast only;
        # inference/decode (self.training False) is untouched.
        if self.training and torch.is_autocast_enabled() and query.dtype not in (torch.float16, torch.bfloat16):
            _ac_dtype = torch.get_autocast_dtype(query.device.type) if hasattr(torch, "get_autocast_dtype") else torch.get_autocast_gpu_dtype()
            query = query.to(_ac_dtype)
            key = key.to(_ac_dtype)
            value = value.to(_ac_dtype)

        use_flash_attn_2 = (
            not is_decode
            and _flash_attn_available
            and self.use_flash_attention
            and attention_mask is None
            and query.dtype in (torch.float16, torch.bfloat16)
            and self.head_dim % 4 == 0
        )
        use_scaled_dot = (
            hasattr(F, "scaled_dot_product_attention")
            and self.use_flash_attention
            and query.dtype in (torch.float16, torch.bfloat16)
            and self.head_dim % 4 == 0
        )

        attn_output = None
        if use_flash_attn_2:
            try:
                flash_dropout = self.attention_dropout if self.training else 0.0
                window = (
                    (self.sliding_window, self.sliding_window)
                    if self.sliding_window is not None
                    else (-1, -1)
                )
                q = query.transpose(1, 2).contiguous()
                k = key.transpose(1, 2).contiguous()
                v = value.transpose(1, 2).contiguous()
                attn_output = flash_attn_func(
                    q,
                    k,
                    v,
                    dropout_p=flash_dropout,
                    softmax_scale=None,
                    causal=True,
                    window_size=window,
                ).transpose(1, 2)
            except RuntimeError:
                attn_output = None

        if attn_output is None and use_scaled_dot:
            try:
                if decode_mask is not None:
                    attn_output = F.scaled_dot_product_attention(
                        query, key, value, attn_mask=decode_mask,
                        dropout_p=0.0, is_causal=False,
                    )
                elif attention_mask is None and self.sliding_window is not None and not is_decode:
                    # argonne4.5 FIX (Finding A): flash-attn-2 is the ONLY path above that honors
                    # `window_size`, and this env ships flash-attn-4 without the 2.x
                    # `flash_attn_interface` module -- so `_flash_attn_available` is False and every
                    # Argonne pretrain since 3.0 silently ran FULL attention on every layer while
                    # pretrain.py advertised LOCAL_ATTENTION_WINDOW=256. experiments/model.py got
                    # this fix during the arch sweep; production never did. Apply the window
                    # explicitly here so the configured layout is the layout that trains.
                    # Token i attends to keys in [i-w, i].
                    qi = torch.arange(seqlen, device=hidden_states.device)
                    ki = torch.arange(kv_len, device=hidden_states.device)
                    win_mask = (ki[None,:] <= qi[:, None]) & (
                        ki[None,:] >= qi[:, None] - self.sliding_window)
                    attn_output = F.scaled_dot_product_attention(
                        query,
                        key,
                        value,
                        attn_mask=win_mask[None, None],
                        dropout_p=self.attention_dropout if self.training else 0.0,
                        is_causal=False,
                    )
                elif attention_mask is None:
                    # Prefill / training: fast causal path.
                    attn_output = F.scaled_dot_product_attention(
                        query,
                        key,
                        value,
                        attn_mask=None,
                        dropout_p=self.attention_dropout if self.training else 0.0,
                        is_causal=True,
                    )
                else:
                    # With attention_mask: need to pass it explicitly (slower but correct)
                    # attention_mask should be 4D: (bsz, 1, seq, seq) or broadcastable
                    attn_output = F.scaled_dot_product_attention(
                        query,
                        key,
                        value,
                        attn_mask=attention_mask,
                        dropout_p=self.attention_dropout if self.training else 0.0,
                        is_causal=False,  # Mask already includes causal component
                    )
            except _OOM_ERRORS:
                # Do NOT degrade to math attention on OOM. torch.OutOfMemoryError subclasses
                # RuntimeError, so the old blanket `except RuntimeError` turned a recoverable SDPA
                # OOM into the seq^2 math path -- which needs strictly MORE memory and then dies
                # for real. At phase-C ctx 65,536 that fallback asked for a single 48 GiB tensor
                # and killed the slice. Surface the original OOM instead; the caller can retry
                # smaller. The fallback below stays for genuine kernel-unavailability.
                raise
            except RuntimeError:
                # Fallback to math attention when kernels are unavailable
                attn_output = None

        if attn_output is None:
            scores = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(self.head_dim)
            mask_value = -65504.0  # Large negative instead of -inf
            if decode_mask is not None:
                scores = scores + decode_mask
            else:
                # Apply causal mask - use large negative instead of -inf for stability
                causal_mask = torch.triu(
                    torch.ones(seqlen, kv_len, dtype=torch.bool, device=hidden_states.device),
                    diagonal=1,
                )
                scores = scores.masked_fill(causal_mask, mask_value)
                # argonne4.5: keep the math fallback consistent with the SDPA path above --
                # it must apply the sliding window too, or the two paths train different models.
                if self.sliding_window is not None:
                    qi = torch.arange(seqlen, device=hidden_states.device)
                    ki = torch.arange(kv_len, device=hidden_states.device)
                    too_old = ki[None,:] < qi[:, None] - self.sliding_window
                    scores = scores.masked_fill(too_old, mask_value)
                if attention_mask is not None:
                    # Bool masks (argonne4.5 document masking) are True == attend; the legacy
                    # float mask is additive.
                    if attention_mask.dtype == torch.bool:
                        scores = scores.masked_fill(~attention_mask, mask_value)
                    else:
                        scores = scores + attention_mask
            attn_weights = torch.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype)
            attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
            attn_output = torch.matmul(attn_weights, value)

        attn_output = (
            attn_output.transpose(1, 2)
            .contiguous()
            .view(bsz, seqlen, self.num_heads * self.head_dim)
        )
        if self.attn_gate:
            attn_output = attn_output * torch.sigmoid(self.gate_proj(hidden_states))
        out = self.o_proj(attn_output)
        if use_cache:
            return out, new_kv
        return out


class SwiGLUMLP(nn.Module):
    def __init__(self, config: ArgonneConfig) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(
            config.hidden_size,
            config.intermediate_size,
            bias=config.mlp_bias,
        )
        self.up_proj = nn.Linear(
            config.hidden_size,
            config.intermediate_size,
            bias=config.mlp_bias,
        )
        self.down_proj = nn.Linear(
            config.intermediate_size,
            config.hidden_size,
            bias=config.mlp_bias,
        )
        self.down_proj._is_residual = True
        self.dropout = nn.Dropout(config.hidden_dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Clamp intermediate values to prevent overflow
        gate = self.gate_proj(x)
        gate = torch.clamp(gate, min=-65504.0, max=65504.0)
        up = self.up_proj(x)
        up = torch.clamp(up, min=-65504.0, max=65504.0)
        return self.dropout(self.down_proj(F.silu(gate) * up))


class ReLU2MLP(nn.Module):
    """Two-matrix ReLU-squared FFN (Nemotron 3.5 Lightning's `mlp_hidden_act: relu2`).

    Drops SwiGLU's `gate_proj` entirely: up -> ReLU(x)^2 -> down. Two matrices instead of three,
    so at equal intermediate width it is 1/3 cheaper in params and FLOPs, and at equal params the
    FFN can be ~50% wider. This is what makes Nemotron's 128 experts/layer affordable. The Argonne
    arch sweep tested SwiGLU *ratios* (2.0x worse, 2.75x optimal, 3.5x neutral) but never the
    activation family, so this axis is unmeasured here.
    """

    def __init__(self, config: ArgonneConfig) -> None:
        super().__init__()
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)
        self.down_proj._is_residual = True
        self.dropout = nn.Dropout(config.hidden_dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.up_proj(x)
        h = torch.clamp(h, min=-65504.0, max=65504.0)
        h = F.relu(h).pow(2)
        # ReLU^2 squares its input, so a large activation overflows bf16 far more easily than
        # SwiGLU's bounded silu gate does. Re-clamp before the down projection.
        h = torch.clamp(h, min=-65504.0, max=65504.0)
        return self.dropout(self.down_proj(h))


def build_mlp(config: ArgonneConfig) -> nn.Module:
    return ReLU2MLP(config) if getattr(config, "mlp_type", "swiglu") == "relu2" else SwiGLUMLP(config)


class Block(nn.Module):
    """Transformer block with GQA attention and SwiGLU feed-forward."""

    def __init__(self, config: ArgonneConfig, layer_idx: int = 0) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.attn = GroupedQueryAttention(config)
        self.input_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.sandwich_norm = config.sandwich_norm
        if self.sandwich_norm:
            self.attn_out_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
            self.mlp_out_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = build_mlp(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor] = None,
        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
    ):
        residual = hidden_states
        hidden_states = self.input_norm(hidden_states)
        attn_out = self.attn(
            hidden_states, position_embeddings, attention_mask,
            past_kv=past_kv, use_cache=use_cache,
        )
        new_kv = None
        if use_cache:
            hidden_states, new_kv = attn_out
        else:
            hidden_states = attn_out
        if self.sandwich_norm:
            hidden_states = self.attn_out_norm(hidden_states)
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_norm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        if self.sandwich_norm:
            hidden_states = self.mlp_out_norm(hidden_states)
        hidden_states = residual + hidden_states
        if use_cache:
            return hidden_states, new_kv

        return hidden_states


class MTPModule(nn.Module):
    """One multi-token-prediction depth, DeepSeek-V3 / Nemotron-3.5-Lightning style.

    This is NOT what `config.mtp_horizon` does. That path re-uses the SAME lm_head on the SAME
    hidden state with shifted labels, so the auxiliary objective competes with t+1 for the trunk's
    capacity -- which is exactly what exp045 measured (horizon 2, weight 0.3, 450 steps: 3.8513 ->
    3.9346, +0.083 LOSS). A real MTP module instead owns its own parameters:

        h'_i = eh_proj([ enorm(emb(t_{i+k})) ; hnorm(h_i) ])   ->  Block  ->  final_norm -> lm_head

    so the trunk's next-token objective is untouched and the extra signal is additive. It also
    yields a speculative-decoding drafter for free, which matters here because the deployable
    reasoning recipe is self-consistency at K=8-32: a 2-3x decode speedup is 2-3x the samples per
    GPU-hour on every downstream experiment.
    """

    def __init__(self, config: ArgonneConfig, layer_idx: int = 0) -> None:
        super().__init__()
        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
        self.block = Block(config, layer_idx)
        self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(self, hidden_states, next_token_embeds, position_embeddings, attention_mask=None):
        merged = torch.cat(
            [self.enorm(next_token_embeds), self.hnorm(hidden_states)], dim=-1
        )
        h = self.eh_proj(merged)
        return self.final_norm(self.block(h, position_embeddings, attention_mask))


class ArgonneModel(PreTrainedModel):
    config_class = ArgonneConfig
    _no_split_modules = ["Block"]
    _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}

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

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList([Block(config, idx) for idx in range(config.num_hidden_layers)])
        # argonne4.5 layout: an explicit 'L'/'G' pattern (e.g. Muse Glimmer's "LLLG" = 3
        # sliding@2048: 1 full) takes precedence over the legacy odd/even interleave. On 'G'
        # layers `nope_global` removes RoPE outright -- Muse sets layer_rope_theta=0 on exactly
        # those layers, so the only layers that see the whole sequence carry no frequency basis
        # to extrapolate incorrectly past the training block.
        pattern = getattr(config, "attn_pattern", None)
        if pattern:
            for idx, block in enumerate(self.blocks):
                is_local = pattern[idx % len(pattern)] == "L"
                block.attn.sliding_window = config.sliding_window_size if is_local else None
                block.attn.use_rope = is_local or not config.nope_global
        elif config.interleaved_local_attention and config.local_attention_window is not None:
            for idx, block in enumerate(self.blocks):
                block.attn.sliding_window = config.local_attention_window if (idx % 2 == 1) else None
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = RotaryEmbedding(
            config.hidden_size // config.num_attention_heads,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta,
        )
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        # Real MTP heads (config.mtp_module_layers > 0). Depth k predicts token t+1+k from the
        # previous depth's state plus the embedding of the token that state predicts. They share
        # lm_head with the trunk (DeepSeek-V3 does the same) so the cost is the blocks only.
        self.mtp_modules = nn.ModuleList(
            [
                MTPModule(config, config.num_hidden_layers + k)
                for k in range(getattr(config, "mtp_module_layers", 0))
            ]
        )

        if config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight

        self.gradient_checkpointing = config.use_gradient_checkpointing
        # Selective activation checkpointing (ported from argonne4.0). Runtime-only --
        # NOT part of config or state_dict, so checkpoint/resume compatibility with every
        # existing argonne3.5 checkpoint is unaffected. checkpoint_stride S:
        #   S == 1: checkpoint ALL layers (default; identical to prior behavior).
        #   S >= 2: checkpoint every layer EXCEPT store (un-checkpoint) every Sth layer
        #            -> store ceil(n_layers/S), recompute the rest. Smaller S stores MORE
        #            (more HBM, less recompute, faster); too-small S OOMs.
        # Numerically identical either way: it only chooses store-vs-recompute.
        self.checkpoint_stride = 1
        self._nan_loss_count = 0
        self.post_init()

        global _attention_path_logged
        rank = os.environ.get("RANK") or os.environ.get("SLURM_PROCID") or "0"
        if not _attention_path_logged and rank == "0":
            _attention_path_logged = True
            # Report what the model ACTUALLY does, by reading the per-block state that was just
            # assigned above -- never by re-deriving it from config. The old version hardcoded
            # "full attention ... IGNORED on this path" whenever flash-attn was missing, which
            # became a lie the moment the argonne4.5 SDPA fix (Finding A) taught the SDPA path to
            # honor the window: a4.5 trained 12 of 24 layers at window 256 while this line
            # insisted attention was full. A banner that can disagree with the model is worse
            # than no banner.
            windows = sorted({b.attn.sliding_window for b in self.blocks
                              if b.attn.sliding_window is not None})
            n_local = sum(1 for b in self.blocks if b.attn.sliding_window is not None)
            if _flash_attn_available and config.use_flash_attention:
                backend = "flash-attn-2"
            else:
                backend = "SDPA/math ({})".format(
                    "flash-attn unavailable" if config.use_flash_attention
                    else "use_flash_attention=False"
                )
            if n_local == 0:
                layout = "full attention on all {} layers".format(len(self.blocks))
            else:
                sizes = ",".join(str(w) for w in windows)
                layout = (
                    f"sliding window {sizes} ACTIVE on {n_local}/{len(self.blocks)} layers"
                    f" (pattern={config.attn_pattern or 'odd-layer legacy'})"
                )
            print(f"[ArgonneModel] attention path: {backend}; {layout}", flush=True)

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

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

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

    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
        self.lm_head = new_embeddings
        if isinstance(new_embeddings, nn.Linear):
            self.config.vocab_size = new_embeddings.out_features

    def tie_weights(self, **kwargs) -> None:
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.embed_tokens.weight

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        """Wrap PreTrainedModel.from_pretrained to self-heal after loading.

        from_pretrained may materialize the model via the meta device
        (low_cpu_mem_usage), which fills the rotary embedding's non-persistent
        buffers with uninitialized memory and can drop the
        embed_tokens/lm_head tie. Callers previously had to remember to call
        init_buffers() and tie_weights() by hand; do it here so every loader
        gets a working model.
        """
        model = super().from_pretrained(*args, **kwargs)
        for module in model.modules():
            if isinstance(module, RotaryEmbedding):
                module.init_buffers()
        model.tie_weights()
        return model

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            std = self.config.hidden_size ** -0.5
            if hasattr(module, "_is_residual"):
                std = (2 * self.config.num_hidden_layers) ** -0.5
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.hidden_size ** -0.5)

    def set_gradient_checkpointing(self, enabled: bool = True) -> None:
        self.gradient_checkpointing = enabled

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None) -> None:
        self.set_gradient_checkpointing(True)

    def gradient_checkpointing_disable(self) -> None:
        self.set_gradient_checkpointing(False)

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[list] = None,
        use_cache: bool = False,
        document_ids: Optional[torch.LongTensor] = None,
        **kwargs,  # Accept extra args from newer transformers (e.g., num_items_in_batch)
    ) -> CausalLMOutput:
        _, seq_length = input_ids.shape

        device = self.embed_tokens.weight.device
        if input_ids.device != device:
            input_ids = input_ids.to(device)
        hidden_states = self.embed_tokens(input_ids)

        # The training path does not use attention masks.
        attention_mask = None
        # argonne4.5 document masking. The production loaders pack many documents into one flat
        # contiguous window, and this line used to be the whole story -- so every token attended
        # freely across document boundaries. That is worst on exactly the target domains (short
        # math problems, short code files), where a block-1024 window spans many documents. When
        # `doc_mask` is on and per-token document ids are supplied, build a block-diagonal mask
        # per DISTINCT window size (at most two: sliding and global) and share it across layers.
        # Costs one (B, T, T) bool per distinct window; without document_ids nothing changes.
        doc_masks = None
        if getattr(self.config, "doc_mask", False) and document_ids is not None and not use_cache:
            doc_masks = self._build_doc_masks(document_ids.to(device), seq_length)

        # RoPE positions are offset by the cached length so incremental decode
        # uses the correct absolute positions.
        past_len = past_key_values[0][0].shape[2] if past_key_values else 0
        cos_full, sin_full = self.rotary_emb(hidden_states, past_len + seq_length)
        rotary = (cos_full[past_len:past_len + seq_length],
                  sin_full[past_len:past_len + seq_length])

        new_cache = [] if use_cache else None
        for i, layer in enumerate(self.blocks):
            _skip_ckpt = (self.checkpoint_stride >= 2) and (i % self.checkpoint_stride == 0)
            layer_mask = (
                doc_masks[layer.attn.sliding_window] if doc_masks is not None else attention_mask
            )
            if self.gradient_checkpointing and self.training and not _skip_ckpt:
                hidden_states = torch.utils.checkpoint.checkpoint(
                    layer,
                    hidden_states,
                    rotary,
                    layer_mask,
                    use_reentrant=False,
                )
            elif use_cache:
                past_kv = past_key_values[i] if past_key_values else None
                hidden_states, layer_kv = layer(
                    hidden_states, rotary, layer_mask,
                    past_kv=past_kv, use_cache=True,
                )
                new_cache.append(layer_kv)
            else:
                hidden_states = layer(hidden_states, rotary, layer_mask)

        hidden_states = self.norm(hidden_states)

        # Chunked cross-entropy (training only, flag-gated via config.loss_chunk_size): avoids
        # materializing the full (batch*seq, vocab) fp32 logit tensor -- the HBM wall at long context.
        # Early-return so the standard full-logit path below stays byte-for-byte unchanged (zero risk to
        # non-chunked / default runs). Only taken when MTP + z-loss are off (both default off).
        _chunk = getattr(self.config, "loss_chunk_size", 0)
        if (self.training and labels is not None and _chunk > 0
                and self.config.mtp_horizon == 1 and self.config.z_loss_weight == 0
                and len(self.mtp_modules) == 0):
            loss = self._chunked_lm_loss(hidden_states, labels, _chunk)
            if torch.isnan(loss):
                self._nan_loss_count += 1
                if self._nan_loss_count <= 5 or self._nan_loss_count % 100 == 0:
                    print(f"WARNING: NaN loss detected (occurrence {self._nan_loss_count}); "
                          "zeroing this step's loss.", flush=True)
                zero = torch.zeros((), device=loss.device, dtype=torch.float32)
                for p in self.parameters():
                    if p.requires_grad:
                        zero = zero + torch.nan_to_num(p).sum().float() * 0.0
                loss = zero.to(loss.dtype)
            return CausalLMOutput(logits=None, loss=loss)

        logits = self.lm_head(hidden_states)
        if self.config.logit_softcap > 0:
            cap = self.config.logit_softcap
            logits = torch.tanh(logits / cap) * cap
        
        # Sanitize logits only at inference. During training the full-tensor
        # scan forces a GPU sync every step over [batch, seq, vocab]; a NaN
        # would surface in the loss check below anyway.
        if not self.training and torch.isnan(logits).any():
            logits = torch.nan_to_num(logits, nan=0.0, posinf=65504.0, neginf=-65504.0)

        loss = None
        if labels is not None:
            if labels.device != logits.device:
                labels = labels.to(logits.device)
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                labels.view(-1),
                ignore_index=-100,
            )
            if self.training and self.config.mtp_horizon > 1:
                mtp_terms = []
                max_horizon = min(self.config.mtp_horizon, labels.shape[1])
                for horizon in range(2, max_horizon + 1):
                    shift = horizon - 1
                    shifted_logits = logits[:,:-shift,:]
                    shifted_labels = labels[:, shift:]
                    if shifted_logits.numel() == 0:
                        continue
                    mtp_loss = F.cross_entropy(
                        shifted_logits.reshape(-1, shifted_logits.size(-1)),
                        shifted_labels.reshape(-1),
                        ignore_index=-100,
                    )
                    mtp_terms.append(mtp_loss / horizon)
                if mtp_terms:
                    loss = loss + (self.config.mtp_loss_weight * torch.stack(mtp_terms).mean()).to(loss.dtype)
            if self.training and len(self.mtp_modules) > 0 and self.config.mtp_loss_weight > 0:
                loss = loss + (
                    self.config.mtp_loss_weight * self._mtp_module_loss(hidden_states, labels, rotary)
                ).to(loss.dtype)
            if self.config.z_loss_weight > 0:
                z = torch.logsumexp(logits.float(), dim=-1)
                loss = loss + (self.config.z_loss_weight * z.pow(2).mean()).to(loss.dtype)

            # A NaN loss becomes a zero loss that is still connected to every
            # parameter: but through the parameters directly, NOT through the
            # network's (possibly NaN) activations. Backward then delivers an
            # exact zero gradient to each parameter: DDP/FSDP gradient hooks
            # all fire (no collective desync/hang) and the optimizer step is a
            # true no-op. Warn (bounded) instead of hiding it.
            if torch.isnan(loss):
                self._nan_loss_count += 1
                if self._nan_loss_count <= 5 or self._nan_loss_count % 100 == 0:
                    print(
                        f"WARNING: NaN loss detected (occurrence {self._nan_loss_count}); "
                        "zeroing this step's loss.",
                        flush=True,
                    )
                zero = torch.zeros((), device=loss.device, dtype=torch.float32)
                for p in self.parameters():
                    if p.requires_grad:
                        zero = zero + torch.nan_to_num(p).sum().float() * 0.0
                loss = zero.to(loss.dtype)

        if use_cache:
            return CausalLMOutputWithPast(logits=logits, loss=loss, past_key_values=new_cache)
        return CausalLMOutput(logits=logits, loss=loss)

    def _build_doc_masks(self, document_ids, seq_length):
        """{sliding_window -> bool mask (B, 1, T, T), True == attend}, one per distinct window.

        Combines causal, the layer's sliding window, and the document boundary, so a token can
        only see earlier tokens of its OWN packed document. Built once per forward and shared
        across every layer with the same window (at most two under an 'LLLG' layout).
        """
        qi = torch.arange(seq_length, device=document_ids.device)
        same_doc = document_ids[:,:, None] == document_ids[:, None,:]   # (B, T, T)
        base = same_doc & (qi[:, None] >= qi[None,:])
        masks = {}
        for window in {block.attn.sliding_window for block in self.blocks}:
            if window is None:
                masks[None] = base.unsqueeze(1)
            else:
                masks[window] = (base & (qi[None,:] >= qi[:, None] - window)).unsqueeze(1)
        return masks

    def _mtp_module_loss(self, hidden_states, labels, rotary):
        """Mean CE over the real MTP depths (see MTPModule).

        Depth k reads the previous depth's state at position i plus the embedding of token t+k
        (the token that state predicts) and predicts t+k+1, so each depth is truncated by one
        more position. `labels[:, i]` is already token i+1 (the loaders hand back y = x shifted),
        hence token t+k at position i is `labels[:, i + k - 1]`.
        """
        terms = []
        h_prev = hidden_states
        cap = float(self.config.logit_softcap)
        seq_len = labels.shape[1]
        batch = labels.shape[0]
        # fp8 `_scaled_mm` requires the flattened row count (batch * length) to be divisible by 16.
        # Truncating by depth k breaks that -- at k=1, 8*2047 = 16,376 is not -- and the lm_head GEMM
        # inside this loss then dies under torchao float8. Round the usable length DOWN so batch*L
        # stays 16-aligned; that discards at most 15 of ~2047 positions, which is negligible against
        # the alignment requirement. Costs nothing when fp8 is off.
        align = 16 // math.gcd(batch, 16)
        for k, module in enumerate(self.mtp_modules, start=1):
            usable = ((seq_len - k) // align) * align
            if usable < 1:
                break
            h_in = h_prev[:,:usable,:]
            tokens_in = labels[:, k - 1: k - 1 + usable]
            targets = labels[:, k: k + usable]
            # Ignored positions (-100) would index the embedding out of range; their targets are
            # ignored by cross_entropy anyway, so any in-range id works as a placeholder.
            embeds_in = self.embed_tokens(tokens_in.clamp_min(0))
            rot_k = (rotary[0][:usable], rotary[1][:usable])
            h_out = module(h_in, embeds_in, rot_k)
            # ⛔ DO NOT "optimise" this with _chunked_lm_loss -- MEASURED WORSE, 2026-08-27.
            # Real MTP (mtp_module_layers=1) is a genuine iso-token quality win (tgt -0.058, code
            # -0.118, both seeds, every domain) but costs -42% throughput (5,348 -> 3,082 tok/s), so
            # iso-compute it LOSES by 0.27-0.33: the same wall buys 1.73x more tokens without it.
            # I hypothesised the -42% was this line materialising [T, 151680] logits per depth, and
            # routed it through _chunked_lm_loss (equivalence proven to 9.5e-7). Result: tok/s went
            # 3,082 -> 2,723, i.e. cost rose to -49.1%, at HBM 0.960 with quality unchanged.
            # ⇒ The transient was NOT the bottleneck; chunking's recompute overhead is pure loss here.
            # The cost is INHERENT: each depth runs a full extra block AND a 151,680-wide lm_head GEMM
            # in forward and backward. MTP is rejected on cost, not on an artifact -- do not retry
            # this optimisation.
            mtp_logits = self.lm_head(h_out)
            if cap > 0:
                mtp_logits = torch.tanh(mtp_logits / cap) * cap
            terms.append(
                F.cross_entropy(
                    mtp_logits.reshape(-1, mtp_logits.size(-1)),
                    targets.reshape(-1),
                    ignore_index=-100,
                )
            )
            h_prev = h_out
        if not terms:
            return hidden_states.new_zeros(())
        return torch.stack(terms).mean()

    @torch.compiler.disable
    def _chunked_lm_loss(self, hidden_states, labels, chunk_size):
        """Cross-entropy without ever materializing the full (batch*seq, vocab) fp32 logit tensor.
        MUST run eager (@torch.compiler.disable): torch.compile/inductor otherwise fuses the chunks
        back into the full (N, vocab) logit tensor, re-creating the OOM this is meant to avoid. The
        transformer body still compiles (+FP8); only this loss runs eager (a small fraction of compute).
        Splits the flattened rows into chunks; torch.utils.checkpoint recomputes each chunk's logits
        in backward, so peak logit memory = chunk_size*vocab (not N*vocab). Numerically identical to
        F.cross_entropy(full_logits, labels.view(-1), ignore_index=-100) with per-chunk softcap
        (mean over non-ignored). Caller guarantees mtp_horizon==1 and z_loss_weight==0.
        Uses self.lm_head(...) per chunk so the fp8 (Float8Linear) GEMM path is preserved."""
        from torch.utils.checkpoint import checkpoint as _ckpt
        cap = float(self.config.logit_softcap)
        H = hidden_states.reshape(-1, hidden_states.size(-1))
        L = labels.reshape(-1).to(H.device)
        N = H.size(0)

        def _seg(hc, lc):
            lg = self.lm_head(hc)                    # (chunk, vocab); Float8Linear under fp8
            if cap > 0:
                lg = torch.tanh(lg / cap) * cap
            return F.cross_entropy(lg, lc, ignore_index=-100, reduction="sum")

        total = H.new_zeros((), dtype=torch.float32)
        ntok = torch.zeros((), device=H.device, dtype=torch.float32)
        cs = max(1, int(chunk_size))
        for i in range(0, N, cs):
            hc = H[i:i + cs]
            lc = L[i:i + cs]
            seg = _ckpt(_seg, hc, lc, use_reentrant=False) if self.training else _seg(hc, lc)
            total = total + seg.float()
            ntok = ntok + (lc != -100).sum().float()
        return total / ntok.clamp(min=1.0)

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_length: int = 1024,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        do_sample: bool = True,
        repetition_penalty: float = 1.0,
        no_repeat_ngram_size: int = 0,
        eos_token_id: Optional[int] = None,
    ) -> torch.Tensor:
        self.eval()
        device = self.embed_tokens.weight.device
        input_ids = input_ids.to(device)
        ctx = self.config.max_position_embeddings
        # Stop at EOS (2026-08-02). This loop previously had NO eos check at all -- its only
        # exit was `max_length`, so config.eos_token_id was inert and every caller following the
        # model card's `.generate()` snippet ran to the token budget and then degenerated into
        # repetition long after the answer was finished. Harmless-looking on a base model (which
        # just continues text) but bad on a chat/think model, where the assistant turn must end
        # at <|im_end|>. Longstanding across the whole Argonne line, including the published
        # 3.0-think bundle -- not introduced by any recent change.
        # Pass eos_token_id=-1 to explicitly opt out and restore run-to-max_length behavior.
        if eos_token_id is None:
            eos_token_id = getattr(self.config, "eos_token_id", None)
        stop_id = None if eos_token_id is None or eos_token_id < 0 else int(eos_token_id)
        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=device)
        # KV-cache decode: prefill the prompt once, then feed a single token per
        # step reusing past_key_values. This yields identical logits to the
        # recompute-the-whole-prefix path (gated by verify_cache.py) but turns the
        # per-step cost from O(seq_len) down to O(1). When the running sequence
        # would exceed the context window we rebuild the cache from the truncated
        # window, matching the original chunk = input_ids[:, -ctx:] behavior.
        past = None
        outputs = None
        while input_ids.shape[1] < max_length:
            if past is None or past[0][0].shape[2] >= ctx:
                chunk = input_ids[:, -ctx:]
                outputs = self.forward(chunk, use_cache=True)
                past = outputs.past_key_values
            logits = outputs.logits[:, -1,:] / temperature

            if repetition_penalty != 1.0:
                for b in range(input_ids.shape[0]):
                    seen_tokens = torch.unique(input_ids[b])
                    seen_logits = logits[b, seen_tokens]
                    adjusted = torch.where(
                        seen_logits < 0,
                        seen_logits * repetition_penalty,
                        seen_logits / repetition_penalty,
                    )
                    logits[b, seen_tokens] = adjusted

            if no_repeat_ngram_size > 0 and input_ids.shape[1] + 1 >= no_repeat_ngram_size:
                n = int(no_repeat_ngram_size)
                for b in range(input_ids.shape[0]):
                    seq = input_ids[b].tolist()
                    ngrams = {}
                    for i in range(len(seq) - n + 1):
                        prefix = tuple(seq[i: i + n - 1]) if n > 1 else tuple()
                        next_token = seq[i + n - 1]
                        if prefix not in ngrams:
                            ngrams[prefix] = set()
                        ngrams[prefix].add(next_token)
                    current_prefix = tuple(seq[-(n - 1):]) if n > 1 else tuple()
                    banned = ngrams.get(current_prefix, set())
                    if banned:
                        logits[b, list(banned)] = float("-inf")

            if do_sample:
                if top_k is not None:
                    top_values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                    logits = logits.masked_fill(logits < top_values[:, [-1]], float("-inf"))
                if top_p is not None:
                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[...,:-1].clone()
                    sorted_indices_to_remove[..., 0] = 0
                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                    logits = logits.masked_fill(indices_to_remove, float("-inf"))
                probs = F.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(logits, dim=-1, keepdim=True)

            next_token = next_token.to(input_ids.device)
            input_ids = torch.cat([input_ids, next_token], dim=-1)
            if input_ids.shape[1] >= max_length:
                break
            # Batch-safe: only stop once EVERY row has emitted EOS, so a short sequence cannot
            # truncate the others. (This arch has no padding support, so batch is normally 1.)
            if stop_id is not None:
                finished |= next_token.squeeze(-1).eq(stop_id)
                if bool(finished.all()):
                    break
            # Advance the cache by one token unless we must rebuild it next loop
            # (cache full); the top-of-loop guard re-prefills in that case.
            if past[0][0].shape[2] < ctx:
                outputs = self.forward(next_token, past_key_values=past, use_cache=True)
                past = outputs.past_key_values
        return input_ids.to(device)


AutoConfig.register("argonne2", ArgonneConfig)
AutoModel.register(ArgonneConfig, ArgonneModel)
AutoModelForCausalLM.register(ArgonneConfig, ArgonneModel)

# Backwards compatibility exports
CausalSelfAttention = GroupedQueryAttention
MLP = SwiGLUMLP