Blame for modeling_ricursive.py

"""HuggingFace modeling implementation for the Ricursive transformer architecture.

Designed for inference with vLLM (via --trust-remote-code). Key differences from
standard Llama/Mistral:
- PeriNorm blocks: 4 layer norms per block (pre/post attention + pre/post FFN)
- Gated attention: elementwise sigmoid gate on attention output
- Mixed local/global layers: local = sliding window + RoPE, global = full attention + no RoPE
- Per-head Q/K RMSNorm

Transformers 5.x compatibility (2026-11-05)
------------------------------------------
`generate()` used to break on the very first prefill because
`prepare_inputs_for_generation` sliced `input_ids` with a `None`
`cache_position`, and the causal/sliding-window mask builders assumed the
transformers 4.x cache API. The fixes are applied inline below:

* `prepare_inputs_for_generation` rebuilds a missing `cache_position`, slices
  only when the cache is warm, derives per-row `position_ids` from the padding
  mask, and moves inputs to the model device.
* `_update_causal_mask` sizes the key axis from `cache_position` instead of the
  removed `get_max_cache_shape()`, emits `0.0` for visible keys, and never
  returns an all-masked single-token decode row.
* `_apply_sliding_window_mask` anchors key positions at the first key the layer
  actually holds (the cache evicts from the left) and is skipped when the
  window cannot bind.
* `RicursiveForCausalLM` inherits `GenerationMixin` explicitly; transformers 5
  no longer injects it into remote-code classes.

DFlash2 inference fixes (2026-11-06)
-----------------------------------
* `RicursiveAttention.forward` builds an end-anchored causal mask for a
  multi-token verify block running against a warm cache. SDPA's implicit mask
  is top-left aligned, so with `q_len < kv_len` it hid valid context.
  Preallocated caches (`StaticCache`) return their whole capacity from
  `update()`, so the mask is now built from the layer's *actual* cache length
  and query positions; the unused tail is trimmed before attention.
* `RicursiveForCausalLM.forward` honours `logits_to_keep` by slicing the LM
  head *input*, so the full `[batch, seq, vocab]` projection is never
  materialised; `prepare_inputs_for_generation` requests one row. Labeled
  calls keep the full sequence so the shifted causal-LM loss stays defined,
  then trim only the returned logits to the requested width.
* `RicursiveMLP.forward` chunks the row dimension
  (`RIPPLE_MLP_CHUNK`, default 2048, `0` disables) to bound the
  prefill activation peak. Every op is row-wise, so the result is exact.

Cache-alignment fixes from review (2026-11-06)
----------------------------------------------
* A mask longer than the key tensor is now aligned to the absolute positions
  the layer actually holds, `[total_kv - kv_len, total_kv)`, instead of to the
  tail of the mask. With a mask sized to the allocated cache *capacity*
  (`StaticCache`), the old `[-kv_len:]` slice selected unwritten slots after
  sliding-window eviction and silently returned wrong attention.
* `RicursiveAttention.forward` forwards `cache_position` to
  `past_key_value.update()`. A transformers 4.x `StaticCache` otherwise falls
  back to `arange(q_len)` and the first decode step overwrites the prompt
  instead of appending after it.
"""

import os
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

try:  # transformers >= 4.26, re-exported from here in 5.x
    from transformers.generation.utils import GenerationMixin
except ImportError:  # pragma: no cover - very old transformers
    from transformers.generation import GenerationMixin

from .configuration_ricursive import RicursiveConfig


# Set when the transformers-5.x `generate()` fixes above are present natively.
# Downstream loaders that monkeypatch older checkpoints check this marker and
# stand down instead of re-installing the equivalent runtime patch.
RICURSIVE_TF5_COMPAT = "2026-11-05"


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

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


