DuoAttention：结合检索头与流式头的高效长上下文大模型推理

[论文] [幻灯片]

演示

https://github.com/user-attachments/assets/b372882b-bf14-4c89-a610-22724d91a415

简要概述

我们在不牺牲长上下文能力的前提下，显著降低了长上下文大模型在预填充和解码阶段的内存占用与延迟。

摘要

部署长上下文大型语言模型至关重要，但同时也带来了巨大的计算和内存挑战。 缓存所有注意力头中的键值状态会消耗大量内存。 现有的键值缓存修剪方法要么会损害大模型的长上下文能力，要么只能带来有限的效率提升。 在本文中，我们发现只有一小部分注意力头——即检索头——对于处理长上下文至关重要，需要对所有 token 进行全注意力计算。 相比之下，其他主要关注近期 token 和注意力汇聚点的头——称为流式头——则不需要进行全注意力计算。 基于这一洞察，我们提出了 DuoAttention 框架：仅对检索头应用完整的键值缓存，而对流式头使用轻量级的固定长度键值缓存。这种方法能够在不降低长上下文能力的情况下，同时减少大模型的解码和预填充阶段的内存占用与延迟。 DuoAttention 使用一种基于优化的轻量级算法，并结合合成数据来准确识别检索头。 我们的方法可将多头注意力（MHA）模型的长上下文推理内存最多减少 2.55 倍，将分组多头注意力（GQA）模型的长上下文推理内存最多减少 1.67 倍；同时，解码速度分别加快至多 2.18 倍和 1.50 倍，预填充速度分别加快至多 1.73 倍和 1.63 倍，且与全注意力相比，精度损失极小。 值得注意的是，结合量化技术后，DuoAttention 使得 Llama-3-8B 在单张 A100 GPU 上即可实现 330 万上下文长度的解码。

安装与使用

环境搭建

训练与评估环境

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

# LongBench 评估
pip install seaborn rouge_score einops pandas

pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/

# 安装 DuoAttention
pip install -e .

# 安装块稀疏流式注意力
git clone https://github.com/mit-han-lab/Block-Sparse-Attention
cd Block-Sparse-Attention
python setup.py install

演示环境

conda create -yn duo_demo python=3.10
conda activate duo_demo

# 安装 DuoAttention
pip install -e .

conda install -y git
conda install -y nvidia/label/cuda-12.4.0::cuda-toolkit
conda install -y nvidia::cuda-cudart-dev

# 安装 QServe
git clone https://github.com:mit-han-lab/qserve
cd qserve
pip install -e .
pip install ninja packaging
pip install flash-attn==2.4.1 --no-build-isolation
cd kernels
python setup.py install

# 安装 FlashInfer
pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3/
pip install tensor_parallel

数据集

下载数据集：

mkdir -p datasets
cd datasets

wget https://huggingface.co/datasets/togethercomputer/Long-Data-Collections/resolve/main/fine-tune/booksum.jsonl.zst

模型

下载 DuoAttention 支持的模型：

mkdir -p models
cd models

# 目前 DuoAttention 支持用于评估的模型
huggingface-cli download togethercomputer/Llama-2-7B-32K-Instruct --local-dir Llama-2-7B-32K-Instruct
huggingface-cli download gradientai/Llama-3-8B-Instruct-Gradient-1048k --local-dir Llama-3-8B-Instruct-Gradient-1048k
huggingface-cli download gradientai/Llama-3-8B-Instruct-Gradient-4194k --local-dir Llama-3-8B-Instruct-Gradient-4194k
huggingface-cli download mistralai/Mistral-7B-Instruct-v0.2 --local-dir Mistral-7B-Instruct-v0.2
huggingface-cli download mistralai/Mistral-7B-Instruct-v0.3 --local-dir Mistral-7B-Instruct-v0.3

# 用于演示的 W8A8KV4 模型，采用 SmoothQuant 和 QServe 技术
huggingface-cli download mit-han-lab/Llama-3-8B-Instruct-Gradient-1048k-w8a8kv4-per-channel --local-dir Llama-3-8B-Instruct-Gradient-1048k-w8a8kv4-per-channel
huggingface-cli download mit-han-lab/Llama-3-8B-Instruct-Gradient-4194k-w8a8kv4-per-channel --local-dir Llama-3-8B-Instruct-Gradient-4194k-w8a8kv4-per-channel

DuoAttention 快速入门

