### Importing dependencies and initializing variables Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/Attention Sinks.ipynb Initial setup for data manipulation, visualization, and tensor operations. ```python import numpy as np from matplotlib import pyplot as plt import seaborn as sns import torch import math from torch import nn head_id = ``` -------------------------------- ### Install RotateKV Environment and Dependencies Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Sets up a Python 3.10 Conda environment, installs PyTorch with CUDA 11.8, and compiles the required fast-hadamard-transform library. ```bash conda create -n RotateKV python==3.10 -y conda activate RotateKV # CUDA 11.8 conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=11.8 -c pytorch -c nvidia -y cd RotateKV pip install -r requirements.txt # Install the fast-hadamard-transform git clone https://github.com/Dao-AILab/fast-hadamard-transform.git cd fast-hadamard-transform pip install -e. ``` -------------------------------- ### Import necessary libraries Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/3D Visualizations of Key Tensors.ipynb Initializes the environment with Matplotlib, NumPy, and PyTorch. ```python import matplotlib.pyplot as plt import numpy as np import torch ``` -------------------------------- ### Evaluate PPL Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Calculate perplexity for FP16 baseline and various RotateKV quantization levels. ```bash # For FP16 baseline. python PPL_evaluation.py --FP16 True # For RotateKV results # INT2 python PPL_evaluation.py --RotateKV2 True # INT3 python PPL_evaluation.py --RotateKV3 True # INT4 python PPL_evaluation.py --RotateKV4 True ``` -------------------------------- ### Generate key states for visualization Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Save pre-RoPE and post-RoPE key states for 2D and 3D visualization tasks. ```bash # generate the key_states # pre_RoPE python PPL_evaluation.py --save_k_pre_rope True # post_RoPE python PPL_evaluation.py --save_k_post_rope True ``` -------------------------------- ### Generate massive activations Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Save massive activation data for visualization. ```bash # generate the massive activations python PPL_evaluation.py --save_massive_activations True ``` -------------------------------- ### Evaluate GSM8K Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Run GSM8K benchmark evaluations for FP16 baseline and RotateKV quantization levels. ```bash # For FP16 baseline. python gsm8k_evaluation.py --FP16 True # For RotateKV results # INT2 python gsm8k_evaluation.py --RotateKV2 True # INT3 python gsm8k_evaluation.py --RotateKV3 True # INT4 python gsm8k_evaluation.py --RotateKV4 True ``` -------------------------------- ### Calibrate reordering indices Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Run calibration scripts to prepare reordering indices for the specified model. ```bash python PPL_evaluation.py --generate_for_calibration True --model llama2_7b python k_reordering_calibration.py --model llama2_7b ``` -------------------------------- ### Import necessary libraries Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/2D Visualizations of Key Tensors.ipynb Imports required modules for data manipulation and visualization. ```python import matplotlib.pyplot as plt import numpy as np import torch from matplotlib.colors import Normalize ``` -------------------------------- ### Load and process key states Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/3D Visualizations of Key Tensors.ipynb Loads a tensor from a file and reshapes it for visualization. ```python key_states = torch.load("../save_tensors/attn_weights_post_rope_layer_11_sample_82.pt") (bsz, num_key_value_heads, q_len, head_dim) = key_states.shape key_states = key_states.transpose(1, 2).view(q_len, num_key_value_heads * head_dim)[2048,:].cpu().numpy() ``` -------------------------------- ### Load and process key states Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/2D Visualizations of Key Tensors.ipynb Loads a tensor from a file, transposes it, and extracts a specific slice for visualization. ```python key_states = torch.load("../save_tensors/key_states_post_rope_layer_11_sample_82.pt") (bsz, num_key_value_heads, q_len, head_dim) = key_states.shape key_states = key_states.transpose(1, 2).view(q_len, num_key_value_heads * head_dim)[2048,:].cpu().numpy() ``` -------------------------------- ### Visualize key states with a bar chart Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/2D Visualizations of Key Tensors.ipynb Generates a bar plot using the coolwarm colormap to represent the key state values. ```python cmap = plt.cm.coolwarm plt.figure(figsize=(8, 6)) plt.ylim(-15, 15) norm = Normalize(vmin=0, vmax=16) plt.bar(x, key_states, color=cmap(np.abs(key_states)/np.max(key_states)), alpha=1, width=50) plt.tick_params(axis='both', which='major', labelsize=15) plt.tight_layout() # plt.savefig("") plt.show() ``` -------------------------------- ### Generate attention scores Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Save attention scores for visualizing attention sinks. ```bash # generate the attetnion scores python PPL_evaluation.py --save_attention_scores True ``` -------------------------------- ### Loading and plotting attention scores Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/Attention Sinks.ipynb Loads a tensor from a file and displays the heatmap visualization. ```python attention_scores = torch.load("../save_tensors/key_states_post_rope_layer_11_sample_82.pt")[0,head_id,:120,:120].cpu().numpy() vis_attention_matrix(attention_scores) plt.title("", fontsize=20) # plt.savefig("") plt.show() ``` -------------------------------- ### Generate 3D surface plot Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/3D Visualizations of Key Tensors.ipynb Creates a 3D surface plot of the key tensor data using Matplotlib's plot_surface. ```python head_id = k = key_states_origin[0, head_id, :2048,:] tokens, channels = k.shape x = np.arange(tokens) y = np.arange(channels) X, Y = np.meshgrid(x, y) fig = plt.figure(figsize=(6, 6)) ax = fig.add_subplot(111, projection='3d') surf = ax.plot_surface(X, Y, k.T, cmap='coolwarm') ax.set_xlabel('Token', labelpad=-5) ax.set_ylabel('Channel', labelpad=-1) plt.tight_layout() # plt.savefig("") plt.show() ``` -------------------------------- ### Cite RotateKV Research Source: https://github.com/zunhaisu/rotatekv/blob/main/README.md Use this BibTeX entry to cite the RotateKV paper in academic publications. ```bibtex @article{su2025rotatekv, title={RotateKV: Accurate and Robust 2-Bit KV Cache Quantization for LLMs via Outlier-Aware Adaptive Rotations}, author={Su, Zunhai and Chen, Zhe and Shen, Wang and Wei, Hanyu and Li, Linge and Yu, Huangqi and Yuan, Kehong}, journal={arXiv preprint arXiv:2501.16383}, year={2025} } ``` -------------------------------- ### Define x-axis range Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/2D Visualizations of Key Tensors.ipynb Creates an array representing the indices for the bar chart. ```python x = np.arange(0, num_key_value_heads * head_dim) ``` -------------------------------- ### Defining the attention matrix visualization function Source: https://github.com/zunhaisu/rotatekv/blob/main/visualize/Attention Sinks.ipynb Creates a heatmap of the attention scores using seaborn. ```python def vis_attention_matrix(attention_scores, index=0, cmap="coolwarm", head_id=head_id): plt.figure(figsize=(5, 5)) ax = sns.heatmap( attention_map, vmin=0, vmax=0.1, cmap=cmap, square=True) # ax.set_xticks([0, 110]) # ax.set_xticklabels(['0', '110']) ax.set_yticks([]) ``` === COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.