# SnapKV
## 论文
[SnapKV: LLM Knows What You are Looking for Before Generation](https://arxiv.org/pdf/2404.14469)

## 模型简介

SnapKV采用近期query与key点积再池化的topk scores进行剪枝。

<div align=center>
    <img src="./doc/SnapKV.png"/>
</div>

## 环境依赖
| 软件 |   版本        |
| :------: |:---------:|
| DTK |      dtk25.04.1        |
| python |   3.10.12     |
| transformers |   5.3.0    |
| torch | 2.5.1+das.opt1.dtk25042 |

环境配置：
```
mv kvpress-SnapKV-Qwen3-8B_pytorch kvpress_pytorch # 去框架名后缀
```

### Docker（方法一）
```
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.5.1-ubuntu22.04-dtk25.04.2-py3.10
# <your IMAGE ID>为以上拉取的docker的镜像ID替换，本镜像为：8ba136d0c0ab
docker run -it --shm-size=64G -v $PWD/kvpress_pytorch:/home/kvpress_pytorch -v /opt/hyhal:/opt/hyhal:ro --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name kvpress <your IMAGE ID> bash

cd /home/kvpress_pytorch/kvpress/
pip install -e . -i https://mirrors.aliyun.com/pypi/simple
```
更多镜像可前往[光源](https://modelzoo.sourcefind.cn/#/service-list)下载使用；

关于本项目DCU显卡所需的特殊深度学习库可从[光合](https://das.sourcefind.cn:55011/portal/#/home)开发者社区下载安装；

## 数据集
`无`

## 训练
`无`

## 推理
预训练权重目录结构：
```
/home/kvpress_pytorch
    └── Qwen/Qwen3-8B

# 设置HF下载镜像：
export HF_ENDPOINT=https://hf-mirror.com
``` 

### 单机多卡
```bash
cd /home/kvpress_pytorch
python test.py # 本介绍以Qwen3-8B为示例进行步骤说明，需提前将Qwen3-8B下载到对应目录。
```

## 效果展示
`输入: `
```
context = "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
question = "美国面积多大？"
```

`输出:`
```
answer:  美国的陆地面积约为 **983.4万平方公里**（约380万平方英里），是世界第三大国家，仅次于俄罗斯和中国。如果包括水域和领海，美国的总面积约为 **998.4万平方公里**（约385万平方英里）。
美国的国土面积广阔，横跨北美洲大部分地区，从阿拉斯加（位于北美洲最西北端）到夏威夷（位于太平洋中部）都有分布。
```

### 精度
DCU与GPU精度一致，推理框架：pytorch。

## 预训练权重
|          模型名称          | 权重大小 | DCU型号  | 最低卡数需求 |下载地址|
|:----------------------:|:----:|:------:|:------:|:----------:|
|  Qwen3-8B   |  8B  | K100AI |   4    | [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B) |

## 源码仓库及问题反馈
- http://developer.sourcefind.cn/codes/modelzoo/kvpress-SnapKV-Qwen3-8B_pytorch.git

## 参考资料
- https://github.com/NVIDIA/kvpress/blob/main/kvpress/presses/snapkv_press.py