论文

SinkRouter:用于在大语言和多模态模型中进行高效长上下文解码的接收器感知路由

SinkRouter: Sink-Aware Routing for Efficient Long-Context Decoding in Large Language and Multimodal Models

模型推理推理加速

摘要

在 LLM 和 LMM 的长上下文解码中,注意力变得越来越受内存限制,因为每个解码步骤都必须从 GPU 内存加载大量 KV 缓存数据。现有的加速策略通常依靠启发式修剪来牺牲效率以换取准确性,这可能会丢弃有用的信息。在更深层次上,他们也倾向于不加区别地保留所有高分标记,将早期标记视为不可或缺的锚点,或者依赖启发式头部路由,反映出对注意力下沉现象的机制理解不足。在本文中,我们表明注意力池现象对应于训练过程中构造的稳定、可达且误差可控的固定点。基于这一见解,我们提出了 SinkRouter,这是一种 无需训练 选择性路由框架,可检测接收器信号并跳过否则会产生接近零输出的计算。为了将此机制转化为现实世界的加速,我们开发了一个具有块级分支和 Split-K 并行性的硬件感知 Triton 内核。我们使用纯文本和多模式主干(例如 Llama-3.1-8B、Llama-3.1-70B、Yi-9B-200K、LLaVA-1.5-7B 和 LLaVA-1.5-13B)对各种长上下文基准测试套件(包括 LongBench、InfiniteBench、CVBench、MileBench 和 MMVP)进行了广泛的评估。在这些设置中,SinkRouter 持续提高解码效率,同时保持有竞争力的准确性,并在 512K 上下文中实现 2.03 倍的加速。