Fix generate() and attention masking under Transformers 5.x

modeling_ricursive.py

@@ -6,6 +6,25 @@ standard Llama/Mistral:
 - 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 @@ from transformers.modeling_outputs import (
     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 @@ class RicursiveAttention(nn.Module):
         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 @@ class RicursiveModel(RicursivePreTrainedModel):
         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)
+
+        # 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.full(
-            (sequence_length, target_length),
-            fill_value=torch.finfo(dtype).min,
+        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)
-
-        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):
+        # `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)
@@ -583,20 +660,70 @@ class RicursiveForCausalLM(RicursivePreTrainedModel):
         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 inputs_embeds is not None and cache_position[0] == 0:
+        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": 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,