我们提供了一个简单的单击补丁，可在 HuggingFace 模型上启用 DuoAttention 优化，包括 Llama 和 Mistral 系列。五个长上下文模型的预训练检索头模式已存储在 attn_patterns 目录中：Llama-2-7B-32K-Instruct、Llama-3-8B-Instruct-Gradient-1048k、Llama-3-8B-Instruct-Gradient-4194k、Mistral-7B-Instruct-v0.2、Mistral-7B-Instruct-v0.3 以及 Meta-Llama-3.1-8B-Instruct。如果您希望训练自己的检索头模式，可以使用 scripts 目录中提供的训练脚本。以下是如何在 Llama-3-8B-Instruct-Gradient-1048k 模型上启用 DuoAttention 的示例。

from duo_attn.utils import load_attn_pattern, sparsify_attention_heads
from duo_attn.patch import enable_duo_attention_eval
import transformers
import torch

# 加载模型
model = transformers.AutoModelForCausalLM.from_pretrained(
    "models/Llama-3-8B-Instruct-Gradient-1048k",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    attn_implementation="eager",
)

# 加载注意力模式
attn_heads, sink_size, recent_size = load_attn_pattern(
    "attn_patterns/Llama-3-8B-Instruct-Gradient-1048k/lr=0.02-reg=0.05-ctx=1000_32000-multi_passkey10"
)

# 稀疏化注意力头
attn_heads, sparsity = sparsify_attention_heads(attn_heads, sparsity=0.5)

# 启用 DuoAttention
enable_duo_attention_eval(
    model,
    attn_heads,
    sink_size=64,
    recent_size=256,
)

# 将模型转移到 GPU
model = model.cuda()

# 准备推理！

演示

完成环境搭建后，您可以运行以下脚本，在 Llama-3-8B-Instruct-Gradient-4194k 模型上执行 W4A8KV4 结合 DuoAttention 的演示。该演示设计为在单张 A100 GPU 上运行，支持高达 330 万 token 的上下文长度。

bash scripts/run_demo.sh

结果

检索头识别

在准备好数据集和模型后，可以运行训练脚本以识别检索头。对于我们评估的模型，相应的注意力模式可在 attn_patterns 目录中找到。

bash scripts/run_train.sh

草堆里的针（NIAH）

DuoAttention 在“草堆里的针”基准测试上，使用 MHA 模型时全注意力比例为 25%，GQA 模型时为 50%，即可达到与全注意力相当的准确率。

bash scripts/run_niah.sh

LongBench

bash scripts/run_longbench.sh

DuoAttention 在 LongBench 基准测试中提供了更好的 KV 预算与准确率之间的权衡。

效率

bash scripts/run_efficiency.sh

• 不同上下文长度下，DuoAttention 相较于全注意力的每 token 解码延迟和内存占用情况。对于 Llama-2-7B（MHA），DuoAttention 使用 25% 的检索头比例；对于 Llama-3-8B（GQA），则使用 50% 的检索头比例。随着上下文长度的增加，DuoAttention 在 MHA 模型上可实现最高 2.45 倍的内存减少和 2.13 倍的延迟降低，在 GQA 模型上则分别达到 1.65 倍和 1.5 倍的减少效果。这些优化效果接近于检索头比例的倒数。内存不足（OOM）的结果基于实测数据进行线性外推。

• 不同预填充块大小下，DuoAttention 相较于全注意力的预填充延迟和内存占用情况。对于 Llama-2-7B（MHA），DuoAttention 使用 25% 的检索头比例，预填充 10 万 token 的上下文；而对于 Llama-3-8B（GQA），则使用 50% 的检索头比例，预填充 32 万 token 的上下文。随着预填充块大小的减小，DuoAttention 在 MHA 模型上可实现最高 1.73 倍的延迟降低，在 GQA 模型上则可达 1.63 倍；同时，内存占用方面，MHA 模型可减少至多 2.38 倍，GQA 模型则可减少至多 1.53 倍。

• DuoAttention 在固定上下文长度下的解码内存和延迟与 KV 预算的关系。当检索头比例降低时，内存和延迟会线性下降。DuoAttention 在 MHA 模型上可实现最高 2.55 倍的内存减少和 2.18 倍的延迟降低，在 GQA 模型上则分别达到 1.67 倍和 1.5 倍的减少效果。

• 结合 8 位权重和 4 位 KV 缓存量化技术，DuoAttention 可使 Llama-3-8B 模型在单张 A100-80G GPU 上容纳 330 万 token。

引用

如果您认为 DuoAttention 对您的项目或研究有用或具有参考价值，请引用我们的论文：

