Transformers 5.x generate() and attention-mask fixes for Ripple

PATCH.diff

@@ -0,0 +1,260 @@
+--- a/modeling_ricursive.py
++++ b/modeling_ricursive.py
+@@ -6,6 +6,25 @@
+ - 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.
+ """
+ 
+ from typing import Optional, Tuple, Union
+@@ -20,9 +39,20 @@
+     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__()
+@@ -140,13 +170,37 @@
+         if self.layer_type != "sliding_attention":
+             return attention_mask
+ 
+-        key_positions = torch.arange(kv_len, device=device)
++        # 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:
++            attention_mask = attention_mask[..., -kv_len:]
++
+         if cache_position is not None:
+-            query_positions = cache_position.to(device=device)
++            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:
++                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))
+@@ -466,43 +520,66 @@
+         dtype, device = input_tensor.dtype, input_tensor.device
+         sequence_length = input_tensor.shape[1]
+ 
+-        if past_key_values is not None:
+-            target_length = past_key_values.get_max_cache_shape()
+-            if target_length is None or target_length < 0:
+-                target_length = (
+-                    cache_position[-1].item() + 1 if cache_position is not None else sequence_length
+-                )
++        if cache_position is None:
++            query_positions = torch.arange(sequence_length, device=device)
+         else:
+-            target_length = sequence_length
++            query_positions = cache_position.to(device=device).reshape(-1)
+ 
+-        causal_mask = torch.full(
+-            (sequence_length, target_length),
+-            fill_value=torch.finfo(dtype).min,
++        # 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,
+         )
+-        if sequence_length != 1:
+-            causal_mask = torch.triu(causal_mask, diagonal=target_length - sequence_length + 1)
+-
+-        causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
++        # `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
+ 
+-        if attention_mask.dim() == 2:
+-            mask_length = attention_mask.shape[-1]
+-            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+-                causal_mask.dtype
+-            ).masked_fill(attention_mask[:, None, None, :] == 0, torch.finfo(dtype).min)
+-            causal_mask = (
+-                torch.cat([padding_mask, causal_mask[:, :, :, mask_length:]], dim=-1)
+-                if mask_length < target_length
+-                else padding_mask
+-            )
+-
+-        return causal_mask
+ 
+-
+-class RicursiveForCausalLM(RicursivePreTrainedModel):
++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)
+@@ -583,20 +660,70 @@
+         inputs_embeds: Optional[torch.FloatTensor] = None,
+         cache_position: Optional[torch.LongTensor] = None,
+         use_cache: Optional[bool] = None,
++        position_ids: Optional[torch.LongTensor] = None,
+         **kwargs,
+     ):
+-        if past_key_values is not None:
+-            if input_ids.shape[1] != 1:
+-                input_ids = input_ids[:, cache_position]
++        # 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[0] == 0:
++        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": cache_position.unsqueeze(0) if cache_position is not None else None,
++                "position_ids": position_ids,
+                 "past_key_values": past_key_values,
+                 "use_cache": use_cache,
+                 "attention_mask": attention_mask,

README.md

@@ -0,0 +1,122 @@
+---
+license: apache-2.0
+library_name: transformers
+tags:
+  - ricursive
+  - ripple
+  - transformers
+  - patch
+---
+
+# Ripple-65B-Preview — Transformers 5.x compatibility patch
+
+Drop-in replacement for `modeling_ricursive.py` from
+[`ricursiveAI/Ripple-65B-Preview`](https://huggingface.co/ricursiveAI/Ripple-65B-Preview).
+
+The upstream custom code was written against the Transformers 4.x cache API.
+Under Transformers 5.x, `generate()` fails on the very first prefill, and even
+after that single line is fixed, padded batches decode incorrectly and
+generations longer than `sliding_window` crash.
+
+This repository distributes the fix only. It contains **no weights**; use it
+with the original checkpoint.
+
+## Apply
+
+Copy the patched file over the upstream one:
+
+```bash
+cp modeling_ricursive.py /path/to/Ripple-65B-Preview/modeling_ricursive.py
+```
+
+Or apply the unified diff from the root of a checkout of the model repository:
+
+```bash
+git apply PATCH.diff
+```
+
+Then load as usual:
+
+```python
+from transformers import AutoModelForCausalLM
+
+model = AutoModelForCausalLM.from_pretrained(
+    "ricursiveAI/Ripple-65B-Preview",
+    trust_remote_code=True,
+)
+```
+
+## What is broken upstream
+
+1. **Spurious axis on prefill.** `prepare_inputs_for_generation` does
+   `input_ids[:, cache_position]` whenever `input_ids.shape[1] != 1`.
+   Transformers 5 pre-creates an empty `DynamicCache` and passes
+   `cache_position=None` for the prefill call, so this becomes
+   `input_ids[:, None]`: `(batch, seq)` turns into `(batch, 1, seq)`. The
+   embedding lookup returns 4-D hidden states and the first decoder layer dies
+   on `bsz, q_len, _ = hidden_states.shape`.
+
+2. **Every decode step fully masked.** `_update_causal_mask` seeds the mask with
+   `finfo.min` and only writes the visible diagonal via `triu` when
+   `sequence_length != 1`. At decode `sequence_length == 1`, so the whole row
+   stays masked and the model emits a constant token. The function also writes
+   `1.0` (not `0.0`) for visible keys, which biases every score, and calls
+   `past_key_values.get_max_cache_shape()`, which Transformers 5 removed.
+
+3. **`position_ids` ignore padding.** They are overwritten with
+   `cache_position` for every row, which is only correct when nothing is padded.
+   With left padding, a short row's real tokens sit at higher indices, so a
+   shared position vector shifts its RoPE phases and the row decodes to garbage.
+
+4. **Sliding window uses the wrong absolute positions.**
+   `_apply_sliding_window_mask` builds `arange(kv_len)` and assumes the keys sit
+   at positions `0..kv_len-1`. Once the hybrid cache evicts from the left they
+   sit at `[cp-kv_len+1 .. cp]`, so the mask excludes progressively more valid
+   keys (e.g. window 16 at position 32 leaves 0 of 16 keys visible).
+
+5. **Mask shape crash past the window.** Transformers 5 uses a hybrid cache:
+   sliding layers hold at most `sliding_window` keys while full-attention layers
+   keep growing. The model builds one mask per forward, so the shape no longer
+   matches the sliding layer's key tensor
+   (`RuntimeError: expanded size (1024) must match existing size (1025)`).
+
+6. **`position_ids` left on CPU.** Transformers derives them from the caller's
+   tensors before the generation hook runs, so a CPU prompt yields CPU positions
+   even after the inputs are moved, and RoPE fails with `mat2 is on cpu`.
+
+7. **Missing `GenerationMixin`.** Transformers 5 no longer injects it into
+   remote-code classes, so `generate` can be missing entirely and
+   `generation_config` is never built.
+
+## What the patch changes
+
+* `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 `input_ids` / `attention_mask` / `inputs_embeds` /
+  `position_ids` onto the model device.
+* `_update_causal_mask` sizes the key axis from `cache_position`, emits `0.0`
+  for visible keys and `finfo.min` elsewhere, and never returns an all-masked
+  single-token decode row.
+* `_apply_sliding_window_mask` trims the shared mask to the keys the layer
+  holds, anchors key positions at the first key the cache still holds, and is
+  skipped when the window cannot bind.
+* `RicursiveForCausalLM` inherits `GenerationMixin` explicitly.
+
+The patch is backward compatible with Transformers 4.x.
+
+## Validation
+
+* `generate()` works under Transformers 5.22.1.
+* Cached and uncached generation are bit-identical, including past
+  `sliding_window` (the eviction path).
+* Left-padded batches decode bit-identically to the same rows generated
+  unpadded.
+* The patched implementation is bit-identical to an independently written
+  reference patch across prefill logits, decode, padded and uncached paths.
+
+## Files
+
+```
+modeling_ricursive.py   patched drop-in replacement
+PATCH.diff              unified diff against the upstream file
+```

modeling_ricursive.py

@@ -0,0 +1,733 @@
+"""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.
+"""
+
+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],
+    ) -> 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:
+            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:
+                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
+        if past_key_value is not None:
+            k, v = past_key_value.update(k, v, self.layer_idx)
+
+        # GQA: expand KV heads
+        k = repeat_kv(k, self.num_kv_groups)
+        v = repeat_kv(v, self.num_kv_groups)
+
+        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,
+        )
+
+        # 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
+
+
+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(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)
+
+
+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
+
+
+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,
+        **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]
+        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 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 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,
+            }
+        )
+        return model_inputs