codon-model 0.0.5a4__tar.gz → 0.0.6a2__tar.gz

{codon_model-0.0.5a4/codon_model.egg-info → codon_model-0.0.6a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: codon-model
-Version: 0.0.5a4
+Version: 0.0.6a2
 Summary: Codon model package
 Author: CodonTeam
 Requires-Python: >=3.8
@@ -13,7 +13,13 @@ Requires-Dist: pandas
 Requires-Dist: numpy
 Requires-Dist: pyarrow
 Requires-Dist: pillow
-Requires-Dist: flash-linear-attention
+Requires-Dist: fastapi
+Requires-Dist: uvicorn
+Requires-Dist: onnx
+Requires-Dist: onnxsim
+Requires-Dist: onnxscript
+Requires-Dist: onnx-simplifier
+Requires-Dist: onnxconverter_common
 Dynamic: author
 Dynamic: license-file
 Dynamic: requires-dist

{codon_model-0.0.5a4 → codon_model-0.0.6a2}/codon/__init__.py

@@ -1,5 +1,5 @@
 from typing import Optional
 
-__version__ = '0.0.5a4'
+__version__ = '0.0.6a2'
 
 __seed__: Optional[int] = None

{codon_model-0.0.5a4 → codon_model-0.0.6a2}/codon/base.py

@@ -2,14 +2,18 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from typing import Callable, Any, Iterator, Union
+from typing import Callable, Any, Iterator, Union, TypeVar
 
 from safetensors.torch import save_model as safe_save_model
 from safetensors.torch import save_file  as safe_save_file
 from safetensors.torch import load_model as safe_load_model
 
+from codon.builtins.mixins.remote import RemoteResourceMixin
 
-class BasicModel(nn.Module):
+
+TBasicModel = TypeVar('TBasicModel', bound='BasicModel')
+
+class BasicModel(nn.Module, RemoteResourceMixin):
     '''
     Base class for all models, providing common functionality like gradient checkpointing and parameter counting.
     '''
@@ -62,6 +66,10 @@ class BasicModel(nn.Module):
             )
         return function(*args, **kwargs)
     
+    @property
+    def trainable_params(self) -> Iterator[torch.nn.Parameter]:
+        return self.get_params(trainable_only=True)
+    
     def get_params(self, trainable_only:bool=False) -> Iterator[torch.nn.Parameter]:
         '''
         Get an iterator over the model parameters.
@@ -134,7 +142,7 @@ class BasicModel(nn.Module):
         
         return total
     
-    def load_pretrained(self, path: str, strict: bool = False) -> 'BasicModel':
+    def load_pretrained(self: TBasicModel, path: str, strict: bool = False) -> TBasicModel:
         '''
         Load a pretrained model from a file.
         Args:
@@ -160,14 +168,14 @@ class BasicModel(nn.Module):
         return self
     
     def save_pretrained(
-            self, 
+            self: TBasicModel, 
             path: str, 
             trainable_only: bool = False, 
             include_buffer: bool = True, 
             exclude_modules: list[Union[type, nn.Module]] = None,
             only: list[str] = None,
             exclude: list[str] = None
-        ) -> 'BasicModel':
+        ) -> TBasicModel:
         '''
         Save the model to a file.
 
@@ -236,7 +244,7 @@ class BasicModel(nn.Module):
             
         return self
     
-    def freeze(self) -> 'BasicModel':
+    def freeze(self: TBasicModel) -> TBasicModel:
         '''
         Freeze all parameters in the model by setting requires_grad to False.
 
@@ -247,7 +255,7 @@ class BasicModel(nn.Module):
             param.requires_grad = False
         return self
 
