modeling_iquestloopcoder.py

"""
Modified MIT License 

Software Copyright© 2025 IQuest Research

Our only modification is that, if the Software (or any derivative works
thereof) is used for any of your commercial products or services, you shall
prominently display "IQuest Coder" on the user interface of such product or
service.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

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