import torch import torch.nn as nn from typing import Dict from transformers import LlamaForCausalLM, LlamaConfig from transformers.generation.utils import GenerationConfig import os import pdb import copy import math import numpy as np from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union import gc import traceback import torch from torch import nn import torch.utils.checkpoint import torch.nn.functional as F from torch.cuda.amp import autocast from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, LlamaAttention, apply_rotary_pos_emb from transformers.cache_utils import DynamicCache class PredictorDynamicCache(DynamicCache): def __init__(self): super().__init__() self.predictor_primary_key: List[Optional[torch.Tensor]] = [] self.predictor_primary_value: List[Optional[torch.Tensor]] = [] self.predictor_importance_key: List[Optional[torch.Tensor]] = [] def update_predictor_primary( self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Append or create the predictor's "primary" K/V states for `layer_idx`. shape for key_states, value_states is typically [batch_size, num_heads, seq_len, head_dim]. """ # Extend the lists so that `predictor_primary_key[layer_idx]` and # `predictor_primary_value[layer_idx]` exist. self._ensure_list_capacity( self.predictor_primary_key, layer_idx, fill=None ) self._ensure_list_capacity( self.predictor_primary_value, layer_idx, fill=None ) # If this is the very first time we are updating that layer's predictor cache, just assign if self.predictor_primary_key[layer_idx] is None: self.predictor_primary_key[layer_idx] = key_states self.predictor_primary_value[layer_idx] = value_states else: # Otherwise, concatenate along the seq_len dimension (=-2 or =2 depending on your shape). self.predictor_primary_key[layer_idx] = torch.cat( [self.predictor_primary_key[layer_idx], key_states], dim=2 ) self.predictor_primary_value[layer_idx] = torch.cat( [self.predictor_primary_value[layer_idx], value_states], dim=2 ) return ( self.predictor_primary_key[layer_idx], self.predictor_primary_value[layer_idx], ) def update_predictor_importance( self, key_states: torch.Tensor, layer_idx: int, ) -> torch.Tensor: """ Append or create the predictor's "importance" key for `layer_idx`. """ self._ensure_list_capacity( self.predictor_importance_key, layer_idx, fill=None ) if self.predictor_importance_key[layer_idx] is None: self.predictor_importance_key[layer_idx] = key_states else: self.predictor_importance_key[layer_idx] = torch.cat( [self.predictor_importance_key[layer_idx], key_states], dim=2 ) return self.predictor_importance_key[layer_idx] def crop(self, max_length: int): super().crop(max_length) # Now also crop predictor caches for idx in range(len(self.predictor_primary_key)): if self.predictor_primary_key[idx] is not None: self.predictor_primary_key[idx] = self.predictor_primary_key[idx][..., :max_length, :] self.predictor_primary_value[idx] = self.predictor_primary_value[idx][..., :max_length, :] for idx in range(len(self.predictor_importance_key)): if self.predictor_importance_key[idx] is not None: self.predictor_importance_key[idx] = self.predictor_importance_key[idx][..., :max_length, :] # Remember to adjust self._seen_tokens accordingly self._seen_tokens = min(self._seen_tokens, max_length) def batch_split( self, full_batch_size: int, split_size: int, num_hidden_layers: int = None ) -> List["PredictorDynamicCache"]: # Use the base split logic for the standard K/V base_splits = super().batch_split(full_batch_size, split_size, num_hidden_layers) # `base_splits` is now a list of new DynamicCache objects. But we *actually* # want them to be PredictorDynamicCache so we can store the predictor states. # Easiest: we can cast and fill them. out: List[PredictorDynamicCache] = [] for split_i, base_split in enumerate(base_splits): # Construct an empty PredictorDynamicCache new_cache = PredictorDynamicCache() # Copy over the underlying fields from base_split new_cache.key_cache = base_split.key_cache new_cache.value_cache = base_split.value_cache new_cache._seen_tokens = base_split._seen_tokens # Now also slice our predictor fields # The slice in batch dim is [i:i+split_size]. b_start = split_i * split_size b_end = min(full_batch_size, b_start + split_size) new_cache.predictor_primary_key = self._slice_list_tensors( self.predictor_primary_key, b_start, b_end ) new_cache.predictor_primary_value = self._slice_list_tensors( self.predictor_primary_value, b_start, b_end ) new_cache.predictor_importance_key = self._slice_list_tensors( self.predictor_importance_key, b_start, b_end ) out.append(new_cache) return out @classmethod def from_batch_splits(cls, splits: List["PredictorDynamicCache"], num_hidden_layers: int = None) -> "PredictorDynamicCache": # Let the base class handle the normal K/V merges base_merged = DynamicCache.from_batch_splits(splits, num_hidden_layers=num_hidden_layers) merged = cls() merged.key_cache = base_merged.key_cache merged.value_cache = base_merged.value_cache merged._seen_tokens = base_merged._seen_tokens # Now unify predictor states by concatenating along batch dim=0 merged.predictor_primary_key = cls._merge_list_tensors( [split.predictor_primary_key for split in splits] ) merged.predictor_primary_value = cls._merge_list_tensors( [split.predictor_primary_value for split in splits] ) merged.predictor_importance_key = cls._merge_list_tensors( [split.predictor_importance_key for split in splits] ) return merged def batch_repeat_interleave(self, repeats: int): super().batch_repeat_interleave(repeats) self.predictor_primary_key = self._repeat_list_tensors( self.predictor_primary_key, repeats ) self.predictor_primary_value = self._repeat_list_tensors( self.predictor_primary_value, repeats ) self.predictor_importance_key = self._repeat_list_tensors( self.predictor_importance_key, repeats ) def batch_select_indices(self, indices: torch.Tensor): super().batch_select_indices(indices) self.predictor_primary_key = self._select_list_tensors( self.predictor_primary_key, indices ) self.predictor_primary_value = self._select_list_tensors( self.predictor_primary_value, indices ) self.predictor_importance_key = self._select_list_tensors( self.predictor_importance_key, indices ) @staticmethod def _ensure_list_capacity(lst: list, idx: int, fill=None): if len(lst) <= idx: lst.extend([fill] * (idx + 1 - len(lst))) @staticmethod def _slice_list_tensors( tensor_list: List[Optional[torch.Tensor]], start: int, end: int ) -> List[Optional[torch.Tensor]]: out = [] for t in tensor_list: if t is None: out.append(None) else: out.append(t[start:end, ...]) return out @classmethod def _merge_list_tensors( cls, list_of_lists: List[List[Optional[torch.Tensor]]] ) -> List[Optional[torch.Tensor]]: # If no splits, return empty if not list_of_lists: return [] # Number of layers is length of the sub-list from the first split max_len = len(list_of_lists[0]) merged = [None] * max_len for layer_idx in range(max_len): # collect that layer_idx from each split chunk_tensors = [] for split in list_of_lists: t = split[layer_idx] if layer_idx < len(split) else None if t is not None: chunk_tensors.append(t) if len(chunk_tensors) == 0: merged[layer_idx] = None else: merged[layer_idx] = torch.cat(chunk_tensors, dim=0) return merged @staticmethod def _repeat_list_tensors( tensor_list: List[Optional[torch.Tensor]], repeats: int ) -> List[Optional[torch.Tensor]]: out = [] for t in tensor_list: if t is None: out.append(None) else: out.append(t.repeat_interleave(repeats, dim=0)) return out @staticmethod def _select_list_tensors( tensor_list: List[Optional[torch.Tensor]], indices: torch.Tensor ) -> List[Optional[torch.Tensor]]: out = [] for t in tensor_list: if t is None: out.append(None) else: out.append(t.index_select(0, indices)) return out class TokenImportancePredictorAttentive(nn.Module): def __init__(self, config, pred_hid_size, num_heads, num_hidden_layers, dDash, intdim, \ attn_reduce_factor, dropout=0.1): """ Optimized Token Importance Predictor with parallel Q-K projections and simplified mapping. Args: config: Configuration object containing model parameters. pred_hid_size (int): Hidden size for the predictor's attention layer. num_heads (int): Number of attention heads. num_hidden_layers (int): Number of transformer layers to predict. dropout (float): Dropout probability. q_downscale (int): Factor to downscale the Q dimension for efficiency. intermediate_dim (int): Intermediate dimension for non-linear transformations in projections. """ super().__init__() self.config = config self.hidden_size = pred_hid_size self.num_heads = num_heads self.num_hidden_layers = num_hidden_layers self.dropout = dropout self.head_dim = pred_hid_size // (num_heads * 4) # Predictor head dimension is not the same as the model head dimension. self.rope_theta = config.rope_theta self.dDash = dDash self.intermediate_dim = intdim self.attn_reduce_factor = attn_reduce_factor self.max_position_embeddings = config.max_position_embeddings self.flash_attn = False assert pred_hid_size % (num_heads * 4) == 0, "pred_hid_size must be divisible by num_heads * 4." # Reduce the hidden size for attention computations self.hidden_size_reduced = self.hidden_size // self.attn_reduce_factor # For example, reduce to 1/4th assert self.hidden_size_reduced % self.num_heads == 0, "Reduced hidden size must be divisible by num_heads" self.attn_head_dim = self.hidden_size_reduced // self.num_heads # Input projection to reduce hidden size self.input_proj = nn.Linear(self.hidden_size, self.hidden_size_reduced, bias=False) # Query, Key, Value projections for attention self.q_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) self.k_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) self.v_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) # Output projection to restore hidden size # self.o_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) self.attn_dropout = nn.Dropout(self.dropout) # LayerNorm and Feed-forward network self.norm1 = nn.LayerNorm(self.hidden_size_reduced) self.norm2 = nn.LayerNorm(self.hidden_size) self.ffn_hidden_size = 2 * self.hidden_size_reduced # Typical FFN hidden size self.ffn = nn.Sequential( nn.Linear(self.hidden_size_reduced, self.ffn_hidden_size), nn.GELU(), nn.Linear(self.ffn_hidden_size, self.hidden_size), nn.Dropout(self.dropout) ) # Add extra LayerNorm for the importance branch when not using the old design. self.norm_importance = nn.LayerNorm(self.hidden_size) # Define Q and K projection layers for all layers in parallel with non-linearity[] # Output shape: [B, L, N * H * D'] self.q_proj_importance = nn.Sequential( nn.Linear(pred_hid_size, self.intermediate_dim, bias=False), nn.SiLU(), nn.Linear(self.intermediate_dim, num_hidden_layers * num_heads * self.dDash, bias=False) ) self.k_proj_importance = nn.Sequential( nn.Linear(pred_hid_size, self.intermediate_dim, bias=False), nn.SiLU(), nn.Linear(self.intermediate_dim, num_hidden_layers * num_heads * self.dDash, bias=False) ) # Initialize rotary positional embeddings self._init_rope() self._initialize_weights() self.device = None def _initialize_weights(self): for name, module in self.named_modules(): if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) # Xavier initialization for linear layers if module.bias is not None: nn.init.constant_(module.bias, 0) elif isinstance(module, nn.LayerNorm): nn.init.constant_(module.weight, 1.0) nn.init.constant_(module.bias, 0.0) elif isinstance(module, nn.MultiheadAttention): # Initialize in_proj_weight nn.init.xavier_uniform_(module.in_proj_weight) if module.in_proj_bias is not None: nn.init.constant_(module.in_proj_bias, 0) # Initialize out_proj nn.init.xavier_uniform_(module.out_proj.weight) if module.out_proj.bias is not None: nn.init.constant_(module.out_proj.bias, 0) def _init_rope(self): # send self.config but after modifying head_dim to be self.head_dim just in the function call config_copy = copy.deepcopy(self.config) config_copy.rope_scaling = { "factor": 32.0, "high_freq_factor": 4.0, "low_freq_factor": 1.0, "original_max_position_embeddings": 8192, "rope_type": "llama3" } config_copy.head_dim = self.attn_head_dim # Rotary embedding for attention layer self.rotary_emb_attn = LlamaRotaryEmbedding( config_copy ) config_copy.head_dim = self.dDash # Rotary embedding for importance projection self.rotary_emb_importance = LlamaRotaryEmbedding( config_copy ) def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False, layer_idx=None): """ Forward pass for the Optimized Token Importance Predictor. Args: hidden_states (torch.Tensor): Input tensor of shape [B, L, HQ]. attention_mask (torch.Tensor, optional): Attention mask of shape [B, 1, 1, L] or [B, 1, L, L]. position_ids (torch.Tensor, optional): Position IDs. past_key_value (tuple, optional): Past key and value states. use_cache (bool, optional): Whether to use cache. Returns: torch.Tensor: Importance scores of shape [B, N, H, L, L]. """ layer_idx = 0 # Guaranteed to be 0, as we only have one predictor! # Set device if not already set if self.device != hidden_states.device: self.device = hidden_states.device self.to(self.device) B, L, E = hidden_states.size() # Reduce hidden size hidden_states = hidden_states.to(self.input_proj.weight.dtype) hidden_states_reduced = self.input_proj(hidden_states) # [B, L, hidden_size_reduced] # Compute q, k, v for attention q = self.q_proj_attn(hidden_states_reduced) # [B, L, hidden_size_reduced] k = self.k_proj_attn(hidden_states_reduced) # [B, L, hidden_size_reduced] v = self.v_proj_attn(hidden_states_reduced) # [B, L, hidden_size_reduced] # Reshape q, k, v to [B, num_heads, L, attn_head_dim] q = q.view(B, L, self.num_heads, self.attn_head_dim).transpose(1, 2) # [B, num_heads, L, attn_head_dim] k = k.view(B, L, self.num_heads, self.attn_head_dim).transpose(1, 2) # [B, num_heads, L, attn_head_dim] v = v.view(B, L, self.num_heads, self.attn_head_dim).transpose(1, 2) # [B, num_heads, L, attn_head_dim] if (past_key_value is not None and layer_idx < len(past_key_value.predictor_primary_key) and past_key_value.predictor_primary_key[layer_idx] is not None): offset = past_key_value.predictor_primary_key[layer_idx].shape[2] # old_k.shape[2] else: offset = 0 # total seq length for new + old kv_seq_len = offset + L # Step 2: build position_ids for just the new chunk [offset..offset+L-1] if position_ids is None: # shape [B, L], e.g. [0..(offset+L-1)] position_ids = torch.arange(offset, offset + L, dtype=torch.long, device=self.device) position_ids = position_ids.unsqueeze(0).expand(B, L) # Step 3: apply rotary to just the new chunk k,v with the correct offset cos, sin = self.rotary_emb_attn(v, position_ids) q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids) # Step 4: ask the cache to append them. Then re‐assign k, v to the full cat if use_cache and past_key_value is not None: k, v = past_key_value.update_predictor_primary(k.detach(), v.detach(), layer_idx) kv_seq_len = k.size(2) # now includes old + new attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=True) attn_output = attn_output.to(q.dtype) attn_output = attn_output.transpose(1, 2).contiguous().view(B, L, self.hidden_size_reduced) attn_output = self.norm1(attn_output) ffn_output = self.ffn(attn_output) # Temporary measure, till old predictor fully deprecated hidden_states = self.norm2(hidden_states + ffn_output) B, L, E = hidden_states.size() # Importance projections H = self.num_heads N = self.num_hidden_layers hidden_states_for_importance = self.norm_importance(hidden_states) q_importance = self.q_proj_importance(hidden_states_for_importance) k_importance = self.k_proj_importance(hidden_states_for_importance) # Reshape and permute to [B, N, H, L, D'] q_importance = q_importance.view(B, L, N, H, self.dDash).permute(0, 2, 3, 1, 4).contiguous() # [B, N, H, L, D'] k_importance = k_importance.view(B, L, N, H, self.dDash).permute(0, 2, 3, 1, 4).contiguous() # [B, N, H, L, D'] # Flatten N and H for efficient computation q_importance = q_importance.view(B * N * H, L, self.dDash) # [BNH, L, D'] k_importance = k_importance.view(B * N * H, L, self.dDash) # [BNH, L, D'] # Apply rotary positional embeddings cos, sin = self.rotary_emb_importance(k_importance, position_ids) q_importance, k_importance = apply_rotary_pos_emb(q_importance, k_importance, cos, sin, position_ids) if use_cache and past_key_value is not None: k_importance = past_key_value.update_predictor_importance(k_importance.detach(), layer_idx) k_importance = k_importance.view(B * H, N, -1, self.dDash) # [BNH, L, D'] q_importance = q_importance.view(B * H, N, -1, self.dDash) # [BH, N, L, D'] return q_importance, k_importance class HeadImportancePredictor(nn.Module): def __init__(self, config, pred_hid_size, num_heads, num_hidden_layers, dDash, intdim, \ attn_reduce_factor, dropout=0.1): """ Optimized Token Importance Predictor with parallel Q-K projections and simplified mapping. Args: config: Configuration object containing model parameters. pred_hid_size (int): Hidden size for the predictor's attention layer. num_heads (int): Number of attention heads. num_hidden_layers (int): Number of transformer layers to predict. dropout (float): Dropout probability. q_downscale (int): Factor to downscale the Q dimension for efficiency. intermediate_dim (int): Intermediate dimension for non-linear transformations in projections. """ super().__init__() self.is_head_predictor = None self.config = config self.hidden_size = pred_hid_size self.num_heads = num_heads self.num_hidden_layers = num_hidden_layers self.dropout = dropout self.head_dim = pred_hid_size // (num_heads * 4) self.rope_theta = config.rope_theta self.dDash = dDash self.intermediate_dim = intdim self.attn_reduce_factor = attn_reduce_factor self.max_position_embeddings = config.max_position_embeddings self.flash_attn = False # Reduce the hidden size for attention computations self.hidden_size_reduced = self.hidden_size // self.attn_reduce_factor # For example, reduce to 1/4th assert self.hidden_size_reduced % self.num_heads == 0, "Reduced hidden size must be divisible by num_heads" self.attn_head_dim = self.hidden_size_reduced // self.num_heads # Input projection to reduce hidden size self.input_proj = nn.Linear(self.hidden_size, self.hidden_size_reduced, bias=False) # Query, Key, Value projections for attention self.q_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) self.k_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) self.v_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) # Output projection to restore hidden size # self.o_proj_attn = nn.Linear(self.hidden_size_reduced, self.hidden_size_reduced, bias=False) self.attn_dropout = nn.Dropout(self.dropout) # LayerNorm and Feed-forward network self.norm1 = nn.LayerNorm(self.hidden_size_reduced) self.norm2 = nn.LayerNorm(self.hidden_size) self.ffn_hidden_size = 4 * self.hidden_size_reduced # Typical FFN hidden size self.ffn = nn.Sequential( nn.Linear(self.hidden_size_reduced, self.ffn_hidden_size), nn.GELU(), nn.Linear(self.ffn_hidden_size, self.num_heads * self.num_hidden_layers), ) # Initialize rotary positional embeddings self._init_rope() self._initialize_weights() self.device = None def _initialize_weights(self): for name, module in self.named_modules(): if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) # Xavier initialization for linear layers if module.bias is not None: nn.init.constant_(module.bias, 0) elif isinstance(module, nn.LayerNorm): nn.init.constant_(module.weight, 1.0) nn.init.constant_(module.bias, 0.0) elif isinstance(module, nn.MultiheadAttention): # Initialize in_proj_weight nn.init.xavier_uniform_(module.in_proj_weight) if module.in_proj_bias is not None: nn.init.constant_(module.in_proj_bias, 0) # Initialize out_proj nn.init.xavier_uniform_(module.out_proj.weight) if module.out_proj.bias is not None: nn.init.constant_(module.out_proj.bias, 0) def _init_rope(self): config_copy = copy.deepcopy(self.config) config_copy.head_dim = self.attn_head_dim # Rotary embedding for attention layer self.rotary_emb_attn = LlamaRotaryEmbedding( config_copy ) # Rotary embedding for importance projection self.rotary_emb_importance = LlamaRotaryEmbedding( config_copy ) def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False): """ Forward pass for the Optimized Token Importance Predictor. Args: hidden_states (torch.Tensor): Input tensor of shape [B, L, HQ]. attention_mask (torch.Tensor, optional): Attention mask of shape [B, 1, 1, L] or [B, 1, L, L]. position_ids (torch.Tensor, optional): Position IDs. past_key_value (tuple, optional): Past key and value states. use_cache (bool, optional): Whether to use cache. Returns: torch.Tensor: Importance scores of shape [B, N, H, L, L]. """ # Set device if not already set if self.device != hidden_states.device: self.device = hidden_states.device self.to(self.device) B, L, E = hidden_states.size() if past_key_value is None: past_key_value = {} # if L == 1: # import pdb; pdb.set_trace() past_primary = past_key_value.get('primary', None) # Reduce hidden size hidden_states = hidden_states.to(self.input_proj.weight.dtype) hidden_states_reduced = self.input_proj(hidden_states) # [B, L, hidden_size_reduced] # Compute q, k, v for attention q = self.q_proj_attn(hidden_states_reduced) # [B, L, hidden_size_reduced] k = self.k_proj_attn(hidden_states_reduced) # [B, L, hidden_size_reduced] v = self.v_proj_attn(hidden_states_reduced) # [B, L, hidden_size_reduced] # Reshape q, k, v to [B, num_heads, L, attn_head_dim] q = q.view(B, L, self.num_heads, self.attn_head_dim).transpose(1, 2) # [B, num_heads, L, attn_head_dim] k = k.view(B, L, self.num_heads, self.attn_head_dim).transpose(1, 2) # [B, num_heads, L, attn_head_dim] v = v.view(B, L, self.num_heads, self.attn_head_dim).transpose(1, 2) # [B, num_heads, L, attn_head_dim] # Compute kv_seq_len before concatenation if past_primary is not None: past_L = past_primary[0].shape[2] kv_seq_len = past_L + L else: kv_seq_len = L # Apply rotary positional embeddings based on kv_seq_len cos, sin = self.rotary_emb_attn(v, position_ids) if position_ids is None: position_ids = torch.arange(kv_seq_len, dtype=torch.long, device=self.device) position_ids = position_ids.unsqueeze(0).expand(B, kv_seq_len) if past_primary is not None: # Concatenate past k and v k = torch.cat([past_primary[0], k], dim=2) # [B, num_heads, past_L + L, attn_head_dim] v = torch.cat([past_primary[1], v], dim=2) # [B, num_heads, past_L + L, attn_head_dim] # Apply rotary embeddings after concatenation q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids) # Update cache if use_cache is True if use_cache: past_key_value['primary'] = (k.detach(), v.detach()) # if self.flash_attn: # sm_scale = 1.0 / math.sqrt(self.attn_head_dim) # attn_output = attention(q.contiguous().to(torch.float16), k.contiguous().to(torch.float16), v.contiguous().to(torch.float16), True, sm_scale).to(q.dtype) # else: # attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=True) attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=True) attn_output = attn_output.to(q.dtype) attn_output = attn_output.transpose(1, 2).contiguous().view(B, L, self.hidden_size_reduced) attn_output = self.norm1(attn_output) head_importances = self.ffn(attn_output) return head_importances, past_key_value def calculate_hit_metrics(estimated_importance: torch.Tensor, true_importance: torch.Tensor, top_k_ratio: float = 0.5) -> Tuple[float, float, float]: """ Calculate hit accuracy, mean, and max rank correlation between estimated and true importance tensors. We compute metrics along the last dimension of the input tensors. Shapes: - 4D token-importance: [B, H, L, L]. We slice the last query (index -1) => [B, H, L]. - 3D head-importance: [B, L, H]. We use all of it as-is => [B, L, H]. Args: estimated_importance (torch.Tensor): [B, H, L, L] or [B, L, H] true_importance (torch.Tensor): [B, H, L, L] or [B, L, H] top_k_ratio (float): Fraction of top-k elements to consider for hit accuracy (default=0.5). Returns: (hit_accuracy, mean_corr, max_corr): hit_accuracy (float): Intersection ratio of top-k sets (0..1). mean_corr (float): Average Spearman rank correlation over all [B, ...]. max_corr (float): Maximum Spearman rank correlation among all [B, ...]. """ # 1) Standardize shapes so the last dimension is what we rank over. if estimated_importance.dim() == 4: # Shape is [B, H, L, L] => slice to keep only the last query => [B, H, L] estimated_importance = estimated_importance[:, :, -1, :] true_importance = true_importance[:, :, -1, :] # after slicing: [B, H, L] # For intersection denominator => top_k * B * H denom_for_hits = estimated_importance.size(0) * estimated_importance.size(1) elif estimated_importance.dim() == 3: # Shape is [B, L, H], the last dimension is H # For intersection denominator => top_k * B * L denom_for_hits = estimated_importance.size(0) * estimated_importance.size(1) else: raise ValueError("Tensors must be either 4D [B,H,L,L] or 3D [B,L,H].") # 2) Compute Spearman rank correlation along the last dimension. # Sort indices in descending order => get 'ranks' for correlation. _, sorted_esti = torch.sort(estimated_importance, dim=-1, descending=True) _, sorted_true = torch.sort(true_importance, dim=-1, descending=True) # Spearman's rho = 1 - 6 * sum(d^2) / [n*(n^2 - 1)] n = sorted_esti.shape[-1] d = sorted_esti.float() - sorted_true.float() d_squared = d ** 2 sum_d_squared = d_squared.sum(dim=-1) rank_corr = 1 - (6 * sum_d_squared) / (n * (n**2 - 1)) # shape: [B,H] or [B,L] mean_corr = rank_corr.mean().item() max_corr = rank_corr.max().item() # 3) Compute top-k hit accuracy along the last dimension. top_k = max(1, int(n * top_k_ratio)) _, top_esti_indices = torch.topk(estimated_importance, top_k, dim=-1) _, top_true_indices = torch.topk(true_importance, top_k, dim=-1) # top_esti_indices => [B,H,top_k] or [B,L,top_k] # top_true_indices => [B,H,top_k] or [B,L,top_k] # matches => [B,H,top_k,top_k] or [B,L,top_k,top_k] matches = (top_esti_indices.unsqueeze(-1) == top_true_indices.unsqueeze(-2)) intersection = matches.any(dim=-1).sum(dim=-1) # => [B,H] or [B,L] # Each [B,H] or [B,L] element can have at most 'top_k' matches, so total is top_k * denom_for_hits. total_possible = top_k * denom_for_hits hit_accuracy = intersection.sum().item() / total_possible # => 0..1 return hit_accuracy, mean_corr, max_corr def threshold_to_mask(unadj_importance_mask, perhead_thresholds, min_sparse_index, bsz, q_len, key_len): """ Create a mask tensor based on per-head thresholds, setting values below the threshold to -inf. Args: - unadj_importance_mask: torch.Tensor of shape [B, H, Lq, Lk]. - perhead_thresholds: torch.Tensor of shape [H], per-head thresholds. - min_sparse_index: Minimum index for sparsity; values below this index will not be masked. - bsz: Batch size. - q_len: Query length (Lq). - key_len: Key length (Lk). Returns: - mask_tensor: torch.Tensor of shape [B, H, Lq, Lk], with values below threshold as -inf. """ # Ensure perhead_thresholds is in the correct shape for broadcasting thresholds_broadcast = perhead_thresholds.view(1, -1, 1, 1) # [1, H, 1, 1] # Compare unadj_importance_mask with thresholds to create a mask mask_tensor = torch.where( unadj_importance_mask >= thresholds_broadcast, torch.zeros_like(unadj_importance_mask), torch.full_like(unadj_importance_mask, float('-inf')) ) # [B, H, Lq, Lk] # Ensure mask_tensor has mask_tensor[:, :, :, :min_sparse_index] = 0 mask_tensor[:, :, :, :min_sparse_index] = 0.0 return mask_tensor class SlidingWindowCache: def __init__(self, max_seq_len, sliding_window, device): self.sliding_window = sliding_window self.device = device if sliding_window is None: self.max_seq_len = 0 self.window = None else: self.max_seq_len = max_seq_len self.window = self._create_window(self.max_seq_len) def _create_window(self, seq_len): idx = torch.arange(seq_len, device=self.device) query = idx.unsqueeze(1) # [seq_len, 1] key = idx.unsqueeze(0) # [1, seq_len] win = (key >= (query - self.sliding_window + 1)) & (key <= query) return win.unsqueeze(0).unsqueeze(0) # [1,1,seq_len,seq_len] def get_window(self, q_len, key_len): if self.sliding_window is None: return None req = max(q_len, key_len) if req > self.max_seq_len: self.max_seq_len = req self.window = self._create_window(self.max_seq_len) return self.window[:, :, :q_len, :key_len] def enforce_sliding_window(mask_tensor, window): if window is None: return mask_tensor return mask_tensor.masked_fill(window, 0.0) def sorted_index_to_mask( sorted_indices, attention_mask, min_sparse_index, bsz, q_len, key_len, sparse_aggression, sliding_window=None ): """ sorted_indices: [B, H, q_len, key_len] attention_mask: [1, 1, q_len, key_len] (True = keep, False = mask out, or vice versa) min_sparse_index: guaranteed front region to keep sliding_window: guaranteed trailing region (for each query) to keep sparse_aggression: float in [0,1], fraction of keys to drop or keep """ device = sorted_indices.device dtype = sorted_indices.dtype # Step 1: Compute base K if q_len == 1: query_positions = torch.arange(q_len, device=device).view(1, 1, q_len, 1).float() query_positions[0] = key_len + 1 else: query_positions = torch.arange(q_len, device=device).view(1, 1, q_len, 1).float() + 1.0 K_original = torch.ceil(query_positions * sparse_aggression).long() # [1,1,q_len,1] K_original = torch.clamp(K_original, max=key_len) # Step 1b: Incorporate guaranteed region guaranteed = min_sparse_index if sliding_window is not None: guaranteed += sliding_window # Subtract guaranteed from the original K K_adjusted = K_original - guaranteed # Ensure K_adjusted is at least 0 K_adjusted = torch.clamp(K_adjusted, min=0, max=key_len) # Step 2: Expand attention_mask to [B,H,q_len,key_len] attention_mask_expanded = attention_mask.expand(bsz, -1, -1, -1) attention_mask_expanded = attention_mask_expanded.expand(-1, sorted_indices.size(1), -1, -1) # Convert True -> 1, False -> 0 attention_mask_expanded = (~attention_mask_expanded.bool()).int() # Step 3: Gather (reorder) mask by sorted_indices gathered_mask = torch.gather(attention_mask_expanded, dim=-1, index=sorted_indices) # Step 4: cumsum along sorted dimension gathered_mask_float = gathered_mask.float() cum_sum = torch.cumsum(gathered_mask_float, dim=-1) # [B,H,q_len,key_len] # Step 5: Compare cumsum <= K_adjusted # Expand K_adjusted to [B,H,q_len,key_len] for broadcast K_broadcast = K_adjusted.view(1, 1, q_len, 1).expand_as(cum_sum) selected_mask = (cum_sum <= K_broadcast) # Step 6: Prepare final mask_tensor with -inf by default mask_tensor = torch.full_like(attention_mask_expanded.float(), float('-inf')) # Step 7: Scatter 0 where selected, -inf otherwise scatter_values = torch.zeros_like(gathered_mask_float) scatter_values = scatter_values.masked_fill(~selected_mask, float('-inf')) mask_tensor.scatter_(-1, sorted_indices, scatter_values) # Step 8: Force the guaranteed front region unmasked mask_tensor[:, :, :, :min_sparse_index] = 0.0 # We do NOT forcibly unmask the trailing `sliding_window` here, # because we typically do it with a separate function that # ensures the last `sliding_window` positions are unmasked for each query. # Replace with self.sliding_window where referenced # Where not referenced, reduce budget in calculation. return mask_tensor class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, config=None): self.scaling_factor = scaling_factor super().__init__(config) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) t = t / self.scaling_factor freqs = torch.einsum("i,j->ij", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False) class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding): """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, config=None): self.scaling_factor = scaling_factor super().__init__(config) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len if seq_len > self.max_position_embeddings: base = self.base * ( (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1) ) ** (self.dim / (self.dim - 2)) inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False) self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class LlamaAttentionExperimental(nn.Module): def __init__(self, config: LlamaConfig, producer=None, layer_idx=0): super().__init__() self.config = config self.hidden_size = config.hidden_size self.num_hidden_layers = config.num_hidden_layers self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.inference_mode = False self.producer = producer self.layer_idx = layer_idx self.token_sparse_method = None self.sparse_aggression = None self.stream_llm_start_size = None self.dDash = None self.intdim = None self.attn_reduce_factor = None self.head_attn_reduce_factor = None self.effective_sparsity = None self.min_sparse_index = None self.pred_hid_size = self.hidden_size self.num_tok_per_page = None self.calc_hitrates = False self.flash_attn = False self.train_headpredictor = False self.calibrate_thresholds = False self.test_with_thresholds = False self.old_predictor = None if self.layer_idx > 0: self.mseloss = MSELoss(reduction='none') self.msemagn_loss = None self.headmseloss = MSELoss(reduction='none') self.headmsemagn_loss = None if self.producer is None: # This is the producer layer self.q_importance = None # Shared mask across layers during inference self.k_importance = None self.head_importances = None self.actmagn_masklist = {} self.available_tokens = {} # Attention setup 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._init_rope() def update_predictor(self): self.sparse_token_predictor = TokenImportancePredictorAttentive( self.config, self.pred_hid_size, self.num_heads, self.num_layers_pred, dropout=0.1, dDash = self.dDash, \ intdim = self.intdim, attn_reduce_factor=self.attn_reduce_factor ).to('cuda:0') self.sparse_token_predictor.flash_attn = self.flash_attn if self.train_headpredictor: self.sparse_head_predictor = HeadImportancePredictor( self.config, self.pred_hid_size, self.num_heads, self.num_layers_pred, dropout=0.1, dDash = self.dDash, \ intdim = self.intdim, attn_reduce_factor=self.head_attn_reduce_factor ).to('cuda:0') self.sparse_head_predictor.flash_attn = self.flash_attn def set_token_sparsity(self): assert self.token_sparse_method is not None, "Set token sparse method first!" if self.token_sparse_method is not None: try: mname = self.config._name_or_path.split("/")[-1] read_path = f"threshold_calibs/{mname}/{self.token_sparse_method}.pkl" threshold_model_dictionary = torch.load(read_path) self.tok_calibration_set = threshold_model_dictionary except: pass if self.token_sparse_method == "LazyLLM": if self.layer_idx <= 9: self.sparse_aggression = 1 elif self.layer_idx <= 19: self.sparse_aggression = 0.7 elif self.layer_idx <= 28: self.sparse_aggression = 0.4 else: self.sparse_aggression = 0.1 elif "fixed" in self.token_sparse_method: if self.layer_idx == 0: self.sparse_aggression = 1 else: self.sparse_aggression = 1 - float(self.token_sparse_method.split("_")[1].split("pc")[0])/100. elif "progressive" in self.token_sparse_method: pc_drop = float(self.token_sparse_method.split("_")[1].split("pc")[0])/100. self.sparse_aggression = (1 - pc_drop) ** (self.layer_idx) # (x% per layer, progressive_xpc style) else: raise ValueError(f"Unknown token sparsity method {self.token_sparse_method}") def _init_rope(self): if self.config.rope_scaling is None: self.rotary_emb = LlamaRotaryEmbedding( self.config ) else: scaling_type = self.config.rope_scaling.get("type") or self.config.rope_scaling.get("rope_type") scaling_factor = self.config.rope_scaling["factor"] if scaling_type == "linear" or scaling_type == 'llama3': self.rotary_emb = LlamaLinearScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope_theta, config=self.config ) elif scaling_type == "dynamic": self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope_theta, config=self.config ) else: raise ValueError(f"Unknown RoPE scaling type {scaling_type}") def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Union[DynamicCache, PredictorDynamicCache]] = None, output_attentions: bool = False, use_cache: bool = False, padding_mask: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[PredictorDynamicCache]]: bsz, q_len, _ = hidden_states.size() Ltrack = hidden_states.size(1) if self.config.pretraining_tp > 1: key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp query_slices = self.q_proj.weight.split( (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0 ) key_slices = self.k_proj.weight.split(key_value_slicing, dim=0) value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)] query_states = torch.cat(query_states, dim=-1) key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)] key_states = torch.cat(key_states, dim=-1) value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)] value_states = torch.cat(value_states, dim=-1) else: query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) evalmode = self.eval_llm_mode num_tokens_to_keep = int(q_len * self.sparse_aggression) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) # cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) # AHMED: Modified this to use the newer version. cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) if use_cache: key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx) kv_seq_len = key_states.shape[-2] final_mask = None key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) key_len = key_states.size(2) bsz, q_len = query_states.size(0), query_states.size(2) if attention_mask is None: # We want a [q_len, kv_seq_len] boolean upper-triangular mask causal_mask_2d = torch.ones(q_len, kv_seq_len, device=hidden_states.device, dtype=torch.bool).triu(diagonal=1) # Then shape it to [bsz, 1, q_len, kv_seq_len] causal_mask_4d = causal_mask_2d.unsqueeze(0).expand(bsz, 1, q_len, kv_seq_len) # Now fill -inf where the mask is True attention_mask = torch.full_like(causal_mask_4d, 0, dtype=hidden_states.dtype) if q_len != 1: attention_mask = attention_mask.masked_fill(causal_mask_4d, float("-inf")) if self.inference_mode: min_sparse_index = self.min_sparse_index with torch.no_grad(): if evalmode == "ExpPred": if self.layer_idx > 0: q_importance_tensor = self.producer.q_importance[:, self.layer_idx % self.producer_frequency, :, :].float().to(query_states.device) # [BH, Lq, D'] k_importance_tensor = self.producer.k_importance[:, self.layer_idx % self.producer_frequency, :, :].float().to(key_states.device) # [BH, Lk, D'] importance_mask = torch.bmm(q_importance_tensor, k_importance_tensor.transpose(-2, -1)) / math.sqrt(self.dDash) # [BH, Lq, Lk] importance_mask = importance_mask.view(bsz, self.num_heads, q_len, key_len) # [B, H, Lq, Lk] attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) / math.sqrt(self.head_dim) if self.calc_hitrates: self.tok_hit_acc, self.tok_mean_rank_corr, self.tok_max_rank_corr = calculate_hit_metrics( estimated_importance=importance_mask, true_importance=attn_weights, top_k_ratio=0.5 ) if self.calibrate_thresholds: ### Threshold variance investigation unadj_importance_mask = importance_mask.clone() importance_mask = torch.softmax(importance_mask + attention_mask, dim=-1) sorted_indices = torch.argsort(importance_mask, dim=-1, descending=True) sorted_indices = sorted_indices[:, :, -q_len:, :] sorted_values, sorted_ix = torch.sort(importance_mask, dim=-1) sorted_true_values, _ = torch.sort(torch.gather(unadj_importance_mask, dim=-1, index=sorted_ix), dim=-1) true_thresholds = sorted_true_values[:, :, :, int(importance_mask.size(-1) * self.sparse_aggression)] thresholds = sorted_values[:, :, :, int(importance_mask.size(-1) * self.sparse_aggression)] self.true_threshmean = true_thresholds self.threshmean = thresholds if self.test_with_thresholds: unadj_importance_mask = importance_mask.clone() perhead_thresholds = self.tok_calibration_set[self.layer_idx - 1].to(unadj_importance_mask.device) # 0 does not have calibration data. mask_tensor = threshold_to_mask(unadj_importance_mask, perhead_thresholds, min_sparse_index, bsz, q_len, key_len) else: importance_mask = torch.softmax(importance_mask + attention_mask, dim=-1) sorted_indices = torch.argsort(importance_mask, dim=-1, descending=True) sorted_indices = sorted_indices[:, :, -q_len:, :] mask_tensor = sorted_index_to_mask(sorted_indices, attention_mask, min_sparse_index, bsz, q_len, key_len, self.sparse_aggression, self.sliding_window) ### Threshold variance investigation if self.sliding_window is not None: if not hasattr(self, "window_cache"): self.window_cache = SlidingWindowCache(max_seq_len=1024, sliding_window=self.sliding_window, device=mask_tensor.device) window = self.window_cache.get_window(q_len, key_len) mask_tensor = enforce_sliding_window(mask_tensor, window) final_mask = mask_tensor self.final_mask_investigate = final_mask attn_weights = attn_weights + mask_tensor + attention_mask else: attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) / math.sqrt(self.head_dim) attn_weights = attn_weights + attention_mask else: raise ValueError(f"Unknown eval mode {evalmode}") attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype) attn_output = torch.matmul(attn_weights, value_states) else: attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) / math.sqrt(self.head_dim) if self.layer_idx > 0: q_importance_tensor = self.producer.q_importance[:, self.layer_idx % self.producer_frequency, :, :].float().to(query_states.device) # [BH, Lq, D'] k_importance_tensor = self.producer.k_importance[:, self.layer_idx % self.producer_frequency, :, :].float().to(key_states.device) # [BH, Lk, D'] importance_mask = torch.bmm(q_importance_tensor, k_importance_tensor.transpose(-2, -1)) / math.sqrt(self.dDash) # [BH, Lq, Lk] importance_mask = importance_mask.view(bsz, self.num_heads, q_len, key_len) # [B, H, Lq, Lk] if self.lookahead == 0: self.msemagn_loss = self.mseloss(attn_weights, importance_mask) else: self.msemagn_loss = self.mseloss(attn_weights[:, :, self.lookahead:, :], importance_mask[:, :, :-self.lookahead, :]) self.msemagn_loss = (self.msemagn_loss).mean(dim=(-1, -2)) self.msemagn_loss = self.msemagn_loss.mean() if self.calc_hitrates: self.tok_hit_acc, self.tok_mean_rank_corr, self.tok_max_rank_corr = calculate_hit_metrics( estimated_importance=importance_mask, true_importance=attn_weights, top_k_ratio=0.5 ) if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(value_states.dtype) attn_output = torch.matmul(attn_weights, value_states) if self.layer_idx > 0 and self.train_headpredictor: head_importance_tensor = self.producer.head_importances[:, :, :, self.layer_idx % self.producer_frequency].float().to(attn_output.device) attn_head_weights = attn_output.mean(dim=-1).permute(0, 2, 1) self.headmsemagn_loss = self.headmseloss(attn_head_weights, head_importance_tensor).mean() if self.calc_hitrates: self.head_hit_acc, self.head_mean_rank_corr, self.head_max_rank_corr = calculate_hit_metrics( estimated_importance=head_importance_tensor, true_importance=attn_head_weights, top_k_ratio=0.5 ) else: self.headmsemagn_loss = 0 if self.calc_hitrates: self.head_hit_acc, self.head_mean_rank_corr, self.head_max_rank_corr = 0, 0, 0 checkeverytime = hasattr(self, 'test_with_thresholds') if checkeverytime: checkeverytime = self.test_with_thresholds if final_mask is not None: if self.effective_sparsity is None or checkeverytime: true_mask = final_mask + attention_mask num_deact = true_mask.bool().sum(dim=-1) # Number of tokens disabled. causally_deact = (attention_mask.bool()).sum(dim=-1).expand_as(num_deact) # Number of tokens disabled causally anyway additional_deact = (num_deact - causally_deact) num_active = (~attention_mask.bool()).sum(dim=-1).expand_as(num_deact) # Number of tokens active at this position if zero-sparsity effective_sparsity = 100 * (additional_deact.float() / num_active.float()).mean().item() self.effective_sparsity = effective_sparsity print("Effective Sparsity:", effective_sparsity, "%\t Sequence Length:", q_len) if self.layer_idx == 0: if self.effective_sparsity is None: self.effective_sparsity = 0.0 attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(bsz, -1, self.hidden_size) if self.config.pretraining_tp > 1: attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2) o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1) attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)]) else: attn_output = self.o_proj(attn_output) if self.producer is None: try: q_importance, k_importance = self.sparse_token_predictor( hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, # the same single cache use_cache=use_cache, layer_idx=self.layer_idx, # or pass 0 ) if self.train_headpredictor: head_importances, past_key_value_hp = self.sparse_head_predictor( hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value_hp, use_cache=use_cache ) head_importances = head_importances.view(bsz, q_len, self.num_heads, self.num_hidden_layers) # [B L H N] q_len = attn_output.size(1) k_len = k_importance.size(-1) except: print(traceback.format_exc()) import pdb; pdb.set_trace() self.q_importance = q_importance self.k_importance = k_importance if self.train_headpredictor: if self.head_importances is None: self.head_importances = head_importances else: self.head_importances = torch.cat([self.head_importances, head_importances], dim=1) # if self.layer_idx == 31: # if q_len == 1: # self.dtok += 1 # print(f"Primary Key-Value Shape: {past_key_value.predictor_primary_key[0].shape}, Importance: {past_key_value.predictor_importance_key[0].shape}, Tok-Decoded: {self.dtok}") # else: # self.dtok = 0 if not output_attentions: attn_weights = None return attn_output, attn_weights def convert_kvcache_experimental(model, config, producer_frequency): producer_layer = None producer_layer_device = None layer_counter = {'idx': 0} def recurse_convert(parent_module): nonlocal producer_layer nonlocal producer_layer_device for name, module in parent_module._modules.items(): if len(list(module.children())) > 0: recurse_convert(module) if isinstance(module, LlamaAttention): device = next(module.parameters()).device dtype = next(module.parameters()).dtype if layer_counter['idx'] % producer_frequency == 0: new_module = LlamaAttentionExperimental(config).to(dtype).to(device) producer_layer = new_module producer_layer_device = device else: new_module = LlamaAttentionExperimental( config, producer=producer_layer, layer_idx=layer_counter['idx'] ).to(dtype).to(device) new_module.load_state_dict(module.state_dict(), strict=False) is_producer = layer_counter['idx'] % producer_frequency == 0 if is_producer: print(f"Converted Producer layer '{name}' to LlamaAttentionExperimental at layer index {layer_counter['idx']}") else: print(f"Converted layer '{name}' to LlamaAttentionExperimental at layer index {layer_counter['idx']}") parent_module._modules[name] = new_module layer_counter['idx'] += 1 recurse_convert(model) producer_layer = producer_layer.to(producer_layer_device) return model # --------------------------------------------------------------------- # 1) Custom Config subclass # --------------------------------------------------------------------- class LlamaButlerConfig(LlamaConfig): """ Extends HF's LlamaConfig to hold optional extra parameters for the "Butler" logic. You can store your custom attributes here, so they can be serialized in config.json. """ model_type = "llama_butler" def __init__( self, eval_llm_mode="ExpPred", token_sparse_method="fixed_50pc", producer_frequency=8, dDash=16, attn_reduce_factor=4, head_attn_reduce_factor=4, intdim=256, flash_attn=False, train_headpredictor=False, min_sparse_index=5, lookahead=0, sliding_window=None, **kwargs ): super().__init__(**kwargs) self.eval_llm_mode = eval_llm_mode self.token_sparse_method = token_sparse_method self.producer_frequency = producer_frequency self.dDash = dDash self.attn_reduce_factor = attn_reduce_factor self.head_attn_reduce_factor = head_attn_reduce_factor self.intdim = intdim self.flash_attn = flash_attn self.train_headpredictor = train_headpredictor self.min_sparse_index = min_sparse_index self.lookahead = lookahead self.sliding_window = sliding_window # --------------------------------------------------------------------- # 2) The main Butler model class # --------------------------------------------------------------------- class LlamaButlerForCausalLM(LlamaForCausalLM): """ A subclass of HF's LlamaForCausalLM that: - Patches each LlamaAttention to your LlamaAttentionExperimental - Sets specialized attributes (eval_llm_mode, etc.) - Overrides _prepare_cache_for_generation to inject PredictorDynamicCache """ # Let HF auto-detect this config class from config.json: config_class = LlamaButlerConfig def __init__(self, config: LlamaButlerConfig): super().__init__(config) """ HF's LlamaForCausalLM initializes: self.model = LlamaModel(config) self.lm_head = nn.Linear(...) """ # 1) Patch the underlying LlamaModel to replace LlamaAttention with LlamaAttentionExperimental self.model = convert_kvcache_experimental( self.model, config, config.producer_frequency ) # 2) Optionally, set per-module attributes so each LlamaAttentionExperimental knows about them: for module in self.model.modules(): if module.__class__.__name__.endswith("AttentionExperimental"): # Set these from your config. Or you can hardcode them if you prefer. module.eval_llm_mode = config.eval_llm_mode module.token_sparse_method = config.token_sparse_method module.set_token_sparsity() # e.g. sets module.sparse_aggression module.producer_frequency = config.producer_frequency module.dDash = config.dDash module.attn_reduce_factor = config.attn_reduce_factor module.head_attn_reduce_factor = config.head_attn_reduce_factor module.intdim = config.intdim module.flash_attn = config.flash_attn module.train_headpredictor = config.train_headpredictor module.min_sparse_index = config.min_sparse_index module.lookahead = config.lookahead module.sliding_window = config.sliding_window module.num_layers_pred = config.producer_frequency # example usage # If this is a "producer layer" (mod.layer_idx % freq == 0), run update_predictor(): if hasattr(module, "layer_idx") and (module.layer_idx % config.producer_frequency == 0): module.update_predictor() # 3) Patch the dynamic cache (past_key_values) creation. For your evaluation modes: if config.eval_llm_mode in ["ExpPred", "ReplAttn"]: self._prepare_cache_for_generation = self._patched_prepare_cache_for_generation.__get__( self, self.__class__ ) # ----------------------------------------------------------------- # 3) The custom `_prepare_cache_for_generation` override # ----------------------------------------------------------------- def _patched_prepare_cache_for_generation( self, generation_config: GenerationConfig, model_kwargs: Dict, *args, **kwargs ): """ This override injects a PredictorDynamicCache in place of the standard 'past_key_values'. """ if "past_key_values" not in model_kwargs or model_kwargs["past_key_values"] is None: model_kwargs["past_key_values"] = PredictorDynamicCache() return model_kwargs