Spaces:
Running on Zero
Running on Zero
Premchandyadav369
feat: deploy interactive Q-TensorFormer demo with trained weights and Gradio 4 UI
67ded87 | """ | |
| 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 | |
| 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, | |
| } | |