Flux Attention：上下文感知混合注意力，实现高效 LLM 推理

标准注意力机制的二次计算复杂度给 LLM 在长上下文场景中带来了严重的可扩展性瓶颈。虽然结合完全注意力（FA）和稀疏注意力（SA）的混合注意力机制提供了潜在的解决方案，但现有方法通常依赖于静态分配比率，无法适应不同任务的可变检索需求。此外，头级动态稀疏性通常会引入严重的计算负载不平衡和同步长尾，这会阻碍自回归解码期间的硬件加速。为了弥补这一差距，我们引入了 Flux Attention，这是一个上下文感知框架，可以在层级别动态优化注意力计算。通过将轻量级层路由器集成到冻结的预训练 LLM 中，所提出的方法根据输入上下文自适应地将每个层路由到 FA 或 SA。这种分层路由保留了高保真信息检索，同时确保连续的内存访问，将理论计算的减少转化为实际的挂钟加速。作为一种参数高效的方法，我们的框架只需要在 8$\times$A800 GPU 上进行 12 小时的训练。跨多个长上下文和数学推理基准的大量实验表明，与基线模型相比，Flux Attention 在性能和推理速度之间实现了卓越的权衡，在预填充和解码阶段的速度提升高达 2.8\times$ 和 $2.0\times$。