通过分层记忆路由提升长文本推理效率，训练快5倍、推理延迟降1.8倍。
MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning
- 设计分层键值缓存，动态路由注意力到局部、会话和长期记忆。
- 在长序列上达到与MLA相当的精度，训练速度提升5倍，推理延迟降低1.8倍。
- 适合需要高效长文本处理的场景，如文档理解、代码生成等。
随着长上下文语言建模的重要性提升，维护和访问大型键值（KV）缓存的成本迅速增加，成为训练和推理的主要瓶颈。尽管先前工作如多查询注意力（MQA）和多潜伏注意力（MLA）通过共享或压缩KV特征来减少内存占用，但往往牺牲表示质量或引入运行时开销。本文提出记忆键注意力（MKA），一种分层注意力机制，整合局部、会话和长期三级KV缓存，并学习动态路由注意力。进一步提出路由融合型MKA（FastMKA），采用广播路由方式在注意力计算前融合记忆源，提升效率。在不同序列长度下的实验表明，FastMKA在保持与MLA相当困惑度的同时，实现最高5倍的训练吞吐量提升和1.8倍的评估延迟降低。结果表明MKA是一种实用且可扩展的高效长上下文注意力框架。
原文摘要 · Abstract (English)
As long-context language modeling becomes increasingly important, the cost of maintaining and attending to large Key/Value (KV) caches grows rapidly, becoming a major bottleneck in both training and inference. While prior works such as Multi-Query Attention (MQA) and Multi-Latent Attention (MLA) reduce memory by sharing or compressing KV features, they often trade off representation quality or incur runtime overhead. We propose Memory-Keyed Attention (MKA), a hierarchical attention mechanism that integrates multi-level KV caches (local, session, and long-term) and learns to route attention across them dynamically. We further introduce Route-Fused MKA (FastMKA), a broadcast-routed variant that fuses memory sources before attention computation for improved efficiency. Experiments on different sequence lengths show that FastMKA achieves a favorable accuracy-efficiency trade-off: comparable perplexity to MLA while achieving up to 5x faster training throughput and 1.8x lower evaluation latency. These results highlight MKA as a practical and extensible framework for efficient long-context attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。