| |
| |
| @@ -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, |
|
|