class RicursiveRotaryEmbedding(nn.Module):
    def __init__(self, dim: int, theta: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.theta = theta
        # NOTE: inv_freq is computed on-the-fly in forward() rather than stored
        # as a buffer, because non-persistent buffers get corrupted during
        # from_pretrained() loading.

    @torch.no_grad()
    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
        inv_freq = 1.0 / (
            self.theta ** (torch.arange(0, self.dim, 2, device=x.device).float() / self.dim)
        )
        inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
        position_ids_expanded = position_ids[:, None, :].float()

        device_type = x.device.type
        with torch.autocast(device_type=device_type, enabled=False):
            freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos()
            sin = emb.sin()
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)


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,
) -> Tuple[torch.Tensor, torch.Tensor]:
    cos = cos.unsqueeze(1)  # (B, 1, T, dim)
    sin = sin.unsqueeze(1)  # (B, 1, T, dim)
    # OLMo-core applies RoPE in full precision (float32) for numerical stability
    orig_dtype = q.dtype
    q_float, k_float = q.float(), k.float()
    cos_float, sin_float = cos.float(), sin.float()
    q_embed = (q_float * cos_float) + (rotate_half(q_float) * sin_float)
    k_embed = (k_float * cos_float) + (rotate_half(k_float) * sin_float)
    return q_embed.to(orig_dtype), k_embed.to(orig_dtype)


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    if n_rep == 1:
        return hidden_states
    batch, num_kv_heads, slen, head_dim = hidden_states.shape
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch, num_kv_heads, n_rep, slen, head_dim
    )
    return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)


