韩松等提出FlashMoBA，比MoBA快7.4倍，序列扩到512K也不会溢出

文章摘要

【关 键 词】 注意力机制、稀疏计算、GPU优化、长上下文处理、模型训练

MoBA（Mixture of Block Attention）是一种将混合专家原理应用于注意力机制的创新方法，通过让模型自主决定关注哪些位置，显著降低了长上下文处理的计算成本。该方法的核心在于Query仅稀疏关注少量Key-Value块，但此前存在GPU实现效率低下的问题。MIT与NVIDIA的研究团队通过建立统计模型发现，路由器的准确性对性能至关重要，其信噪比直接影响检索效果，进而提出两条改进路径：缩小块尺寸和在Key上应用短卷积。

理论分析表明，较小的块尺寸（如从512缩减至128）能提升路由精度，使340M参数模型的WikiText困惑度降低5.7%，RULER准确率提升44.6%。Key卷积（kconv3/kconv5）则通过聚集块内相关信号，进一步将语言建模准确率提高0.5-1个百分点。然而，小块尺寸导致内存访问碎片化，使原始实现速度甚至慢于稠密注意力。为此，研究者开发了硬件友好的FlashMoBA内核，采用三级优化策略：分块Top-K选择避免显式生成分数矩阵，”收集并致密化”策略通过两级分块机制优化内存访问，反向传播则通过重计算注意力分数保持线性复杂度。

实验结果显示，FlashMoBA在64K序列长度下比原始实现快7.4倍，内存占用减少83%，且能扩展到512K长度。其融合内核设计将前向传播时间压缩至FlashAttention-2的49%，在小块配置下实现最高14.7倍加速。质量评估证实，优化后的MoBA在语言建模、64K长度检索等任务中达到与密集注意力相当的性能，1B参数模型在RULER基准上准确率提升至68.2%。该研究为稀疏注意力机制的实际部署提供了关键解决方案，通过协同优化算法设计与硬件实现，突破了长序列处理的效率瓶颈。

原文和模型

【原文链接】 阅读原文 [ 2634字 | 11分钟 ]
【原文作者】 机器之心
【摘要模型】 deepseek/deepseek-v3-0324
【摘要评分】 ★★★☆☆

派欧算力云 – AI 云端一体化解决方案: 模型 API、Serverless、GPU 租赁

派欧算力云提供全面的 AI 云端平台解决方案,包括模型 API、Serverless 和高性能 GPU 实例等,帮助您节省高达 50% 的成本,并同时加速 AI 产品的构建和业务发展,助力企业实现创新突破,赢得市场竞争优势。