@article{xiao2024duo,
        title={DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads},
        author={Xiao, Guangxuan and Tang, Jiaming and Zuo, Jingwei and Guo, Junxian and Yang, Shang and Tang, Haotian and Fu, Yao and Han, Song},
        journal={arXiv},
        year={2024}
}

DuoAttention 快速上手指南

DuoAttention 是一个用于长上下文大语言模型（LLM）高效推理的框架。它通过区分“检索头”（Retrieval Heads）和“流式头”（Streaming Heads），仅对关键头部保留完整 KV 缓存，从而显著降低显存占用并提升推理速度，同时保持长上下文能力。

1. 环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu)
• GPU: NVIDIA GPU (支持 CUDA 12.4)
• Python: 3.10
• 显存: 建议 24GB+ (根据模型大小和上下文长度而定)

前置依赖

确保已安装 conda 和 git。本指南使用 Conda 管理环境以避免依赖冲突。

2. 安装步骤

以下命令将创建一个名为 duo 的独立环境并安装所有必要依赖。

# 1. 创建并激活 Conda 环境
conda create -yn duo python=3.10
conda activate duo

# 2. 安装基础工具与 CUDA  toolkit (版本 12.4)
conda install -y git
conda install -y nvidia/label/cuda-12.4.0::cuda-toolkit
conda install -y nvidia::cuda-cudart-dev

# 3. 安装 PyTorch (CUDA 12.4 版本)
# 国内用户可使用清华源加速：-i https://pypi.tuna.tsinghua.edu.cn/simple
pip install torch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia

# 4. 安装 Python 依赖包
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

# 5. 安装 Flash Attention (需无隔离构建)
pip install flash-attn==2.6.3 --no-build-isolation

# 6. 安装评估工具 (可选，如需运行 LongBench)
pip install seaborn rouge_score einops pandas

# 7. 安装 FlashInfer (注意 CUDA 和 Torch 版本匹配)
pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/

# 8. 安装 DuoAttention
# 先克隆仓库或下载源码后，在根目录执行：
pip install -e .

# 9. 安装 Block Sparse Streaming Attention 组件
git clone https://github.com/mit-han-lab/Block-Sparse-Attention
cd Block-Sparse-Attention
python setup.py install
cd ..

提示：如果在国内网络环境下安装 pip 包较慢，可在每条 pip install 命令后添加 -i https://pypi.tuna.tsinghua.edu.cn/simple。

3. 基本使用

DuoAttention 提供了一键补丁功能，可快速应用于 HuggingFace 模型（如 Llama-3, Mistral 等）。以下是启用优化的最小代码示例。

前置准备：下载模型与注意力模式

在使用前，请确保已下载支持的模型和预训练的注意力头模式文件。

mkdir -p models attn_patterns
cd models

# 示例：下载 Llama-3-8B-Instruct-Gradient-1048k 模型
huggingface-cli download gradientai/Llama-3-8B-Instruct-Gradient-1048k --local-dir Llama-3-8B-Instruct-Gradient-1048k

# 注意：attn_patterns 目录需包含对应的模式文件，可从项目 releases 或仓库获取
# 此处假设已存在文件：attn_patterns/Llama-3-8B-Instruct-Gradient-1048k/lr=0.02-reg=0.05-ctx=1000_32000-multi_passkey10

Python 推理示例

from duo_attn.utils import load_attn_pattern, sparsify_attention_heads
from duo_attn.patch import enable_duo_attention_eval
import transformers
import torch

# 1. 加载模型
model = transformers.AutoModelForCausalLM.from_pretrained(
    "models/Llama-3-8B-Instruct-Gradient-1048k",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    attn_implementation="eager",
)

# 2. 加载预训练的注意力头模式
# 路径需对应具体模型和训练配置
attn_heads, sink_size, recent_size = load_attn_pattern(
    "attn_patterns/Llama-3-8B-Instruct-Gradient-1048k/lr=0.02-reg=0.05-ctx=1000_32000-multi_passkey10"
)

# 3. 稀疏化注意力头 (例如设置 50% 的检索头比例)
attn_heads, sparsity = sparsify_attention_heads(attn_heads, sparsity=0.5)

# 4. 启用 DuoAttention 优化
enable_duo_attention_eval(
    model,
    attn_heads,
    sink_size=64,      # 注意力汇大小
    recent_size=256,   # 近期 token 窗口大小
)

# 5. 将模型移至 GPU
model = model.cuda()

# 6. 开始推理
inputs = tokenizer("你的长文本输入...", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0]))

运行官方 Demo

若需体验极致长上下文（如 3.3M tokens），可运行官方提供的 Demo 脚本（需单卡 A100）：

bash scripts/run_demo.sh