class RicursiveAttention(nn.Module):
    def __init__(self, config: RicursiveConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.layer_type = config.layer_types[layer_idx]

        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 = config.head_dim
        self.num_kv_groups = self.num_heads // self.num_kv_heads

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
        if config.use_valres:
            self.value_res_lambda = nn.Parameter(torch.zeros(1))

        # Per-head Q/K norms (RMSNorm on head_dim)
        self.q_norm = RicursiveRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RicursiveRMSNorm(self.head_dim, eps=config.rms_norm_eps)

        # Gated attention
        if config.use_gated_attention:
            self.g_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)

        # RoPE only for local (sliding_attention) layers
        if self.layer_type == "sliding_attention":
            self.rotary_emb = RicursiveRotaryEmbedding(self.head_dim, theta=config.rope_theta)

    def _apply_sliding_window_mask(
        self,
        attention_mask: Optional[torch.Tensor],
        q_len: int,
        kv_len: int,
        dtype: torch.dtype,
        device: torch.device,
        cache_position: Optional[torch.LongTensor],
        total_kv: Optional[int] = None,
    ) -> Optional[torch.Tensor]:
        if self.layer_type != "sliding_attention":
            return attention_mask

        # transformers >= 5 builds a hybrid cache: a sliding layer holds at most
        # `sliding_window` keys while a full-attention layer keeps growing, so the
        # shared mask produced by `_update_causal_mask` can be longer than this
        # layer's key tensor. Keep only the keys this layer actually holds.
        if attention_mask is not None and attention_mask.shape[-1] > kv_len:
            mask_len = attention_mask.shape[-1]
            if total_kv is not None:
                # The retained keys sit at `[total_kv - kv_len, total_kv)`.
                start = max(0, min(total_kv - kv_len, mask_len - kv_len))
                attention_mask = attention_mask[..., start : start + kv_len]
            else:
                attention_mask = attention_mask[..., -kv_len:]

        if cache_position is not None:
            query_positions = cache_position.to(device=device).reshape(-1)
        else:
            query_positions = torch.arange(kv_len - q_len, kv_len, device=device)

        window = self.config.sliding_window
        if window is None:
            return attention_mask
        if kv_len <= window:
            # The window cannot exclude any key, so the result is either the
            # unchanged (already causal) mask or, without one, plain causal
            # attention -- which SDPA applies through `is_causal`.
            if attention_mask is None:
                # SDPA's implicit causal mask is top-left aligned. It only
                # matches this query block when it is the whole sequence
                # (`q_len == kv_len`) or a single decode position
                # (`q_len == 1`). A multi-token block verified against a warm
                # cache sits at the *end* of the key span, so fall through and
                # build the end-anchored mask below.
                if q_len == 1 or q_len == kv_len:
                    return None
            if q_len == 1:
                return attention_mask

        # Keys are contiguous and end at the newest query position. After the cache
        # evicts from the left they no longer start at position 0, so anchor the
        # range at the first key the layer holds instead of at zero.
        key_positions = torch.arange(kv_len, device=device) + (
            query_positions[-1] - kv_len + 1
        )

        allowed = (
            (key_positions[None, :] <= query_positions[:, None])
            & (key_positions[None, :] >= query_positions[:, None] - (window - 1))
        )[None, None, :, :]

        if attention_mask is None:
            return allowed
        if attention_mask.dtype == torch.bool:
            return attention_mask & allowed
        return attention_mask.masked_fill(~allowed, torch.finfo(dtype).min)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        value_residual: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        bsz, q_len, _ = hidden_states.shape

        # Project Q, K, V
        q = self.q_proj(hidden_states)
        k = self.k_proj(hidden_states)
        value_states = hidden_states
        if value_residual is not None:
            if not self.config.use_valres:
                raise RuntimeError("value_residual was supplied but config.use_valres is false")
            blend = torch.sigmoid(self.value_res_lambda[0]).to(dtype=hidden_states.dtype)
            value_states = torch.lerp(
                value_residual.to(dtype=hidden_states.dtype), hidden_states, blend
            )
        v = self.v_proj(value_states)

        # Reshape to (B, T, n_heads, head_dim) then transpose to (B, n_heads, T, head_dim)
        q = q.view(bsz, q_len, self.num_heads, self.head_dim)
        k = k.view(bsz, q_len, self.num_kv_heads, self.head_dim)
        v = v.view(bsz, q_len, self.num_kv_heads, self.head_dim)

        # Per-head Q/K norm (applied before transpose, on head_dim)
        q = self.q_norm(q)
        k = self.k_norm(k)

        # Transpose to (B, n_heads, T, head_dim) for attention
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # RoPE only for local layers
        if self.layer_type == "sliding_attention":
            cos, sin = self.rotary_emb(v, position_ids)
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # KV cache
        total_kv = None
        if past_key_value is not None:
            # Forward the absolute positions of the tokens being written. A
            # transformers 4.x `StaticCache` falls back to `arange(q_len)` when
            # `cache_position` is absent, so the first decode step would write
            # at slots 0..q_len-1 and overwrite the prompt instead of appending
            # after it. `DynamicCache`/5.x caches ignore the extra kwarg.
            cache_kwargs = (
                {"cache_position": cache_position} if cache_position is not None else None
            )
            if cache_kwargs is not None:
                try:
                    k, v = past_key_value.update(k, v, self.layer_idx, cache_kwargs)
                except TypeError:  # cache without a `cache_kwargs` parameter
                    k, v = past_key_value.update(k, v, self.layer_idx)
            else:
                k, v = past_key_value.update(k, v, self.layer_idx)
            # Number of keys this layer has actually written. A preallocated
            # (StaticCache) layer returns its full capacity from `update()`,
            # including slots that have not been filled yet; an evicting
            # (sliding-window) layer reports more than the trailing window it
            # returns. The mask must use the real positions, not the tensor size.
            total_kv = _cache_layer_seq_length(past_key_value, self.layer_idx)

        # GQA: expand KV heads
        k = repeat_kv(k, self.num_kv_groups)
        v = repeat_kv(v, self.num_kv_groups)

        # Drop the unwritten tail of a preallocated cache. Without this a
        # six-token prefill into a sixteen-slot StaticCache looks like a
        # ten-token prefix, and the end-anchored verify mask below would let
        # each query attend slots belonging to a different absolute position.
        if total_kv is not None and 0 < total_kv < k.shape[2]:
            k = k[:, :, :total_kv]
            v = v[:, :, :total_kv]

        # Align an incoming mask that is longer than the keys this layer holds.
        # The shared mask is built over absolute key positions
        # `0..target_length-1`; keep the slice that matches the keys present.
        #
        # The keys returned by `update()` occupy the absolute positions
        # `[total_kv - kv_len, total_kv)`. Anchoring on the tensor tail
        # (`[-kv_len:]`) is only correct when the mask length equals
        # `total_kv`; with a mask sized to the *allocated cache capacity* it
        # selects the unwritten tail instead of the retained window and
        # silently produces wrong attention.
        if attention_mask is not None and attention_mask.shape[-1] > k.shape[2]:
            kv_len_now = k.shape[2]
            mask_len = attention_mask.shape[-1]
            if total_kv is not None:
                start = max(0, min(total_kv - kv_len_now, mask_len - kv_len_now))
                attention_mask = attention_mask[..., start : start + kv_len_now]
            else:
                # Unknown cache: the keys are the trailing window.
                attention_mask = attention_mask[..., -kv_len_now:]

        attention_mask = self._apply_sliding_window_mask(
            attention_mask=attention_mask,
            q_len=q.shape[2],
            kv_len=k.shape[2],
            dtype=q.dtype,
            device=q.device,
            cache_position=cache_position,
            total_kv=total_kv,
        )

        if attention_mask is None and q.shape[2] > 1 and k.shape[2] > q.shape[2]:
            # Multi-token verify block against a warm cache. SDPA's implicit
            # causal mask is top-left aligned, so when `q_len < kv_len` it hides
            # every key before the block. Build the mask from the actual query
            # positions (and the absolute positions of the returned keys) rather
            # than assuming the key tensor starts at position 0.
            if cache_position is not None and cache_position.numel() == q.shape[2]:
                query_positions = cache_position.to(
                    device=q.device, dtype=torch.long
                ).reshape(-1)
            else:
                query_positions = torch.arange(
                    k.shape[2] - q.shape[2], k.shape[2], device=q.device
                )
            if total_kv is not None and total_kv > k.shape[2]:
                # Evicting cache: the returned keys are the trailing window that
                # ends at the newest query, not positions 0..kv_len-1.
                last_query = int(query_positions[-1].item())
                key_positions = torch.arange(
                    last_query - k.shape[2] + 1, last_query + 1, device=q.device
                )
            else:
                key_positions = torch.arange(k.shape[2], device=q.device)
            attention_mask = (
                key_positions[None, :] <= query_positions[:, None]
            )[None, None]

        # Attention (use is_causal when no explicit mask for correct causal behavior)
        # Cast to plain bool to avoid SymBool issues with torch.compile
        is_causal = bool(attention_mask is None and q.shape[2] > 1)
        attn_output = F.scaled_dot_product_attention(
            q, k, v, attn_mask=attention_mask, dropout_p=0.0, is_causal=is_causal
        )

        # (B, n_heads, T, head_dim) -> (B, T, n_heads * head_dim)
        attn_output = attn_output.transpose(1, 2).contiguous()

        # Gated attention: output * sigmoid(g_proj(input))
        if self.config.use_gated_attention:
            gate = self.g_proj(hidden_states)
            if self.config.gate_full_precision:
                gate = gate.float()
            gate = torch.sigmoid(gate).to(attn_output.dtype)
            attn_output = attn_output.view(bsz, q_len, -1) * gate
        else:
            attn_output = attn_output.view(bsz, q_len, -1)

        # Output projection
        attn_output = self.o_proj(attn_output)

        return attn_output, None, past_key_value


