epfl_final/SparseAtt/sparseatt/attention.py

"""
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)