MISA：用于长上下文 LLM 推理的混合索引稀疏注意力
机器学习
2026-05-11 v1 人工智能
摘要
DeepSeek 稀疏注意力(DSA)通过引入学习的 token 级索引器对每个前缀 token 进行评分并选择最相关的几个，用于主注意力，树立了细粒度推理时稀疏注意力的状态最佳纪录。为保持表达性，索引器使用大量查询头(例如 DeepSeek-V3.2 中为64)共享相同的所选 token 集合；这种多头设计正是使索引器成为长上下文中主要成本的原因。我们提出 MISA(混合索引稀疏注意力)，作为 DSA 索引器的即插即用替代方案，将索引器头视为混合专家池。轻量级路由器使用低成本的块级统计信息选取少量活跃头子集，仅这些头运行重型的 token 级评分。这既保留了原始索引器池的多样性，又将每个前缀 token 的评分成本从使用每个头降低到仅使用少数路由头，外加在小型聚合 key 集合上计算的可忽略路由项。我们进一步引入 MISA 的层次化变体，利用路由 pass 保留扩大的候选集，然后用原始 DSA 索引器重新排序以几乎精确恢复最终选中的 token。仅需8个活跃头且无需额外训练，MISA 在 DeepSeek-V3.2 和 GLM-5 上长文本基准测试中匹配密集 DSA 索引器性能，分别减少8倍和4倍的索引头数量，平均超越 HISA。它在最高达128K token 上文中完全保留绿色 Needle-in-a-Haystack 热图，并在每层恢复超过92%的 DSA 索引器所选 token。我们的 TileLang kernel 在单张 NVIDIA H200 GPU 上相较于 DSA 原索引器 kernel 速度提升约3.82倍。
引用
@article{arxiv.2605.07363,
title = {MISA: Mixture of Indexer Sparse Attention for Long-Context LLM Inference},
author = {Ruijie Zhou and Fanxu Meng and Yufei Xu and Tongxuan Liu and Guangming Lu and Muhan Zhang and Wenjie Pei},
journal= {arXiv preprint arXiv:2605.07363},
year = {2026}
}
备注
https://github.com/MuLabPKU/TransArch