def _mlp_chunk_size() -> int:
    """Row budget for the chunked MLP forward. `0` (or negative) disables it."""
    try:
        return int(os.environ.get("RIPPLE_MLP_CHUNK", "2048"))
    except ValueError:
        return 2048


class RicursiveMLP(nn.Module):
    def __init__(self, config: RicursiveConfig):
        super().__init__()
        self.clamp_threshold = config.ffn_clamp_threshold
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def _forward_rows(self, x: torch.Tensor) -> torch.Tensor:
        gate = self.gate_proj(x)
        up = self.up_proj(x)
        if self.clamp_threshold:
            gate = gate.clamp(min=None, max=self.clamp_threshold)
            up = up.clamp(min=-self.clamp_threshold, max=self.clamp_threshold)
        return self.down_proj(F.silu(gate) * up)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Ripple's MLP expands hidden_size -> intermediate_size (5120 -> 40960
        # for the 65B checkpoint), so a long prefill materialises `gate`, `up`
        # and `silu(gate) * up` at once. Every op here is row-wise, so splitting
        # the row dimension is exact: each row still uses the same kernels with
        # the same reduction order, and the pieces are concatenated back. Only
        # the transient peak shrinks. `RIPPLE_MLP_CHUNK` overrides the
        # row budget; `0` disables chunking.
        chunk = _mlp_chunk_size()
        if chunk <= 0 or not torch.is_tensor(x) or x.dim() not in (2, 3):
            return self._forward_rows(x)

        shape = x.shape
        hidden = shape[-1]
        rows = 1
        for size in shape[:-1]:
            rows *= size
        if rows <= chunk:
            return self._forward_rows(x)

        flat = x.reshape(rows, hidden)
        pieces = [
            self._forward_rows(flat[start : start + chunk])
            for start in range(0, rows, chunk)
        ]
        out = torch.cat(pieces, dim=0) if len(pieces) > 1 else pieces[0]
        return out.reshape(*shape[:-1], out.shape[-1])


