DuoAttention：高效长上下文LLM推理框架

DuoAttention：高效长上下文LLM推理框架

duo-attention [ICLR 2025] DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads 项目地址: https://gitcode.com/gh_mirrors/du/duo-attention

1. 项目介绍

DuoAttention 是一个针对长上下文的大型语言模型（LLM）推理的高效框架。该框架通过区分处理长上下文的检索头和处理近期信息的流式头，大幅减少了LLM的解码和预填充内存及延迟，同时不牺牲模型处理长上下文的能力。

2. 项目快速启动

环境搭建

首先，创建一个名为 duo 的虚拟环境，并安装必要的依赖：

conda create -yn duo python=3.10
conda activate duo
conda install -y git
conda install -y nvidia/label/cuda-12.4.0::cuda-toolkit
conda install -y nvidia::cuda-cudart-dev
conda install -y pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia
pip install transformers==4.45.2 accelerate sentencepiece datasets wandb zstandard matplotlib huggingface_hub==0.25.2
pip install tensor_parallel==2.0.0
pip install ninja packaging
pip install flash-attn==2.6.3 --no-build-isolation

安装 DuoAttention

将 DuoAttention 项目克隆到本地，并安装：

git clone https://github.com/mit-han-lab/duo-attention.git
cd duo-attention
pip install -e .

Demo 环境

如果你想运行Demo，需要创建一个名为 duo_demo 的虚拟环境，并安装相应的依赖：

conda create -yn duo_demo python=3.10
conda activate duo_demo
pip install -e ./duo-attention
conda install -y git
conda install -y nvidia/label/cuda-12.4.0::cuda-toolkit
conda install -y nvidia::cuda-cudart-dev
# 安装 QServe 和 FlashInfer
# ...

运行 Demo

Demo 脚本会执行带有 DuoAttention 的 W4A8KV4 模型推理：

bash scripts/run_demo.sh

3. 应用案例和最佳实践

在应用 DuoAttention 时，以下是一些最佳实践：

• 确保在推理前已经正确地标识了检索头。
• 在长上下文推理中使用 DuoAttention 可以显著减少内存使用和推理时间。
• 结合量化技术，可以进一步提高模型效率。

4. 典型生态项目

DuoAttention 可以与以下项目配合使用，以构建更高效的LLM应用：

• Block Sparse Attention：用于实现稀疏注意力机制的库。
• QServe：提供模型量化服务的框架。
• FlashInfer：一个推理加速工具。

通过上述介绍和教程，您应该能够开始使用 DuoAttention，并在长上下文LLM推理任务中实现效率提升。

duo-attention [ICLR 2025] DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads 项目地址: https://gitcode.com/gh_mirrors/du/duo-attention