RotateKV

https://github.com/zunhaisu/rotatekv
Admin
RotateKV is an accurate and robust 2-bit KV cache quantization framework for large language models...
Tokens:1,837
Snippets:23
Trust Score:7
Update:1 month ago
Show doc for...
Tokens:
Raw
### 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([])
```
Tokens:
Raw