class RicursiveDecoderLayer(nn.Module):
    """PeriNorm decoder layer with 4 layer norms."""

    def __init__(self, config: RicursiveConfig, layer_idx: int):
        super().__init__()
        self.self_attn = RicursiveAttention(config, layer_idx)
        self.mlp = RicursiveMLP(config)

        # PeriNorm: 4 norms
        self.input_layernorm = RicursiveRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RicursiveRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.pre_feedforward_layernorm = RicursiveRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )
        self.post_feedforward_layernorm = RicursiveRMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        value_residual: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, ...]:
        # Attention sub-block (PeriNorm: pre-norm -> attn -> post-norm -> residual)
        normed = self.input_layernorm(hidden_states)
        attn_output, attn_weights, present_key_value = self.self_attn(
            normed,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            cache_position=cache_position,
            value_residual=value_residual,
            **kwargs,
        )
        hidden_states = hidden_states + self.post_attention_layernorm(attn_output)

        # FFN sub-block (PeriNorm: pre-norm -> ffn -> post-norm -> residual)
        normed = self.pre_feedforward_layernorm(hidden_states)
        ff_output = self.mlp(normed)
        hidden_states = hidden_states + self.post_feedforward_layernorm(ff_output)

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (attn_weights,)
        if use_cache:
            outputs += (present_key_value,)
        return outputs


class RicursivePreTrainedModel(PreTrainedModel):
    config_class = RicursiveConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["RicursiveDecoderLayer"]
    _supports_sdpa = True
    _supports_cache_class = True


class RicursiveModel(RicursivePreTrainedModel):
    def __init__(self, config: RicursiveConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [RicursiveDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
        )
        self.norm = RicursiveRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = (
            output_attentions if output_attentions is not None else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        # embed_scale: non-folded checkpoints carry the scale (e.g. sqrt(hidden_size)) in the
        # config and apply it here at runtime; folded checkpoints set embed_scale=None and bake
        # the scale into embed_tokens.weight, so this is a no-op. This mirrors the SGLang native
        # port (sglang_ricursive) so served and trained forwards agree. It must run BEFORE the
        # value-residual seed below, whose source is the (scaled) layer-0 embeddings.
        embed_scale = getattr(self.config, "embed_scale", None)
        if embed_scale is not None:
            inputs_embeds = inputs_embeds * embed_scale

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        if cache_position is None:
            past_seen_tokens = (
                past_key_values.get_seq_length() if past_key_values is not None else 0
            )
            cache_position = torch.arange(
                past_seen_tokens,
                past_seen_tokens + inputs_embeds.shape[1],
                device=inputs_embeds.device,
            )

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        # Build causal mask
        causal_mask = self._update_causal_mask(
            attention_mask, inputs_embeds, cache_position, past_key_values
        )

        hidden_states = inputs_embeds
        value_residual = hidden_states if self.config.use_valres else None
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None

        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            layer_outputs = layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=position_ids,
                past_key_value=past_key_values,
                output_attentions=output_attentions,
                use_cache=use_cache,
                cache_position=cache_position,
                value_residual=value_residual,
                **kwargs,
            )
            hidden_states = layer_outputs[0]

            if output_attentions:
                all_self_attns += (layer_outputs[1],)

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = past_key_values if use_cache else None

        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
                if v is not None
            )
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )

    def _update_causal_mask(
        self,
        attention_mask: Optional[torch.Tensor],
        input_tensor: torch.Tensor,
        cache_position: torch.Tensor,
        past_key_values: Optional[Cache],
    ) -> Optional[torch.Tensor]:
        # When there's no explicit attention_mask (no padding) and we're doing
        # standard prefill or cached decoding, return None and let SDPA use
        # its efficient is_causal path. This also avoids numerical issues
        # from float min mask values with large activations.
        if attention_mask is None:
            return None
        if torch.all(attention_mask == 1):
            return None

        # Only build an explicit mask when we have a padding mask
        dtype, device = input_tensor.dtype, input_tensor.device
        sequence_length = input_tensor.shape[1]

        if cache_position is None:
            query_positions = torch.arange(sequence_length, device=device)
        else:
            query_positions = cache_position.to(device=device).reshape(-1)

        # Size the key axis from `cache_position`, which is exactly what the
        # attention layer uses to size its key tensor.
        # `past_key_values.get_max_cache_shape()` was removed in transformers 5,
        # and its replacement reports the *capacity* (e.g. `sliding_window`),
        # which disagrees with the live key count whenever the cache is not full.
        target_length = int(query_positions[-1].item()) + 1
        target_length = max(target_length, attention_mask.shape[-1])

        # A shorter padding mask leaves the trailing keys unspecified; treat them as
        # visible -- the causal term below still applies.
        if attention_mask.shape[-1] < target_length:
            pad = attention_mask.new_ones(
                attention_mask.shape[0], target_length - attention_mask.shape[-1]
            )
            attention_mask = torch.cat([pad, attention_mask], dim=-1)

        key_positions = torch.arange(target_length, device=device)
        allowed = key_positions[None, :] <= query_positions[:, None]
        allowed = allowed[None, None, :, :] & attention_mask[:, None, None, :].to(torch.bool)

        causal_mask = torch.zeros(
            (attention_mask.shape[0], 1, sequence_length, target_length),
            dtype=dtype,
            device=device,
        )
        # `0.0` for visible keys (not `1.0`, which would bias every score) and the
        # dtype minimum elsewhere. Building the mask from `allowed` also keeps a
        # single-token decode step from coming out fully masked.
        return causal_mask.masked_fill(~allowed, torch.finfo(dtype).min)


