""" Attention modules for Q-TensorFormer. Features: - Scaled Dot-Product Attention (SDPA / FlashAttention compatible) - Quantum Kernel Self-Attention (QKSAM) with fidelity kernel - First-class Adaptive KV Cache integration (FP16/INT8/INT4/eviction) - Pre-token Attention Entropy & Importance calculation for Information State - RoPE (Rotary Position Embeddings) """ import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Optional, Tuple, Dict, Union from .kv_cache import AdaptiveKVCache, KVPrecision from .quantum_backend import QuantumBackend, BackendType class MultiHeadAttention(nn.Module): """ Hardware-efficient Multi-Head Attention with Adaptive KV Cache and RoPE. """ def __init__( self, d_model: int = 128, n_heads: int = 4, n_kv_heads: Optional[int] = None, dropout: float = 0.1, max_seq_len: int = 128, use_quantum_kernel: bool = False, n_qubits: int = 4, backend_type: Union[str, BackendType] = BackendType.CLASSICAL_SURROGATE, ): super().__init__() assert d_model % n_heads == 0, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})" self.d_model = d_model self.n_heads = n_heads self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads assert self.n_heads % self.n_kv_heads == 0, f"n_heads ({n_heads}) must be divisible by n_kv_heads ({self.n_kv_heads})" self.num_queries_per_kv = self.n_heads // self.n_kv_heads self.head_dim = d_model // n_heads self.max_seq_len = max_seq_len self.use_quantum_kernel = use_quantum_kernel self.scale = 1.0 / math.sqrt(self.head_dim) if self.n_kv_heads == self.n_heads: self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) self.q_proj = None self.k_proj = None self.v_proj = None else: self.qkv = None self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(d_model, self.n_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(d_model, self.n_kv_heads * self.head_dim, bias=False) self.out_proj = nn.Linear(d_model, d_model, bias=False) self.dropout_p = dropout # Quantum backend for QKSAM self.quantum_backend = QuantumBackend( backend_type=backend_type, n_qubits=n_qubits, n_layers=2, d_model=d_model, ) # Precompute RoPE buffers self.register_buffer("rope_cos", None, persistent=False) self.register_buffer("rope_sin", None, persistent=False) def _init_rope(self, device: torch.device, required_len: int = 0): target_len = max(self.max_seq_len, required_len) if ( self.rope_cos is not None and self.rope_cos.device == device and self.rope_cos.shape[0] >= target_len ): return pos = torch.arange(target_len, device=device, dtype=torch.float32) dim = torch.arange(0, self.head_dim // 2, device=device, dtype=torch.float32) dim = dim / max(1, (self.head_dim // 2)) freqs = 1.0 / (10000.0 ** dim) angles = torch.outer(pos, freqs) self.rope_cos = torch.cos(angles) self.rope_sin = torch.sin(angles) def _apply_rope(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor: B, H, T, D = x.shape self._init_rope(x.device, required_len=offset + T) cos = self.rope_cos[offset:offset + T, :].unsqueeze(0).unsqueeze(0) # (1, 1, T, D/2) sin = self.rope_sin[offset:offset + T, :].unsqueeze(0).unsqueeze(0) x_rot = x.reshape(B, H, T, D // 2, 2) x1, x2 = x_rot[..., 0], x_rot[..., 1] x_rot1 = x1 * cos - x2 * sin x_rot2 = x1 * sin + x2 * cos return torch.stack([x_rot1, x_rot2], dim=-1).reshape(B, H, T, D) def _build_hide_mask( self, B: int, T: int, T_k: int, mask: Optional[torch.Tensor], device: torch.device, grouped_5d: bool = False, ) -> Optional[torch.Tensor]: """ Build a boolean hide_mask where True indicates positions to mask out (-inf). Combines causal masking (when T > 1) and user-provided padding/attention masks without ever multiplying 0.0 * float('-inf'). """ hide_mask: Optional[torch.Tensor] = None # 1. Causal mask whenever prefilling / attending over multiple query tokens if T > 1: offset = max(0, T_k - T) causal_hide = torch.triu( torch.ones(T, T_k, device=device, dtype=torch.bool), diagonal=1 + offset, ) # (T, T_k) hide_mask = causal_hide.view(1, 1, T, T_k) # 2. User-provided attention/padding mask if mask is not None: m = mask.to(device=device) if m.dtype == torch.bool: # Standard attention_mask: True = valid token, False = pad token pad_hide = ~m elif m.is_floating_point(): if (m < 0).any(): # Additive mask containing negative / -inf values pad_hide = m < -1e4 else: # 0.0 / 1.0 mask: 1.0 = valid, 0.0 = pad pad_hide = (m == 0.0) else: # Integer 0/1 mask: 1 = valid, 0 = pad pad_hide = (m == 0) if pad_hide.dim() == 2: # Shape (B, T_k) or (T, T_k) if pad_hide.shape == (B, T_k): pad_hide = pad_hide[:, None, None, :] elif pad_hide.shape[0] == B and pad_hide.shape[1] > T_k: pad_hide = pad_hide[:, -T_k:][:, None, None, :] elif pad_hide.shape[0] == B and pad_hide.shape[1] < T_k: padded = torch.zeros(B, T_k, device=device, dtype=torch.bool) padded[:, -pad_hide.shape[1]:] = pad_hide pad_hide = padded[:, None, None, :] elif pad_hide.shape == (T, T_k): pad_hide = pad_hide[None, None, :, :] else: pad_hide = pad_hide.view(B, 1, 1, -1)[..., :T_k] elif pad_hide.dim() == 3: pad_hide = pad_hide.unsqueeze(1) hide_mask = pad_hide if hide_mask is None else (hide_mask | pad_hide) if hide_mask is not None and grouped_5d: hide_mask = hide_mask.unsqueeze(2) # (B or 1, 1, 1, T, T_k) return hide_mask @staticmethod def _masked_softmax(scores: torch.Tensor, hide_mask: Optional[torch.Tensor]) -> torch.Tensor: """ Compute softmax over scores while masking `hide_mask == True` positions using `masked_fill` (avoiding 0 * -inf = NaN) and guarding against completely masked query rows (e.g. left-padded query positions). """ if hide_mask is None: return F.softmax(scores, dim=-1) row_all_masked = hide_mask.all(dim=-1, keepdim=True) safe_hide = hide_mask & ~row_all_masked masked_scores = scores.masked_fill(safe_hide, float("-inf")) weights = F.softmax(masked_scores, dim=-1) return weights.masked_fill(row_all_masked, 0.0) def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, kv_cache: Optional[AdaptiveKVCache] = None, use_quantum: bool = False, return_attn_weights: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """ Forward attention computation. Args: x: (batch, seq_len, d_model) mask: optional attention mask (1/True = valid token, 0/False = masked padding) kv_cache: optional AdaptiveKVCache instance for decode use_quantum: whether to activate QKSAM quantum fidelity kernel return_attn_weights: if True, returns attention probability weights Returns: out: (batch, seq_len, d_model) attn_weights: (batch, n_heads, seq_len_q, seq_len_k) or None """ B, T, C = x.shape if self.qkv is not None: qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim) q, k, v = qkv.unbind(dim=2) q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) else: q = self.q_proj(x).reshape(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) offset = kv_cache.seq_len if kv_cache is not None else 0 q = self._apply_rope(q, offset=offset) k = self._apply_rope(k, offset=offset) # Append to KV cache if provided (stores compact n_kv_heads) if kv_cache is not None: k, v = kv_cache.update(k, v) # Zero-copy GQA/MQA view broadcasting: keep K and V in (B, n_kv_heads, 1, T_k, head_dim) # without materializing full (B, n_heads, T_k, head_dim) copies. T_k = k.shape[-2] grouped_5d = self.num_queries_per_kv > 1 if grouped_5d: q_for_attn = q.reshape(B, self.n_kv_heads, self.num_queries_per_kv, T, self.head_dim) k_for_attn = k.unsqueeze(2) # Zero-copy view: (B, n_kv_heads, 1, T_k, head_dim) v_for_attn = v.unsqueeze(2) # Zero-copy view: (B, n_kv_heads, 1, T_k, head_dim) else: q_for_attn = q k_for_attn = k v_for_attn = v # Track KV storage footprint for diagnostic verification self.last_kv_storage_elements = k.numel() self.last_kv_shared_storage = ( k_for_attn.untyped_storage().data_ptr() == k.untyped_storage().data_ptr() and v_for_attn.untyped_storage().data_ptr() == v.untyped_storage().data_ptr() ) hide_mask = self._build_hide_mask( B=B, T=T, T_k=T_k, mask=mask, device=x.device, grouped_5d=grouped_5d ) if use_quantum or self.use_quantum_kernel: # Genuine Quantum Kernel Self-Attention (QKSAM): K_fidelity q_scaled = q_for_attn * self.scale kernel_mat = self.quantum_backend.compute_kernel_matrix(q_scaled, k_for_attn) attn_weights = self._masked_softmax(kernel_mat, hide_mask) if self.dropout_p > 0 and self.training: attn_weights = F.dropout(attn_weights, p=self.dropout_p) out = torch.matmul(attn_weights, v_for_attn) if grouped_5d: out = out.reshape(B, self.n_heads, T, self.head_dim) attn_weights = attn_weights.reshape(B, self.n_heads, T, T_k) else: # High-performance classical Scaled Dot-Product Attention (SDPA) is_causal_fast = (T == T_k and mask is None) if not return_attn_weights and hasattr(F, "scaled_dot_product_attention") and mask is None: out = F.scaled_dot_product_attention( q_for_attn, k_for_attn, v_for_attn, attn_mask=None, dropout_p=self.dropout_p if self.training else 0.0, is_causal=is_causal_fast, ) if grouped_5d: out = out.reshape(B, self.n_heads, T, self.head_dim) attn_weights = None else: scores = torch.matmul(q_for_attn, k_for_attn.transpose(-2, -1)) * self.scale attn_weights = self._masked_softmax(scores, hide_mask) if self.dropout_p > 0 and self.training: attn_weights = F.dropout(attn_weights, p=self.dropout_p) out = torch.matmul(attn_weights, v_for_attn) if grouped_5d: out = out.reshape(B, self.n_heads, T, self.head_dim) attn_weights = attn_weights.reshape(B, self.n_heads, T, T_k) out = out.transpose(1, 2).reshape(B, T, C) out = self.out_proj(out) return out, attn_weights def flops(self, batch_size: int = 1, seq_len: int = 1) -> Dict[str, int]: """Compute attention FLOPs.""" D = self.d_model H = self.n_heads hd = self.head_dim T = seq_len qkv_flops = 2 * batch_size * T * D * 3 * D attn_flops = 2 * batch_size * H * T * T * hd out_flops = 2 * batch_size * T * D * D return { "qkv_proj": qkv_flops, "attention": attn_flops, "out_proj": out_flops, "total": qkv_flops + attn_flops + out_flops, }