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