SalesforceAIResearch/Random-Attention

来源：https://github.com/SalesforceAIResearch/Random-Attention

Random Attention：重新思考高效推理的 KV 缓存驱逐策略

论文 Random Attention: Rethinking KV Cache Eviction for Efficient Reasoning（https://arxiv.org/abs/2609.03430）(arXiv:2609.03430) 的代码实现。

📄 论文 (arXiv:2609.03430) (https://arxiv.org/abs/2609.03430) · 🌐 项目页面 (https://arthur-heng.github.io/Random-Attention-page/) · 🧵 讨论帖 (https://x.com/HengWang_uiuc/status/2095685764302475563)

Random Attention (代码中为 random_pp) 是一种无需信号的推理模型 KV 缓存驱逐策略：保留提示（prompt），然后在生成的 token 中，按每个 KV 头 均匀随机 的方式保留一个子集直至达到预算（再加上一个短暂的近期窗口）。它不读取任何注意力分数、值统计数据或校准数据，因此每一轮驱逐仅需执行压缩操作本身。在 MATH-500、GPQA-Diamond、AIME、HMMT 和 LiveCodeBench 基准测试上，在匹配的预算下，它在 Qwen3-4B/14B/32B 和 Phi-4-reasoning 模型上达到或超越了学习型选择器（SnapKV、R-KV、VaSE、TriAttention）的表现，并且无论是在 Hugging Face 测试环境还是 vLLM 服务栈中，它都是最快的驱逐器。

本仓库包含驱逐引擎、评估框架、显著性检验、效率基准测试、vLLM 移植版本以及用于生成论文中所有数据的机制研究工具。

目录结构

kvcompress/engine/    KV驱逐引擎：论文中的每种方法对应一种驱逐模式（cache_utils.py），  
                      用于 Qwen3 / phi3 / llama 的 evict-attention 前向传播，  
                      忠实的 TriAttention 评分器，引擎测试  
kvcompress/harness/   评估框架（改编自 VaSE，见 THIRD_PARTY_NOTICES.md）：  
                      eval_hf.py，分片多工作者启动器 parallel_run_hf_mw.py，  
                      评分工具，TriAttention 校准与统计  
kvcompress/eval/      评分（数学/科学，LiveCodeBench），单元格完整性审计，  
                      配对显著性检验，分片修复，LaTeX 表格生成器  
kvcompress/analysis/  机制研究：保留日志，分支重放/尸检，  
                      载体质量，驱逐时间基准测试  
kvcompress/synth/     受控合成检索研究（注册协议，表格）  
scripts/run_cell.sh   为一个（模型，任务，方法，预算）准确率单元格的标准启动脚本  
scripts/grade_cell.sh 单元格的完整性检查与评分  
scripts/efficiency/   Hugging Face 吞吐量协议（VaSE 固定批次，等内存，最大批次，驱逐轮次计时）  
scripts/vllm_rp_bench/ Random Attention 在 TriAttention 的 vLLM 0.19 运行时内的实现  
                      （包含独立 README 与运行手册）  
scripts/mechanism/    保留日志面板与分支重放启动器  
figures/              图表脚本（读取已评分的 TSV 文件）  
data/                 基准布局及我们的 LiveCodeBench 难度子集（data/README.md）  

环境设置

git clone https://github.com/SalesforceAIResearch/Random-Attention && cd Random-Attention  
bash setup.sh          # Python 3.10 虚拟环境：torch 2.4.0 (cu121)，flash-attn 2.7.3，transformers 5.0.0，等等  
. env.sh               # 设置 RA_ROOT / RA_ENGINE / RA_DATA_DIR / RA_MODELS_DIR / PYTHONPATH  

• 模型：位于 $RA_MODELS_DIR/（默认为 models/）下的 Hugging Face 检查点：Qwen3-4B、Qwen3-14B、Qwen3-32B、phi-4-reasoning，可选 DeepSeek-R1-Distill-Llama-8B。
• 数据：$RA_DATA_DIR//test.jsonl – 格式和来源请参阅 data/README.md。已包含我们的 LiveCodeBench 子集。
• 硬件：论文中的实验在 8x H200 (141 GB) 上运行。

运行准确率单元格

scripts/run_cell.sh Qwen3-4B math random_pp  
# K 默认为任务约 4 倍点数（此处为 1024）  
scripts/run_cell.sh Qwen3-4B math vase  
# VaSE 使用忠实的 n_large = K/4  
scripts/run_cell.sh Qwen3-4B math triattn  
# TriAttention 使用每模型校准统计  
scripts/run_cell.sh phi-4-reasoning gpqa snapkv 2048  
scripts/run_cell.sh Qwen3-32B aime25 rkv && scripts/run_cell.sh Qwen3-32B aime26 rkv  

添加 --dry-run 可打印确切的 parallel_run_hf_mw.py 命令。单元格在重新启动时会恢复；完成结果位于 results//_K// 下。scripts/run_cell.sh 的头部列出了每种方法及其标志。

评分、检验、制表

scripts/grade_cell.sh Qwen3-4B math_K1024 dense,random_pp,attn,attn_rkv_l05,vase_faithful,triattn_ph_memofix  
python kvcompress/eval/stats_paired.py --base results/Qwen3-4B/math_K1024 --data_name math \  
  --method_a random_pp --methods_b attn,attn_rkv_l05,vase_faithful,triattn_ph_memofix  
# 配对 bootstrap + 符号检验  
python kvcompress/eval/gen_paper_tables.py --results --lcb --out tables/  

grade_cell.sh 首先运行 cell_integrity.py：评分是基于位置的（问题 = 分片偏移量 + 行），因此混合、参差不齐或重叠的单元格会被拒绝，而不是被错误评分。准确率是 flag_acc（答案正确）；acc_strict 还要求在 32k 上限内终止，并作为诊断信息报告。

效率

• Hugging Face 测试环境：scripts/efficiency/*.sh（固定批次 VaSE 协议，等内存最大批次服务，驱逐轮次计时）。请勿从批量准确率运行中报告 tokens/s——它们共享 GPU。
• vLLM 服务：scripts/vllm_rp_bench/ – Random Attention 作为选择器集成在 TriAttention 的 vLLM 0.19 运行时内（使用其分页-KV压缩机制，采用我们的选择规则），包含精度转移检查和运行时补丁，详见其 README。

机制与合成研究

通过 kvcompress/engine/cache_utils.py 头部文档记录的引擎环境开关来驱动：
KEEPLOG=1 开启保留日志记录，FORCE_KEEP_RANGE 开启分支重放（kvcompress/analysis/fork_replay.py、fork_autopsy.py），载载体头混合模式，以及注册的合成检索协议（kvcompress/synth/）。

许可证

本项目采用 Apache License 2.0 发布（见 LICENSE.txt）；这是伴随学术论文的研究性发布——请阅读 AI_ETHICS.md。
贡献指南：CONTRIBUTING.md；安全报告：SECURITY.md。

致谢与第三方代码

评估框架和引擎骨架源自 VaSE (https://github.com/terarachang/VaSE) (MIT)；TriAttention 基准和 vLLM 服务基准基于 TriAttention (https://github.com/WeianMao/triattention) (Apache-2.0)。
精确的文件级来源和许可证文本请见 THIRD_PARTY_NOTICES.md。

引用

@article{randomattention2026,  
  title = {Random Attention: Rethinking KV Cache Eviction for Efficient Reasoning},  
  author = {Heng Wang and Jielin Qiu and Wenting Zhao and Cheng Qian and Liangwei Yang and Jiawei Han and Heng Ji and Silvio Savarese and Shelby Heinecke and Huan Wang},  
  journal = {arXiv preprint arXiv:2609.03430},  
  year = {2026},  
  eprint = {2609.03430},  
  archivePrefix = {arXiv},  
  primaryClass = {cs.CL},  
  url = {https://arxiv.org/abs/2609.03430}  
}