""" Qwen3 Attention Module. This file is kept separate for easy modification during token pruning experiments. """ from typing import Optional, Tuple import math import torch import torch.nn as nn import torch.nn.functional as F try: from torch.nn.attention.flex_attention import create_block_mask, flex_attention except ImportError: create_block_mask = flex_attention = None from .config import Qwen3Config def alibi_slopes(num_heads: int) -> torch.Tensor: """Press et al. slopes: 1/2, 1/4, ..., 2**(-8) for eight heads.""" return torch.pow(2.0, -8.0 * torch.arange(1, num_heads + 1, dtype=torch.float32) / num_heads) def _alibi_causal_mask(b, h, q_idx, kv_idx): return q_idx >= kv_idx def _alibi_score_mod(num_heads: int): # q/k indices and h are scalar symbolic tensors under FlexAttention. This # produces no LxL bias tensor; tau scales Q only and therefore never this term. def score_mod(score, b, h, q_idx, kv_idx): slope = 2.0 ** (-8.0 * (h + 1) / num_heads) return score + slope * (kv_idx - q_idx) return score_mod class Qwen3RMSNorm(nn.Module): """RMS Normalization used in Qwen3.""" def __init__(self, hidden_size: int, eps: float = 1e-6, weight_fp32: bool = False): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.eps = eps self.weight_fp32 = weight_fp32 self._eval_temperature: Optional[float] = None def set_eval_temperature(self, tau: float) -> None: """Apply an eval-only Q gain scale before its activation-dtype cast.""" self._eval_temperature = float(tau) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.eps) # QK gains can be deliberate FP32 master parameters while all model # activations stay bf16. Casting here preserves fp32 optimizer updates # without promoting the attention activations. weight = self.weight if self._eval_temperature is None else self.weight / self._eval_temperature return weight.to(input_dtype) * hidden_states.to(input_dtype) class Qwen3RotaryEmbedding(nn.Module): """Rotary Position Embedding for Qwen3.""" def __init__( self, dim: int, max_position_embeddings: int = 32768, base: float = 1000000.0, train_position_scale_set: Optional[list[float]] = None, device: Optional[torch.device] = None, ): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base raw_scales = train_position_scale_set or [] scales = [float(scale) for scale in raw_scales] if any(not torch.isfinite(torch.tensor(scale)) or scale <= 0.0 for scale in scales): raise ValueError("rope_train_position_scale_set must contain finite positive values") self.register_buffer( "train_position_scales", torch.tensor(scales, dtype=torch.float32, device=device), persistent=False, ) inv_freq = 1.0 / ( self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim) ) self.register_buffer("inv_freq", inv_freq, persistent=False) @torch.no_grad() def forward( self, x: torch.Tensor, position_ids: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: # Mixed sparse training resamples a mask (and therefore its realised # key-distance distribution) independently in every layer. This dense # control mirrors only that positional exposure: each layer samples one # scalar for the whole micro-batch, while all attention links remain # present. Eval is deliberately untouched. if self.training and self.train_position_scales.numel() > 0: scale_idx = torch.randint( self.train_position_scales.numel(), (1,), device=position_ids.device, ) position_ids = position_ids.float() * self.train_position_scales[scale_idx] inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() device_type = x.device.type with torch.autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() sin = emb.sin() return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) def rotate_half(x: torch.Tensor) -> torch.Tensor: """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1, ) -> Tuple[torch.Tensor, torch.Tensor]: """Applies Rotary Position Embedding to the query and key tensors.""" cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed class Qwen3Attention(nn.Module): """Multi-head attention with Grouped Query Attention (GQA) support.""" def __init__(self, config: Qwen3Config, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = config.head_dim self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.attention_dropout = config.attention_dropout self.positional = config.positional self.qk_norm = str(getattr(config, "qk_norm", "on")).lower() self.qk_gain_fp32 = bool(getattr(config, "qk_gain_fp32", False)) if self.qk_norm not in ("on", "off"): raise ValueError(f"unsupported qk_norm mode: {self.qk_norm}") if self.positional not in ("rope", "alibi"): raise ValueError(f"unsupported positional mode: {self.positional}") # Attention output dimension is num_heads * head_dim self.attention_hidden_size = self.num_heads * self.head_dim self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) self.rotary_emb = Qwen3RotaryEmbedding( self.head_dim, max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, train_position_scale_set=config.rope_train_position_scale_set, ) self.register_buffer("alibi_slopes", alibi_slopes(self.num_heads), persistent=False) self._alibi_masks = {} self._alibi_score_mod = _alibi_score_mod(self.num_heads) if self.positional == "alibi" else None # For Qwen3, q_norm and k_norm are per layer. The explicit ``off`` # regime removes the modules entirely, rather than leaving identity # modules in the state dict. The default remains the legacy path. if self.qk_norm == "on": self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps, weight_fp32=self.qk_gain_fp32) self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps, weight_fp32=self.qk_gain_fp32) # Evaluation-only temperature for qk_norm=off. It is a non-persistent # buffer so no checkpoint weights are touched and compiled graphs can # reuse the same shape/static graph across temperatures. self.register_buffer("_query_temperature_scale", torch.ones((), dtype=torch.float32), persistent=False) # Eval-only diagnostic. When positive, sample this many causal query # positions and retain head-averaged attention entropy. The default is # disabled and therefore leaves training/evaluation behavior unchanged. self.entropy_probe_queries = 0 self.last_entropy_probe = None self._use_query_temperature = False @torch.no_grad() def set_attention_temperature(self, tau: float) -> None: """Set the query-side temperature without changing learned weights. This is the temperature path for qk_norm=off and is also exposed for equivalence checks on qk_norm=on checkpoints. """ tau = float(tau) if not math.isfinite(tau) or tau <= 0.0: raise ValueError("attention temperature must be finite and > 0") if self.qk_norm == "on": # This is algebraically query-side scaling, but applying it before # the bf16 activation cast matches the learned-gain division path # exactly without mutating a checkpoint parameter. self.q_norm.set_eval_temperature(tau) self._use_query_temperature = False else: self._query_temperature_scale.fill_(1.0 / tau) self._use_query_temperature = True @torch.no_grad() def _measure_attention_entropy(self, query_states, key_states): seq_len = query_states.shape[2] count = min(int(self.entropy_probe_queries), seq_len) # Include short and long prefixes while avoiding the degenerate first token. positions = torch.linspace(1, seq_len - 1, count, device=query_states.device).round().long().unique() kt = key_states.transpose(-1, -2) records = [] for chunk in positions.split(4): q = query_states.index_select(2, chunk) scores = torch.matmul(q, kt) / math.sqrt(self.head_dim) keys = torch.arange(seq_len, device=scores.device) allowed = keys[None, None, None, :] <= chunk[None, None, :, None] scores = scores.float().masked_fill(~allowed, float('-inf')) logz = torch.logsumexp(scores, dim=-1) probs = torch.softmax(scores, dim=-1) entropy = logz - torch.sum(probs * scores.masked_fill(~allowed, 0.0), dim=-1) mean_heads = entropy.mean(dim=(0, 1)).cpu().tolist() records.extend((int(pos), int(pos)+1, float(ent)) for pos,ent in zip(chunk.cpu().tolist(),mean_heads)) del q,scores,allowed,logz,probs,entropy self.last_entropy_probe = records def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, use_cache: bool = False, output_attentions: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]: batch_size, seq_len, _ = hidden_states.size() # Project Q, K, V query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) # Reshape for multi-head attention query_states = query_states.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) # Apply Q/K normalisation when enabled. In the no-QK-norm regime the # projections go straight to RoPE/logits, with only the evaluation-time # query-side temperature scale below. if self.qk_norm == "on": query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) if self.qk_norm == "off" or self._use_query_temperature: query_states = query_states * self._query_temperature_scale.to(dtype=query_states.dtype) # RoPE and ALiBi are mutually exclusive. ALiBi is added by FlexAttention # after QK scaling, so changing q_norm for temperature rescales content only. if self.positional == "rope": cos, sin = self.rotary_emb(value_states, position_ids) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) # Handle KV cache - store keys/values BEFORE repeating for GQA if past_key_value is not None: past_key, past_value = past_key_value key_states = torch.cat([past_key, key_states], dim=2) value_states = torch.cat([past_value, value_states], dim=2) if use_cache: present_key_value = (key_states, value_states) else: present_key_value = None key_states = self._repeat_kv(key_states, self.num_key_value_groups) value_states = self._repeat_kv(value_states, self.num_key_value_groups) if self.entropy_probe_queries: self._measure_attention_entropy(query_states, key_states) if self.positional == "alibi": if flex_attention is None or create_block_mask is None: raise ImportError("ALiBi requires torch FlexAttention support") if attention_mask is not None or past_key_value is not None: raise ValueError("ALiBi currently requires packed causal batches without a KV cache") key = (seq_len, query_states.device) block_mask = self._alibi_masks.get(key) if block_mask is None: block_mask = create_block_mask(_alibi_causal_mask, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=query_states.device) self._alibi_masks[key] = block_mask attn_output = flex_attention(query_states, key_states, value_states, score_mod=self._alibi_score_mod, block_mask=block_mask, enable_gqa=False, kernel_options={"BACKEND": "TRITON"}) else: attn_output = F.scaled_dot_product_attention( query_states, key_states, value_states, attn_mask=attention_mask, dropout_p=self.attention_dropout if self.training else 0.0, is_causal=attention_mask is None) attn_weights = None # SDPA doesn't materialize weights, saving massive memory attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(batch_size, seq_len, self.attention_hidden_size) # Output projection attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, present_key_value @staticmethod def _repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """Repeat KV heads to match the number of query heads for GQA.""" if n_rep == 1: return hidden_states batch, num_kv_heads, slen, head_dim = hidden_states.shape hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)