def _cache_seq_length(past_key_values: Optional[Cache]) -> int:
    """Number of cached keys, across transformers 4.x and 5.x cache objects."""
    if past_key_values is None:
        return 0
    getter = getattr(past_key_values, "get_seq_length", None)
    if callable(getter):
        return int(getter())
    try:
        return int(past_key_values[0][0].shape[-2])
    except Exception:  # pragma: no cover - unknown cache implementations
        return 0


def _cache_layer_seq_length(
    past_key_values: Optional[Cache], layer_idx: int
) -> Optional[int]:
    """Actual number of keys written for `layer_idx`, or `None` if unknown.

    `Cache.get_seq_length(layer_idx)` counts *written* tokens. For a
    preallocated static layer that is smaller than the tensor `update()`
    returns until the cache fills; for an evicting sliding-window layer it is
    larger, because the layer only returns a trailing window.
    """
    if past_key_values is None:
        return None
    getter = getattr(past_key_values, "get_seq_length", None)
    if not callable(getter):
        return None
    try:
        return int(getter(layer_idx))
    except TypeError:  # transformers 4.x cache without a layer argument
        pass
    except Exception:  # pragma: no cover - unknown cache implementations
        return None
    try:
        return int(getter())
    except Exception:  # pragma: no cover - unknown cache implementations
        return None


