The increasing demand for long-context modeling in large language models (LLMs) is bottlenecked by the quadratic complexity of the standard self-attention mechanism. The community has proposed sparse attention to mitigate this issue. However, position-aware sparse attention methods rely on static sparse structures that lack adaptability to diverse query contexts, while content-aware sparse attention methods depend on heuristic key-value selection, hindering full differentiability. We introduce a trainable dynamic mask sparse attention mechanism, a method that merges the advantages of both position-aware and content-aware approaches. Dynamic Mask Attention (DMA) achieves this through three key innovations: First, it leverages value vector representations to generate content-aware dynamic masks, enabling the model to adaptively identify and attend to critical information. Second, it computes position-aware sparse weights in a hardware-friendly manner, efficiently skipping unnecessary computational regions. Finally, we demonstrate that the introduced dynamic mask and sparse weights do not obstruct gradients, supporting end-to-end training. We have validated the performance of DMA through comprehensive experiments. A large body of experimental evidence shows that DMA consistently holds a Pareto advantage over state-of-the-art sparse attention baselines in tasks including scaling laws, multi-query associative recall, standard benchmarks, and needle in a haystack tests, while also delivering up to a 10x overall speedup. These results highlight its ability to effectively balance model efficiency with long-context modeling capabilities. Our computational kernel code is now open-source at https://github.com/SmallDoges/flash-dmattn to encourage further research and application by the community.

翻译：大型语言模型（LLMs）对长上下文建模的需求日益增长，但标准自注意力机制的二次复杂度成为其瓶颈。学界已提出稀疏注意力以缓解此问题。然而，基于位置的稀疏注意力方法依赖于静态稀疏结构，缺乏对不同查询上下文的适应性；而基于内容的稀疏注意力方法则依赖启发式的键值选择，阻碍了完全可微性。本文提出一种可训练动态掩码稀疏注意力机制，该方法融合了基于位置和基于内容两种方法的优势。动态掩码注意力（DMA）通过三项关键创新实现这一目标：首先，它利用值向量表示生成内容感知的动态掩码，使模型能够自适应地识别并关注关键信息；其次，它以硬件友好的方式计算位置感知的稀疏权重，高效跳过不必要的计算区域；最后，我们证明所引入的动态掩码和稀疏权重不会阻碍梯度传播，支持端到端训练。我们通过全面实验验证了DMA的性能。大量实验证据表明，在包括缩放定律、多查询关联召回、标准基准测试以及“大海捞针”测试等任务中，DMA相较于最先进的稀疏注意力基线始终具有帕累托优势，同时还能实现高达10倍的整体加速。这些结果突显了其在模型效率与长上下文建模能力之间有效平衡的能力。我们的计算内核代码已在 https://github.com/SmallDoges/flash-dmattn 开源，以促进学界的进一步研究和应用。