"""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. DFlash2 inference fixes (2026-11-06) ----------------------------------- * `RicursiveAttention.forward` builds an end-anchored causal mask for a multi-token verify block running against a warm cache. SDPA's implicit mask is top-left aligned, so with `q_len < kv_len` it hid valid context. Preallocated caches (`StaticCache`) return their whole capacity from `update()`, so the mask is now built from the layer's *actual* cache length and query positions; the unused tail is trimmed before attention. * `RicursiveForCausalLM.forward` honours `logits_to_keep` by slicing the LM head *input*, so the full `[batch, seq, vocab]` projection is never materialised; `prepare_inputs_for_generation` requests one row. Labeled calls keep the full sequence so the shifted causal-LM loss stays defined, then trim only the returned logits to the requested width. * `RicursiveMLP.forward` chunks the row dimension (`RIPPLE_MLP_CHUNK`, default 2048, `0` disables) to bound the prefill activation peak. Every op is row-wise, so the result is exact. Cache-alignment fixes from review (2026-11-06) ---------------------------------------------- * A mask longer than the key tensor is now aligned to the absolute positions the layer actually holds, `[total_kv - kv_len, total_kv)`, instead of to the tail of the mask. With a mask sized to the allocated cache *capacity* (`StaticCache`), the old `[-kv_len:]` slice selected unwritten slots after sliding-window eviction and silently returned wrong attention. * `RicursiveAttention.forward` forwards `cache_position` to `past_key_value.update()`. A transformers 4.x `StaticCache` otherwise falls back to `arange(q_len)` and the first decode step overwrites the prompt instead of appending after it. """ import os 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], total_kv: Optional[int] = None, ) -> 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: mask_len = attention_mask.shape[-1] if total_kv is not None: # The retained keys sit at `[total_kv - kv_len, total_kv)`. start = max(0, min(total_kv - kv_len, mask_len - kv_len)) attention_mask = attention_mask[..., start : start + kv_len] else: 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: # SDPA's implicit causal mask is top-left aligned. It only # matches this query block when it is the whole sequence # (`q_len == kv_len`) or a single decode position # (`q_len == 1`). A multi-token block verified against a warm # cache sits at the *end* of the key span, so fall through and # build the end-anchored mask below. if q_len == 1 or q_len == kv_len: 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 total_kv = None if past_key_value is not None: # Forward the absolute positions of the tokens being written. A # transformers 4.x `StaticCache` falls back to `arange(q_len)` when # `cache_position` is absent, so the first decode step would write # at slots 0..q_len-1 and overwrite the prompt instead of appending # after it. `DynamicCache`/5.x caches ignore the extra kwarg. cache_kwargs = ( {"cache_position": cache_position} if cache_position is not None else None ) if cache_kwargs is not None: try: k, v = past_key_value.update(k, v, self.layer_idx, cache_kwargs) except TypeError: # cache without a `cache_kwargs` parameter k, v = past_key_value.update(k, v, self.layer_idx) else: k, v = past_key_value.update(k, v, self.layer_idx) # Number of keys this layer has actually written. A preallocated # (StaticCache) layer returns its full capacity from `update()`, # including slots that have not been filled yet; an evicting # (sliding-window) layer reports more than the trailing window it # returns. The mask must use the real positions, not the tensor size. total_kv = _cache_layer_seq_length(past_key_value, self.layer_idx) # GQA: expand KV heads k = repeat_kv(k, self.num_kv_groups) v = repeat_kv(v, self.num_kv_groups) # Drop the unwritten tail of a preallocated cache. Without this a # six-token prefill into a sixteen-slot StaticCache looks like a # ten-token prefix, and the end-anchored verify mask below would let # each query attend slots belonging to a different absolute position. if total_kv is not None and 0 < total_kv < k.shape[2]: k = k[:, :, :total_kv] v = v[:, :, :total_kv] # Align an incoming mask that is longer than the keys this layer holds. # The shared mask is built over absolute key positions # `0..target_length-1`; keep the slice that matches the keys present. # # The keys returned by `update()` occupy the absolute positions # `[total_kv - kv_len, total_kv)`. Anchoring on the tensor tail # (`[-kv_len:]`) is only correct when the mask length equals # `total_kv`; with a mask sized to the *allocated cache capacity* it # selects the unwritten tail instead of the retained window and # silently produces wrong attention. if attention_mask is not None and attention_mask.shape[-1] > k.shape[2]: kv_len_now = k.shape[2] mask_len = attention_mask.shape[-1] if total_kv is not None: start = max(0, min(total_kv - kv_len_now, mask_len - kv_len_now)) attention_mask = attention_mask[..., start : start + kv_len_now] else: # Unknown cache: the keys are the trailing window. attention_mask = attention_mask[..., -kv_len_now:] 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, total_kv=total_kv, ) if attention_mask is None and q.shape[2] > 1 and k.shape[2] > q.shape[2]: # Multi-token verify block against a warm cache. SDPA's implicit # causal mask is top-left aligned, so when `q_len < kv_len` it hides # every key before the block. Build the mask from the actual query # positions (and the absolute positions of the returned keys) rather # than assuming the key tensor starts at position 0. if cache_position is not None and cache_position.numel() == q.shape[2]: query_positions = cache_position.to( device=q.device, dtype=torch.long ).reshape(-1) else: query_positions = torch.arange( k.shape[2] - q.shape[2], k.shape[2], device=q.device ) if total_kv is not None and total_kv > k.shape[2]: # Evicting cache: the returned keys are the trailing window that # ends at the newest query, not positions 0..kv_len-1. last_query = int(query_positions[-1].item()) key_positions = torch.arange( last_query - k.shape[2] + 1, last_query + 1, device=q.device ) else: key_positions = torch.arange(k.shape[2], device=q.device) attention_mask = ( key_positions[None, :] <= query_positions[:, None] )[None, None] # 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 def _mlp_chunk_size() -> int: """Row budget for the chunked MLP forward. `0` (or negative) disables it.""" try: return int(os.environ.get("RIPPLE_MLP_CHUNK", "2048")) except ValueError: return 2048 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_rows(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) def forward(self, x: torch.Tensor) -> torch.Tensor: # Ripple's MLP expands hidden_size -> intermediate_size (5120 -> 40960 # for the 65B checkpoint), so a long prefill materialises `gate`, `up` # and `silu(gate) * up` at once. Every op here is row-wise, so splitting # the row dimension is exact: each row still uses the same kernels with # the same reduction order, and the pieces are concatenated back. Only # the transient peak shrinks. `RIPPLE_MLP_CHUNK` overrides the # row budget; `0` disables chunking. chunk = _mlp_chunk_size() if chunk <= 0 or not torch.is_tensor(x) or x.dim() not in (2, 3): return self._forward_rows(x) shape = x.shape hidden = shape[-1] rows = 1 for size in shape[:-1]: rows *= size if rows <= chunk: return self._forward_rows(x) flat = x.reshape(rows, hidden) pieces = [ self._forward_rows(flat[start : start + chunk]) for start in range(0, rows, chunk) ] out = torch.cat(pieces, dim=0) if len(pieces) > 1 else pieces[0] return out.reshape(*shape[:-1], out.shape[-1]) 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 def _cache_layer_seq_length( past_key_values: Optional[Cache], layer_idx: int ) -> Optional[int]: """Actual number of keys written for `layer_idx`, or `None` if unknown. `Cache.get_seq_length(layer_idx)` counts *written* tokens. For a preallocated static layer that is smaller than the tensor `update()` returns until the cache fills; for an evicting sliding-window layer it is larger, because the layer only returns a trailing window. """ if past_key_values is None: return None getter = getattr(past_key_values, "get_seq_length", None) if not callable(getter): return None try: return int(getter(layer_idx)) except TypeError: # transformers 4.x cache without a layer argument pass except Exception: # pragma: no cover - unknown cache implementations return None try: return int(getter()) except Exception: # pragma: no cover - unknown cache implementations return None 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, logits_to_keep: Optional[int] = 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] # Honour `logits_to_keep` by slicing the LM head *input*. Projecting the # full `[batch, seq, vocab]` tensor only to discard all but the last row # costs several GiB of transient memory at prefill batch sizes. # # When `labels` are supplied the full sequence must be kept: the causal-LM # loss shifts logits and labels by one position, so trimming both to the # same short window leaves the shift empty (and the loss NaN). Labeled # calls therefore project the full sequence to get a defined loss, then # trim only the *returned* logits so the output shape still honours # `logits_to_keep`. The generation path, which never passes labels, keeps # the memory win from slicing the LM head input. keep = ( int(logits_to_keep) if logits_to_keep is not None and logits_to_keep > 0 else 0 ) if keep and labels is None and hidden_states.shape[1] > keep: hidden_states = hidden_states[:, -keep:, :] 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 keep: # Loss used the full sequence; only the returned tensor is trimmed. logits = logits[..., -keep:, :] 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 _update_model_kwargs_for_generation( self, outputs, model_kwargs, is_encoder_decoder=False, num_new_tokens=1 ): """Advance explicit positions when the generation backend leaves them unchanged.""" position_ids = model_kwargs.get("position_ids") model_kwargs = super()._update_model_kwargs_for_generation( outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder, num_new_tokens=num_new_tokens, ) # Transformers 4 leaves positions unchanged; Transformers 5 already extends them. if position_ids is not None and model_kwargs.get("position_ids") is position_ids: offsets = torch.arange( 1, num_new_tokens + 1, device=position_ids.device, dtype=position_ids.dtype ) model_kwargs["position_ids"] = torch.cat( [position_ids, position_ids[:, -1:] + offsets], dim=-1 ) return model_kwargs 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, } ) if use_cache is not False: # Generation only reads the final position. Requesting one row keeps # the 100k-way LM head from projecting the whole prompt. model_inputs["logits_to_keep"] = 1 return model_inputs