class RicursiveForCausalLM(RicursivePreTrainedModel, GenerationMixin):
    _tied_weights_keys = []

    # transformers 5 no longer injects `GenerationMixin` into remote-code classes,
    # so inherit it explicitly above. Without it `generate` is missing entirely and
    # `PreTrainedModel.__init__` never builds `generation_config`.
    #
    # This attribute marks the native transformers-5 compatibility fixes so that
    # downstream loaders can skip their runtime monkeypatch.
    _ricursive_tf5_native_fix = RICURSIVE_TF5_COMPAT

    def __init__(self, config: RicursiveConfig):
        super().__init__(config)
        self.model = RicursiveModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Optional[int] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            cache_position=cache_position,
            **kwargs,
        )
        hidden_states = outputs[0]
        # Honour `logits_to_keep` by slicing the LM head *input*. Projecting the
        # full `[batch, seq, vocab]` tensor only to discard all but the last row
        # costs several GiB of transient memory at prefill batch sizes.
        #
        # When `labels` are supplied the full sequence must be kept: the causal-LM
        # loss shifts logits and labels by one position, so trimming both to the
        # same short window leaves the shift empty (and the loss NaN). Labeled
        # calls therefore project the full sequence to get a defined loss, then
        # trim only the *returned* logits so the output shape still honours
        # `logits_to_keep`. The generation path, which never passes labels, keeps
        # the memory win from slicing the LM head input.
        keep = (
            int(logits_to_keep)
            if logits_to_keep is not None and logits_to_keep > 0
            else 0
        )
        if keep and labels is None and hidden_states.shape[1] > keep:
            hidden_states = hidden_states[:, -keep:, :]
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
            )
            if keep:
                # Loss used the full sequence; only the returned tensor is trimmed.
                logits = logits[..., -keep:, :]

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def _update_model_kwargs_for_generation(
        self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1
    ):
        """Advance explicit positions when the generation backend leaves them unchanged."""
        position_ids = model_kwargs.get("position_ids")
        model_kwargs = super()._update_model_kwargs_for_generation(
            outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder,
            num_new_tokens=num_new_tokens,
        )
        # Transformers 4 leaves positions unchanged; Transformers 5 already extends them.
        if position_ids is not None and model_kwargs.get("position_ids") is position_ids:
            offsets = torch.arange(
                1, num_new_tokens + 1, device=position_ids.device, dtype=position_ids.dtype
            )
            model_kwargs["position_ids"] = torch.cat(
                [position_ids, position_ids[:, -1:] + offsets], dim=-1
            )
        return model_kwargs

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.LongTensor,
        past_key_values: Optional[Cache] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        position_ids: Optional[torch.LongTensor] = None,
        **kwargs,
    ):
        # transformers 5.x does not move remote-code inputs onto the model device for
        # this model, so a plain CPU `input_ids` reaches the embedding lookup and
        # raises. Its own generation helper performs this move, so mirror it here.
        device = getattr(self, "device", None)
        if device is not None and getattr(device, "type", None) != "meta":
            if input_ids is not None and input_ids.device != device:
                input_ids = input_ids.to(device)
            if attention_mask is not None and attention_mask.device != device:
                attention_mask = attention_mask.to(device)
            if inputs_embeds is not None and inputs_embeds.device != device:
                inputs_embeds = inputs_embeds.to(device)
            # transformers derives `position_ids` from the *caller's* tensors before
            # this hook runs, so a CPU prompt yields CPU positions even after the
            # move above. RoPE then multiplies them against CUDA frequencies and fails.
            if position_ids is not None and position_ids.device != device:
                position_ids = position_ids.to(device)

        past_length = _cache_seq_length(past_key_values)

        if cache_position is None:
            # With a non-empty cache and the full sequence supplied, only the trailing
            # (uncached) tokens are new. Otherwise every token in `input_ids` is new.
            if past_length > 0 and input_ids.shape[1] > 1:
                new_len = input_ids.shape[1] - past_length
                cache_position = torch.arange(
                    past_length, past_length + new_len, device=input_ids.device
                )
            else:
                cache_position = torch.arange(
                    past_length, past_length + input_ids.shape[1], device=input_ids.device
                )

        # Only slice when the cache is warm and `input_ids` still carries cached
        # tokens. Slicing unconditionally turned a `None` cache position into
        # `input_ids[:, None]` and injected a bogus axis.
        sliced = past_length > 0 and input_ids.shape[1] > cache_position.numel()
        if sliced:
            input_ids = input_ids[:, cache_position]

        if inputs_embeds is not None and cache_position.numel() and int(cache_position[0]) == 0:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        # Derive per-row positions from the padding mask. Using `cache_position` for
        # every row is only correct when nothing is padded: with left padding the real
        # tokens of a short row sit at higher indices, so a shared position vector
        # shifts its RoPE phases and the row decodes to garbage.
        if position_ids is not None:
            if sliced:
                position_ids = position_ids[:, cache_position]
        elif attention_mask is not None and attention_mask.dim() == 2:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids = position_ids.masked_fill(attention_mask == 0, 0)
            position_ids = position_ids[:, -input_ids.shape[1]:]
        else:
            position_ids = cache_position.unsqueeze(0)

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": use_cache,
                "attention_mask": attention_mask,
                "cache_position": cache_position,
            }
        )
        if use_cache is not False:
            # Generation only reads the final position. Requesting one row keeps
            # the 100k-way LM head from projecting the whole prompt.
            model_inputs["logits_to_keep"] = 1
        return model_inputs