"""Grouped Query Attention (GQA). GQA reduces KV cache memory by sharing key-value heads across groups of query heads. With G query heads per KV head, the KV cache is G× smaller than standard multi-head attention (MHA). This is the primary motivation for its use in Llama 3 at 128K context: 8 KV heads shared across 32 query heads gives a 4× cache reduction. Setting num_key_value_heads == num_attention_heads recovers standard MHA. Setting num_key_value_heads == 1 gives multi-query attention (MQA). Attention is computed via torch.nn.functional.scaled_dot_product_attention (SDPA), which selects FlashAttention when hardware and dtype allow, falling back to standard attention otherwise. No custom kernel or additional dependency required. KV caching is handled via HuggingFace's Cache protocol. The cache owns K/V storage and accumulation; attention only calls cache.update() to store new projections and retrieve the full accumulated history. This cleanly separates attention computation from cache management: different Cache subclasses (DynamicCache, StaticCache, custom research variants) can be dropped in without touching the attention logic. No bias on any projection — a fixed architectural constant of this model. """ import math import torch import torch.nn as nn import torch.nn.functional as F from transformers import PretrainedConfig from transformers.cache_utils import Cache from .rope import RotaryEmbedding class GroupedQueryAttention(nn.Module): """Grouped Query Attention with RoPE, causal masking, and KV cache support. Implements GQA as used in Llama 3: Q heads are split into groups, each group sharing a single KV head. Before attention is computed, K and V are expanded by repeating each KV head across its group of query heads. The forward pass is strictly causal. An optional pre-built boolean attention mask can be threaded in from the caller; when absent, SDPA's native ``is_causal`` mode applies — correct for full-sequence training. Args: config: Model config. Must expose ``num_attention_heads``, ``num_key_value_heads``, ``head_dim``, ``hidden_size``, and ``attention_dropout``. Raises: ValueError: If ``num_attention_heads`` is not divisible by ``num_key_value_heads``. """ def __init__(self, config: PretrainedConfig) -> None: super().__init__() self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.head_dim = config.head_dim self.num_groups = self.num_heads // self.num_kv_heads self.attention_dropout = config.attention_dropout if self.num_heads % self.num_kv_heads != 0: raise ValueError( f"num_attention_heads ({self.num_heads}) must be divisible by " f"num_key_value_heads ({self.num_kv_heads})." ) # No bias on any projection — architectural constant. self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) self.rope = RotaryEmbedding(config) def forward( self, x: torch.Tensor, position_ids: torch.Tensor, cache: Cache | None = None, layer_idx: int = 0, causal_mask: torch.Tensor | None = None, ) -> torch.Tensor: """Apply grouped query attention to the input. Args: x: Input of shape (batch, seq_len, hidden_size). position_ids: Absolute positions of shape (batch, seq_len). Used by RoPE to rotate Q and K at the correct frequencies. cache: HuggingFace Cache object for KV accumulation, or None when caching is disabled (``use_cache=False``). When provided, ``cache.update(k, v, layer_idx)`` stores the new K/V and returns the full accumulated key and value tensors for this layer. layer_idx: Which slot in the cache to read and write. Each decoder layer has its own index so they accumulate independently. causal_mask: Optional boolean attention mask of shape (1, 1, seq_len, kv_len), where True indicates a position that should be attended to. When None, SDPA's built-in ``is_causal`` mode is used, which is correct for full-sequence training (square Q×K matrix). When provided, ``is_causal`` is disabled and the explicit mask governs attention — required for any generation pattern where Q and K lengths differ. Returns: Output tensor of shape (batch, seq_len, hidden_size). """ batch, seq_len, _ = x.shape # Project and reshape: (batch, seq_len, heads * head_dim) # → (batch, heads, seq_len, head_dim) q = self.q_proj(x).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(batch, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) # Apply RoPE. attention_scaling is 1.0 for default/linear; YaRN returns a # value != 1.0 that corrects attention magnitude after frequency manipulation. q, k, attention_scaling = self.rope(q, k, position_ids) if cache is not None: k_full, v_full = cache.update(k, v, layer_idx) else: k_full, v_full = k, v # Expand KV heads to align with query heads for GQA. # Each KV head is repeated num_groups times so SDPA sees matching head counts. if self.num_groups > 1: k_full = k_full.repeat_interleave(self.num_groups, dim=1) v_full = v_full.repeat_interleave(self.num_groups, dim=1) attn_output = F.scaled_dot_product_attention( q, k_full, v_full, attn_mask=causal_mask, dropout_p=self.attention_dropout if self.training else 0.0, is_causal=causal_mask is None, scale=attention_scaling / math.sqrt(self.head_dim), ) # Merge heads and project back to hidden_size. attn_output = ( attn_output .transpose(1, 2) .contiguous() .view(batch, seq_len, self.num_heads * self.head_dim) ) return self.o_proj(attn_output)