# Copyright 2024 IQuestLoopCoder Authors # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. """ IQuestLoopCoder Model Implementation for HuggingFace. Loop model passes hidden states through the decoder multiple times: - Loop 1: Standard attention, stores K1, V1 for each layer - Loop 2+: Mixed attention with gated combination of: - A: Full attention with Loop1's KV (global context) - B: Sliding window attention with Loop2's KV (local, high-precision context) - Gate g = sigmoid(linear(Q)), per-head - Output = g * A + (1 - g) * B """ import math from typing import Any, List, Optional, Tuple, Union import torch import torch.nn.functional as F import torch.utils.checkpoint from torch import nn from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.modeling_utils import PreTrainedModel from transformers.generation.utils import GenerationMixin from transformers.utils import ( add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings, ) from .configuration_iquestloopcoder import IQuestLoopCoderConfig logger = logging.get_logger(__name__) _CONFIG_FOR_DOC = "IQuestLoopCoderConfig" class IQuestLoopCoderCache(Cache): """Cache implementation for IQuestLoopCoder that manages shared and local KV caches. - shared_key_cache/shared_value_cache: Stores KV from Loop 1 (global context) - local_key_cache/local_value_cache: Stores KV from Loop 2+ (local window, only window_size tokens) """ def __init__(self, window_size: int, num_layers: int): # We intentionally don't call super().__init__ because the parent assumes static cache sizes. self.window_size = window_size self.num_layers = num_layers # Shared cache: stores Loop 1 KV (global context) self.shared_key_cache: List[Optional[torch.Tensor]] = [None] * num_layers self.shared_value_cache: List[Optional[torch.Tensor]] = [None] * num_layers # Local cache: stores Loop 2+ KV (sliding window, only window_size tokens) self.local_key_cache: List[Optional[torch.Tensor]] = [None] * num_layers self.local_value_cache: List[Optional[torch.Tensor]] = [None] * num_layers self.layers: List[Any] = [] # attribute expected by HF Cache utilities self._seen_tokens = 0 def update_shared( self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Update shared cache (Loop 1 KV).""" if layer_idx < 0 or layer_idx >= self.num_layers: raise ValueError(f"layer_idx must be in [0, {self.num_layers}), got {layer_idx}") cached_key = self.shared_key_cache[layer_idx] cached_value = self.shared_value_cache[layer_idx] if cached_key is None: self.shared_key_cache[layer_idx] = key_states self.shared_value_cache[layer_idx] = value_states else: if ( key_states.shape[0] != cached_key.shape[0] or key_states.shape[1] != cached_key.shape[1] or key_states.shape[3] != cached_key.shape[3] ): raise ValueError( "Cached and incoming key/value tensors must match on batch, head, and head_dim dimensions." ) assert cached_value is not None self.shared_key_cache[layer_idx] = torch.cat([cached_key, key_states], dim=2) self.shared_value_cache[layer_idx] = torch.cat([cached_value, value_states], dim=2) result_key = self.shared_key_cache[layer_idx] result_value = self.shared_value_cache[layer_idx] assert result_key is not None and result_value is not None # Track sequence length self._seen_tokens = result_key.shape[2] return result_key, result_value def update_local( self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Update local cache (Loop 2+ KV) with sliding window management. If the cache is full (window_size tokens), remove the oldest token and add the new one. """ if layer_idx < 0 or layer_idx >= self.num_layers: raise ValueError(f"layer_idx must be in [0, {self.num_layers}), got {layer_idx}") cached_key = self.local_key_cache[layer_idx] cached_value = self.local_value_cache[layer_idx] if cached_key is None: # First token in local cache self.local_key_cache[layer_idx] = key_states self.local_value_cache[layer_idx] = value_states else: if ( key_states.shape[0] != cached_key.shape[0] or key_states.shape[1] != cached_key.shape[1] or key_states.shape[3] != cached_key.shape[3] ): raise ValueError( "Cached and incoming key/value tensors must match on batch, head, and head_dim dimensions." ) assert cached_value is not None # Check if we need to remove the oldest token current_len = cached_key.shape[2] if current_len >= self.window_size: # Remove the first token (oldest) and add the new one self.local_key_cache[layer_idx] = torch.cat([cached_key[:, :, 1:, :], key_states], dim=2) self.local_value_cache[layer_idx] = torch.cat([cached_value[:, :, 1:, :], value_states], dim=2) else: # Just append self.local_key_cache[layer_idx] = torch.cat([cached_key, key_states], dim=2) self.local_value_cache[layer_idx] = torch.cat([cached_value, value_states], dim=2) result_key = self.local_key_cache[layer_idx] result_value = self.local_value_cache[layer_idx] assert result_key is not None and result_value is not None return result_key, result_value def get_shared(self, layer_idx: int) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """Get shared cache for a layer.""" if layer_idx < 0 or layer_idx >= self.num_layers: return None, None return self.shared_key_cache[layer_idx], self.shared_value_cache[layer_idx] def get_local(self, layer_idx: int) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """Get local cache for a layer.""" if layer_idx < 0 or layer_idx >= self.num_layers: return None, None return self.local_key_cache[layer_idx], self.local_value_cache[layer_idx] def update( self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, cache_kwargs: Optional[dict] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Default update method (for compatibility, updates shared cache).""" return self.update_shared(key_states, value_states, layer_idx, cache_kwargs) def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: """Get sequence length from shared cache.""" if layer_idx is None: layer_idx = 0 if layer_idx < 0 or layer_idx >= len(self.shared_key_cache): return 0 cached = self.shared_key_cache[layer_idx] if cached is None: return 0 return cached.shape[2] def get_max_length(self) -> Optional[int]: return None def get_usable_length( self, new_seq_length: int, layer_idx: Optional[int] = 0 ) -> int: return self.get_seq_length(layer_idx) def reorder_cache(self, beam_idx: torch.LongTensor) -> None: """Reorder cache for beam search.""" for layer_idx in range(self.num_layers): if self.shared_key_cache[layer_idx] is not None: device = self.shared_key_cache[layer_idx].device self.shared_key_cache[layer_idx] = self.shared_key_cache[layer_idx].index_select(0, beam_idx.to(device)) self.shared_value_cache[layer_idx] = self.shared_value_cache[layer_idx].index_select(0, beam_idx.to(device)) if self.local_key_cache[layer_idx] is not None: device = self.local_key_cache[layer_idx].device self.local_key_cache[layer_idx] = self.local_key_cache[layer_idx].index_select(0, beam_idx.to(device)) self.local_value_cache[layer_idx] = self.local_value_cache[layer_idx].index_select(0, beam_idx.to(device)) @property def is_compileable(self) -> bool: return False def clear(self) -> None: """Clear all caches.""" logger.debug("Clearing IQuestLoopCoderCache") self.shared_key_cache = [None] * self.num_layers self.shared_value_cache = [None] * self.num_layers self.local_key_cache = [None] * self.num_layers self.local_value_cache = [None] * self.num_layers self._seen_tokens = 0 class IQuestLoopCoderRMSNorm(nn.Module): """RMS Normalization layer.""" def __init__(self, hidden_size, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) class IQuestLoopCoderRotaryEmbedding(nn.Module): """Rotary Position Embedding (RoPE).""" def __init__(self, dim, max_position_embeddings=8192, base=500000.0, device=None, scaling_factor=1.0): super().__init__() self.scaling_factor = scaling_factor self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self.max_seq_len_cached = max_position_embeddings @torch.no_grad() def forward(self, x, position_ids): # x: [batch_size, num_heads, seq_len, head_dim] inv_freq_expanded = self.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.float() @ position_ids_expanded.float()).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): """Rotates half the hidden dims of the input.""" 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, k, cos, sin, position_ids=None, unsqueeze_dim=1): """Applies Rotary Position Embedding to the query and key tensors.""" cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """Expand KV heads to match query heads for GQA.""" batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class IQuestLoopCoderMLP(nn.Module): """MLP with SwiGLU activation.""" def __init__(self, config): super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class LoopGateProjection(nn.Module): """Gate projection for mixed attention in Loop 2+. Computes: g = sigmoid(linear(Q)) for each head independently. This gate determines how much to use Loop1's KV (global) vs current loop's KV (local). """ def __init__(self, num_heads: int, head_dim: int): super().__init__() self.num_heads = num_heads self.head_dim = head_dim # Each head has its own gate: Linear(head_dim -> 1) per head # Implemented as [num_heads, head_dim] weight + [num_heads] bias self.weight = nn.Parameter(torch.zeros(num_heads, head_dim)) self.bias = nn.Parameter(torch.zeros(num_heads)) def forward(self, query: torch.Tensor) -> torch.Tensor: """Compute gate values from query tensor. Args: query: [batch, num_heads, seq_len, head_dim] Returns: gate: [batch, num_heads, seq_len, 1] """ # query: [batch, num_heads, seq_len, head_dim] # weight: [num_heads, head_dim] # For each head h: gate_h = query[:, h, :, :] @ weight[h, :].T + bias[h] # Using einsum: gate = einsum('bhsd,hd->bhs', query, weight) + bias gate_logits = torch.einsum('bhsd,hd->bhs', query, self.weight) # [batch, num_heads, seq_len] gate_logits = gate_logits + self.bias[None, :, None] # broadcast bias gate = torch.sigmoid(gate_logits) return gate.unsqueeze(-1) # [batch, num_heads, seq_len, 1] class IQuestLoopCoderAttention(nn.Module): """Multi-head attention with GQA support.""" def __init__(self, config: IQuestLoopCoderConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = config.head_dim self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.attention_dropout = config.attention_dropout self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) self.rotary_emb = IQuestLoopCoderRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, ) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = self.rotary_emb(value_states, position_ids) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) # Repeat KV for GQA key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.o_proj(attn_output) return attn_output, attn_weights if output_attentions else None, past_key_value def forward_with_external_kv( self, hidden_states: torch.Tensor, external_key: torch.Tensor, external_value: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, sliding_window: Optional[int] = None, ) -> torch.Tensor: """Forward pass using external K, V (for Loop 2+ mixed attention). Args: hidden_states: Input for computing Q external_key: Pre-computed K (already with RoPE applied) external_value: Pre-computed V attention_mask: Causal attention mask position_ids: Position IDs sliding_window: If set, apply sliding window attention Returns: Attention output [batch, seq_len, num_heads, head_dim] """ bsz, q_len, _ = hidden_states.size() # Compute Q from current hidden states query_states = self.q_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) # Apply RoPE to Q cos, sin = self.rotary_emb(query_states, position_ids) query_states = (query_states * cos.unsqueeze(1)) + (rotate_half(query_states) * sin.unsqueeze(1)) # Use external K, V (already have RoPE for K) key_states = external_key value_states = external_value # Repeat KV for GQA key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) # Compute attention attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) # Apply attention mask (causal) if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask # Apply sliding window mask if needed if sliding_window is not None and q_len > sliding_window: # Create sliding window mask # For each position i, can only attend to [i-window+1, i] seq_len = key_states.shape[2] row_idx = torch.arange(q_len, device=query_states.device).unsqueeze(1) col_idx = torch.arange(seq_len, device=query_states.device).unsqueeze(0) window_mask = (col_idx > row_idx) | (col_idx < row_idx - sliding_window + 1) window_mask = window_mask.unsqueeze(0).unsqueeze(0) # [1, 1, q_len, seq_len] attn_weights = attn_weights.masked_fill(window_mask, float('-inf')) attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states) # Don't apply o_proj here - return raw attention output attn_output = attn_output.transpose(1, 2).contiguous() return attn_output # [batch, seq_len, num_heads, head_dim] def get_qkv( self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Get Q, K, V tensors with RoPE applied. Returns: query: [batch, num_heads, seq_len, head_dim] key: [batch, num_kv_heads, seq_len, head_dim] value: [batch, num_kv_heads, seq_len, head_dim] """ bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = self.rotary_emb(value_states, position_ids) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) return query_states, key_states, value_states def forward_decode_loop1( self, hidden_states: torch.Tensor, past_shared_key: Optional[torch.Tensor], past_shared_value: Optional[torch.Tensor], attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Forward pass for Loop 1 in decode stage. Args: hidden_states: Current hidden states [batch, 1, hidden_size] past_shared_key: Past shared keys from cache [batch, num_kv_heads, past_len, head_dim] past_shared_value: Past shared values from cache [batch, num_kv_heads, past_len, head_dim] attention_mask: Causal attention mask position_ids: Position IDs cache_position: Cache position Returns: output: Attention output [batch, 1, hidden_size] k1: Current key [batch, num_kv_heads, 1, head_dim] (only current token) v1: Current value [batch, num_kv_heads, 1, head_dim] (only current token) """ bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) cos, sin = self.rotary_emb(value_states, position_ids) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) # Store current token's k1, v1 for return (before concatenation) k1_current = key_states # [batch, num_kv_heads, 1, head_dim] v1_current = value_states # [batch, num_kv_heads, 1, head_dim] # Concatenate with past shared KV cache for attention computation if past_shared_key is not None and past_shared_value is not None: key_states = torch.cat([past_shared_key, key_states], dim=2) value_states = torch.cat([past_shared_value, value_states], dim=2) # Repeat KV for GQA key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) attn_output = self.o_proj(attn_output) return attn_output, k1_current, v1_current def forward_decode_loop2( self, hidden_states: torch.Tensor, k1: torch.Tensor, v1: torch.Tensor, past_shared_key: Optional[torch.Tensor], past_shared_value: Optional[torch.Tensor], past_local_key: Optional[torch.Tensor], past_local_value: Optional[torch.Tensor], gate_proj: LoopGateProjection, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, loop_window_size: int = 64, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Forward pass for Loop 2 in decode stage with mixed attention. Args: hidden_states: Current hidden states [batch, 1, hidden_size] k1: Key from Loop 1 (current token) [batch, num_kv_heads, 1, head_dim] v1: Value from Loop 1 (current token) [batch, num_kv_heads, 1, head_dim] past_shared_key: Past shared keys from cache [batch, num_kv_heads, past_len, head_dim] past_shared_value: Past shared values from cache [batch, num_kv_heads, past_len, head_dim] past_local_key: Past local keys from cache [batch, num_kv_heads, window_len, head_dim] past_local_value: Past local values from cache [batch, num_kv_heads, window_len, head_dim] gate_proj: Gate projection module attention_mask: Causal attention mask position_ids: Position IDs loop_window_size: Window size for sliding window attention Returns: output: Attention output [batch, 1, hidden_size] k2: Current key [batch, num_kv_heads, 1, head_dim] v2: Current value [batch, num_kv_heads, 1, head_dim] """ bsz, q_len, _ = hidden_states.size() # Get Q2, K2, V2 for current loop q2, k2, v2 = self.get_qkv(hidden_states, position_ids) # Compute gate: g = sigmoid(linear(Q2)) gate = gate_proj(q2) # [batch, num_heads, 1, 1] # For attention A: concatenate past shared KV with current k1, v1 (full global context) if past_shared_key is not None and past_shared_value is not None: k1_full = torch.cat([past_shared_key, k1], dim=2) v1_full = torch.cat([past_shared_value, v1], dim=2) else: k1_full = k1 v1_full = v1 # For attention B: concatenate past local KV with current k2, v2 (sliding window) if past_local_key is not None and past_local_value is not None: k2_full = torch.cat([past_local_key, k2], dim=2) v2_full = torch.cat([past_local_value, v2], dim=2) else: k2_full = k2 v2_full = v2 # Repeat KV for GQA k1_expanded = repeat_kv(k1_full, self.num_key_value_groups) v1_expanded = repeat_kv(v1_full, self.num_key_value_groups) k2_expanded = repeat_kv(k2_full, self.num_key_value_groups) v2_expanded = repeat_kv(v2_full, self.num_key_value_groups) # Attention A: Q2 @ K1_full, V1_full (global, full sequence) head_dim = q2.shape[-1] attn_weights_A = torch.matmul(q2, k1_expanded.transpose(2, 3)) / math.sqrt(head_dim) if attention_mask is not None: causal_mask = attention_mask[:, :, :, : k1_expanded.shape[-2]] attn_weights_A = attn_weights_A + causal_mask attn_weights_A = nn.functional.softmax(attn_weights_A, dim=-1, dtype=torch.float32).to(q2.dtype) attn_A = torch.matmul(attn_weights_A, v1_expanded) # Attention B: Q2 @ K2_full, V2_full (local sliding window) attn_weights_B = torch.matmul(q2, k2_expanded.transpose(2, 3)) / math.sqrt(head_dim) if attention_mask is not None: causal_mask = attention_mask[:, :, :, : k2_expanded.shape[-2]] attn_weights_B = attn_weights_B + causal_mask # Apply sliding window mask q_len_attn = q2.shape[2] k_len_attn = k2_expanded.shape[2] if q_len_attn <= loop_window_size: # If sequence fits in window, use standard attention attn_weights_B = nn.functional.softmax(attn_weights_B, dim=-1, dtype=torch.float32).to(q2.dtype) else: # Apply sliding window mask row_idx = torch.arange(q_len_attn, device=q2.device).unsqueeze(1) col_idx = torch.arange(k_len_attn, device=q2.device).unsqueeze(0) window_mask = (col_idx > row_idx) | (col_idx < row_idx - loop_window_size + 1) window_mask = window_mask.unsqueeze(0).unsqueeze(0) attn_weights_B = attn_weights_B.masked_fill(window_mask, float('-inf')) attn_weights_B = nn.functional.softmax(attn_weights_B, dim=-1, dtype=torch.float32).to(q2.dtype) attn_B = torch.matmul(attn_weights_B, v2_expanded) # Mixed attention: gate * A + (1 - gate) * B mixed_attn = gate * attn_A + (1 - gate) * attn_B # Reshape and apply output projection bsz, num_heads, seq_len, head_dim = mixed_attn.shape mixed_attn = mixed_attn.transpose(1, 2).contiguous().reshape(bsz, seq_len, -1) attn_output = self.o_proj(mixed_attn) return attn_output, k2, v2 class IQuestLoopCoderDecoderLayer(nn.Module): """Transformer decoder layer.""" def __init__(self, config: IQuestLoopCoderConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = IQuestLoopCoderAttention(config=config, layer_idx=layer_idx) self.mlp = IQuestLoopCoderMLP(config) self.input_layernorm = IQuestLoopCoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = IQuestLoopCoderRMSNorm(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.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, 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, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) return outputs def forward_loop2_mixed( self, hidden_states: torch.Tensor, k1: torch.Tensor, v1: torch.Tensor, gate_proj: LoopGateProjection, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, loop_window_size: int = 64, ) -> Tuple[torch.Tensor, float]: """Forward pass for Loop 2+ with mixed attention. Args: hidden_states: Current hidden states k1: Key from Loop 1 [batch, num_kv_heads, seq_len, head_dim] v1: Value from Loop 1 [batch, num_kv_heads, seq_len, head_dim] gate_proj: Gate projection module for this layer attention_mask: Causal attention mask position_ids: Position IDs loop_window_size: Window size for sliding window attention Returns: output hidden states, gate mean value """ residual = hidden_states hidden_states_normed = self.input_layernorm(hidden_states) # Get Q2, K2, V2 for current loop q2, k2, v2 = self.self_attn.get_qkv(hidden_states_normed, position_ids) # Compute gate: g = sigmoid(linear(Q2)) # q2: [batch, num_heads, seq_len, head_dim] gate = gate_proj(q2) # [batch, num_heads, seq_len, 1] gate_mean = gate.detach().mean().item() # Repeat K1, V1 for GQA k1_expanded = repeat_kv(k1, self.self_attn.num_key_value_groups) v1_expanded = repeat_kv(v1, self.self_attn.num_key_value_groups) k2_expanded = repeat_kv(k2, self.self_attn.num_key_value_groups) v2_expanded = repeat_kv(v2, self.self_attn.num_key_value_groups) # Attention A: Q2 @ K1, V1 (global, full sequence) attn_A = self._compute_attention(q2, k1_expanded, v1_expanded, attention_mask) # Attention B: Q2 @ K2, V2 (local sliding window) attn_B = self._compute_attention_with_window(q2, k2_expanded, v2_expanded, attention_mask, loop_window_size) # Mixed attention: gate * A + (1 - gate) * B # attn_A, attn_B: [batch, num_heads, seq_len, head_dim] mixed_attn = gate * attn_A + (1 - gate) * attn_B # Reshape and apply output projection bsz, num_heads, seq_len, head_dim = mixed_attn.shape mixed_attn = mixed_attn.transpose(1, 2).contiguous().reshape(bsz, seq_len, -1) hidden_states = self.self_attn.o_proj(mixed_attn) hidden_states = residual + hidden_states # MLP residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states, gate_mean def _compute_attention( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], ) -> torch.Tensor: """Standard attention computation.""" head_dim = query.shape[-1] attn_weights = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(head_dim) if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key.shape[-2]] attn_weights = attn_weights + causal_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_output = torch.matmul(attn_weights, value) return attn_output def _compute_attention_with_window( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: Optional[torch.Tensor], window_size: int, ) -> torch.Tensor: """Attention with sliding window.""" q_len = query.shape[2] k_len = key.shape[2] head_dim = query.shape[-1] # If sequence fits in window, use standard attention if q_len <= window_size: return self._compute_attention(query, key, value, attention_mask) attn_weights = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(head_dim) # Apply causal mask if attention_mask is not None: causal_mask = attention_mask[:, :, :, : key.shape[-2]] attn_weights = attn_weights + causal_mask # Apply sliding window mask row_idx = torch.arange(q_len, device=query.device).unsqueeze(1) col_idx = torch.arange(k_len, device=query.device).unsqueeze(0) # Can only attend to positions in [i - window_size + 1, i] window_mask = (col_idx > row_idx) | (col_idx < row_idx - window_size + 1) window_mask = window_mask.unsqueeze(0).unsqueeze(0) attn_weights = attn_weights.masked_fill(window_mask, float('-inf')) attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_output = torch.matmul(attn_weights, value) return attn_output class IQuestLoopCoderPreTrainedModel(PreTrainedModel): """Base class for IQuestLoopCoder models.""" config_class = IQuestLoopCoderConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["IQuestLoopCoderDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_cache_class = True _supports_static_cache = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() class IQuestLoopCoderModel(IQuestLoopCoderPreTrainedModel): """IQuestLoopCoder Transformer decoder model.""" def __init__(self, config: IQuestLoopCoderConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList([ IQuestLoopCoderDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ]) self.norm = IQuestLoopCoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps) # Gate projections for Loop 2+ (one per layer) self.gate_projections = nn.ModuleList([ LoopGateProjection(config.num_attention_heads, config.head_dim) for _ in range(config.num_hidden_layers) ]) # Loop configuration self.loop_num = config.loop_num self.loop_window_size = config.loop_window_size self.gradient_checkpointing = False self.post_init() def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def forward( self, input_ids: 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, ) -> 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) seq_length = inputs_embeds.shape[1] # Determine which forward path to use: # 1. If past_key_values exists and seq_length == 1: autoregressive generation step # -> Use standard attention with KV cache (no loop needed for single token) # 2. Otherwise (prefill or training): use loop mechanism is_generation_step = past_key_values is not None and seq_length == 1 # import pdb; pdb.set_trace() if is_generation_step: # Autoregressive generation: single token, use KV cache return self._forward_with_cache( inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, ) # Prefill or training: use loop mechanism return self._forward_loop( inputs_embeds=inputs_embeds, attention_mask=attention_mask, position_ids=position_ids, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, use_cache=use_cache, cache_position=cache_position, ) def _forward_loop( self, inputs_embeds: torch.Tensor, attention_mask: Optional[torch.Tensor], position_ids: Optional[torch.LongTensor], output_attentions: bool, output_hidden_states: bool, return_dict: bool, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: """Forward with loop mechanism (for training and prefill). This implements the Loop mechanism: - Loop 1: Standard attention, stores K1, V1 for each layer - Loop 2+: Mixed attention with gated combination of global (K1,V1) and local (K2,V2) """ batch_size, seq_length, _ = inputs_embeds.shape if position_ids is None: device = inputs_embeds.device position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0) if cache_position is None: cache_position = torch.arange(seq_length, device=inputs_embeds.device) # Create causal mask causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, None, output_attentions) hidden_states = inputs_embeds all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None # For KV cache during prefill - use IQuestLoopCoderCache # In prefill, past_key_values should be None, so we create a new cache if use_cache: next_decoder_cache = IQuestLoopCoderCache(self.loop_window_size, len(self.layers)) else: next_decoder_cache = None # ============ Loop 1: Standard forward, store K1, V1 in shared cache ============ for layer_idx, decoder_layer in enumerate(self.layers): if output_hidden_states: all_hidden_states += (hidden_states,) # Get K1, V1 before standard forward (from original hidden_states, after layernorm) hidden_states_normed = decoder_layer.input_layernorm(hidden_states) q1, k1, v1 = decoder_layer.self_attn.get_qkv(hidden_states_normed, position_ids) # Store K1, V1 in shared cache if use_cache: next_decoder_cache.update_shared(k1, v1, layer_idx) # Standard forward layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=None, output_attentions=output_attentions, use_cache=False, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) # ============ Loop 2 to loop_num: Mixed attention, store in local cache ============ for loop_idx in range(2, self.loop_num + 1): for layer_idx, decoder_layer in enumerate(self.layers): # Get K1, V1 from shared cache k1, v1 = next_decoder_cache.get_shared(layer_idx) if use_cache else (None, None) if k1 is None or v1 is None: # Fallback: compute K1, V1 if not in cache (shouldn't happen in prefill) hidden_states_normed = decoder_layer.input_layernorm(hidden_states) _, k1, v1 = decoder_layer.self_attn.get_qkv(hidden_states_normed, position_ids) gate_proj = self.gate_projections[layer_idx] hidden_states, gate_mean = decoder_layer.forward_loop2_mixed( hidden_states, k1=k1, v1=v1, gate_proj=gate_proj, attention_mask=causal_mask, position_ids=position_ids, loop_window_size=self.loop_window_size, ) # Store Loop 2+ KV in local cache (only for loop_idx == 2) if use_cache and loop_idx == 2: hidden_states_normed = decoder_layer.input_layernorm(hidden_states) _, k2, v2 = decoder_layer.self_attn.get_qkv(hidden_states_normed, position_ids) next_decoder_cache.update_local(k2, v2, layer_idx) hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) if not return_dict: return tuple(v for v in [hidden_states, next_decoder_cache, all_hidden_states, all_self_attns] if v is not None) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_decoder_cache, hidden_states=all_hidden_states, attentions=all_self_attns, ) def _forward_with_cache( self, inputs_embeds: torch.Tensor, attention_mask: Optional[torch.Tensor], position_ids: Optional[torch.LongTensor], past_key_values: Optional[Cache], use_cache: bool, output_attentions: bool, output_hidden_states: bool, return_dict: bool, cache_position: Optional[torch.LongTensor], ) -> Union[Tuple, BaseModelOutputWithPast]: """Forward with KV cache using loop mechanism (for inference generation). Loop 1: Standard attention, uses shared KV cache (previous tokens + current token) Loop 2+: Mixed attention, uses local KV cache (sliding window) """ batch_size, seq_length, _ = inputs_embeds.shape 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 + seq_length, device=inputs_embeds.device) if position_ids is None: position_ids = cache_position.unsqueeze(0) causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions) # Ensure we're using IQuestLoopCoderCache if use_cache: if not isinstance(past_key_values, IQuestLoopCoderCache): # Convert to IQuestLoopCoderCache if needed next_decoder_cache = IQuestLoopCoderCache(self.loop_window_size, len(self.layers)) # Copy existing cache if possible if past_key_values is not None: for layer_idx in range(len(self.layers)): try: past_k = past_key_values.key_cache[layer_idx] if hasattr(past_key_values, 'key_cache') else None past_v = past_key_values.value_cache[layer_idx] if hasattr(past_key_values, 'value_cache') else None if past_k is not None and past_v is not None: next_decoder_cache.update_shared(past_k, past_v, layer_idx) except: pass else: next_decoder_cache = past_key_values else: next_decoder_cache = None hidden_states = inputs_embeds all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None # ============ Loop 1: Standard attention, store in shared cache ============ for layer_idx, decoder_layer in enumerate(self.layers): if output_hidden_states: all_hidden_states += (hidden_states,) # Get past shared KV cache past_shared_key, past_shared_value = None, None if next_decoder_cache is not None: past_shared_key, past_shared_value = next_decoder_cache.get_shared(layer_idx) # Forward Loop 1 attn_output, k1, v1 = decoder_layer.self_attn.forward_decode_loop1( hidden_states=decoder_layer.input_layernorm(hidden_states), past_shared_key=past_shared_key, past_shared_value=past_shared_value, attention_mask=causal_mask, position_ids=position_ids, cache_position=cache_position, ) # Update shared cache with current token's Loop 1 KV if use_cache: next_decoder_cache.update_shared(k1, v1, layer_idx) hidden_states = hidden_states + attn_output # MLP residual = hidden_states hidden_states = decoder_layer.post_attention_layernorm(hidden_states) hidden_states = decoder_layer.mlp(hidden_states) hidden_states = residual + hidden_states if output_attentions: all_self_attns += (None,) # We don't return attention weights in decode loop # ============ Loop 2 to loop_num: Mixed attention, store in local cache ============ # Store k1, v1 from Loop 1 for use in Loop 2+ loop1_kv = [] for layer_idx in range(len(self.layers)): if next_decoder_cache is not None: k1_full, v1_full = next_decoder_cache.get_shared(layer_idx) if k1_full is not None and v1_full is not None: # Get only the last token (current token) loop1_kv.append((k1_full[:, :, -1:, :], v1_full[:, :, -1:, :], k1_full, v1_full)) else: loop1_kv.append((None, None, None, None)) else: loop1_kv.append((None, None, None, None)) for loop_idx in range(2, self.loop_num + 1): for layer_idx, decoder_layer in enumerate(self.layers): # Get k1, v1 (current token's Loop 1 KV) and full shared cache k1_current, v1_current, k1_full, v1_full = loop1_kv[layer_idx] if k1_current is None or v1_current is None: continue # Get past local KV cache past_local_key, past_local_value = None, None if next_decoder_cache is not None: past_local_key, past_local_value = next_decoder_cache.get_local(layer_idx) gate_proj = self.gate_projections[layer_idx] # Forward Loop 2+ attn_output, k2, v2 = decoder_layer.self_attn.forward_decode_loop2( hidden_states=decoder_layer.input_layernorm(hidden_states), k1=k1_current, v1=v1_current, past_shared_key=k1_full[:, :, :-1, :] if k1_full is not None and k1_full.shape[2] > 1 else None, past_shared_value=v1_full[:, :, :-1, :] if v1_full is not None and v1_full.shape[2] > 1 else None, past_local_key=past_local_key, past_local_value=past_local_value, gate_proj=gate_proj, attention_mask=causal_mask, position_ids=position_ids, loop_window_size=self.loop_window_size, ) # Update local cache with current token's Loop 2+ KV if use_cache and loop_idx == 2: next_decoder_cache.update_local(k2, v2, layer_idx) hidden_states = hidden_states + attn_output # MLP residual = hidden_states hidden_states = decoder_layer.post_attention_layernorm(hidden_states) hidden_states = decoder_layer.mlp(hidden_states) hidden_states = residual + hidden_states hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) next_cache = next_decoder_cache 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: torch.Tensor, input_tensor: torch.Tensor, cache_position: torch.Tensor, past_key_values: Cache, output_attentions: bool, ): """Create causal attention mask.""" dtype, device = input_tensor.dtype, input_tensor.device min_dtype = torch.finfo(dtype).min sequence_length = input_tensor.shape[1] # Determine target length for attention if past_key_values is not None: # For DynamicCache: use get_seq_length() to get cached length # target_length = cached_length + current_sequence_length past_length = past_key_values.get_seq_length() target_length = past_length + sequence_length elif attention_mask is not None: target_length = attention_mask.shape[-1] else: target_length = sequence_length # Create causal mask causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device) if sequence_length != 1: # For prefill: standard causal mask causal_mask = torch.triu(causal_mask, diagonal=1) # Adjust for cache position (for generation steps after prefill) causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) if attention_mask is not None: causal_mask = causal_mask.clone() mask_length = attention_mask.shape[-1] if mask_length <= target_length: padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] padding_mask = padding_mask == 0 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype) return causal_mask class IQuestLoopCoderForCausalLM(IQuestLoopCoderPreTrainedModel, GenerationMixin): """IQuestLoopCoder model with a causal language modeling head.""" _tied_weights_keys = ["lm_head.weight"] def __init__(self, config): super().__init__(config) self.model = IQuestLoopCoderModel(config) self.vocab_size = config.vocab_size 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 set_decoder(self, decoder): self.model = decoder def get_decoder(self): return self.model def forward( self, input_ids: 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, ) -> Union[Tuple, CausalLMOutputWithPast]: 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 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, ) hidden_states = outputs[0] logits = self.lm_head(hidden_states) logits = logits.float() loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = nn.CrossEntropyLoss() shift_logits = shift_logits.view(-1, self.config.vocab_size) shift_labels = shift_labels.view(-1) shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) 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 prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, use_cache=True, **kwargs, ): past_length = 0 if past_key_values is not None: past_length = past_key_values.get_seq_length() if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :] elif past_length < input_ids.shape[1]: input_ids = input_ids[:, past_length:] if cache_position is None: cache_position = torch.arange(past_length, past_length + input_ids.shape[1], device=input_ids.device) elif use_cache: cache_position = cache_position[-input_ids.shape[1]:] position_ids = cache_position.unsqueeze(0) if inputs_embeds is not None and past_key_values is None: model_inputs = {"inputs_embeds": inputs_embeds} else: model_inputs = {"input_ids": input_ids.contiguous()} model_inputs.update( { "position_ids": position_ids, "cache_position": cache_position, "past_key_values": past_key_values, "use_cache": use_cache, "attention_mask": attention_mask, } ) return model_inputs