Add final Butler model

config.json

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "LlamaButlerForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attn_reduce_factor": 8,
+  "auto_map": {
+    "AutoConfig": "modeling_llama_butler:LlamaButlerConfig",
+    "AutoModel": "modeling_llama_butler:LlamaButlerForCausalLM",
+    "AutoModelForCausalLM": "modeling_llama_butler:LlamaButlerForCausalLM"
+  },
+  "bos_token_id": 128000,
+  "dDash": 16,
+  "eos_token_id": 128001,
+  "eval_llm_mode": "ExpPred",
+  "flash_attn": false,
+  "head_attn_reduce_factor": 2,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intdim": 512,
+  "intermediate_size": 8192,
+  "lookahead": 0,
+  "max_position_embeddings": 131072,
+  "min_sparse_index": 8,
+  "mlp_bias": false,
+  "model_type": "llama_butler",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 8,
+  "pretraining_tp": 1,
+  "producer_frequency": 16,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "factor": 32.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "llama3"
+  },
+  "rope_theta": 500000.0,
+  "sliding_window": 128,
+  "tie_word_embeddings": true,
+  "token_sparse_method": "fixed_50pc",
+  "torch_dtype": "float32",
+  "train_headpredictor": false,
+  "transformers_version": "4.48.3",
+  "use_cache": true,
+  "vocab_size": 128256
+}

conversion.py

@@ -0,0 +1,81 @@
+from transformers import LlamaForCausalLM, LlamaConfig, AutoTokenizer
+import torch
+import os
+
+question = "A $y$-intercept is a point on the graph that lies on the $y$-axis, so $x = 0$. Hence, the number $y$-intercepts corresponds to the number of real solutions of the quadratic equation $y^2 - 4y - 1 = 0$. The discriminant of this quadratic equation is $(-4)^2 + 4 \cdot 1 \cdot (-1) = 20$, which is positive, so the quadratic has two distinct real roots. Therefore, the number of $y$-intercepts is $\boxed{2}$. \n  \n [asy] \n size(150); \n real ticklen=3; \n real tickspace=2; \n  \n real ticklength=0.1cm; \n real axisarrowsize=0.14cm; \n pen axispen=black+1.3bp; \n real vectorarrowsize=0.2cm; \n real tickdown=-0.5; \n real tickdownlength=-0.15inch; \n real tickdownbase=0.3; \n real wholetickdown=tickdown; \n void rr_cartesian_axes(real xleft, real xright, real ybottom, real ytop, real xstep=1, real ystep=1, bool \n  \n useticks=false, bool complexplane=false, bool usegrid=true) { \n  \n import graph; \n  \n real i; \n  \n if(complexplane) { \n  \n label('$\textnormal{Re}$',(xright,0),SE); \n  \n label('$\textnormal{Im}$',(0,ytop),NW); \n  \n } else { \n  \n label('$x$',(xright+0.4,-0.5)); \n  \n label('$y$',(-0.5,ytop+0.2)); \n  \n } \n  \n ylimits(ybottom,ytop); \n  \n xlimits( xleft, xright); \n  \n real[] TicksArrx,TicksArry; \n  \n for(i=xleft+xstep; i<xright; i+=xstep) { \n  \n if(abs(i) >0.1) { \n  \n TicksArrx.push(i); \n  \n } \n  \n } \n  \n for(i=ybottom+ystep; i<ytop; i+=ystep) { \n  \n if(abs(i) >0.1) { \n  \n TicksArry.push(i); \n  \n } \n  \n } \n  \n if(usegrid) {"    
+
+def get_producer_layers(model):
+    """
+    Traverses the model to find the producer layer (layer_idx=0).cc
+    """
+    producer_modules = []
+    for module in model.modules():
+        if module.__class__.__name__.endswith("AttentionExperimental") and module.layer_idx == 0:
+            producer_modules.append(module)
+    return producer_modules
+    
+# 1) Load the base model from HF
+base_model_name = "meta-llama/Llama-3.2-1B"
+base_model = LlamaForCausalLM.from_pretrained(base_model_name, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+inputs = tokenizer(question, return_tensors="pt")
+inputs = {k: v.to(base_model.device) for k, v in inputs.items()}
+question_length = inputs['attention_mask'].shape[1]
+
+with torch.no_grad():
+    base_output_ids = base_model.generate(
+        **inputs,
+        max_new_tokens=200,
+        do_sample=True,
+        top_p=0.95,
+        temperature=0.7,
+    )
+base_output_text = tokenizer.decode(base_output_ids[0][question_length:], skip_special_tokens=True)
+
+
+from modeling_llama_butler import LlamaButlerConfig, LlamaButlerForCausalLM
+butler_config = LlamaButlerConfig.from_pretrained('config.json')
+
+
+butler_model = LlamaButlerForCausalLM(butler_config)
+butler_model.load_state_dict(base_model.state_dict(), strict=False)
+
+predictor_load_path = "/home/ya255/projects/TokenButler/expt_model/TrainTokenButler_42_finetune_None_None_500_llama_meta-llama_Llama-3.2-1B_L3_1B_2k.csv_L3_1B_2k_False_False_2000_False_redpajama_1024_1_1_20_0.001_1024/16_False_4_1000_ExpPred_fixed_40pc_True_False_0_None_False_False_4_8_2_16_512_False_False_True_16_0.37500000000000006__best.pt"
+
+model_producer_layers = get_producer_layers(butler_model)
+producer_layer_weights = torch.load(predictor_load_path)
+for idx, producer_layer_weight in enumerate(producer_layer_weights):
+    try:
+        model_producer_layers[idx].load_state_dict(producer_layer_weight, strict=False)
+    except Exception as e:
+        print(f"Error loading producer layer {idx}: {e}")
+        print("\n\nContinuing... !! Bad Perf If Unintentional !!\n\n")
+
+
+butler_model.to(base_model.device)
+butler_model.eval()
+
+with torch.no_grad():
+    butler_output_ids = butler_model.generate(
+        **inputs,
+        max_new_tokens=200,
+        do_sample=True,
+        top_p=0.95,
+        temperature=0.7,
+    )
+
+butler_output_text = tokenizer.decode(butler_output_ids[0][question_length:], skip_special_tokens=True)
+
+print("\n=== Base Model Output (Newlines Removed For Brevity) ===\n")
+print(base_output_text.replace("\n", ""))
+print("\n")
+print("=== Butler Model Output (Newlines Removed For Brevity) ===\n")
+print(butler_output_text.replace("\n", ""))
+print("\n")
+
+OUTPUT_DIR = "."
+print(f"\nSaving final merged model to: {OUTPUT_DIR}")
+butler_model.save_pretrained(OUTPUT_DIR, safe_serialization=False)
+
+# tokenizer.save_pretrained(OUTPUT_DIR)
+print("\nAll done! The folder should now have `pytorch_model.bin` and the updated `config.json`.\n")

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 128000,
+  "eos_token_id": 128001,
+  "transformers_version": "4.48.3"
+}

modeling_llama_butler.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+from typing import Dict
+from transformers import LlamaForCausalLM, LlamaConfig
+from transformers.generation.utils import GenerationConfig
+
+# Import your custom patching logic & custom modules:
+from modify_llama import convert_kvcache_experimental
+from predictor import TokenImportancePredictorAttentive, HeadImportancePredictor, PredictorDynamicCache
+from modify_llama import LlamaAttentionExperimental
+
+
+
+# ---------------------------------------------------------------------
+# 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

modify_llama.py

@@ -0,0 +1,464 @@
+import os
+import pdb
+import copy
+import math
+import numpy as np 
+from dataclasses import dataclass
+from typing import 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 utils import LlamaLinearScalingRotaryEmbedding, LlamaDynamicNTKScalingRotaryEmbedding, repeat_kv, sorted_index_to_mask
+from utils import calculate_hit_metrics, calculate_effective_sparsity, threshold_to_mask, SlidingWindowCache, enforce_sliding_window
+from transformers.cache_utils import DynamicCache
+from predictor import TokenImportancePredictorAttentive, PredictorDynamicCache, HeadImportancePredictor, attention_mse_loss, attention
+
+from triton_kernels.flash_attn import attention
+from triton_kernels.flash_attn_mse_loss import attention_mse_loss
+
+# torch.backends.cuda.enable_flash_sdp(enabled=True)
+# torch.backends.cuda.enable_mem_efficient_sdp(enabled=True)
+
+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:
+            if self.flash_attn:
+                if self.layer_idx > 0:
+                    # Token hit-rates cannot be calculated if using flash attention.
+                    self.tok_hit_acc = 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']
+                    q_importance_tensor = q_importance_tensor.view(bsz, self.num_heads, q_len, self.dDash)
+                    k_importance_tensor = k_importance_tensor.view(bsz, self.num_heads, key_len, self.dDash)
+                    device_index = query_states.device.index
+                    assert self.lookahead == 0, "Lookahead not supported with flash attention yet. Please disable --flash_attn"
+                    with torch.cuda.device(device_index):
+                        attn_output, mse_loss = attention_mse_loss(query_states.contiguous().to(torch.float16),
+                                                                    key_states.contiguous().to(torch.float16),
+                                                                    value_states.contiguous().to(torch.float16),
+                                                                    q_importance_tensor.contiguous().to(torch.float16),
+                                                                    k_importance_tensor.contiguous().to(torch.float16), 
+                                                                    True
+                                                                    )
+                    self.tok_hit_acc, self.tok_mean_rank_corr, self.tok_max_rank_corr = 0, 0, 0
+                    attn_output = attn_output.to(query_states.dtype)
+                    if not torch.isnan(mse_loss):
+                        self.msemagn_loss = mse_loss
+                    else:
+                        raise ValueError(f"NaN loss detected: {mse_loss}")
+                else:
+                    attn_output = torch.nn.functional.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=None, is_causal=True)
+            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

predictor.py

@@ -0,0 +1,643 @@
+import os
+import pdb
+import copy
+import math
+import numpy as np 
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+import gc
+
+from typing import Any, Dict, List, Optional, Tuple
+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 utils import LlamaLinearScalingRotaryEmbedding, LlamaDynamicNTKScalingRotaryEmbedding, repeat_kv, sorted_index_to_mask
+from transformers.cache_utils import DynamicCache
+
+from triton_kernels.flash_attn import attention
+from triton_kernels.flash_attn_mse_loss import attention_mse_loss
+
+
+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

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b6022d8733580ff4744a5e7002122b704a1e6426980a07db95e717d917ec553
+size 4992957478

triton_kernels/flash_attn.py

