File size: 12,685 Bytes
ede5883 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 | --- 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,
|