# PyramidKV
## 论文
[PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling](https://arxiv.org/pdf/2406.02069)

## 模型简介

PyramidKV构建金字塔形缓存结构，低层（浅层特征弥散）分配更多缓存，高层（深层特征聚焦）分配较少缓存，从而减少显存占用。

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

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

环境配置：
```
mv kvpress-PyramidKV-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:  美国的面积大约是963万平方千米（约3.84百万平方英亩）。它是世界上面积第三大的国家，仅次于俄罗斯和中国。美国国土面积广阔，包括多个州和联邦地区，从阿拉斯克到夏威夷，从墨西哥湾到加拿大边境，地理环境多样，涵盖从沙漠到森林，从山脉到平原等多种地貌。
```

### 精度
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-PyramidKV-Qwen3-8B_pytorch.git

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