@@ -0,0 +1,543 @@
+import torch
+
+import triton
+import triton.language as tl
+
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+
+@triton.jit
+def _attn_fwd_inner(acc, l_i, m_i, q,  #
+                    K_block_ptr, V_block_ptr,  #
+                    start_m, qk_scale,  #
+                    BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr,  #
+                    STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr,  #
+                    N_CTX: tl.constexpr, fp8_v: tl.constexpr):
+    # range of values handled by this stage
+    if STAGE == 1:
+        lo, hi = 0, start_m * BLOCK_M
+    elif STAGE == 2:
+        lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
+        lo = tl.multiple_of(lo, BLOCK_M)
+    # causal = False
+    else:
+        lo, hi = 0, N_CTX
+    K_block_ptr = tl.advance(K_block_ptr, (0, lo))
+    V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
+    # loop over k, v and update accumulator
+    for start_n in range(lo, hi, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        # -- compute qk ----
+        k = tl.load(K_block_ptr)
+        qk = tl.dot(q, k)
+        if STAGE == 2:
+            mask = offs_m[:, None] >= (start_n + offs_n[None, :])
+            qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
+            m_ij = tl.maximum(m_i, tl.max(qk, 1))
+            qk -= m_ij[:, None]
+        else:
+            m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
+            qk = qk * qk_scale - m_ij[:, None]
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        # -- update output accumulator --
+        acc = acc * alpha[:, None]
+        # update acc
+        v = tl.load(V_block_ptr)
+        if fp8_v:
+            p = p.to(tl.float8e5)
+        else:
+            p = p.to(tl.float16)
+        acc = tl.dot(p, v, acc)
+        # update m_i and l_i
+        m_i = m_ij
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+    return acc, l_i, m_i
+
+
+# We don't run auto-tuning every time to keep the tutorial fast. Keeping
+# the code below and commenting out the equivalent parameters is convenient for
+# re-tuning.
+configs = [
+    triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+    for BM in [16, 32, 64, 128]\
+    for BN in [16, 32]\
+    for s in ([1] if is_hip() else [3, 4, 7])\
+    for w in [4, 8]\
+]
+
+
+def keep(conf):
+    BLOCK_M = conf.kwargs["BLOCK_M"]
+    BLOCK_N = conf.kwargs["BLOCK_N"]
+    if BLOCK_M * BLOCK_N < 128 * 128 and conf.num_warps == 8:
+        return False
+    return True
+
+
+@triton.autotune(list(filter(keep, configs)), key=["N_CTX", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd(Q, K, V, sm_scale, M, Out,  #
+              stride_qz, stride_qh, stride_qm, stride_qk,  #
+              stride_kz, stride_kh, stride_kn, stride_kk,  #
+              stride_vz, stride_vh, stride_vk, stride_vn,  #
+              stride_oz, stride_oh, stride_om, stride_on,  #
+              Z, H, N_CTX,  #
+              HEAD_DIM: tl.constexpr,  #
+              BLOCK_M: tl.constexpr,  #
+              BLOCK_N: tl.constexpr,  #
+              STAGE: tl.constexpr  #
+              ):
+    tl.static_assert(BLOCK_N <= HEAD_DIM)
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+
+    # block pointers
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_block_ptr = tl.make_block_ptr(
+        base=V + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=v_order,
+    )
+    K_block_ptr = tl.make_block_ptr(
+        base=K + qvk_offset,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    # load scales
+    qk_scale = sm_scale
+    qk_scale *= 1.44269504  # 1/log(2)
+    # load q: it will stay in SRAM throughout
+    q = tl.load(Q_block_ptr)
+    # stage 1: off-band
+    # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE
+    # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE
+    if STAGE & 1:
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,  #
+                                        start_m, qk_scale,  #
+                                        BLOCK_M, HEAD_DIM, BLOCK_N,  #
+                                        4 - STAGE, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5  #
+                                        )
+    # stage 2: on-band
+    if STAGE & 2:
+        # barrier makes it easier for compielr to schedule the
+        # two loops independently
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, K_block_ptr, V_block_ptr,  #
+                                        start_m, qk_scale,  #
+                                        BLOCK_M, HEAD_DIM, BLOCK_N,  #
+                                        2, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5  #
+                                        )
+    # epilogue
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    m_ptrs = M + off_hz * N_CTX + offs_m
+    tl.store(m_ptrs, m_i)
+    tl.store(O_block_ptr, acc.to(Out.type.element_ty))
+
+
+@triton.jit
+def _attn_bwd_preprocess(O, DO,  #
+                         Delta,  #
+                         Z, H, N_CTX,  #
+                         BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr  #
+                         ):
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv(dk, dv,  #
+                   Q, k, v, sm_scale,  #
+                   DO,  #
+                   M, D,  #
+                   # shared by Q/K/V/DO.
+                   stride_tok, stride_d,  #
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
+                   BLOCK_N1: tl.constexpr,  #
+                   HEAD_DIM: tl.constexpr,  #
+                   # Filled in by the wrapper.
+                   start_n, start_m, num_steps,  #
+                   MASK: tl.constexpr):
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    curr_m = start_m
+    step_m = BLOCK_M1
+    for blk_idx in range(num_steps):
+        qT = tl.load(qT_ptrs)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = curr_m + tl.arange(0, BLOCK_M1)
+        m = tl.load(M + offs_m)
+        qkT = tl.dot(k, qT)
+        pT = tl.math.exp2(qkT - m[None, :])
+        # Autoregressive masking.
+        if MASK:
+            mask = (offs_m[None, :] >= offs_n[:, None])
+            pT = tl.where(mask, pT, 0.0)
+        do = tl.load(do_ptrs)
+        # Compute dV.
+        ppT = pT
+        ppT = ppT.to(tl.float16)
+        dv += tl.dot(ppT, do)
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(tl.float16)
+        dk += tl.dot(dsT, tl.trans(qT))
+        # Increment pointers.
+        curr_m += step_m
+        qT_ptrs += step_m * stride_tok
+        do_ptrs += step_m * stride_tok
+    return dk, dv
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq(dq, q, K, V,  #
+                 do, m, D,
+                 # shared by Q/K/V/DO.
+                 stride_tok, stride_d,  #
+                 H, N_CTX,  #
+                 BLOCK_M2: tl.constexpr,  #
+                 BLOCK_N2: tl.constexpr,  #
+                 HEAD_DIM: tl.constexpr,
+                 # Filled in by the wrapper.
+                 start_m, start_n, num_steps,  #
+                 MASK: tl.constexpr):
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    # D (= delta) is pre-divided by ds_scale.
+    Di = tl.load(D + offs_m)
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    curr_n = start_n
+    step_n = BLOCK_N2
+    for blk_idx in range(num_steps):
+        kT = tl.load(kT_ptrs)
+        vT = tl.load(vT_ptrs)
+        qk = tl.dot(q, kT)
+        p = tl.math.exp2(qk - m)
+        # Autoregressive masking.
+        if MASK:
+            offs_n = curr_n + tl.arange(0, BLOCK_N2)
+            mask = (offs_m[:, None] >= offs_n[None, :])
+            p = tl.where(mask, p, 0.0)
+        # Compute dP and dS.
+        dp = tl.dot(do, vT).to(tl.float32)
+        ds = p * (dp - Di[:, None])
+        ds = ds.to(tl.float16)
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        dq += tl.dot(ds, tl.trans(kT))
+        # Increment pointers.
+        curr_n += step_n
+        kT_ptrs += step_n * stride_tok
+        vT_ptrs += step_n * stride_tok
+    return dq
+
+
+@triton.jit
+def _attn_bwd(Q, K, V, sm_scale,  #
+              DO,  #
+              DQ, DK, DV,  #
+              M, D,
+              # shared by Q/K/V/DO.
+              stride_z, stride_h, stride_tok, stride_d,  #
+              H, N_CTX,  #
+              BLOCK_M1: tl.constexpr,  #
+              BLOCK_N1: tl.constexpr,  #
+              BLOCK_M2: tl.constexpr,  #
+              BLOCK_N2: tl.constexpr,  #
+              BLK_SLICE_FACTOR: tl.constexpr,  #
+              HEAD_DIM: tl.constexpr):
+    LN2: tl.constexpr = 0.6931471824645996  # = ln(2)
+
+    bhid = tl.program_id(2)
+    off_chz = (bhid * N_CTX).to(tl.int64)
+    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
+    pid = tl.program_id(0)
+
+    # offset pointers for batch/head
+    Q += adj
+    K += adj
+    V += adj
+    DO += adj
+    DQ += adj
+    DK += adj
+    DV += adj
+    M += off_chz
+    D += off_chz
+
+    # load scales
+    offs_k = tl.arange(0, HEAD_DIM)
+
+    start_n = pid * BLOCK_N1
+    start_m = start_n
+
+    MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 // BLK_SLICE_FACTOR
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    num_steps = BLOCK_N1 // MASK_BLOCK_M1
+
+    dk, dv = _attn_bwd_dkdv(dk, dv,  #
+                            Q, k, v, sm_scale,  #
+                            DO,  #
+                            M, D,  #
+                            stride_tok, stride_d,  #
+                            H, N_CTX,  #
+                            MASK_BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
+                            start_n, start_m, num_steps,  #
+                            MASK=True  #
+                            )
+
+    start_m += num_steps * MASK_BLOCK_M1
+    num_steps = (N_CTX - start_m) // BLOCK_M1
+
+    # Compute dK and dV for non-masked blocks.
+    dk, dv = _attn_bwd_dkdv(  #
+        dk, dv,  #
+        Q, k, v, sm_scale,  #
+        DO,  #
+        M, D,  #
+        stride_tok, stride_d,  #
+        H, N_CTX,  #
+        BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
+        start_n, start_m, num_steps,  #
+        MASK=False  #
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv)
+
+    # Write back dK.
+    dk *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk)
+
+    # THIS BLOCK DOES DQ:
+    start_m = pid * BLOCK_M2
+    end_n = start_m + BLOCK_M2
+
+    MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 // BLK_SLICE_FACTOR
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    m = tl.load(M + offs_m)
+    m = m[:, None]
+
+    # Compute dQ for masked (diagonal) blocks.
+    # NOTE: This code scans each row of QK^T backward (from right to left,
+    # but inside each call to _attn_bwd_dq, from left to right), but that's
+    # not due to anything important.  I just wanted to reuse the loop
+    # structure for dK & dV above as much as possible.
+    num_steps = BLOCK_M2 // MASK_BLOCK_N2
+    dq = _attn_bwd_dq(dq, q, K, V,  #
+                      do, m, D,  #
+                      stride_tok, stride_d,  #
+                      H, N_CTX,  #
+                      BLOCK_M2, MASK_BLOCK_N2, HEAD_DIM,  #
+                      start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps,  #
+                      MASK=True  #
+                      )
+    end_n -= num_steps * MASK_BLOCK_N2
+    # stage 2
+    num_steps = end_n // BLOCK_N2
+    dq = _attn_bwd_dq(dq, q, K, V,  #
+                      do, m, D,  #
+                      stride_tok, stride_d,  #
+                      H, N_CTX,  #
+                      BLOCK_M2, BLOCK_N2, HEAD_DIM,  #
+                      start_m, end_n - num_steps * BLOCK_N2, num_steps,  #
+                      MASK=False  #
+                      )
+    # Write back dQ.
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq *= LN2
+    tl.store(dq_ptrs, dq)
+
+
+class _attention(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v, causal, sm_scale):
+        # shape constraints
+        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
+        # when v is in float8_e5m2 it is transposed.
+        HEAD_DIM_V = v.shape[-1]
+        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
+        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
+        o = torch.empty_like(q)
+        stage = 3 if causal else 1
+        extra_kern_args = {}
+        # Tuning for AMD target
+        if is_hip():
+            waves_per_eu = 3 if HEAD_DIM_K <= 64 else 2
+            extra_kern_args = {"waves_per_eu": waves_per_eu, "allow_flush_denorm": True}
+
+        grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
+        M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+        _attn_fwd[grid](
+            q, k, v, sm_scale, M, o,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),  #
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
+            q.shape[0], q.shape[1],  #
+            N_CTX=q.shape[2],  #
+            HEAD_DIM=HEAD_DIM_K,  #
+            STAGE=stage,  #
+            **extra_kern_args)
+
+        ctx.save_for_backward(q, k, v, o, M)
+        ctx.grid = grid
+        ctx.sm_scale = sm_scale
+        ctx.HEAD_DIM = HEAD_DIM_K
+        ctx.causal = causal
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, M = ctx.saved_tensors
+        do = do.contiguous()
+        assert do.is_contiguous()
+        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
+        dq = torch.empty_like(q)
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+        BATCH, N_HEAD, N_CTX = q.shape[:3]
+        PRE_BLOCK = 128
+        NUM_WARPS, NUM_STAGES = 4, 5
+        BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 128, 128, 32
+        BLK_SLICE_FACTOR = 2
+        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+        arg_k = k
+        arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
+        PRE_BLOCK = 128
+        assert N_CTX % PRE_BLOCK == 0
+        pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
+        delta = torch.empty_like(M)
+        _attn_bwd_preprocess[pre_grid](
+            o, do,  #
+            delta,  #
+            BATCH, N_HEAD, N_CTX,  #
+            BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM  #
+        )
+        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+        _attn_bwd[grid](
+            q, arg_k, v, ctx.sm_scale, do, dq, dk, dv,  #
+            M, delta,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            N_HEAD, N_CTX,  #
+            BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
+            BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2,  #
+            BLK_SLICE_FACTOR=BLK_SLICE_FACTOR,  #
+            HEAD_DIM=ctx.HEAD_DIM,  #
+            num_warps=NUM_WARPS,  #
+            num_stages=NUM_STAGES  #
+        )
+
+        return dq, dk, dv, None, None
+
+
+attention = _attention.apply
+
+from torch import nn
+from torch.nn import MSELoss
+import math
+if __name__ == '__main__':
+    B, H, N_CTX, D = 1, 32, 84, 16
+    causal = True
+    DTYPE = torch.float32
+    torch.manual_seed(20)
+    gain_q = math.sqrt(5.0 / D)  
+    gain_k = math.sqrt(5.0 / D)
+    gain_v = math.sqrt(3.0 / D)    
+
+    q = (torch.randn((B, H, N_CTX, D), dtype=DTYPE, device=DEVICE) * gain_q).requires_grad_()
+    k = (torch.randn((B, H, N_CTX, D), dtype=DTYPE, device=DEVICE) * gain_k).requires_grad_()
+    v = (torch.randn((B, H, N_CTX, D), dtype=DTYPE, device=DEVICE) * gain_v).requires_grad_()
+
+    sm_scale = 1.0 / math.sqrt(D)
+    att_output_triton = attention(q.to(torch.float16), k.to(torch.float16), v.to(torch.float16), causal, sm_scale)
+
+    M = torch.tril(torch.ones((N_CTX, N_CTX), device=DEVICE))
+
+    attn_weights = torch.matmul(q, k.transpose(-2, -1)) * sm_scale
+    if causal:
+        causal_mask = torch.triu(torch.ones(N_CTX, N_CTX, device=attn_weights.device), diagonal=1).bool()
+        attn_weights = attn_weights.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1).to(v.dtype)
+    attn_output_original = torch.matmul(attn_weights, v)
+
+    print(f'att_output_triton: {att_output_triton}')
+    print(f'attn_output_original: {attn_output_original}')
+    # assert torch.allclose(attn_output_original, att_output_triton, atol=1e-2, rtol=0), f"Error: {torch.mean(torch.abs(attn_output_original - att_output_triton))}"
+    perc_diff = 100 * torch.abs(att_output_triton - attn_output_original).mean() / torch.abs(attn_output_original).mean()
+    print(f'passed test for mse loss with {att_output_triton.mean()} vs {attn_output_original.mean()}, \t\t\t\tpercentage diff: {perc_diff.item()}%')

triton_kernels/flash_attn_mse_loss.py

