MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning

• URL: http://arxiv.org/abs/2603.20586v2
• Date: Tue, 24 Mar 2026 11:05:44 GMT
• Title: MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning
• Abstract summary: We propose a hierarchical attention mechanism that integrates multi-level KV caches (local, session, and long-term) and learns to route attention across them dynamically.<n>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.
• Score: 37.38708392928324
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration [67.17709877327232]
  Cate KV is a hybrid KV cache method that retains only critical token information for consistent heads.<n>Cate KV reduces memory usage by up to $2.72times$ and accelerates decoding by $2.18times$ in single-sample inputs.
  arXiv  Detail & Related papers  (2026-08-31T06:02:37Z)
• ZoomR: Memory Efficient Reasoning through Multi-Granularity Key Value Retrieval [58.575695990976136]
  Large language models (LLMs) have shown great performance on complex reasoning tasks.<n>LLMs often require generating long intermediate thoughts before reaching a final answer.<n>We introduce ZoomR, a novel approach that enables LLMs to adaptively compress verbose reasoning thoughts into summaries.
  arXiv  Detail & Related papers  (2026-04-13T02:00:35Z)
• KEEP: A KV-Cache-Centric Memory Management System for Efficient Embodied Planning [8.216400469571084]
  We propose KEEP, a KV-cache-centric memory management system for efficient embodied planning.<n>KEEP features 3 key innovations: (1) a Static-Dynamic Memory Construction algorithm that reduces KV cache recomputation by mixed-granularity memory group; (2) a Multi-hop Memory Re-computation algorithm that dynamically identifies important cross-attention among different memory groups; and (3) a Layer-balanced Memory Loading that eliminates unbalanced KV cache loading and cross-attention across different layers.
  arXiv  Detail & Related papers  (2026-02-27T01:48:07Z)
• Multi-head Temporal Latent Attention [27.475917680869657]
  Multi-head latent attention was recently developed to compress the Key-Value cache into a low-rank latent space.<n>This paper proposes Multi-head Temporal Latent Attention (MTLA), which further reduces the KV cache size along the temporal dimension.<n>Experiments across tasks, including speech translation, speech recognition, speech understanding and text summarisation, demonstrate that MTLA achieves competitive performance.
  arXiv  Detail & Related papers  (2025-05-19T02:09:41Z)
• Sparse Attention Remapping with Clustering for Efficient LLM Decoding on PIM [7.651654889371008]
  Transformer-based models are the foundation of modern machine learning, but their execution places significant pressure on memory systems.<n> processing-in-memory (PIM) architectures are a promising solution, offering high internal bandwidth and compute parallelism near memory.<n>Current PIM designs are primarily optimized for dense attention and struggle with the dynamic, irregular access patterns introduced by modern KV cache sparsity techniques.
  arXiv  Detail & Related papers  (2025-05-09T04:17:05Z)
• ZSMerge: Zero-Shot KV Cache Compression for Memory-Efficient Long-Context LLMs [15.76582272387931]
  We propose ZSMerge, a dynamic KV cache compression framework for efficient cache management.<n>ZSMerge significantly enhances memory efficiency and inference speed with negligible performance degradation.
  arXiv  Detail & Related papers  (2025-03-13T03:36:03Z)
• Inference-Friendly Models With MixAttention [7.103010772135246]
  MixAttention combines sliding window attention, where only a small subset of recent tokens is stored in the KV cache, with KV cache sharing across layers. Our experiments demonstrate that MixAttention significantly reduces memory usage and improves inference speed without sacrificing model performance in both short and long-context tasks.
  arXiv  Detail & Related papers  (2024-09-23T13:37:25Z)
• ThinK: Thinner Key Cache by Query-Driven Pruning [63.13363917871414]
  Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications.<n>This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference.<n>We propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels.
  arXiv  Detail & Related papers  (2024-07-30T17:59:08Z)
• Training-Free Exponential Context Extension via Cascading KV Cache [49.608367376911694]
  We introduce a novel mechanism that leverages cascading sub-cache buffers to selectively retain the most relevant tokens.<n>Our method reduces prefill stage latency by a factor of 6.8 when compared to flash attention on 1M tokens.
  arXiv  Detail & Related papers  (2024-06-24T03:59:17Z)
• CORM: Cache Optimization with Recent Message for Large Language Model Inference [57.109354287786154]
  We introduce an innovative method for optimizing the KV cache, which considerably minimizes its memory footprint. CORM, a KV cache eviction policy, dynamically retains essential key-value pairs for inference without the need for model fine-tuning. Our validation shows that CORM reduces the inference memory usage of KV cache by up to 70% with negligible performance degradation across six tasks in LongBench.
  arXiv  Detail & Related papers  (2024-04-24T16:11:54Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.