-    def unfreeze(self) -> 'BasicModel':
+    def unfreeze(self: TBasicModel) -> TBasicModel:
         '''
         Unfreeze all parameters in the model by setting requires_grad to True.
 
@@ -257,3 +265,6 @@ class BasicModel(nn.Module):
         for param in self.parameters():
             param.requires_grad = True
         return self
+    
+    def compile(self: TBasicModel) -> TBasicModel:
+        return torch.compile(self)

{codon_model-0.0.5a4 → codon_model-0.0.6a2}/codon/block/__init__.py

@@ -36,6 +36,7 @@ from .manifold import (
     BasicManifoldLinear, RiemannianManifoldLinear,
     BasicManifoldConv2d, RiemannianManifoldConv2d
 )
+from .fourier import MultiHeadFourier
 
 __all__ = [
     # attention
@@ -90,5 +91,7 @@ __all__ = [
     'BasicManifoldLinear',
     'RiemannianManifoldLinear',
     'BasicManifoldConv2d',
-    'RiemannianManifoldConv2d'
+    'RiemannianManifoldConv2d',
+    # Fourier
+    'MultiHeadFourier',
 ]

{codon_model-0.0.5a4 → codon_model-0.0.6a2}/codon/block/attention.py

@@ -1,358 +1,178 @@
-from codon.base  import *
-from codon.block.embedding import BasicEmbedding
-from codon.ops.attention   import (
-    AttentionOutput,
-    apply_attention, chunk_linear_attn, chunk_gla
-)
-
-
-class MultiHeadAttention(BasicModel):
-    ''' 
-    Multi-Head Attention module.
-    Supports Grouped Query Attention (GQA), QK Normalization, and Gating mechanisms.
-
-    Attributes:
-        q_proj (nn.Linear): Linear layer for query projection.
-        k_proj (nn.Linear): Linear layer for key projection.
-        v_proj (nn.Linear): Linear layer for value projection.
-        o_proj (nn.Linear): Linear layer for output projection.
-        q_norm (nn.RMSNorm, optional): Normalization layer for queries.
-        k_norm (nn.RMSNorm, optional): Normalization layer for keys.
-        g_proj (nn.Linear, optional): Linear layer for gating mechanism.
-    '''
-    def __init__(
-        self,
-        hidden_size,
-        num_heads,
-        num_kv_heads=None,
-        use_qk_norm=True,
-        use_gate=False,
-        dropout=0.1,
-        bias: bool=True,
-        is_causal=True
-    ):
-        '''
-        Initialize the Multi-Head Attention module.
-
-        Args:
-            hidden_size (int): Size of the hidden layer.
-            num_heads (int): Number of attention heads.
-            num_kv_heads (int, optional): Number of key/value heads for GQA. 
-                                          If None, defaults to num_heads.
-            use_qk_norm (bool, optional): Whether to apply RMSNorm to queries and keys. 
-                                          Defaults to True.
-            use_gate (bool, optional): Whether to apply a gating mechanism. Defaults to False.
-            dropout (float, optional): Dropout probability. Defaults to 0.1.
-            is_causal (bool, optional): Whether to apply a causal mask. 
-                                        Defaults to True (for Decoder architectures).
-        '''
-        super(MultiHeadAttention, self).__init__()
-
-        if num_kv_heads is None: num_kv_heads = num_heads
-
-        assert hidden_size % num_heads == 0
-        assert num_heads % num_kv_heads == 0
-
-        self.hidden_size = hidden_size
-        self.num_heads  = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.num_kv_queries = num_heads // num_kv_heads
-        self.head_dim  = hidden_size // num_heads
-        self.kv_dim = self.num_kv_heads * self.head_dim
-        self.use_qk_norm = use_qk_norm
-        self.use_gate = use_gate
-        self.dropout = dropout
-        self.is_causal = is_causal
-        
-        if use_qk_norm:
-            self.q_norm = nn.RMSNorm(self.head_dim)
-            self.k_norm = nn.RMSNorm(self.head_dim)
-        
-        if use_gate:
-            self.g_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-
-        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-        self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
-        self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
-        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-        
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        kv_states: torch.Tensor = None,
-        attention_mask: torch.Tensor = None,
-        output_attentions: bool = False,
-        position_emb: BasicEmbedding = None,
-        embedding_start: int = 0,
-        embedding_pos: torch.Tensor = None,
-        past_key_value: tuple[torch.Tensor, torch.Tensor] = None,
-        use_cache: bool = False
-    ) -> AttentionOutput:
-        ''' 
-        Perform forward pass of Multi-Head Attention.
-
-        Args:
-            hidden_states (torch.Tensor): Input hidden states.
-            kv_states (torch.Tensor, optional): Hidden states for keys/values. 
-                                                If None, uses hidden_states. Defaults to None.
-            attention_mask (torch.Tensor, optional): Attention mask. 
-                                                     Defaults to None.
-            output_attentions (bool, optional): Whether to output attention weights. 
-                                                Defaults to False.
-            position_emb (BasicEmbedding, optional): Positional embedding module. 
-                                                    Defaults to None.
-            embedding_start (int, optional): Starting position for embedding. Defaults to 0.
-            embedding_pos (torch.Tensor, optional): Explicit position indices for positional embedding. 
-                                                    Defaults to None.
-            past_key_value (tuple[torch.Tensor, torch.Tensor], optional): Past key-value cache. 
-                                                                          Defaults to None.
-            use_cache (bool, optional): Whether to use KV cache. Defaults to False.
-        
-        Returns:
-            AttentionOutput: Object containing output, attention weights, and KV cache.
-        '''
-        
-        if kv_states is None:
-            kv_states = hidden_states
-
-        batch_size, q_len, _ = hidden_states.shape
-        kv_len_input = kv_states.shape[1]
-
-        if self.use_gate:
-            G = torch.sigmoid(self.g_proj(hidden_states))
-        
-        Q = self.q_proj(hidden_states)
-        K = self.k_proj(kv_states)
-        V = self.v_proj(kv_states)
-
-        Q = Q.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        K = K.view(batch_size, kv_len_input, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        V = V.view(batch_size, kv_len_input, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        
-        if self.use_qk_norm:
-            Q = self.q_norm(Q)
-            K = self.k_norm(K)
-        
-        if position_emb is not None:
-            Q = position_emb(Q, start_pos=embedding_start, positions=embedding_pos)
-            K = position_emb(K, start_pos=embedding_start, positions=embedding_pos)
-        
-        current_key_value = None
-        if use_cache:
-            if past_key_value is not None:
-                past_k, past_v = past_key_value
-                K = torch.cat((past_k, K), dim=2)
-                V = torch.cat((past_v, V), dim=2)
-            current_key_value = (K, V)
-
-        kv_seq_len_total = K.shape[2]
-
-        if self.num_kv_queries > 1:
-            # [B, H_kv, 1, L, D] -> [B, H_kv, G, L, D]
-            K = K[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_seq_len_total, self.head_dim)
-            V = V[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_seq_len_total, self.head_dim)
-            
-            K = K.reshape(batch_size, self.num_heads, kv_seq_len_total, self.head_dim)
-            V = V.reshape(batch_size, self.num_heads, kv_seq_len_total, self.head_dim)
-            
-        attn_output = apply_attention(
-            Q, K, V, 
-            attention_mask=attention_mask, 
-            output_attentions=output_attentions,
-            is_causal=self.is_causal,
-            dropout=self.dropout if self.training else 0.0
-        )
-
-        output = attn_output.output
-        attention_weights = attn_output.attention_weights
-        output = output.transpose(1, 2).contiguous().view(batch_size, q_len, self.hidden_size)
-        output = self.o_proj(output)
-
-        if self.use_gate: output = output * G
-
-        return AttentionOutput(
-            output=output,
-            attention_weights=attention_weights,
-            past_key_value=current_key_value
-        )
-
-
-class LinearAttention(BasicModel):
-    def __init__(self, hidden_size, num_heads, num_kv_heads=None, use_qk_norm=True, use_gate=False, dropout=0.1, is_causal=True, bias: bool=True):
-        super().__init__()
-        if num_kv_heads is None: num_kv_heads = num_heads
-        assert hidden_size % num_heads == 0
-        assert num_heads % num_kv_heads == 0
-
-        self.hidden_size = hidden_size
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.num_kv_queries = num_heads // num_kv_heads
-        self.head_dim = hidden_size // num_heads
-        self.kv_dim = self.num_kv_heads * self.head_dim
-        self.use_qk_norm = use_qk_norm
-        self.use_gate = use_gate
-        self.dropout = dropout
-        self.is_causal = is_causal
-
-        if use_qk_norm:
-            self.q_norm = nn.RMSNorm(self.head_dim)
-            self.k_norm = nn.RMSNorm(self.head_dim)
-            
-        if use_gate:
-            self.g_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-
-        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-        self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
-        self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
-        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-
-    def forward(
-            self,
-            hidden_states,
-            kv_states=None,
-            position_emb=None,
-            embedding_start=0,
-            embedding_pos=None,
-            past_key_value=None,
-            use_cache=False
-        ):
-        if kv_states is None: kv_states = hidden_states
-        batch_size, q_len, _ = hidden_states.shape
-        kv_len = kv_states.shape[1]
-
-        if self.use_gate:
-            G_out = torch.sigmoid(self.g_proj(hidden_states))
-
-        Q = self.q_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        K = self.k_proj(kv_states).view(batch_size, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        V = self.v_proj(kv_states).view(batch_size, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-
-        if self.use_qk_norm:
-            Q, K = self.q_norm(Q), self.k_norm(K)
-        if position_emb is not None:
-            Q = position_emb(Q, start_pos=embedding_start, positions=embedding_pos)
-            K = position_emb(K, start_pos=embedding_start, positions=embedding_pos)
-
-        if self.num_kv_queries > 1:
-            K = K[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_len, self.head_dim)
-            V = V[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_len, self.head_dim)
-            K = K.reshape(batch_size, self.num_heads, kv_len, self.head_dim)
-            V = V.reshape(batch_size, self.num_heads, kv_len, self.head_dim)
-
-        Q = F.elu(Q, alpha=1.0) + 1.0
-        K = F.elu(K, alpha=1.0) + 1.0
-
-        Q = Q.contiguous()
-        K = K.contiguous()
-        V = V.contiguous()
-
-        scale = self.head_dim ** -0.5  
-        output, current_key_value = chunk_linear_attn(
-            q=Q,
-            k=K,
-            v=V,
-            scale=scale,
-            initial_state=past_key_value,
-            output_final_state=use_cache
-        )
-
-        output = output.transpose(1, 2).contiguous().view(batch_size, q_len, self.hidden_size)
-        output = self.o_proj(output)
-        
-        if self.use_gate: output = output * G_out
-
-        return AttentionOutput(
-            output=output,
-            attention_weights=None,
-            past_key_value=current_key_value
-        )
-
-
-class GatedLinearAttention(BasicModel):
-    def __init__(self, hidden_size, num_heads, num_kv_heads=None, use_qk_norm=True, dropout=0.1, is_causal=True, bias: bool=True):
-        super().__init__()
-        if num_kv_heads is None: num_kv_heads = num_heads
-        assert hidden_size % num_heads == 0
-        assert num_heads % num_kv_heads == 0
-
-        self.hidden_size = hidden_size
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.num_kv_queries = num_heads // num_kv_heads
-        self.head_dim = hidden_size // num_heads
-        self.kv_dim = self.num_kv_heads * self.head_dim
-        self.use_qk_norm = use_qk_norm
-        self.dropout = dropout
-        self.is_causal = is_causal
-
-        if use_qk_norm:
-            self.q_norm = nn.RMSNorm(self.head_dim)
-            self.k_norm = nn.RMSNorm(self.head_dim)
-
-        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-        self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
-        self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
-        
-        self.g_proj = nn.Linear(hidden_size, num_heads, bias=False)
-        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
-
-    def forward(
-            self,
-            hidden_states,
-            kv_states=None,
-            position_emb=None,
-            embedding_start=0,
-            embedding_pos=None,
-            past_key_value=None,
-            use_cache=False
-        ):
-        if kv_states is None: kv_states = hidden_states
-        batch_size, q_len, _ = hidden_states.shape
-        kv_len = kv_states.shape[1]
-
-        Q = self.q_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        K = self.k_proj(kv_states).view(batch_size, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-        V = self.v_proj(kv_states).view(batch_size, kv_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
-
-        if self.use_qk_norm:
-            Q, K = self.q_norm(Q), self.k_norm(K)
-        if position_emb is not None:
-            Q = position_emb(Q, start_pos=embedding_start, positions=embedding_pos)
-            K = position_emb(K, start_pos=embedding_start, positions=embedding_pos)
-
-        if self.num_kv_queries > 1:
-            K = K[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_len, self.head_dim)
-            V = V[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_len, self.head_dim)
-            K = K.reshape(batch_size, self.num_heads, kv_len, self.head_dim)
-            V = V.reshape(batch_size, self.num_heads, kv_len, self.head_dim)
-
-        Q = F.elu(Q, alpha=1.0) + 1.0
-        K = F.elu(K, alpha=1.0) + 1.0
-
-        Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
-
-        gate_logits = self.g_proj(hidden_states)
-        gate_logits = gate_logits.transpose(1, 2).contiguous()
-
-        log_g = F.logsigmoid(gate_logits) 
-        
-        scale = self.head_dim ** -0.5
-        output, current_key_value = chunk_gla(
-            q=Q,
-            k=K,
-            v=V,
-            g=log_g,
-            scale=scale,
-            initial_state=past_key_value,
-            output_final_state=use_cache
-        )
-
-        output = output.transpose(1, 2).contiguous().view(batch_size, q_len, self.hidden_size)
-        output = self.o_proj(output)
-
-        return AttentionOutput(
-            output=output,
-            attention_weights=None,
-            past_key_value=current_key_value
-        )
+from codon.base  import *
+from codon.block.embedding import BasicEmbedding
+from codon.ops.attention   import (
+    AttentionOutput,
+    apply_attention
+)
+
+
+class MultiHeadAttention(BasicModel):
+    ''' 
+    Multi-Head Attention module.
+    Supports Grouped Query Attention (GQA), QK Normalization, and Gating mechanisms.
+
+    Attributes:
+        q_proj (nn.Linear): Linear layer for query projection.
+        k_proj (nn.Linear): Linear layer for key projection.
+        v_proj (nn.Linear): Linear layer for value projection.
+        o_proj (nn.Linear): Linear layer for output projection.
+        q_norm (nn.RMSNorm, optional): Normalization layer for queries.
+        k_norm (nn.RMSNorm, optional): Normalization layer for keys.
+        g_proj (nn.Linear, optional): Linear layer for gating mechanism.
+    '''
+    def __init__(
+        self,
+        hidden_size,
+        num_heads,
+        num_kv_heads=None,
+        use_qk_norm=True,
+        use_gate=False,
+        dropout=0.1,
+        bias: bool=True,
+        is_causal=True
+    ):
+        '''
+        Initialize the Multi-Head Attention module.
+
+        Args:
+            hidden_size (int): Size of the hidden layer.
+            num_heads (int): Number of attention heads.
+            num_kv_heads (int, optional): Number of key/value heads for GQA. 
+                                          If None, defaults to num_heads.
+            use_qk_norm (bool, optional): Whether to apply RMSNorm to queries and keys. 
+                                          Defaults to True.
+            use_gate (bool, optional): Whether to apply a gating mechanism. Defaults to False.
+            dropout (float, optional): Dropout probability. Defaults to 0.1.
+            is_causal (bool, optional): Whether to apply a causal mask. 
+                                        Defaults to True (for Decoder architectures).
+        '''
+        super(MultiHeadAttention, self).__init__()
+
+        if num_kv_heads is None: num_kv_heads = num_heads
+
+        assert hidden_size % num_heads == 0
+        assert num_heads % num_kv_heads == 0
+
+        self.hidden_size = hidden_size
+        self.num_heads  = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.num_kv_queries = num_heads // num_kv_heads
+        self.head_dim  = hidden_size // num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.use_qk_norm = use_qk_norm
+        self.use_gate = use_gate
+        self.dropout = dropout
+        self.is_causal = is_causal
+        
+        if use_qk_norm:
+            self.q_norm = nn.RMSNorm(self.head_dim)
+            self.k_norm = nn.RMSNorm(self.head_dim)
+        
+        if use_gate:
+            self.g_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
+        self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=bias)
+        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        kv_states: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        output_attentions: bool = False,
+        position_emb: BasicEmbedding = None,
+        embedding_start: int = 0,
+        embedding_pos: torch.Tensor = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] = None,
+        use_cache: bool = False
+    ) -> AttentionOutput:
+        ''' 
+        Perform forward pass of Multi-Head Attention.
+
+        Args:
+            hidden_states (torch.Tensor): Input hidden states.
+            kv_states (torch.Tensor, optional): Hidden states for keys/values. 
+                                                If None, uses hidden_states. Defaults to None.
+            attention_mask (torch.Tensor, optional): Attention mask. 
+                                                     Defaults to None.
+            output_attentions (bool, optional): Whether to output attention weights. 
+                                                Defaults to False.
+            position_emb (BasicEmbedding, optional): Positional embedding module. 
+                                                    Defaults to None.
+            embedding_start (int, optional): Starting position for embedding. Defaults to 0.
+            embedding_pos (torch.Tensor, optional): Explicit position indices for positional embedding. 
+                                                    Defaults to None.
+            past_key_value (tuple[torch.Tensor, torch.Tensor], optional): Past key-value cache. 
+                                                                          Defaults to None.
+            use_cache (bool, optional): Whether to use KV cache. Defaults to False.
+        
+        Returns:
+            AttentionOutput: Object containing output, attention weights, and KV cache.
+        '''
+        
+        if kv_states is None:
+            kv_states = hidden_states
+
+        batch_size, q_len, _ = hidden_states.shape
+        kv_len_input = kv_states.shape[1]
+
+        if self.use_gate:
+            G = torch.sigmoid(self.g_proj(hidden_states))
+        
+        Q = self.q_proj(hidden_states)
+        K = self.k_proj(kv_states)
+        V = self.v_proj(kv_states)
+
+        Q = Q.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        K = K.view(batch_size, kv_len_input, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, kv_len_input, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        
+        if self.use_qk_norm:
+            Q = self.q_norm(Q)
+            K = self.k_norm(K)
+        
+        if position_emb is not None:
+            Q = position_emb(Q, start_pos=embedding_start, positions=embedding_pos)
+            K = position_emb(K, start_pos=embedding_start, positions=embedding_pos)
+        
+        current_key_value = None
+        if use_cache:
+            if past_key_value is not None:
+                past_k, past_v = past_key_value
+                K = torch.cat((past_k, K), dim=2)
+                V = torch.cat((past_v, V), dim=2)
+            current_key_value = (K, V)
+
+        kv_seq_len_total = K.shape[2]
+
+        if self.num_kv_queries > 1:
+            # [B, H_kv, 1, L, D] -> [B, H_kv, G, L, D]
+            K = K[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_seq_len_total, self.head_dim)
+            V = V[:, :, None, :, :].expand(batch_size, self.num_kv_heads, self.num_kv_queries, kv_seq_len_total, self.head_dim)
+            
+            K = K.reshape(batch_size, self.num_heads, kv_seq_len_total, self.head_dim)
+            V = V.reshape(batch_size, self.num_heads, kv_seq_len_total, self.head_dim)
+            
+        attn_output = apply_attention(
+            Q, K, V, 
+            attention_mask=attention_mask, 
+            output_attentions=output_attentions,
+            is_causal=self.is_causal,
+            dropout=self.dropout if self.training else 0.0
+        )
+
+        output = attn_output.output
+        attention_weights = attn_output.attention_weights
+        output = output.transpose(1, 2).contiguous().view(batch_size, q_len, self.hidden_size)
+        output = self.o_proj(output)
+
+        if self.use_gate: output = output * G
+
+        return AttentionOutput(
+            output=output,
+            attention_weights=attention_weights,
+            past_key_value=current_key_value
+        )
+