MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning
This addresses efficiency issues in long-context reasoning for AI practitioners, offering a practical improvement over existing methods.
The paper tackled the bottleneck of high memory and computational costs in long-context language modeling by proposing Memory-Keyed Attention (MKA) and its variant FastMKA, which achieved comparable perplexity to prior methods while delivering up to 5x faster training throughput and 1.8x lower evaluation latency.
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.