@@ -0,0 +1,733 @@
+import torch
+
+import triton
+import triton.language as tl
+
+import torch
+import math
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+@triton.jit
+def _attn_fwd_inner(acc, l_i, m_i, q, q_imp,  #
+                    K_block_ptr, K_imp_block_ptr, V_block_ptr,  #
+                    mse_loss,  #
+                    start_m, qk_sqrt, qk_scale, qk_sqrt_imp, qk_scale_imp, #
+                    BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, NUM_ELEMENTS: tl.constexpr, BLOCK_N: tl.constexpr,  #
+                    STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr,  #
+                    N_CTX: tl.constexpr, fp8_v: tl.constexpr):
+    # range of values handled by this stage
+    if STAGE == 1:
+        lo, hi = 0, start_m * BLOCK_M
+    elif STAGE == 2:
+        lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M
+        lo = tl.multiple_of(lo, BLOCK_M)
+    # causal = False
+    else:
+        lo, hi = 0, N_CTX
+    K_block_ptr = tl.advance(K_block_ptr, (0, lo))
+    K_imp_block_ptr = tl.advance(K_imp_block_ptr, (0, lo))
+    V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
+    mse_contrib_total = 0.0
+    # loop over k, v and update accumulator
+    for start_n in range(lo, hi, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        # -- compute qk ----
+        k = tl.load(K_block_ptr).to(tl.float32)
+        k_imp = tl.load(K_imp_block_ptr).to(tl.float32)
+        qf32 = q.to(tl.float32)
+        qf32_imp = q_imp.to(tl.float32)
+        qk = tl.dot(qf32, k)
+        qk_imp = tl.dot(qf32_imp, k_imp)
+        diff = (qk * qk_sqrt - qk_imp * qk_sqrt_imp)
+        diff_sqr = diff * diff
+        mse_contrib = tl.sum(diff_sqr)
+        # tl.atomic_add(mse_loss, mse_contrib)
+        mse_contrib_total += mse_contrib
+        if STAGE == 2:
+            mask = offs_m[:, None] >= (start_n + offs_n[None, :])
+            qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
+            m_ij = tl.maximum(m_i, tl.max(qk, 1))
+            qk -= m_ij[:, None]
+        else:
+            m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
+            qk = qk * qk_scale - m_ij[:, None]
+        p = tl.math.exp2(qk)
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        # -- update output accumulator --
+        acc = acc * alpha[:, None]
+        # update acc
+        v = tl.load(V_block_ptr)
+        p = p.to(v.type.element_ty)
+        acc = tl.dot(p, v, acc)
+        # update m_i and l_i
+        m_i = m_ij
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+        K_imp_block_ptr = tl.advance(K_imp_block_ptr, (0, BLOCK_N))
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+
+    if STAGE == 2:
+        for start_n in range(hi, N_CTX, BLOCK_N):
+            start_n = tl.multiple_of(start_n, BLOCK_N)
+            # -- compute qk ----
+            k = tl.load(K_block_ptr).to(tl.float32)
+            k_imp = tl.load(K_imp_block_ptr).to(tl.float32)
+            qf32 = q.to(tl.float32)
+            qf32_imp = q_imp.to(tl.float32)
+            qk = tl.dot(qf32, k)
+            qk_imp = tl.dot(qf32_imp, k_imp)
+            diff = (qk * qk_sqrt - qk_imp * qk_sqrt_imp)
+            diff_sqr = diff * diff
+            mse_contrib = tl.sum(diff_sqr)
+            # tl.atomic_add(mse_loss, mse_contrib)
+            mse_contrib_total += mse_contrib
+            K_imp_block_ptr = tl.advance(K_imp_block_ptr, (0, BLOCK_N))
+            K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+    mse_contrib_total /= NUM_ELEMENTS
+    tl.atomic_add(mse_loss, mse_contrib_total)
+    return acc, l_i, m_i
+
+
+# We don't run auto-tuning every time to keep the tutorial fast. Keeping
+# the code below and commenting out the equivalent parameters is convenient for
+# re-tuning.
+# configs = [
+#     triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+#     for BM in [32]\
+#     for BN in [32]\
+#     for s in ([1] if is_hip() else [3, 4, 7])\
+#     for w in [4, 8]\
+# ]
+
+fixed_config = triton.Config(
+    {'BLOCK_M': 16, 'BLOCK_N': 16},
+    num_stages=4,
+    num_warps=8
+)
+
+def keep(conf):
+    BLOCK_M = conf.kwargs["BLOCK_M"]
+    BLOCK_N = conf.kwargs["BLOCK_N"]
+    if BLOCK_M * BLOCK_N < 128 * 128 and conf.num_warps == 8:
+        return False
+    return True
+
+
+# @triton.autotune(list(filter(keep, configs)), key=["N_CTX", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd(Q, K, V, sm_scale, sm_scale_imp, M, Out,  #
+              Q_importance, K_importance, mse_loss,  
+              stride_qz, stride_qh, stride_qm, stride_qk,  #
+              stride_kz, stride_kh, stride_kn, stride_kk,  #
+              stride_vz, stride_vh, stride_vk, stride_vn,  #
+              stride_oz, stride_oh, stride_om, stride_on,  #
+              stride_qz_imp, stride_qh_imp, stride_qm_imp, stride_qk_imp,  #
+              stride_kz_imp, stride_kh_imp, stride_kn_imp, stride_kk_imp,  #
+              Z, H, N_CTX,  #
+              HEAD_DIM: tl.constexpr,  #
+              D_DASH: tl.constexpr,  #
+              NUM_ELEMENTS: tl.constexpr,  #
+              BLOCK_M: tl.constexpr,  #
+              BLOCK_N: tl.constexpr,  #
+              STAGE: tl.constexpr  #
+              ):
+    tl.static_assert(BLOCK_N <= HEAD_DIM)
+    tl.static_assert(BLOCK_N <= D_DASH)
+    start_m = tl.program_id(0)
+    off_hz = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh
+    qv_imp_offset = off_z.to(tl.int64) * stride_qz_imp + off_h.to(tl.int64) * stride_qh_imp
+
+    # block pointers
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_block_ptr = tl.make_block_ptr(
+        base=V + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, HEAD_DIM),
+        order=v_order,
+    )
+    K_block_ptr = tl.make_block_ptr(
+        base=K + qvk_offset,
+        shape=(HEAD_DIM, N_CTX),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(HEAD_DIM, BLOCK_N),
+        order=(0, 1),
+    )
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + qvk_offset,
+        shape=(N_CTX, HEAD_DIM),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, HEAD_DIM),
+        order=(1, 0),
+    )
+    Q_imp_block_ptr = tl.make_block_ptr(
+        base=Q_importance + qv_imp_offset,
+        shape=(N_CTX, D_DASH),
+        strides=(stride_qm_imp, stride_qk_imp),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, D_DASH),
+        order=(1, 0),
+    )
+    K_imp_block_ptr = tl.make_block_ptr(
+        base=K_importance + qv_imp_offset,
+        shape=(D_DASH, N_CTX),
+        strides=(stride_kk_imp, stride_kn_imp),
+        offsets=(0, 0),
+        block_shape=(D_DASH, BLOCK_N),
+        order=(0, 1),
+    )
+    # initialize offsets
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    # load scales
+    qk_sqrt = sm_scale
+    qk_scale = qk_sqrt * 1.44269504  # 1/log(2)
+    qk_sqrt_imp = sm_scale_imp
+    qk_scale_imp = qk_sqrt_imp * 1.44269504
+    # load q: it will stay in SRAM throughout
+    q = tl.load(Q_block_ptr)
+    q_imp = tl.load(Q_imp_block_ptr)
+    # stage 1: off-band
+    # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE
+    # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE
+    # _attn_fwd_get_loss(q, q_imp, K_block_ptr, K_imp_block_ptr, V_block_ptr, mse_loss, start_m, qk_scale, qk_scale_imp, BLOCK_M, HEAD_DIM, BLOCK_N, STAGE, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5)
+    if STAGE & 1:
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, q_imp, K_block_ptr, K_imp_block_ptr, V_block_ptr,  #
+                                        mse_loss,  #
+                                        start_m, qk_sqrt, qk_scale, qk_sqrt_imp, qk_scale_imp,  #
+                                        BLOCK_M, HEAD_DIM, NUM_ELEMENTS, BLOCK_N,  #
+                                        4 - STAGE, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5  #
+                                        )
+    # stage 2: on-band
+    if STAGE & 2:
+        # barrier makes it easier for compielr to schedule the
+        # two loops independently
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, q_imp, K_block_ptr, K_imp_block_ptr, V_block_ptr,  #
+                                        mse_loss,  #
+                                        start_m, qk_sqrt, qk_scale, qk_sqrt_imp, qk_scale_imp,  #
+                                        BLOCK_M, HEAD_DIM, NUM_ELEMENTS, BLOCK_N,  #
+                                        2, offs_m, offs_n, N_CTX, V.dtype.element_ty == tl.float8e5  #
+                                        )
+    # epilogue
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    m_ptrs = M + off_hz * N_CTX + offs_m
+    tl.store(m_ptrs, m_i)
+    tl.store(O_block_ptr, acc.to(Out.type.element_ty))
+
+
+@triton.jit
+def _attn_bwd_preprocess(O, DO,  #
+                         Delta,  #
+                         Z, H, N_CTX,  #
+                         BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr  #
+                         ):
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv(dk, dv,  #
+                   Q, k, v, sm_scale,  #
+                   DO,  #
+                   M, D,  #
+                   # shared by Q/K/V/DO.
+                   stride_tok, stride_d,  #
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
+                   BLOCK_N1: tl.constexpr,  #
+                   HEAD_DIM: tl.constexpr,  #
+                   # Filled in by the wrapper.
+                   start_n, start_m, num_steps,  #
+                   MASK: tl.constexpr):
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    curr_m = start_m
+    step_m = BLOCK_M1
+    for blk_idx in range(num_steps):
+        qT = tl.load(qT_ptrs)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = curr_m + tl.arange(0, BLOCK_M1)
+        m = tl.load(M + offs_m)
+        qkT = tl.dot(k, qT)
+        pT = tl.math.exp2(qkT - m[None, :])
+        # Autoregressive masking.
+        if MASK:
+            mask = (offs_m[None, :] >= offs_n[:, None])
+            pT = tl.where(mask, pT, 0.0)
+        do = tl.load(do_ptrs)
+        # Compute dV.
+        ppT = pT
+        ppT = ppT.to(tl.float16)
+        dv += tl.dot(ppT, do)
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(tl.float16)
+        dk += tl.dot(dsT, tl.trans(qT))
+        # Increment pointers.
+        curr_m += step_m
+        qT_ptrs += step_m * stride_tok
+        do_ptrs += step_m * stride_tok
+    return dk, dv
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq(dq, q, K, V,  #
+                 do, m, D,
+                 # shared by Q/K/V/DO.
+                 stride_tok, stride_d,  #
+                 H, N_CTX,  #
+                 BLOCK_M2: tl.constexpr,  #
+                 BLOCK_N2: tl.constexpr,  #
+                 HEAD_DIM: tl.constexpr,
+                 # Filled in by the wrapper.
+                 start_m, start_n, num_steps,  #
+                 MASK: tl.constexpr):
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    # D (= delta) is pre-divided by ds_scale.
+    Di = tl.load(D + offs_m)
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    curr_n = start_n
+    step_n = BLOCK_N2
+    for blk_idx in range(num_steps):
+        kT = tl.load(kT_ptrs)
+        vT = tl.load(vT_ptrs)
+        qk = tl.dot(q, kT)
+        p = tl.math.exp2(qk - m)
+        # Autoregressive masking.
+        if MASK:
+            offs_n = curr_n + tl.arange(0, BLOCK_N2)
+            mask = (offs_m[:, None] >= offs_n[None, :])
+            p = tl.where(mask, p, 0.0)
+        # Compute dP and dS.
+        dp = tl.dot(do, vT).to(tl.float32)
+        ds = p * (dp - Di[:, None])
+        ds = ds.to(tl.float16)
+        # Compute dQ.
+        # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
+        dq += tl.dot(ds, tl.trans(kT))
+        # Increment pointers.
+        curr_n += step_n
+        kT_ptrs += step_n * stride_tok
+        vT_ptrs += step_n * stride_tok
+    return dq
+
+@triton.jit
+def _attn_bwd_dk_imp(
+                   Q, Q_imp, k, k_imp, sm_scale, sm_scale_imp, num_elements, #
+                   DMSE,  #
+                   # shared by Q/K/V/DO.
+                   stride_tok, stride_d,  #
+                   stride_tok_imp, stride_d_imp,  #
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
+                   BLOCK_N1: tl.constexpr,  #
+                   HEAD_DIM: tl.constexpr,  #
+                   D_DASH: tl.constexpr,  #
+                   # Filled in by the wrapper.
+                   start_n, start_m, num_steps, is_float16):
+    dk_imp = tl.zeros([BLOCK_N1, D_DASH], dtype=tl.float32)
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    offs_k_imp = tl.arange(0, D_DASH)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    qT_imp_ptrs = Q_imp + offs_m[None, :] * stride_tok_imp + offs_k_imp[:, None] * stride_d_imp
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    curr_m = start_m
+    step_m = BLOCK_M1
+    dmse_eval = tl.load(DMSE)
+    for blk_idx in range(num_steps):
+        qT = tl.load(qT_ptrs)
+        qT_imp = tl.load(qT_imp_ptrs)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = curr_m + tl.arange(0, BLOCK_M1)
+        qkT = tl.dot(k, qT)
+        qkT_imp = tl.dot(k_imp, qT_imp)
+        diff = (qkT_imp * sm_scale_imp - qkT * sm_scale)
+        tmp = dmse_eval * 2.0 * (1 / num_elements) * sm_scale_imp
+        diff = diff.to(tl.float16)
+        qT_imp = qT_imp.to(tl.float16)
+        dk_imp += tmp * tl.dot(diff, tl.trans(qT_imp))
+        if not is_float16:
+            dk_imp = dk_imp.to(tl.float32)
+
+        # Increment pointers.
+        curr_m += step_m
+        qT_ptrs += step_m * stride_tok
+        qT_imp_ptrs += step_m * stride_tok_imp
+    return dk_imp
+
+@triton.jit
+def _attn_bwd_dq_imp(
+                   q, q_imp, K, K_imp, sm_scale, sm_scale_imp, num_elements, #
+                   DMSE,  #
+                   # shared by Q/K/V/DO.
+                   stride_tok, stride_d,  #
+                   stride_tok_imp, stride_d_imp,  #
+                   H, N_CTX, BLOCK_M1: tl.constexpr,  #
+                   BLOCK_N1: tl.constexpr,  #
+                   HEAD_DIM: tl.constexpr,  #
+                   D_DASH: tl.constexpr,  #
+                   # Filled in by the wrapper.
+                   start_n, start_m, num_steps, is_float16):
+    dq_imp = tl.zeros([BLOCK_N1, D_DASH], dtype=tl.float32)
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    offs_k_imp = tl.arange(0, D_DASH)
+    kT_ptrs = K + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    kT_imp_ptrs = K_imp + offs_m[None, :] * stride_tok_imp + offs_k_imp[:, None] * stride_d_imp
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    curr_m = start_m
+    step_m = BLOCK_M1
+    dmse_eval = tl.load(DMSE)
+    for blk_idx in range(num_steps):
+        kT = tl.load(kT_ptrs)
+        kT_imp = tl.load(kT_imp_ptrs)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = curr_m + tl.arange(0, BLOCK_M1)
+        qkT = tl.dot(q, kT)
+        qkT_imp = tl.dot(q_imp, kT_imp)
+        diff = (qkT_imp * sm_scale_imp - qkT * sm_scale)
+        tmp = dmse_eval * 2.0 * (1 / num_elements) * sm_scale_imp
+        diff = diff.to(tl.float16)
+        kT_imp = kT_imp.to(tl.float16)
+        dq_imp += tmp * tl.dot(diff, tl.trans(kT_imp))
+        if not is_float16:
+            dq_imp = dq_imp.to(tl.float32)
+        # Increment pointers.
+        curr_m += step_m
+        kT_ptrs += step_m * stride_tok
+        kT_imp_ptrs += step_m * stride_tok_imp
+    return dq_imp
+
+
+@triton.jit
+def _attn_bwd(Q, Q_imp, K, K_imp, sm_scale, sm_scale_imp, num_elements, #
+              DQ_imp, DK_imp, DMSE, #
+              # shared by Q/K/V/DO.
+              stride_z, stride_h, stride_tok, stride_d,  #
+              stride_z_imp, stride_h_imp, stride_tok_imp, stride_d_imp,  #
+              H, N_CTX,  #
+              BLOCK_M1: tl.constexpr,  #
+              BLOCK_N1: tl.constexpr,  #
+              HEAD_DIM: tl.constexpr,
+              D_DASH: tl.constexpr):
+    LN2: tl.constexpr = 0.6931471824645996  # = ln(2)
+
+    bhid = tl.program_id(2)
+    off_chz = (bhid * N_CTX).to(tl.int64)
+    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
+    adj_imp = (stride_h_imp * (bhid % H) + stride_z_imp * (bhid // H)).to(tl.int64)
+    pid = tl.program_id(0)
+
+    # offset pointers for batch/head
+    Q += adj
+    K += adj
+    Q_imp += adj_imp
+    K_imp += adj_imp
+    DQ_imp += adj_imp
+    DK_imp += adj_imp
+
+    # load scales
+    offs_k = tl.arange(0, HEAD_DIM)
+    offs_k_imp = tl.arange(0, D_DASH)
+
+    start_n = pid * BLOCK_N1
+    start_m = 0
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    num_steps = N_CTX // BLOCK_M1
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    k_imp = tl.load(K_imp + offs_n[:, None] * stride_tok_imp + offs_k_imp[None, :] * stride_d_imp)
+    dk_imp = _attn_bwd_dk_imp(
+                            Q, Q_imp, k, k_imp, sm_scale, sm_scale_imp, num_elements,  #
+                            DMSE,  #
+                            stride_tok, stride_d,  #
+                            stride_tok_imp, stride_d_imp,  #
+                            H, N_CTX, BLOCK_M1,  #
+                            BLOCK_N1, HEAD_DIM, D_DASH,  #
+                            start_n, start_m, num_steps, k.dtype == tl.float16 #
+                            )
+    dk_imp_ptrs = DK_imp + offs_n[:, None] * stride_tok_imp + offs_k_imp[None, :] * stride_d_imp
+    tl.store(dk_imp_ptrs, dk_imp)
+
+    start_n = pid * BLOCK_N1
+    start_m = 0
+    num_steps = N_CTX // BLOCK_M1
+    q = tl.load(Q + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    q_imp = tl.load(Q_imp + offs_n[:, None] * stride_tok_imp + offs_k_imp[None, :] * stride_d_imp)
+    dq_imp = _attn_bwd_dq_imp(
+                            q, q_imp, K, K_imp, sm_scale, sm_scale_imp, num_elements,  #
+                            DMSE,  #
+                            stride_tok, stride_d,  #
+                            stride_tok_imp, stride_d_imp,  #
+                            H, N_CTX, BLOCK_M1,  #
+                            BLOCK_N1, HEAD_DIM, D_DASH,  #
+                            start_n, start_m, num_steps, q.dtype == tl.float16 #
+                            )
+    dq_imp_ptrs = DQ_imp + offs_n[:, None] * stride_tok_imp + offs_k_imp[None, :] * stride_d_imp
+    tl.store(dq_imp_ptrs, dq_imp)
+    
+
+class _attention_mse_loss(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, q, k, v, q_importance, k_importance, causal):
+        # shape constraints
+        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
+        # when v is in float8_e5m2 it is transposed.
+        HEAD_DIM_V = v.shape[-1]
+        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
+        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
+        D_DASH = q_importance.shape[-1]
+        assert D_DASH == k_importance.shape[-1], "q_importance and k_importance must have the same last dimension"
+        sm_scale = 1.0 / math.sqrt(HEAD_DIM_Q)
+        sm_scale_imp = 1.0 / math.sqrt(D_DASH)
+        o = torch.empty_like(q)
+        stage = 3 if causal else 1
+        extra_kern_args = {}
+        # Tuning for AMD target
+        if is_hip():
+            waves_per_eu = 3 if HEAD_DIM_K <= 64 else 2
+            extra_kern_args = {"waves_per_eu": waves_per_eu, "allow_flush_denorm": True}
+
+        grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
+        M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+        mse_loss = torch.zeros(1, device=q.device, dtype=torch.float32)
+        num_elements = (q.shape[0] * q.shape[1] * q.shape[2] * k.shape[2])
+        _attn_fwd[grid](
+            q, k, v, sm_scale, sm_scale_imp, M, o,  #
+            q_importance, k_importance, mse_loss,  #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            k.stride(0), k.stride(1), k.stride(2), k.stride(3),  #
+            v.stride(0), v.stride(1), v.stride(2), v.stride(3),  #
+            o.stride(0), o.stride(1), o.stride(2), o.stride(3),  #
+            q_importance.stride(0), q_importance.stride(1), q_importance.stride(2), q_importance.stride(3),  #
+            k_importance.stride(0), k_importance.stride(1), k_importance.stride(2), k_importance.stride(3),  #
+            Z=q.shape[0], H=q.shape[1],  #
+            N_CTX=q.shape[2],  #
+            HEAD_DIM=HEAD_DIM_K,  #
+            D_DASH=D_DASH,  #
+            NUM_ELEMENTS=num_elements,  #
+            STAGE=stage,  #
+            **fixed_config.kwargs)
+        ctx.save_for_backward(q, q_importance, k, k_importance, v, o, M)
+        ctx.grid = grid
+        ctx.sm_scale = sm_scale
+        ctx.sm_scale_imp = sm_scale_imp
+        ctx.HEAD_DIM = HEAD_DIM_K
+        ctx.D_DASH = D_DASH
+        ctx.num_elements = num_elements
+        ctx.causal = causal
+        return o, mse_loss
+
+    @staticmethod
+    def backward(ctx, do, dmse):
+        q, q_importance, k, k_importance, v, o, M = ctx.saved_tensors
+        do = do.contiguous()
+        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
+        dq = torch.empty_like(q)
+        dq_imp = torch.empty_like(q_importance)
+        dk = torch.empty_like(k)
+        dk_imp = torch.empty_like(k_importance)
+        dv = torch.empty_like(v)
+        BATCH, N_HEAD, N_CTX = q.shape[:3]
+        NUM_WARPS, NUM_STAGES = 4, 5
+        BLOCK_M1, BLOCK_N1 = 16, 16
+        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
+
+        # PRE_BLOCK = 128
+        # assert N_CTX % PRE_BLOCK == 0
+        # pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
+        # delta = torch.empty_like(M)
+        # _attn_bwd_preprocess[pre_grid](
+        #     o, do,  #
+        #     delta,  #
+        #     BATCH, N_HEAD, N_CTX,  #
+        #     BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM  #
+        # )
+        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
+        _attn_bwd[grid]( 
+            q, q_importance, k, k_importance, ctx.sm_scale, ctx.sm_scale_imp, ctx.num_elements,  #
+            dq_imp, dk_imp, dmse, #
+            q.stride(0), q.stride(1), q.stride(2), q.stride(3),  #
+            q_importance.stride(0), q_importance.stride(1), q_importance.stride(2), q_importance.stride(3),  #
+            N_HEAD, N_CTX,  #
+            BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1,  #
+            HEAD_DIM=ctx.HEAD_DIM,  #
+            D_DASH=ctx.D_DASH,  #
+            num_warps=NUM_WARPS,  #
+            num_stages=NUM_STAGES  #
+        )
+
+        return None, None, None, dq_imp, dk_imp, None
+
+
+attention_mse_loss = _attention_mse_loss.apply
+
+from torch import nn
+from torch.nn import MSELoss
+import time
+if __name__ == '__main__':
+    B, H, N_CTX, D = 1, 32, 512, 128
+    D_DASH = 32
+    causal = True
+    DTYPE = torch.float32
+    LOAD_WEIGHT = False
+    import os
+    if LOAD_WEIGHT and os.path.exists("export_params.pt"):
+        print("[Info] Detected export_params.pt, loading saved tensors...")
+        debug_tensors = torch.load("export_params.pt", map_location=DEVICE)
+        
+        q = debug_tensors["q"].detach().clone().requires_grad_().contiguous()
+        k = debug_tensors["k"].detach().clone().requires_grad_().contiguous()
+        v = debug_tensors["v"].detach().clone().requires_grad_().contiguous()
+        q_importance = debug_tensors["q_importance"].detach().clone().requires_grad_().contiguous()
+        k_importance = debug_tensors["k_importance"].detach().clone().requires_grad_().contiguous()
+
+        print("[Success] Tensors loaded successfully!")
+        print(f"q.shape: {q.shape}, k.shape: {k.shape}, v.shape: {v.shape}")
+        print(f"q_importance.shape: {q_importance.shape}, k_importance.shape: {k_importance.shape}")
+    else:
+        print("[Info] No export_params.pt found, initializing random tensors...")
+        DTYPE = torch.float32
+
+        gain_q = math.sqrt(5.0 / D)  
+        gain_k = math.sqrt(5.0 / D)
+        gain_v = math.sqrt(3.0 / D)    
+
+        q = (torch.randn((B, H, N_CTX, D), dtype=DTYPE, device=DEVICE) * gain_q).requires_grad_()
+        k = (torch.randn((B, H, N_CTX, D), dtype=DTYPE, device=DEVICE) * gain_k).requires_grad_()
+        v = (torch.randn((B, H, N_CTX, D), dtype=DTYPE, device=DEVICE) * gain_v).requires_grad_()
+
+        gain_q_imp = math.sqrt(5.0 / D_DASH)
+        gain_k_imp = math.sqrt(5.0 / D_DASH)
+        q_importance = (torch.randn((B, H, N_CTX, D_DASH), dtype=DTYPE, device=DEVICE) * gain_q_imp + 0.1).requires_grad_()
+        k_importance = (torch.randn((B, H, N_CTX, D_DASH), dtype=DTYPE, device=DEVICE) * gain_k_imp - 0.1).requires_grad_()
+        print("[Info] Random tensors initialized.")
+
+    # warm up for the triton implementation
+    attn_output, mse_loss_triton = attention_mse_loss(q.to(torch.float16),
+                                                                    k.to(torch.float16),
+                                                                    v.to(torch.float16),
+                                                                    q_importance.to(torch.float16),
+                                                                    k_importance.to(torch.float16), True)
+
+
+    mse_loss_triton.backward()
+    q_importance.grad = None
+    k_importance.grad = None
+
+    # warm up for the original implementation
+    attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(D)
+    importance_mask = torch.matmul(q_importance, k_importance.transpose(-2, -1)) / math.sqrt(D_DASH) # [B, H, Lq, Lk]
+    mse_func = MSELoss(reduction='none')
+    mse_loss_original = mse_func(attn_weights, importance_mask)
+    mse_loss_original = mse_loss_original.mean()
+    if causal:
+        causal_mask = torch.triu(torch.ones(N_CTX, N_CTX, device=attn_weights.device), diagonal=1).bool()
+        attn_weights = attn_weights.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+    attn_weights = nn.functional.softmax(attn_weights.float(), dim=-1).to(v.dtype)
+    attn_output_original = torch.matmul(attn_weights, v)
+    mse_loss_original.backward()
+    q_importance.grad = None
+    k_importance.grad = None
+
+    # triton implementation
+    tri_start = time.time()
+    att_output_triton, mse_loss_triton = attention_mse_loss(q, k, v, q_importance, k_importance, causal)
+    mse_loss_triton.backward()
+    tri_dq_imp, q_importance.grad = q_importance.grad.clone(), None
+    tri_dk_imp, k_importance.grad = k_importance.grad.clone(), None
+    tri_end = time.time()
+
+    # original implementation
+    ori_start = time.time()
+    attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(D)
+    importance_mask = torch.matmul(q_importance, k_importance.transpose(-2, -1)) / math.sqrt(D_DASH) # [B, H, Lq, Lk]
+    mse_func = MSELoss(reduction='none')
+    mse_loss_original = mse_func(attn_weights, importance_mask)
+    mse_loss_original = mse_loss_original.mean()
+    mse_loss_original = mse_loss_original
+    if causal:
+        causal_mask = torch.triu(torch.ones(N_CTX, N_CTX, device=attn_weights.device), diagonal=1).bool()
+        attn_weights = attn_weights.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(v.dtype)
+    attn_output_original = torch.matmul(attn_weights, v)
+    
+    mse_loss_original.backward()
+    ref_dk_imp, k_importance.grad = k_importance.grad.clone(), None
+    ref_dq_imp, q_importance.grad = q_importance.grad.clone(), None
+    ori_end = time.time()
+    
+    print(f'mse_loss_triton: {mse_loss_triton}')
+    print(f'mse_loss_original: {mse_loss_original}')
+    print(f'mean of ref_dk_imp: {ref_dk_imp.mean()}')
+    print(f'mean of tri_dk_imp: {tri_dk_imp.mean()}')
+    # print(f'error: {torch.abs(ref_dk_imp - tri_dk_imp).mean()}')
+    print(f'ref_dq_imp: {ref_dq_imp[0][0]}')
+    print(f'tri_dq_imp: {tri_dq_imp[0][0]}')
+    print(f'mean of ref_dq_imp: {ref_dq_imp.mean()}')
+    print(f'mean of tri_dq_imp: {tri_dq_imp.mean()}')
+    # print(f'error: {torch.abs(ref_dq_imp - tri_dq_imp).mean()}')
+
+    # assert torch.allclose(attn_output_original, att_output_triton, atol=1e-2, rtol=0), f'{attn_output_original.mean()} vs {att_output_triton.mean()}'
+    perc_diff = 100 * torch.abs(attn_output_original - att_output_triton).mean() / torch.abs(attn_output_original).mean()
+    print(f'passed test for attention output with {attn_output_original.mean()} vs {att_output_triton.mean()}, \t\t\tpercentage diff: {perc_diff}%')
+    # assert torch.allclose(mse_loss_triton, mse_loss_original, atol=1e-1, rtol=0), f'{mse_loss_triton} vs {mse_loss_original}'
+    perc_diff = 100 * torch.abs(mse_loss_triton - mse_loss_original) / torch.abs(mse_loss_original)
+    print(f'passed test for mse loss with {mse_loss_triton.item()} vs {mse_loss_original.item()}, \t\t\t\tpercentage diff: {perc_diff.item()}%')
+    # assert torch.allclose(ref_dk_imp, tri_dk_imp, atol=1e-1, rtol=0), f'{ref_dk_imp.mean()} vs {tri_dk_imp.mean()}'
+    perc_diff = 100 * torch.abs(ref_dk_imp - tri_dk_imp).mean() / torch.abs(ref_dk_imp).mean()
+    print(f'passed test for dk_imp with {ref_dk_imp.mean()} vs {tri_dk_imp.mean()}, \t\tpercentage diff: {perc_diff}%')
+    # assert torch.allclose(ref_dq_imp, tri_dq_imp, atol=1e-1, [rtol=0), f'{ref_dq_imp.mean()} vs {tri_dq_imp.mean()}'
+    perc_diff = 100 * torch.abs(ref_dq_imp - tri_dq_imp).mean() / torch.abs(ref_dq_imp).mean()
+    print(f'passed test for dq_imp with {ref_dq_imp.mean()} vs {tri_dq_imp.mean()}, \t\tpercentage diff: {perc_diff}%')
+    print(f'original time: {ori_end - ori_start}')
+    print(f'triton time: {tri_end - tri_start}')

utils.py

@@ -0,0 +1,1521 @@
+import os
+import pdb
+import copy
+import math
+import numpy as np 
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+import gc
+import matplotlib.pyplot as plt
+
+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 scipy.stats import spearmanr
+from transformers.models.llama.configuration_llama import LlamaConfig
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, LlamaAttention, LlamaRMSNorm, apply_rotary_pos_emb
+from torch.utils.data import Dataset, DataLoader
+import torch
+from typing import Tuple
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+import re
+import time
+import matplotlib.cm as cm
+from scipy.spatial.distance import cosine
+from tqdm import tqdm
+
+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 sanitize_filename(name):
+    return re.sub(r'[<>:"/\\|?*\'\[\]]', '_', name)
+
+def args_to_name(args, timestamp=True):
+    args_dict = vars(args)
+    args_dict = args_dict.copy()
+    # remove longbench_datasets, task_list from args_dict
+    args_dict.pop("longbench_datasets", None)
+    args_dict.pop("task_list", None)
+    model_descr = list(args_dict.values())
+    # Split the model description into two parts
+    split_point = len(model_descr) // 2
+    folder_part = model_descr[:split_point]
+    file_part = model_descr[split_point:]
+    # Create a sanitized folder name from the first part
+    folder_name = "_".join([str(elem) for elem in folder_part])
+    folder_name = sanitize_filename(folder_name)
+    # Create a sanitized file name from the second part
+    file_name = "_".join([str(elem) for elem in file_part])
+    file_name = sanitize_filename(file_name)
+    # Add timestamp to file name
+    if timestamp:
+        timestamp = time.strftime("%Y%m%d-%H%M%S")
+        file_name = file_name + "_" + timestamp
+    return folder_name, file_name
+
+    
+def snapkv_mask_only(self, query_states, key_states, value_states, attention_mask=None):
+    """
+    'Mask-only' version of SnapKV that does not gather/slice the actual key_states:
+    - If q_len < max_capacity_prompt, do nothing.
+    - Else, we compute the 'top prefix tokens' using the last window_size queries, 
+      plus the last window_size tokens themselves.
+    - Then we create a single-step mask that is -inf for all other tokens.
+
+    We store that single-step mask in self.snapkv_cache so that 
+    on the next decode step (q_len=1) we can re-apply it.
+    """
+    bsz, num_heads, q_len, head_dim = query_states.shape
+    # Ensure prefix-phase
+    assert key_states.shape[-2] == query_states.shape[-2], "Prefix shape mismatch"
+
+    # If no compression is needed, just return the normal outputs
+    if q_len < self.max_capacity_prompt:
+        return None  # signals: no special mask built
+
+    # 1) Compute local attention (like SnapKV: last window_size queries vs entire prefix)
+    obs = self.window_size
+    if obs > q_len:
+        obs = q_len  # if the prompt is shorter than window_size
+    attn_logits = torch.matmul(query_states[..., -obs:, :],
+                               key_states.transpose(-2, -1)) / math.sqrt(head_dim)
+    # shape [bsz, num_heads, obs, kv_seq_len]
+
+    # 2) Build a local triangular mask of shape (obs, obs) for the last window_size queries
+    mask = torch.full((obs, obs), float('-inf'), device=attn_logits.device, dtype=attn_logits.dtype)
+    idxs = torch.arange(obs, device=mask.device)
+    mask.masked_fill_(idxs < idxs.unsqueeze(-1), 0)  # lower-tri (including diagonal)=0, above diag=-inf
+    local_mask = mask.unsqueeze(0).unsqueeze(0)  # [1,1,obs,obs]
+
+    # Apply it to the last obs block in attn_logits
+    attn_logits[:, :, -obs:, -obs:] += local_mask
+
+    # 3) Softmax
+    attn_probs = F.softmax(attn_logits, dim=-1, dtype=torch.float32).to(query_states.dtype)
+    # shape [bsz, num_heads, obs, kv_seq_len]
+
+    # 4) Sum across the obs dimension => [bsz, num_heads, kv_seq_len]
+    attn_sum = attn_probs.sum(dim=-2)
+
+    # 5) Optional pooling => must pass [N, C, L] to max_pool1d/avg_pool1d
+    # attn_sum is shape [bsz, num_heads, kv_seq_len]. We can flatten bsz*num_heads
+    bnh = bsz * num_heads
+    L = key_states.shape[-2]  # kv_seq_len
+    x = attn_sum.view(bnh, 1, L)  # [bnh, 1, kv_seq_len]
+
+    if self.pooling == 'avgpool':
+        pooled = F.avg_pool1d(
+            x,
+            kernel_size=self.kernel_size,
+            stride=1,
+            padding=self.kernel_size // 2
+        )
+    elif self.pooling == 'maxpool':
+        pooled = F.max_pool1d(
+            x,
+            kernel_size=self.kernel_size,
+            stride=1,
+            padding=self.kernel_size // 2
+        )
+    else:
+        raise ValueError("Unsupported pooling method")
+
+    # Now pooled is shape [bnh, 1, L']
+    # Usually, L' = L if stride=1, but let's just keep it so it lines up with kv_seq_len
+    pooled = pooled.view(bsz, num_heads, -1)  # [bsz, num_heads, kv_seq_len]
+
+    # 6) topk
+    top_prefix_to_keep = self.max_capacity_prompt - obs
+    prefix_indices = pooled.topk(top_prefix_to_keep, dim=-1).indices  # [bsz, num_heads, top_prefix_to_keep]
+
+    # 7) Build single-step mask => shape [bsz, num_heads, 1, kv_seq_len]
+    single_mask = torch.full(
+        (bsz, num_heads, 1, L),
+        float('-inf'),
+        device=query_states.device,
+        dtype=query_states.dtype
+    )
+
+    # unmask the top prefix positions
+    row_idx = torch.arange(bsz, device=query_states.device).view(bsz, 1, 1)
+    head_idx = torch.arange(num_heads, device=query_states.device).view(1, num_heads, 1)
+    row_idx = row_idx.expand(bsz, num_heads, top_prefix_to_keep)
+    head_idx = head_idx.expand(bsz, num_heads, top_prefix_to_keep)
+    single_mask[row_idx, head_idx, 0, prefix_indices] = 0.0
+
+    # unmask the last obs tokens
+    single_mask[:, :, 0, -obs:] = 0.0
+
+    return single_mask
+
+def calculate_effective_sparsity(final_mask, attention_mask):
+    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()
+    return effective_sparsity
+
+# def sorted_index_to_mask(sorted_indices, attention_mask, min_sparse_index, bsz, q_len, key_len, sparse_aggression):
+#     device = sorted_indices.device
+#     dtype = sorted_indices.dtype
+#     # Step 1: Calculate the number of keys to keep for each query position
+#     # K = ceil((query_position + 1) * sparse_aggression)
+#     # Shape: [1, 1, Lq, 1]
+#     query_positions = torch.arange(q_len, device=device).view(1, 1, q_len, 1).float() + 1.0  # 1-based indexing
+#     K = torch.ceil(query_positions * sparse_aggression).to(dtype=torch.long)  # [1, 1, Lq, 1]
+#     # Ensure K does not exceed key_len
+#     K = torch.clamp(K, max=key_len)
+#     # Step 2: Expand attention_mask to match sorted_indices shape
+#     # attention_mask: [1, 1, Lq, Lk] -> [B, H, Lq, Lk]
+#     attention_mask_expanded = attention_mask.expand(bsz, -1, -1, -1)  # [B, 1, Lq, Lk]
+#     attention_mask_expanded = attention_mask_expanded.expand(-1, sorted_indices.size(1), -1, -1)  # [B, H, Lq, Lk]
+#     attention_mask_expanded = (~attention_mask_expanded.bool()).int()
+#     # Step 3: Gather attention_mask based on sorted_indices
+#     # This rearranges the attention_mask to follow the sorted order of keys
+#     gathered_mask = torch.gather(attention_mask_expanded, dim=-1, index=sorted_indices)  # [B, H, Lq, Lk]
+#     # Step 4: Convert gathered_mask to float for cumulative sum
+#     gathered_mask_float = gathered_mask.float()
+#     # Step 5: Compute cumulative sum along the sorted key dimension
+#     cum_sum = torch.cumsum(gathered_mask_float, dim=-1)  # [B, H, Lq, Lk]
+#     # Step 6: Create a mask where cumulative sum <= K_i
+#     # K has shape [1, 1, Lq, 1], expand it to [B, H, Lq, Lk] for comparison
+#     K_broadcast = K.view(1, 1, q_len, 1).expand(bsz, sorted_indices.size(1), q_len, key_len)  # [B, H, Lq, Lk]
+#     selected_mask = cum_sum <= K_broadcast  # [B, H, Lq, Lk], boolean
+#     # Step 7: Initialize mask_tensor with -inf
+#     mask_tensor = torch.full_like(attention_mask_expanded.to(torch.float32), float('-inf'))  # [B, H, Lq, Lk]
+#     # Step 8: Prepare values to scatter: 0 where selected_mask is True, else -inf
+#     # This ensures that only the selected positions are allowed (0), others remain disallowed (-inf)
+#     scatter_values = torch.zeros_like(gathered_mask_float)
+#     scatter_values = scatter_values.masked_fill(~selected_mask, float('-inf'))  # [B, H, Lq, Lk]
+#     # Step 9: Scatter the values back to the original key positions
+#     # For each sorted key position, place the corresponding scatter_value at the original key index
+#     mask_tensor.scatter_(-1, sorted_indices, scatter_values)
+#     # Step 10: Ensure mask_tensor has mask_tensor[:, :, :, :min_sparse_index] = 0 
+#     mask_tensor[:, :, :, :min_sparse_index] = 0.0
+#     return mask_tensor
+
+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
+
+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
+
+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 plot_thresholds(threshold_tensor, true_threshold_tensor, fpath_base, fpath_specific):
+    """
+    Plots mean and error regions for random layers and heads, showing threshold changes across decode steps.
+    
+    Args:
+    - threshold_tensor: torch.Tensor of shape [163, 31, 32, 1024].
+    - true_threshold_tensor: torch.Tensor of shape [163, 31, 32, 1024].
+    """
+    def create_plot(tensor, title, filename):
+        """
+        Helper function to generate the plot.
+        """
+        # Choose 3 random layers
+        layers = np.random.choice(tensor.shape[1], 3, replace=False)
+        # layers = np.array([0, 15, 30])
+        
+        # Create subplots
+        fig, axs = plt.subplots(1, 3, figsize=(18, 5), sharey=True)
+        x = np.arange(tensor.shape[3])  # Decode steps (1024)
+        
+        for i, layer in enumerate(layers):
+            # Choose 5 random heads for this layer
+            heads = np.random.choice(tensor.shape[2], 5, replace=False)
+            
+            for head in heads:
+                try:
+                    # Extract data for the selected layer and head
+                    data = tensor[:, layer, head, :].numpy()  # Shape [163, 1024]
+                    
+                    # Compute mean and standard deviation across samples (dim=0)
+                    mean = np.mean(data, axis=0)
+                    std = np.std(data, axis=0)
+                    
+                    # Plot mean and shaded error region for the head
+                    axs[i].plot(x, mean, label=f"Head {head}")
+                    axs[i].fill_between(x, mean - std, mean + std, alpha=0.3)
+                except:
+                    import pdb; pdb.set_trace()
+            
+            # Customize subplot
+            axs[i].set_title(f"Layer {layer}")
+            axs[i].set_xlabel("Decode Step")
+            axs[i].grid(True)
+            axs[i].legend(fontsize=8)  # Adjust legend size for multiple heads
+        
+        # Common Y-axis label and adjustments
+        axs[0].set_ylabel("Threshold")
+        fig.suptitle(title, fontsize=16)
+        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
+        
+        # Save the plot
+        plt.savefig(filename)
+        plt.close()
+
+    def compute_mean_threshold(tensor):
+        """
+        Computes the mean threshold value for each head and layer, excluding the first 32 tokens.
+        """
+        # Exclude the first 32 tokens (dimension 1024)
+        tensor_excluded = tensor[:, :, :, 32:]  # Shape [163, 31, 32, 992]
+        
+        # Compute the mean along the first (samples) and last (remaining tokens) dimensions
+        mean_threshold = tensor_excluded.mean(dim=(0, -1))  # Shape [31, 32]
+        
+        return mean_threshold
+
+    # create folder fpath_base if it does not exist
+    if not os.path.exists(f"threshold_plots"):
+        os.makedirs(f"threshold_plots")
+    if not os.path.exists(f"threshold_plots/{fpath_base}"):
+        os.makedirs(f"threshold_plots/{fpath_base}")
+    # Plot for threshold_tensor
+    create_plot(threshold_tensor, "Post-Attention Thresholds", f"threshold_plots/{fpath_base}/{fpath_specific}_postattn_threshold.pdf")
+    
+    # Plot for true_threshold_tensor
+    create_plot(true_threshold_tensor, "Predicted Pre-SM Thresholds", f"threshold_plots/{fpath_base}/{fpath_specific}_pred_presm_threshold.pdf")
+
+    # Compute mean thresholds
+    mean_threshold_postattn = compute_mean_threshold(threshold_tensor)
+    mean_threshold_predpresm = compute_mean_threshold(true_threshold_tensor)
+    
+    return mean_threshold_postattn, mean_threshold_predpresm
+
+
+
+# def plot_thresholds(threshold_tensor, true_threshold_tensor):
+#     """
+#     Plots mean and error regions for random layers and heads, showing threshold changes across decode steps.
+    
+#     Args:
+#     - threshold_tensor: torch.Tensor of shape [163, 31, 32, 1024].
+#     - true_threshold_tensor: torch.Tensor of shape [163, 31, 32, 1024].
+#     """
+#     def create_plot(tensor, title, filename):
+#         """
+#         Helper function to generate the plot.
+#         """
+#         # Choose 3 random layers and heads
+#         layers = np.random.choice(tensor.shape[1], 3, replace=False)
+#         heads = np.random.choice(tensor.shape[2], 3, replace=False)
+        
+#         # Create subplots
+#         fig, axs = plt.subplots(1, 3, figsize=(15, 5), sharey=True)
+#         x = np.arange(tensor.shape[3])  # Decode steps (1024)
+        
+#         for i, (layer, head) in enumerate(zip(layers, heads)):
+#             # Extract data for the selected layer and head
+#             data = tensor[:, layer, head, :].numpy()  # Shape [163, 1024]
+            
+#             # Compute mean and standard deviation across samples (dim=0)
+#             mean = np.mean(data, axis=0)
+#             std = np.std(data, axis=0)
+            
+#             # Plot with shaded error regions
+#             axs[i].plot(x, mean, label=f"Layer {layer}, Head {head}")
+#             axs[i].fill_between(x, mean - std, mean + std, alpha=0.3)
+#             axs[i].set_title(f"Layer {layer}, Head {head}")
+#             axs[i].set_xlabel("Decode Step")
+#             axs[i].grid(True)
+        
+#         # Common Y-axis label and adjustments
+#         axs[0].set_ylabel("Threshold")
+#         fig.suptitle(title, fontsize=16)
+#         plt.tight_layout(rect=[0, 0.03, 1, 0.95])
+        
+#         # Save the plot
+#         plt.savefig(filename)
+#         plt.close()
+
+#     def compute_mean_threshold(tensor):
+#         """
+#         Computes the mean threshold value for each head and layer, excluding the first 32 tokens.
+#         """
+#         # Exclude the first 32 tokens (dimension 1024)
+#         tensor_excluded = tensor[:, :, :, 32:]  # Shape [163, 31, 32, 992]
+        
+#         # Compute the mean along the first (samples) and last (remaining tokens) dimensions
+#         mean_threshold = tensor_excluded.mean(dim=(0, -1))  # Shape [31, 32]
+        
+#         return mean_threshold
+
+#     # Plot for threshold_tensor
+#     create_plot(threshold_tensor, "Post-Attention Thresholds", "postattn_threshold.pdf")
+    
+#     # Plot for true_threshold_tensor
+#     create_plot(true_threshold_tensor, "Predicted Pre-SM Thresholds", "pred_presm_threshold.pdf")
+
+    
+#     # Compute mean thresholds
+#     mean_threshold_postattn = compute_mean_threshold(threshold_tensor)
+#     mean_threshold_predpresm = compute_mean_threshold(true_threshold_tensor)
+    
+#     return mean_threshold_postattn, mean_threshold_predpresm
+
+
+# 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.
+    
+#     Args:
+#         estimated_importance (torch.Tensor): Tensor of estimated importance values [B, H, L, L] OR [B, L, H] (Token / Head Prediction)
+#         true_importance (torch.Tensor): Tensor of true importance values [B, H, L, L] OR [B, L, H] (Token / Head Prediction)
+#         top_k_ratio (float): Fraction of top-k elements to consider for hit accuracy (default: 0.5).
+
+#     Returns:
+#         Tuple[float, float, float]: Hit accuracy, mean rank correlation, max rank correlation.
+#     """
+#     if len(estimated_importance.shape) == 4:
+#         # Only take the final, fully unmasked next-word prediction problem.
+#         import pdb; pdb.set_trace()
+#         estimated_importance = estimated_importance[:, :, -1, :].unsqueeze(-2)
+#         true_importance = true_importance[:, :, -1, :].unsqueeze(-2)
+#         estisize = estimated_importance.size(0) * estimated_importance.size(1) * estimated_importance.size(2)
+#     elif len(estimated_importance.shape) == 3:
+#         # Take all next-word head magnitudes, but note that its not L2
+#         import pdb; pdb.set_trace()
+#         estisize = estimated_importance.size(0) * estimated_importance.size(1)
+#     else:
+#         raise ValueError("Invalid shape for estimated_importance tensor. Must be 3 or 4 dimensional.")
+#     # Sort indices to get ranks
+#     _, sorted_esti_indices = torch.sort(estimated_importance, dim=-1, descending=True)
+#     _, sorted_true_indices = torch.sort(true_importance, dim=-1, descending=True)
+    
+#     # Compute rank correlation
+#     n = sorted_esti_indices.shape[-1]
+#     d = sorted_esti_indices - sorted_true_indices
+#     d_squared = d ** 2
+#     sum_d_squared = torch.sum(d_squared, dim=-1)
+#     rank_correlation = 1 - (6 * sum_d_squared) / (n * (n**2 - 1))
+    
+#     # Compute top-k hit accuracy
+#     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)
+#     intersection = (top_esti_indices.unsqueeze(-1) == top_true_indices.unsqueeze(-2)).any(dim=-1).sum(dim=-1)
+#     hit_accuracy = intersection.sum().item() / (top_k * estisize)
+
+#     # Return metrics
+#     mean_rank_corr = rank_correlation.mean().item()
+#     max_rank_corr = rank_correlation.max().item()
+#     return hit_accuracy, mean_rank_corr, max_rank_corr
+
+    
+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 FlattenedDataset(Dataset):
+    def __init__(self, dataset, max_seq_len, max_repeat_fraction=0.7):
+        self.max_seq_len = max_seq_len
+        self.max_repeat_fraction = max_repeat_fraction
+
+        # Extract and flatten the input_ids column
+        all_tokens = torch.cat([torch.tensor(ids) for ids in dataset["input_ids"]], dim=0)
+
+        # Calculate the number of full chunks
+        num_full_chunks = len(all_tokens) // self.max_seq_len
+        all_chunks = all_tokens[:num_full_chunks * self.max_seq_len].view(-1, self.max_seq_len)
+
+        # Filter out chunks with excessive repeated tokens
+        self.chunks = []
+        for chunk in all_chunks:
+            unique_tokens, counts = torch.unique(chunk, return_counts=True)
+            max_repeats = counts.max().item()
+            if max_repeats <= self.max_repeat_fraction * chunk.numel():
+                self.chunks.append(chunk)
+        
+        self.chunks = torch.stack(self.chunks)  # Stack the remaining chunks
+
+    def __len__(self):
+        return len(self.chunks)
+
+    def __getitem__(self, idx):
+        return self.chunks[idx]
+
+
+def compute_js_divergence(p, q, epsilon=1e-12):
+    """
+    Compute the Jensen-Shannon Divergence between two probability distributions.
+    
+    Args:
+        p (torch.Tensor): Shape [..., D]
+        q (torch.Tensor): Shape [..., D]
+        epsilon (float): Small value to avoid log(0)
+    
+    Returns:
+        torch.Tensor: JS Divergence values per pair (Shape: [...])
+    """
+    # Add epsilon to avoid log(0)
+    p = p + epsilon
+    q = q + epsilon
+    
+    # Normalize to ensure they are valid probability distributions
+    p = p / p.sum(dim=-1, keepdim=True)
+    q = q / q.sum(dim=-1, keepdim=True)
+    
+    # Compute the average distribution
+    m = 0.5 * (p + q)
+    
+    # Compute KL Divergences
+    kl_p = F.kl_div(m.log(), p, reduction='none').sum(dim=-1)
+    kl_q = F.kl_div(m.log(), q, reduction='none').sum(dim=-1)
+    
+    # Compute JS Divergence
+    js = 0.5 * kl_p + 0.5 * kl_q
+    
+    return js
+
+def compute_head_consistency_js(head_data):
+    """
+    Compute the consistency of a head's probability distributions across examples using JS Divergence.
+    
+    Args:
+        head_data (torch.Tensor): Shape [163, 1024], probability distributions for one head.
+    
+    Returns:
+        float: Mean pairwise JS Divergence across examples.
+    """
+    num_examples = head_data.size(0)
+    
+    # Ensure head_data is float for division and normalization
+    head_data = head_data.float()
+    
+    # Normalize each distribution to ensure it sums to 1
+    head_data = head_data / head_data.sum(dim=-1, keepdim=True)
+    
+    # Initialize variables to accumulate JS divergences
+    total_js = 0.0
+    count = 0
+    
+    # Iterate over all unique pairs without redundancy
+    for i in tqdm(range(num_examples), desc="Computing JS Divergence"):
+        # Select the i-th distribution and expand its dimensions
+        p = head_data[i].unsqueeze(0)  # Shape: [1, D]
+        
+        # Select all distributions after the i-th to avoid duplicate pairs
+        q = head_data[i+1:]  # Shape: [N-i-1, D]
+        
+        if q.size(0) == 0:
+            continue  # No more pairs left
+        
+        # Compute JS Divergence between p and q
+        js = compute_js_divergence(p.repeat(q.size(0), 1), q)  # Shape: [N-i-1]
+        
+        # Accumulate the sum of JS Divergence
+        total_js += js.sum().item()
+        count += js.size(0)
+    
+    # Compute the mean JS Divergence
+    mean_js = total_js / count if count > 0 else 0.0
+    
+    return mean_js
+    
+    
+def compute_token_consistency_js(head_data):
+    """
+    Compute token consistency for all heads in a layer using JS Divergence.
+    
+    Args:
+        head_data (torch.Tensor): Shape [163, 24, 1024], layer's head data.
+
+    Returns:
+        np.ndarray: Consistency values for all 24 heads.
+    """
+    num_heads = head_data.shape[1]
+    consistency_metrics = []
+    
+    for head in tqdm(range(num_heads), desc="Processing Heads"):
+        head_consistency = compute_head_consistency_js(head_data[:, head, :])  # Consistency for one head
+        consistency_metrics.append(head_consistency)
+    
+    return np.array(consistency_metrics)
+
+
+def graph_headtok_pos_affinity(head_tokpos_affinity, args):
+    """
+    Generate a violin plot for Token Access Consistency Across Layers using JS Divergence.
+    
+    Args:
+        head_tokpos_affinity (dict): Dictionary where keys are layer identifiers and values are 
+                                     torch.Tensor of shape [163, 24, 1024].
+        args (argparse.Namespace): Arguments containing at least 'model_path'.
+    """
+    # Process the data into a format suitable for plotting
+    layer_ids = []
+    consistency_values = []
+    
+    for layer, tensor in tqdm(head_tokpos_affinity.items(), desc="Processing Layers"):
+        layer_consistency = compute_token_consistency_js(tensor)  # Shape: [24]
+        layer_ids.extend([layer] * len(layer_consistency))
+        consistency_values.extend(layer_consistency)
+
+    # Create directory structure: ablation_plots/traces/tok_js_div
+    trace_dir = f"ablation_plots/traces/tok_js_div"
+    os.makedirs(trace_dir, exist_ok=True)
+    mpath = args.model_path.replace("/", "_")
+    # Save data to a .npy file (NumPy array format)
+    trace_path = os.path.join(trace_dir, f"layer_consistency_{mpath}.npy")
+    os.makedirs(os.path.dirname(trace_path), exist_ok=True)
+    np.save(trace_path, {"Layer": layer_ids, "JS_Divergence": consistency_values})
+    print(f"Consistency data saved to {trace_path}")
+
+    
+    # Prepare data for Seaborn violin plot
+    data = {"Layer": layer_ids, "JS_Divergence": consistency_values}
+    
+    # Create the violin plot
+    plt.figure(figsize=(10, 6))
+    sns.violinplot(x=data["Layer"], y=data["JS_Divergence"], scale="width", inner="quartile", palette="viridis")
+    
+    # Formatting the plot
+    plt.title(f"Token Access Consistency Across Layers for {args.model_path}", fontsize=16)
+    plt.xlabel("Layer", fontsize=14)
+    plt.ylabel("Token Consistency Metric (Mean JS Divergence)", fontsize=14)
+    plt.xticks(fontsize=12)
+    plt.yticks(fontsize=12)
+    
+    # Enhance layout
+    plt.tight_layout()
+    
+    # Create ablation_plots directory if it doesn't exist
+    os.makedirs("ablation_plots", exist_ok=True)
+    
+    # Construct the full file path
+    file_path = f"ablation_plots/{mpath}_headtok_consistency_js_divergence.pdf"
+    
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+    # Save the plot
+    plt.savefig(file_path)
+    plt.close()
+
+
+def compute_head_agreement_js(head_data):
+    """
+    Compute head agreement for a single example using JS Divergence.
+    
+    Args:
+        head_data (torch.Tensor): Shape [num_heads, num_tokens], token distributions for all heads.
+    
+    Returns:
+        float: Mean pairwise JS Divergence for the heads.
+    """
+    num_heads = head_data.size(0)
+    js_divergences = []
+
+    for i in range(num_heads):
+        p = head_data[i].unsqueeze(0)  # Shape: [1, num_tokens]
+        q = head_data[i + 1 :]  # Remaining heads, Shape: [num_heads - i - 1, num_tokens]
+        
+        if q.size(0) == 0:
+            continue
+        
+        js = compute_js_divergence(p.repeat(q.size(0), 1), q)  # Pairwise JS Div
+        js_divergences.extend(js.tolist())
+
+    # Return the mean of the upper triangular JS Divergence matrix
+    return np.mean(js_divergences)
+
+def compute_head_agreement_all_examples(head_tokpos_affinity):
+    """
+    Compute head agreement for all examples across all layers.
+    
+    Args:
+        head_tokpos_affinity (dict): Dictionary where keys are layer identifiers and values are 
+                                     torch.Tensor of shape [num_examples, num_heads, num_tokens].
+
+    Returns:
+        np.ndarray: Head agreement values for all examples.
+    """
+    agreement_values = []
+
+    for layer, tensor in tqdm(head_tokpos_affinity.items(), desc="Processing Layers"):
+        for example_idx in range(tensor.shape[0]):  # Iterate over examples
+            head_data = tensor[example_idx]  # Shape: [num_heads, num_tokens]
+            agreement = compute_head_agreement_js(head_data)
+            agreement_values.append(agreement)
+
+    return np.array(agreement_values)
+
+def plot_and_save_head_agreement(agreement_values, args):
+    """
+    Save head agreement values and plot them as a violin plot.
+    
+    Args:
+        agreement_values (np.ndarray): Head agreement values for all examples.
+        args (argparse.Namespace): Arguments containing at least 'model_path'.
+    """
+    # Save agreement values to a .npy file
+    trace_dir = "ablation_plots/traces/headagreement_js_div"
+    os.makedirs(trace_dir, exist_ok=True)
+    mpath = args.model_path.replace("/", "_")
+    trace_path = os.path.join(trace_dir, f"head_agreement_{mpath}.npy")
+    np.save(trace_path, {"HeadAgreement": agreement_values})
+    print(f"Head agreement values saved to {trace_path}")
+
+    # Plot the violin plot
+    plt.figure(figsize=(10, 6))
+    sns.violinplot(data=[agreement_values], scale="width", inner="quartile", palette="viridis")
+
+    # Formatting
+    plt.title(f"Head Agreement Across Examples for {mpath}", fontsize=16)
+    plt.xlabel("Model", fontsize=14)
+    plt.ylabel("Mean JS Divergence Between Heads", fontsize=14)
+    plt.xticks([0], [mpath], fontsize=12)  # Single violin
+    plt.yticks(fontsize=12)
+
+    # Enhance layout
+    plt.tight_layout()
+
+    # Save the plot
+    plot_path = f"ablation_plots/{mpath}_head_agreement_js_divergence.pdf"
+    plt.savefig(plot_path)
+    print(f"Violin plot saved to {plot_path}")
+    plt.close()
+
+
+def compute_jsd_over_decode_steps(decode_probs):
+    """
+    Compute average JSD over decode steps for a single head.
+    
+    Args:
+        decode_probs (torch.Tensor): Shape [50, 974], softmaxed token importances for 50 decode steps.
+    
+    Returns:
+        float: Mean JSD across the upper diagonal of the pairwise JSD matrix.
+    """
+    # Expand dims to create pairwise matrices
+    p = decode_probs.unsqueeze(0)  # Shape: [1, 50, 974]
+    q = decode_probs.unsqueeze(1)  # Shape: [50, 1, 974]
+    
+    # Compute pairwise JSD (broadcasting handles the pairwise combinations)
+    jsd_matrix = compute_js_divergence(p, q)  # Shape: [50, 50]
+    
+    # Extract upper diagonal without the diagonal itself
+    triu_indices = torch.triu_indices(jsd_matrix.size(0), jsd_matrix.size(1), offset=1)
+    jsd_upper = jsd_matrix[triu_indices[0], triu_indices[1]]  # Shape: [N*(N-1)/2]
+    
+    return jsd_upper.mean().item()
+
+def compute_layer_jsd(decode_tokpos_affinity):
+    """
+    Compute average JSD over decode steps for all heads and layers.
+    
+    Args:
+        decode_tokpos_affinity (dict): Dictionary where keys are layer indices and values are 
+                                       torch.Tensor of shape [163, 24, 50, 974].
+    
+    Returns:
+        dict: Average JSD values for each head in each layer, keyed by layer index.
+    """
+    layer_jsd = {}
+
+    for layer, tensor in tqdm(decode_tokpos_affinity.items(), desc="Processing Layers"):
+        # tensor shape: [163, 24, 50, 974]
+        num_examples, num_heads, num_decode_steps, num_tokens = tensor.shape
+        
+        # Reshape for batch processing
+        decode_probs = tensor.view(-1, num_decode_steps, num_tokens)  # Shape: [163 * 24, 50, 974]
+        
+        # Compute pairwise JSD for all heads and examples
+        jsd_values = []
+        for decode_head in decode_probs:
+            jsd_values.append(compute_jsd_over_decode_steps(decode_head))
+
+        # Reshape back to per-head per-layer values
+        jsd_values = torch.tensor(jsd_values).view(num_examples, num_heads)  # Shape: [163, 24]
+        layer_jsd[layer] = jsd_values.mean(dim=0).tolist()  # Average across examples per head
+
+    return layer_jsd
+
+def plot_decode_jsd_violin(layer_jsd, args):
+    """
+    Plot per-layer JSD values as violins.
+    
+    Args:
+        layer_jsd (dict): Dictionary of average JSD values for each head in each layer.
+        args (argparse.Namespace): Arguments containing at least 'model_path'.
+    """
+    # Prepare data for per-layer violin plot
+    layers = []
+    jsd_values = []
+
+    for layer, values in layer_jsd.items():
+        layers.extend([layer] * len(values))
+        jsd_values.extend(values)
+
+    # Save JSD values to a file
+    trace_dir = "ablation_plots/traces/decode_jsd"
+    os.makedirs(trace_dir, exist_ok=True)
+    mpath = args.model_path.replace("/", "_")
+    trace_path = os.path.join(trace_dir, f"decode_jsd_{mpath}.npy")
+    np.save(trace_path, {"Layer": layers, "JSD": jsd_values})
+    print(f"Decode JSD data saved to {trace_path}")
+
+    # Plot the violin plot
+    plt.figure(figsize=(10, 6))
+    sns.violinplot(x=layers, y=jsd_values, scale="width", inner="quartile", palette="viridis")
+
+    # Formatting
+    plt.title(f"Per-Layer Decode JSD for {args.model_path}", fontsize=16)
+    plt.xlabel("Layer", fontsize=14)
+    plt.ylabel("Average JSD Over Decode Steps", fontsize=14)
+    plt.xticks(fontsize=12)
+    plt.yticks(fontsize=12)
+
+    plt.tight_layout()
+
+    # Save the plot
+    plot_path = f"ablation_plots/{mpath}_decode_jsd_per_layer.pdf"
+    plt.savefig(plot_path)
+    print(f"Violin plot saved to {plot_path}")
+    plt.close()
+
+def compute_percentage_match_vectorized(decode_probs, top_k=0.1):
+    """
+    Compute the average percentage match of top-k token indices across 50 decode steps for a single head.
+    
+    Args:
+        decode_probs (torch.Tensor): Shape [50, 974], softmaxed token importances for 50 decode steps.
+        top_k (float): Percentage of top tokens to consider (e.g., 0.1 for top 10%).
+    
+    Returns:
+        float: Average percentage match of token indices across the 50 decode steps.
+    """
+    num_steps, num_tokens = decode_probs.shape
+    k = int(num_tokens * top_k)  # Number of top tokens to consider
+
+    # Get top-k indices for all steps
+    top_indices = torch.topk(decode_probs, k, dim=-1).indices  # Shape: [50, k]
+
+    # Create a binary mask for top-k tokens
+    binary_mask = torch.zeros(num_steps, num_tokens, device=decode_probs.device)
+    binary_mask.scatter_(1, top_indices, 1)  # Shape: [50, 974]
+
+    # Compute overlap between all pairs of steps
+    overlap_matrix = torch.matmul(binary_mask, binary_mask.T)  # Shape: [50, 50]
+    overlap_per_pair = overlap_matrix / k  # Normalize by k to get match percentages
+
+    # Extract upper triangular without diagonal
+    triu_indices = torch.triu_indices(num_steps, num_steps, offset=1)
+    upper_triangle = overlap_per_pair[triu_indices[0], triu_indices[1]]  # Shape: [N*(N-1)/2]
+
+    return torch.tensor([upper_triangle.mean().item()], device=decode_probs.device)
+
+
+
+def compute_layer_percentage_match_vectorized(decode_tokpos_affinity, top_k=0.1):
+    """
+    Compute average percentage match for top-k token indices across decode steps for all heads and layers.
+    
+    Args:
+        decode_tokpos_affinity (dict): Dictionary where keys are layer indices and values are 
+                                       torch.Tensor of shape [163, 24, 50, 974].
+        top_k (float): Percentage of top tokens to consider (e.g., 0.1 for top 10%).
+    
+    Returns:
+        dict: Average percentage match values for each head in each layer, keyed by layer index.
+    """
+    layer_match = {}
+
+    for layer, tensor in tqdm(decode_tokpos_affinity.items(), desc="Processing Layers"):
+        # tensor shape: [163, 24, 50, 974]
+        num_examples, num_heads, num_decode_steps, num_tokens = tensor.shape
+
+        # Flatten examples and heads for batch processing
+        decode_probs = tensor.view(-1, num_decode_steps, num_tokens)  # Shape: [163 * 24, 50, 974]
+
+        # Vectorized computation for all decode heads
+        match_values = torch.cat([
+            compute_percentage_match_vectorized(decode_head, top_k=top_k) for decode_head in decode_probs
+        ])
+
+        # Reshape back to per-head per-layer values
+        match_values = match_values.view(num_examples, num_heads)  # Shape: [163, 24]
+        layer_match[layer] = match_values.mean(dim=0).tolist()  # Average across examples per head
+
+    return layer_match
+
+def plot_decode_percdrift_vectorized(layer_match, args):
+    """
+    Plot per-layer average percentage match of top-k token indices as violins.
+    
+    Args:
+        layer_match (dict): Dictionary of average percentage match values for each head in each layer.
+        args (argparse.Namespace): Arguments containing at least 'model_path'.
+    """
+    # Prepare data for per-layer violin plot
+    layers = []
+    match_values = []
+
+    for layer, values in layer_match.items():
+        layers.extend([layer] * len(values))
+        match_values.extend(values)
+
+    # Save match values to a file
+    trace_dir = "ablation_plots/traces/percdrift"
+    os.makedirs(trace_dir, exist_ok=True)
+    mpath = args.model_path.replace("/", "_")
+    trace_path = os.path.join(trace_dir, f"decode_percdrift_{mpath}.npy")
+    np.save(trace_path, {"Layer": layers, "Match": match_values})
+    print(f"Decode percentage drift data saved to {trace_path}")
+
+    # Plot the violin plot
+    plt.figure(figsize=(10, 6))
+    sns.violinplot(x=layers, y=match_values, scale="width", inner="quartile", palette="viridis")
+
+    # Formatting
+    plt.title(f"Per-Layer Percentage Match for {args.model_path}", fontsize=16)
+    plt.xlabel("Layer", fontsize=14)
+    plt.ylabel("Average Percentage Match (Top 10%)", fontsize=14)
+    plt.xticks(fontsize=12)
+    plt.yticks(fontsize=12)
+
+    plt.tight_layout()
+
+    # Save the plot
+    plot_path = f"ablation_plots/{mpath}_decode_percdrift_per_layer.pdf"
+    plt.savefig(plot_path)
+    print(f"Violin plot saved to {plot_path}")
+    plt.close()
+
+
+def plot_decode_drift_trajectory(decode_tokpos_affinity, top_k=0.1, args=None):
+    """
+    Plot the trajectory of top-k token overlaps for each decode step, compared to the first decode step.
+
+    Args:
+        decode_tokpos_affinity (dict): Dictionary where keys are layer indices and values are 
+                                       torch.Tensor of shape [num_examples, num_heads, num_decode_steps, num_tokens].
+        top_k (float): Percentage of top tokens to consider (e.g., 0.1 for top 10%).
+        args (argparse.Namespace): Arguments containing at least 'model_path'.
+    """
+
+
+    trajectories = []  # To store trajectories for all layers and heads
+    num_decode_steps = None  # To infer decode step count from the first layer
+    for layer, tensor in tqdm(decode_tokpos_affinity.items(), desc="Processing Layers"):
+        num_examples, num_heads, num_decode_steps, num_tokens = tensor.shape
+        k = int(num_tokens * top_k)  # Top-k token count
+
+        # Flatten examples and heads for batch processing
+        decode_probs = tensor.view(-1, num_decode_steps, num_tokens)  # Shape: [163 * 24, 50, 974]
+
+        # Get top-k indices for the first decode step
+        initial_top_k = torch.topk(decode_probs[:, 0, :], k, dim=-1).indices  # Shape: [163 * 24, k]
+
+        # Create binary masks for the first decode step
+        initial_masks = torch.zeros(decode_probs.size(0), num_tokens, device=decode_probs.device)
+        initial_masks.scatter_(1, initial_top_k, 1)  # Shape: [163 * 24, num_tokens]
+
+        # Get top-k indices for all decode steps
+        top_k_indices = torch.topk(decode_probs, k, dim=-1).indices  # Shape: [163 * 24, 50, k]
+
+        # Create binary masks for all decode steps
+        step_masks = torch.zeros(decode_probs.size(0), num_decode_steps, num_tokens, device=decode_probs.device)
+        step_masks.scatter_(2, top_k_indices, 1)  # Shape: [163 * 24, 50, num_tokens]
+
+        # Compute overlap with the first decode step for all steps
+        overlaps = (step_masks * initial_masks.unsqueeze(1)).sum(dim=-1) / k  # Shape: [163 * 24, 50]
+
+        # Append mean overlap trajectory for this layer
+        # import pdb; pdb.set_trace()
+        # Append all trajectories for this layer
+        trajectories.extend(overlaps.cpu().numpy())  # Shape: [163 * 24, 50]
+
+    # Plot all trajectories with a colormap
+    # Convert trajectories to NumPy for easier processing
+    trajectories = np.array(trajectories)  # Shape: [672, 50]
+    plt.figure(figsize=(10, 6))
+    colormap = cm.get_cmap("viridis", trajectories.shape[0])  # Viridis colormap
+    for i, trajectory in enumerate(trajectories):
+        plt.plot(range(num_decode_steps), trajectory, color=colormap(i), alpha=0.5, linewidth=0.8)
+
+    # Add labels and grid
+    plt.title("Decode Drift Trajectories for All Heads", fontsize=16)
+    plt.xlabel("Decode Step", fontsize=14)
+    plt.ylabel("Top-k Overlap with Initial Step", fontsize=14)
+    plt.grid(True)
+
+    # Save the plot
+    mpath = args.model_path.replace("/", "_")
+    output_path = f"ablation_plots/{mpath}_drift_trajectories.png"
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=600)
+    plt.close()
+    bins = 10
+    # Create Y-axis edges for histogram bins
+    y_edges = np.linspace(0, 1, bins + 1)  # Divide Y-axis into bins (0 to 1 overlap values)
+
+    # Digitize trajectories to find bin indices for all overlap values
+    bin_indices = np.digitize(trajectories, y_edges, right=True) - 1  # Shape: [109536, 50]
+
+    # Clip bin indices to avoid out-of-bound indices
+    bin_indices = np.clip(bin_indices, 0, bins - 1)  # Ensure indices are within [0, bins-1]
+
+    # Create the density map using bincount for each step
+    density_map = np.zeros((bins, num_decode_steps), dtype=np.float32)
+    for step in tqdm(range(num_decode_steps)):
+        # Count occurrences in each bin for the current step
+        counts = np.bincount(bin_indices[:, step], minlength=bins)
+        density_map[:, step] = counts
+    # Normalize density map for better visualization
+    density_map /= density_map.max()
+    # # Create a 2D density histogram (heatmap data)
+    # density_map = np.zeros((bins, num_decode_steps))
+    # y_edges = np.linspace(0, 1, bins + 1)  # Divide Y-axis into bins (0 to 1 overlap values)
+
+    # for step in range(num_decode_steps):
+    #     # Histogram for overlap values at each decode step
+    #     hist, _ = np.histogram(trajectories[:, step], bins=y_edges)
+    #     density_map[:, step] = hist
+
+    # # Normalize density map for better visualization
+    # density_map /= density_map.max()
+
+    # Plot the density heatmap
+    plt.figure(figsize=(12, 8))
+    sns.heatmap(
+        density_map,
+        cmap="viridis",
+        xticklabels=10,  # Optional: Adjust frequency of X-ticks
+        yticklabels=np.round(np.linspace(0, 1, bins), decimals=2),  # Show overlap bins
+        cbar_kws={"label": "Density"}
+    )
+    plt.title("Density of Decode Drift Trajectories", fontsize=16)
+    plt.xlabel("Decode Step", fontsize=14)
+    plt.ylabel("Top-k Overlap with Initial Step", fontsize=14)
+    plt.tight_layout()
+
+    # Save the plot
+    mpath = args.model_path.replace("/", "_")
+    output_path = f"ablation_plots/{mpath}_drift_density_heatmap.png"
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    plt.savefig(output_path, dpi=600)
+    plt.close()
+
+    print(f"Drift trajectory plot saved to {output_path}")
+    # Convert trajectories to NumPy for easier processing
+    trajectories = np.array(trajectories)  # Shape: [672, 50]
+    # average trajectories on dim=0
+    trajectories_to_save = np.mean(trajectories, axis=0)  # Shape: [50]
+
+    # Save rank agreement values to a .npy file
+    trace_dir = "ablation_plots/traces/decode_drift_trajectory"
+    os.makedirs(trace_dir, exist_ok=True)
+    mpath = args.model_path.replace("/", "_")
+    trace_path = os.path.join(trace_dir, f"drift_traj_{mpath}.npy")
+    np.save(trace_path, {"Trajectory": trajectories_to_save})
+
+    # Compute mean and standard deviation across trajectories
+    mean_trajectory = np.mean(trajectories, axis=0)  # Shape: [50]
+    std_trajectory = np.std(trajectories, axis=0)  # Shape: [50]
+
+    # Plot the mean trajectory with shaded error region
+    plt.figure(figsize=(10, 6))
+    plt.plot(range(num_decode_steps), mean_trajectory, label="Mean Drift Trajectory", color="blue")
+    plt.fill_between(
+        range(num_decode_steps),
+        mean_trajectory - std_trajectory,
+        mean_trajectory + std_trajectory,
+        color="blue",
+        alpha=0.2,
+        label="±1 Std Dev"
+    )
+    plt.axhline(y=1.0, color="red", linestyle="--", label="Initial (100% Overlap)")
+    plt.title("Decode Drift Trajectory with Error Region", fontsize=16)
+    plt.xlabel("Decode Step", fontsize=14)
+    plt.ylabel("Top-k Overlap with Initial Step", fontsize=14)
+    plt.legend(fontsize=12)
+    plt.grid(True)
+
+    # Save the plot
+    mpath = args.model_path.replace("/", "_")
+    output_path = f"ablation_plots/{mpath}_drift_trajectory.png"
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=600)
+    plt.close()
+
+    print(f"Drift trajectory plot saved to {output_path}")
+
+
+def compute_rank_agreement_all_examples(head_tokpos_affinity, args):
+    """
+    Compute rank agreement (mean, min, max) for all examples across all heads and layers.
+
+    Args:
+        head_tokpos_affinity (dict): Keys are layers; values are torch.Tensor of shape [num_examples, num_heads, num_tokens].
+
+    Returns:
+        np.ndarray: Shape [num_examples, 3], where 3 corresponds to mean, min, and max rank correlation per example.
+    """
+    num_examples = next(iter(head_tokpos_affinity.values())).shape[0]
+
+    # Flatten heads across layers
+    all_heads = []
+    for layer, tensor in head_tokpos_affinity.items():
+        all_heads.append(tensor)  # Shape: [num_examples, num_heads, num_tokens]
+    all_heads = torch.cat(all_heads, dim=1)  # Shape: [num_examples, total_heads, num_tokens]
+
+    # Rank tokens per head (Spearman's correlation requires ranks)
+    ranks = torch.argsort(all_heads, dim=-1).float()  # Shape: [num_examples, total_heads, num_tokens]
+
+    # Compute rank correlation metrics
+    results = []
+    total_corr_matrix = None
+
+    for example_idx in tqdm(range(num_examples), desc="Computing Rank Correlations"):
+        example_ranks = ranks[example_idx]  # Shape: [total_heads, num_tokens]
+
+        # Compute rank correlation matrix (Spearman)
+        corr_matrix = np.corrcoef(example_ranks.numpy())  # Shape: [total_heads, total_heads]
+
+        # Accumulate the correlation matrix
+        if total_corr_matrix is None:
+            total_corr_matrix = corr_matrix
+        else:
+            total_corr_matrix += corr_matrix
+        # Extract upper triangle
+        triu_indices = np.triu_indices(corr_matrix.shape[0], k=1)
+        upper_triangle = corr_matrix[triu_indices]
+
+        # Compute mean, min, and max of rank correlations
+        mean_corr = np.mean(upper_triangle)
+        min_corr = np.min(upper_triangle)
+        max_corr = np.max(upper_triangle)
+
+        results.append([mean_corr, min_corr, max_corr])
+
+    mean_corr_matrix = total_corr_matrix / num_examples
+    # Plot the heatmap
+    plt.figure(figsize=(12, 10))  # Adjust size as needed
+    sns.heatmap(corr_matrix, square=True, cbar=True, xticklabels=False, yticklabels=False, cmap="viridis")
+
+    # Set title
+    plt.title("Mean Rank Correlation Matrix", fontsize=16)
+
+    # Construct the file path
+    mpath = args.model_path.replace("/", "_")
+    heatmap_path = f"ablation_plots/{mpath}_rankcorr_heatmap.png"
+
+    # Save the heatmap
+    plt.tight_layout()
+    plt.savefig(heatmap_path, dpi=600)
+    plt.close()
+
+    print(f"Mean rank correlation heatmap saved to {heatmap_path}")
+    return np.array(results)  # Shape: [num_examples, 3]
+
+def plot_and_save_rank_agreement(rank_agreement, args):
+    """
+    Save rank agreement values and plot their distribution as a violin plot.
+
+    Args:
+        rank_agreement (np.ndarray): Shape [num_examples, 3], where columns represent mean, min, and max rank correlations.
+        args (argparse.Namespace): Arguments containing at least 'model_path'.
+    """
+    # Save rank agreement values to a .npy file
+    trace_dir = "ablation_plots/traces/rankagreement_allheads"
+    os.makedirs(trace_dir, exist_ok=True)
+    mpath = args.model_path.replace("/", "_")
+    trace_path = os.path.join(trace_dir, f"rank_agreement_{mpath}.npy")
+    np.save(trace_path, {"RankAgreement": rank_agreement})
+    print(f"Rank agreement data saved to {trace_path}")
+
+    # Prepare data for violin plot
+    categories = ["Mean", "Min", "Max"]
+    values = [rank_agreement[:, i] for i in range(3)]  # Separate columns
+
+    # Create the violin plot
+    plt.figure(figsize=(10, 6))
+    sns.violinplot(data=values, scale="width", inner="quartile", palette="viridis")
+
+    # Formatting
+    plt.title(f"Rank Agreement Distribution for {mpath}", fontsize=16)
+    plt.xlabel("Metric", fontsize=14)
+    plt.ylabel("Rank Correlation", fontsize=14)
+    plt.xticks(range(3), categories, fontsize=12)
+    plt.yticks(fontsize=12)
+
+    # Enhance layout and save the plot
+    plt.tight_layout()
+    # plot_path = os.path.join(trace_dir, f"rank_agreement_violin_{mpath}.pdf")
+
+    plot_path = f"ablation_plots/{mpath}_rank_agreement_violin.pdf"
+    plt.savefig(plot_path)
+    print(f"Violin plot saved to {plot_path}")
+    plt.close()
+# def compute_head_consistency(head_data):
+#     """
+#     Compute the consistency of a head's probability distributions across examples.
+
+#     Args:
+#         head_data (torch.Tensor): Shape [163, 1024], probability distributions for one head.
+
+#     Returns:
+#         float: Mean pairwise cosine similarity across examples.
+#     """
+#     # Ensure head_data is float for division and normalization
+#     head_data = head_data.float()
+
+#     # Normalize each distribution to ensure it sums to 1
+#     head_data = head_data / head_data.sum(dim=-1, keepdim=True)
+
+#     # Normalize vectors to unit norm to prepare for cosine similarity
+#     head_data = head_data / head_data.norm(dim=-1, keepdim=True)
+
+#     # Compute cosine similarity matrix: [163, 163]
+#     similarity_matrix = torch.matmul(head_data, head_data.T)
+
+#     # Extract the upper triangular part of the similarity matrix, excluding the diagonal
+#     num_examples = head_data.size(0)
+#     triu_indices = torch.triu_indices(num_examples, num_examples, offset=1)
+#     pairwise_similarities = similarity_matrix[triu_indices[0], triu_indices[1]]
+
+#     # Compute and return the mean similarity
+#     return pairwise_similarities.mean().item()
+
+# def compute_token_consistency(head_data):
+#     """
+#     Compute token consistency for all heads in a layer.
+    
+#     Args:
+#         head_data (torch.Tensor): Shape [163, 24, 1024], layer's head data.
+
+#     Returns:
+#         np.ndarray: Consistency values for all 24 heads.
+#     """
+#     # Ensure head_data is float for division and normalization
+#     head_data = head_data.float()
+
+#     # Normalize each distribution to ensure it sums to 1
+#     head_data = head_data / head_data.sum(dim=-1, keepdim=True)
+
+#     # Normalize vectors to unit norm to prepare for cosine similarity
+#     head_data = head_data / head_data.norm(dim=-1, keepdim=True)
+
+#     # Reshape to [24, 163, 1024] for batch processing
+#     head_data = head_data.permute(1, 0, 2)  # Now shape is [24, 163, 1024]
+
+#     # Compute cosine similarity matrices for all heads: [24, 163, 163]
+#     similarity_matrices = torch.bmm(head_data, head_data.transpose(1, 2))
+
+#     # Extract upper triangular indices, excluding the diagonal
+#     num_examples = head_data.size(1)
+#     triu_indices = torch.triu_indices(num_examples, num_examples, offset=1)
+
+#     # Gather pairwise similarities for all heads: [24, num_pairs]
+#     pairwise_similarities = similarity_matrices[:, triu_indices[0], triu_indices[1]]
+
+#     # Compute the mean similarity for each head: [24]
+#     consistency_metrics = pairwise_similarities.mean(dim=1).cpu().numpy()
+
+#     return consistency_metrics
+
+
+
+# def graph_headtok_pos_affinity(head_tokpos_affinity, args):
+#     # Process the data into a format suitable for plotting
+#     layer_ids = []
+#     consistency_values = []
+
+#     for layer, tensor in tqdm(head_tokpos_affinity.items()):
+#         layer_consistency = compute_token_consistency(tensor)  # Shape: [24]
+#         layer_ids.extend([layer] * len(layer_consistency))
+#         consistency_values.extend(layer_consistency)
+
+#     # Prepare data for Seaborn violin plot
+#     data = {"Layer": layer_ids, "Consistency": consistency_values}
+
+#     # Create the violin plot
+#     plt.figure(figsize=(10, 6))
+#     sns.violinplot(x=data["Layer"], y=data["Consistency"], scale="width", inner="quartile", palette="viridis")
+
+#     # Formatting the plot
+#     plt.title(f"Token Access Consistency Across Layers For {args.model_path}", fontsize=16)
+#     plt.xlabel("Layer", fontsize=14)
+#     plt.ylabel("Token Consistency Metric (Mean Cosine Similarity)", fontsize=14)
+#     plt.xticks(fontsize=12)
+#     plt.yticks(fontsize=12)
+
+#     # Show the plot
+#     plt.tight_layout()
+    
+#     # Create ablation_plots directory if it doesn't exist
+#     os.makedirs("ablation_plots", exist_ok=True)
+    
+#     # Construct the full file path
+#     file_path = f"ablation_plots/{args.model_path}_headtok_consistency.pdf"
+
+#     # Create the ablation_plots and intermediate directories if they don't exist
+#     os.makedirs(os.path.dirname(file_path), exist_ok=True)
+#     # Save the plot
+#     plt.savefig(file_path)
+#     plt.close()