静态稀疏已死：Flux Attention 开启长上下文 LLM 自适应推理新时代

论文http://arxiv.org/abs/2604.07394代码https://github.com/qqtang-code/FluxAttention长上下文 LLM 高效推理的核心矛盾从来都不是 “要不要做稀疏”而是 “怎么让稀疏策略适配千变万化的输入”。大语言模型的上下文窗口正在以惊人的速度扩张从最初的4K、8K到如今128K、256K甚至百万级token的上下文已成为行业主流。无论是长文档解析、多轮对话还是代码库理解超长上下文都让LLM的能力边界不断拓展但一个核心瓶颈始终横亘在所有人面前标准Full Attention全注意力机制的二次计算复杂度。序列长度翻一倍注意力的计算和内存开销就翻四倍这让超长上下文的LLM推理始终面临着性能与效率的两难选择。稀疏注意力Sparse Attention无疑是解决这个问题的核心路径——通过只关注序列中最关键的部分token大幅降低计算开销但如何在削减算力的同时不损失模型的生成效果成了整个领域的核心难题。我们的第一步探索Elastic Attention打破静态稀疏的桎梏在之前的工作中我们提出了Elastic Attentionhttps://arxiv.org/abs/2601.17367第一次系统性地解决了混合注意力机制的静态配比难题。彼时行业内的主流方案大多是在模型中固定好全注意力FA和稀疏注意力SA的比例比如70%的头用稀疏注意力30%的头用全注意力推理全程不再变化。但我们通过实验发现下游任务对稀疏度的敏感度天差地别摘要、代码生成这类任务属于稀疏鲁棒任务即便稀疏度拉满性能也几乎不会下降而单文档/多文档问答这类检索密集型任务属于稀疏敏感任务一旦稀疏度超过阈值性能会出现断崖式下跌。基于这个核心洞察我们在Elastic Attention中设计了轻量级的Attention Router实现了头级别head-wise的动态路由模型会根据输入的内容在推理时为每个注意力头动态分配FA或SA模式对稀疏敏感的问答任务自动降低稀疏度对稀疏鲁棒的摘要任务自动拉高稀疏度。最终仅用8张A800训练12小时我们就在多个长上下文基准上实现了性能与效率的平衡效果超越了DuoAttention、PruLong等一众静态head-wise的混合注意力方案。但当我们尝试将Elastic Attention推向实际工程落地时一个无法回避的硬件瓶颈出现了头级别的动态路由虽然在算法层面实现了最优的稀疏配比却很难在decode阶段将理论上的FLOPs减少转化为实际的wall-clock速度提升。原因在于GPU的decode过程是内存带宽受限的当同一个Transformer层里有的注意力头跑FA、有的跑SA会出现两个致命问题一是层内计算负载严重不均衡稀疏头很快完成计算必须长时间等待计算量大的全注意力头产生严重的同步长尾效应二是内存访问不连续不同模式的头需要读取不同的KV缓存无法利用GPU的连续内存读取优化最终导致即便理论算力砍半实际decode速度也只有微弱提升。全新迭代Flux Attention打通算法优化到硬件落地的最后一公里针对Elastic Attention在落地中暴露的核心问题我们完成了全新的工作Flux Attention将动态稀疏注意力的范式从头级别路由升级为层级别layer-wise路由解决了算法灵活性与硬件执行效率之间的矛盾。我们的核心洞察是Transformer的层本身就存在显著的功能异质性。模型的部分层是核心检索层负责从超长上下文中定位关键信息必须依赖全注意力才能保证效果而大部分层是语义聚合层负责对上下文信息进行整体理解即便用稀疏注意力性能也几乎不受影响。与其在同一个层里让不同的头“各行其是”不如直接以层为单位为整个层动态分配FA或SA模式。我们通过实验完整验证了这一点不同类型的下游任务对Transformer不同层的注意力模式需求呈现出了极强的规律性差异。不同下游任务的层级别路由激活热力图颜色越深代表该层被分配全注意力FA的频率越高。上半部分为检索密集型任务对全注意力的需求显著更高下半部分为上下文整体型任务可适配更高的稀疏度。基于这个思路我们设计了Flux Attention的核心架构上下文感知的Layer Router层路由模块。整个模块的设计遵循极致轻量化与硬件友好的原则核心亮点分为三个部分高效的任务特征提取路由模块只接收Query张量作为输入通过前缀后缀池化操作仅提取序列前后100个token的特征来判断任务类型——这是因为长上下文prompt中任务指令和用户核心需求几乎都分布在序列首尾中间的文档内容对任务类型判断而言是噪声。这种设计让路由模块的计算量几乎不随序列长度增长而增加最终单步路由的平均开销仅0.2ms且从512到1M token的序列长度下延迟完全保持稳定。可微分的训练与稳定的推理为了解决离散路由决策的不可微问题我们沿用并优化了Elastic Attention中经过验证的训练策略采用Gumbel-Softmax连续松弛方案配合温度退火在训练中让路由分布逐渐向离散决策收敛同时用直通估计器STE解决梯度回传问题保证训练与推理的一致性。训练目标上我们构建了语言建模损失与稀疏正则的联合优化目标为不同类型的任务设置非紧致的稀疏度约束让模型自动学习最优的路由策略。整个训练过程中我们冻结了LLM主干模型的所有参数只需更新Layer Router的轻量级参数最终在8张A800显卡上仅需12小时就能完成对Qwen3、Llama3.1等主流开源模型的适配训练成本极低。极致的工程落地优化在部署阶段Flux Attention解决了Elastic Attention的硬件效率痛点。我们在prefill阶段只执行一次路由决策为每个层确定好FA/SA模式后将结果缓存下来在后续所有decode步骤中复用完全消除了decode中额外的路由开销。同时层级别的路由让每个层的所有注意力头都采用同一种计算模式完美适配FlashAttention等主流高性能注意力内核无需开发复杂的融合内核内存访问上稀疏层只需要加载和保存最小规模的KV缓存稠密层则保留完整KV实现了完全连续的内存访问彻底消除了层内同步长尾让理论上的算力节省100%转化为实际的推理速度提升。从Elastic到Flux我们实现了哪些核心突破很多人会问相比于head-wise路由的Elastic Attentionlayer-wise路由的Flux Attention究竟带来了什么实质性的飞跃答案集中在四个维度推理效率的质变Elastic Attention的加速主要集中在prefill阶段decode阶段受限于硬件瓶颈加速效果有限而Flux Attention在256K超长上下文下实现了prefill阶段最高2.8倍端到端加速decode阶段最高2.0倍内核加速且加速比随序列长度增长持续提升真正实现了长上下文场景下的全流程加速。模型能力的泛化拓展Elastic Attention主要聚焦于长上下文理解任务而Flux Attention不仅在LongBench-E、RULER、LongBench-V2三大主流长上下文基准上实现了与全注意力基线持平甚至超越的效果还在GSM8K、AIME24等数学推理任务上实现了性能提升证明了层级别动态路由的范式不仅适配长文本场景对复杂逻辑推理任务同样有效。工程落地的成本骤降Elastic Attention需要定制化的融合内核才能实现不同模式头的并行计算而Flux Attention的层级路由完全兼容现有开源生态中的高性能注意力内核无需复杂的底层开发就能实现开箱即用的加速效果大幅降低了工业界的落地门槛。内存占用的进一步优化层级别的稀疏决策让我们可以在decode阶段完全跳过稀疏层的全量KV缓存加载相比头级别路由KV缓存的内存占用进一步降低让256K甚至更长上下文的推理在消费级显卡上也成为可能。写在最后从Elastic Attention到Flux Attention我们始终在探索同一个问题如何让长上下文大模型在真实的硬件环境中实现真正可用的高效推理。静态稀疏配比的方案无法适配千变万化的用户输入头级别的动态路由解决了算法层面的自适应问题却受限于硬件架构无法完全释放潜力而Flux Attention的层级别上下文感知路由终于在算法灵活性与硬件执行效率之间找到了最优的平衡点。在未来的工作中我们将持续沿着这条技术路线深耕一方面完成 Flux Attention 在更大模型规模、更极端长上下文与全硬件场景的 Scaling 验证另一方面将这套动态路由范式向多模态大模型领域全面拓展最终还会把路由模块原生融入预训练全流程实现从后训练适配到原生稀疏架构的跨越持续探索大模型高效推理的最优解。目前Flux Attention的论文发布在http://arxiv.org/abs/2604.07394代码与模型权重也完全开源在GitHubhttps://github.com/qqtang-code/FluxAttention。我们希望这个工作能为长上下文大模型的落地提供一个更高效、更易用、更具泛化性的解决方案也期待和行业内的同仁一起继续推动大模型高效推理技术的发展。