""" Attention mechanisms for Transformer architecture. This module implements the core attention mechanisms used in the Transformer model: - ScaledDotProductAttention: Fundamental attention operation - MultiHeadAttention: Parallel attention with learned projections - T5RelativePositionBias: Relative position bias for T5-style attention Doing this first for Bottom-Up implementation of the Transformer Author: Oliver Perrin Date: 2025-10-23 """ import math from typing import Optional, Tuple, cast import torch import torch.nn as nn import torch.nn.functional as F class T5RelativePositionBias(nn.Module): """ T5-style relative position bias for attention. T5 uses a learned embedding table to encode relative positions between tokens. Positions are bucketed to handle arbitrary sequence lengths efficiently. This is added to attention scores BEFORE softmax, not to embeddings. """ def __init__( self, num_heads: int, num_buckets: int = 32, max_distance: int = 128, is_decoder: bool = False, ): super().__init__() self.num_heads = num_heads self.num_buckets = num_buckets self.max_distance = max_distance self.is_decoder = is_decoder # Learned embedding table: (num_buckets, num_heads) self.relative_attention_bias = nn.Embedding(num_buckets, num_heads) @staticmethod def _relative_position_bucket( relative_position: torch.Tensor, bidirectional: bool = True, num_buckets: int = 32, max_distance: int = 128, ) -> torch.Tensor: """ Translate relative position to a bucket index. T5 uses a combination of exact positions (for nearby tokens) and logarithmically-spaced buckets (for distant tokens). """ relative_buckets = torch.zeros_like(relative_position, dtype=torch.long) if bidirectional: num_buckets //= 2 relative_buckets += (relative_position > 0).long() * num_buckets relative_position = torch.abs(relative_position) else: relative_position = -torch.min(relative_position, torch.zeros_like(relative_position)) # Half buckets for exact positions max_exact = num_buckets // 2 is_small = relative_position < max_exact # Other half for logarithmically-spaced buckets relative_position_if_large = ( max_exact + ( torch.log(relative_position.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact) ).long() ) relative_position_if_large = torch.min( relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1) ) relative_buckets += torch.where(is_small, relative_position, relative_position_if_large) return relative_buckets def compute_bias( self, query_length: int, key_length: int, device: torch.device, query_position_offset: int = 0, ) -> torch.Tensor: """ Compute relative position bias for attention. Args: query_length: Number of query positions key_length: Number of key positions device: Device to create tensors on query_position_offset: Offset for query positions (for incremental decoding) When decoding step-by-step, query_length=1 but the actual position is past_len, so query_position_offset=past_len. Returns: (1, num_heads, query_length, key_length) """ # Create position indices context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] context_position = ( context_position + query_position_offset ) # Apply offset for incremental decoding memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] # Relative position: (query_length, key_length) relative_position = memory_position - context_position # Convert to bucket indices relative_position_bucket = self._relative_position_bucket( relative_position, bidirectional=(not self.is_decoder), num_buckets=self.num_buckets, max_distance=self.max_distance, ) # Look up bias values: (query_length, key_length, num_heads) values = self.relative_attention_bias(relative_position_bucket) # Reshape to (1, num_heads, query_length, key_length) values = values.permute([2, 0, 1]).unsqueeze(0) return values def forward( self, query_length: int, key_length: int, device: torch.device, query_position_offset: int = 0, ) -> torch.Tensor: return self.compute_bias(query_length, key_length, device, query_position_offset) class ScaledDotProductAttention(nn.Module): """ Scaled Dot-Product Attention using PyTorch's optimized backend. Uses F.scaled_dot_product_attention which automatically selects the best available kernel (FlashAttention v2, Memory-Efficient Attention, or math fallback) based on hardware and input shapes. On CUDA GPUs with appropriate compute capability, this will use FlashAttention for significantly improved speed and memory efficiency. See: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html """ def __init__(self, scale_scores: bool = True): """ Args: scale_scores: Whether to scale attention scores by sqrt(d_k). T5 does NOT scale scores, so set this to False for T5. Standard transformers (BERT, GPT, etc.) use scaling. """ super().__init__() self.scale_scores = scale_scores def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None, return_attn_weights: bool = False, position_bias: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """ Args: query: (batch, num_heads, seq_q, d_k) key: (batch, num_heads, seq_k, d_k) value: (batch, num_heads, seq_k, d_v) mask: Optional boolean mask, True = attend, False = mask position_bias: Optional (1, num_heads, seq_q, seq_k) T5-style relative position bias Returns: output: (batch, num_heads, seq_q, d_v) attention_weights: Optional (batch, num_heads, seq_q, seq_k) """ d_k = query.size(-1) scale_factor = 1.0 / math.sqrt(d_k) if self.scale_scores else 1.0 # If we need attention weights, must use manual path if return_attn_weights: # Manual implementation with float32 softmax for numerical stability scores = torch.matmul(query, key.transpose(-2, -1)) * scale_factor if position_bias is not None: scores = scores + position_bias if mask is not None: mask_bool = mask.to(dtype=torch.bool, device=scores.device) if mask_bool.dim() == 2: mask_bool = mask_bool.unsqueeze(1).unsqueeze(2) elif mask_bool.dim() == 3: mask_bool = mask_bool.unsqueeze(1) scores = scores.masked_fill(~mask_bool, -1e4) p_attn = F.softmax(scores.float(), dim=-1).type_as(scores) p_attn = torch.nan_to_num(p_attn, nan=0.0, posinf=0.0, neginf=0.0) output = torch.matmul(p_attn, value) return output, p_attn # Use optimized SDPA path - torch.compile friendly version # Pre-scale query instead of using SDPA's scale parameter for better compile compatibility # This avoids issues with inductor and custom scale values if self.scale_scores: query = query * scale_factor # Build combined attention mask (float tensor added to scores) attn_mask = None if position_bias is not None or mask is not None: # Start with position bias if provided if position_bias is not None: # Clamp position bias to prevent overflow attn_mask = position_bias.to(dtype=query.dtype).clamp(-100, 100) # Add mask (convert bool mask to additive float mask) if mask is not None: mask_bool = mask.to(dtype=torch.bool, device=query.device) if mask_bool.dim() == 2: mask_bool = mask_bool.unsqueeze(1).unsqueeze(2) elif mask_bool.dim() == 3: mask_bool = mask_bool.unsqueeze(1) mask_float = torch.zeros(mask_bool.shape, dtype=query.dtype, device=query.device) mask_float = mask_float.masked_fill(~mask_bool, -1e4) if attn_mask is not None: attn_mask = attn_mask + mask_float else: attn_mask = mask_float # Use SDPA without custom scale (scale=None uses default 1/sqrt(d_k)) # For T5 (scale_scores=False), we already didn't scale query above, so default scale is wrong # But we pre-scaled query for scaled attention, so we need scale=1.0 here # Actually simpler: always use scale=1.0 since we handle scaling ourselves output = F.scaled_dot_product_attention( query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False, scale=1.0, # We handle scaling manually above ) return output, None # --------------- Rotary Positional Embeddings --------------- class RotaryEmbedding(nn.Module): """ Rotary Positional Embeddings (RoPE). Encodes relative positions by rotating the query and key vectors. Reference: https://arxiv.org/abs/2104.09864 """ def __init__(self, dim, max_seq_len=2048): super().__init__() inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) t = torch.arange(max_seq_len).type_as(inv_freq) freqs = torch.einsum("i,j->ij", t, inv_freq) emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos", emb.cos()) self.register_buffer("sin", emb.sin()) def forward(self, x): # x shape: (batch, num_heads, seq_len, dim) seq_len = x.shape[2] # Slice cos/sin to current sequence length # unsqueeze to broadcast over batch and heads: (1, 1, seq_len, dim) cos_buf = cast(torch.Tensor, self.cos) sin_buf = cast(torch.Tensor, self.sin) cos = cos_buf[:seq_len, :].unsqueeze(0).unsqueeze(0) sin = sin_buf[:seq_len, :].unsqueeze(0).unsqueeze(0) return (x * cos) + (self._rotate_half(x) * sin) def _rotate_half(self, x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) # --------------- Multi-Head Attention --------------- class MultiHeadAttention(nn.Module): """ Multi-Head Attention mechanism. Allows the model to jointly attend to information from different representation subspaces at different positions. Transforming the input into query, key, and value representations Args: d_model: Dimension of model (default: 512) num_heads: Number of attention heads (default: 8) dropout: Dropout probability (default: 0.1) use_rope: Whether to use Rotary Positional Embeddings (default: False) max_len: Maximum sequence length for RoPE (default: 2048) use_lora: Whether to use LoRA (Low-Rank Adaptation) (default: False) lora_rank: Rank of LoRA matrices (default: 8) lora_alpha: Scaling factor for LoRA (default: 16) lora_dropout: Dropout probability for LoRA (default: 0.1) scale_scores: Whether to scale attention scores by sqrt(d_k). T5 does NOT scale. """ def __init__( self, d_model: int = 512, num_heads: int = 8, dropout: float = 0.1, use_rope: bool = False, max_len: int = 2048, use_lora: bool = False, lora_rank: int = 8, lora_alpha: int = 16, lora_dropout: float = 0.1, quantization: Optional[str] = None, scale_scores: bool = True, # T5 uses scale_scores=False ): super().__init__() # Assert that d_model is divisible by num_heads # Why? Because d_k = d_model // num_heads must be an integer assert d_model % num_heads == 0 # Assume d_v always equals d_k self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads # Select Linear layer type based on quantization Linear = nn.Linear kwargs = {} if quantization == "4bit": try: import bitsandbytes as bnb Linear = bnb.nn.Linear4bit # type: ignore kwargs = {"compute_dtype": torch.bfloat16, "quant_type": "nf4"} except (ImportError, AttributeError): print("bitsandbytes not installed or incompatible, falling back to nn.Linear") elif quantization == "8bit": try: import bitsandbytes as bnb Linear = bnb.nn.Linear8bitLt # type: ignore except (ImportError, AttributeError): print("bitsandbytes not installed or incompatible, falling back to nn.Linear") # Create 4 linear layers (W_Q, W_K, W_V, W_O) # All should be nn.Linear(d_model, d_model) self.W_Q = Linear(d_model, d_model, **kwargs) self.W_K = Linear(d_model, d_model, **kwargs) self.W_V = Linear(d_model, d_model, **kwargs) self.W_O = Linear(d_model, d_model, **kwargs) # Create ScaledDotProductAttention instance # Note: T5 does NOT scale attention scores by sqrt(d_k) self.attention = ScaledDotProductAttention(scale_scores=scale_scores) # Create dropout layer self.dropout = nn.Dropout(p=dropout) # RoPE self.use_rope = use_rope if use_rope: self.rope = RotaryEmbedding(self.d_k, max_seq_len=max_len) # LoRA (Low-Rank Adaptation) self.use_lora = use_lora if use_lora: self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.lora_scaling = lora_alpha / lora_rank self.lora_dropout = nn.Dropout(p=lora_dropout) # LoRA for Query: W_Q' = W_Q + B_q @ A_q * scaling self.lora_q_A = nn.Linear(d_model, lora_rank, bias=False) self.lora_q_B = nn.Linear(lora_rank, d_model, bias=False) # LoRA for Value: W_V' = W_V + B_v @ A_v * scaling self.lora_v_A = nn.Linear(d_model, lora_rank, bias=False) self.lora_v_B = nn.Linear(lora_rank, d_model, bias=False) # Initialize LoRA parameters # A: Kaiming uniform, B: Zeros (so training starts with original behavior) nn.init.kaiming_uniform_(self.lora_q_A.weight, a=math.sqrt(5)) nn.init.zeros_(self.lora_q_B.weight) nn.init.kaiming_uniform_(self.lora_v_A.weight, a=math.sqrt(5)) nn.init.zeros_(self.lora_v_B.weight) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None, return_attn_weights: bool = False, position_bias: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """ Args: query: (batch, seq_len, d_model) key: (batch, seq_len, d_model) value: (batch, seq_len, d_model) mask: Optional (batch, seq_len, seq_len) or (batch, 1, seq_len, seq_len) position_bias: Optional (1, num_heads, seq_q, seq_k) T5-style relative position bias Returns: output: (batch, seq_len, d_model) attention_weights: (batch, num_heads, seq_len, seq_len) """ batch_size = query.size(0) # Linear projections Q = self.W_Q(query) # (batch, seq_len, d_model) K = self.W_K(key) V = self.W_V(value) # Apply LoRA if enabled if self.use_lora: # Q += (query @ A^T @ B^T) * scaling # Note: nn.Linear(x) computes x @ weight.T # So lora_q_A(x) is x @ A.T # lora_q_B(lora_q_A(x)) is (x @ A.T) @ B.T = x @ A.T @ B.T lora_q = self.lora_q_B(self.lora_q_A(self.lora_dropout(query))) * self.lora_scaling Q = Q + lora_q # V += (value @ A^T @ B^T) * scaling lora_v = self.lora_v_B(self.lora_v_A(self.lora_dropout(value))) * self.lora_scaling V = V + lora_v # Split into heads # Reshape from (batch, seq_len, d_model) to (batch, num_heads, seq_len, d_k), Apply to Q, K, V Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # Now: (batch, num_heads, seq_len, d_k) # Now all are: (batch=2, num_heads=8, seq_len=10, d_k=64) # Apply RoPE if enabled if self.use_rope: Q = self.rope(Q) K = self.rope(K) # Handle mask broadcasting for multi-head attention if mask is not None: # If mask is 3D (batch, seq, seq), add head dimension if mask.dim() == 3: mask = mask.unsqueeze(1) # (batch, 1, seq, seq) # Now mask broadcasts across all heads: (batch, 1, seq, seq) → (batch, 8, seq, seq) # Apply attention with optional position bias output, attn_weights = self.attention( Q, K, V, mask, return_attn_weights=return_attn_weights, position_bias=position_bias ) # output: (batch, num_heads, seq_len, d_k) # attn_weights: (batch, num_heads, seq_len, seq_len) # Concatenate heads # (batch, num_heads, seq_len, d_k) → (batch, seq_len, num_heads, d_k) → (batch, seq_len, d_model) output = output.transpose(1, 2).contiguous() output = output.view( batch_size, -1, self.d_model ) # -1 in view means 'infer this dimension' # After transpose, the tensor's memory layout # is "scattered", contiguous() just reorganizes it in memory # Final linear projection output = self.W_O(output) # Apply dropout output = self.dropout(output) return output, attn_weights