Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Dong Liu, Yanxuan Yu, Ben Lengerich, Ying Nian Wu

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.

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

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.
Paper Structure (36 sections, 1 theorem, 10 equations, 2 figures, 11 tables, 3 algorithms)

This paper contains 36 sections, 1 theorem, 10 equations, 2 figures, 11 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Long-context Attention Mechanisms
  4. Hierarchical and External Memory Systems
  5. Dynamic Routing and Query-aware Attention
  6. Motivation: Beyond MLA and MHA
  7. Methodology
  8. Route-Fused MKA (FastMKA): Causal Route-Fusion with KV-Cache
  9. Block-Memory Keyed Attention (Block-MKA) Design
  10. Hierarchical Block-wise MKA Algorithm
  11. Algorithm and Pseudocode
  12. Theoretical Formulation: Recursive MKA with Online Softmax
  13. Recursive Reformulation
  14. Numerical Stability via Max-Shift
  15. Local vs. Global MKA Modes
  16. ...and 21 more sections

Key Result

Theorem 5.1

The recursive formulation defined by Equations (6)-(8) computes the gated mixture attention:

Figures (2)

  • Figure 1: MKA hierarchical memory design with three memory levels (L1: local, L2: session, L3: long-term) and dynamic routing gates.
  • Figure 2: Hierarchical Memory-Keyed Attention (MKA) with Multi-Level Routing: illustrates the three-tier memory architecture (L1/L2/L3) and query-based routing mechanism.

Theorems & Definitions (1)

  • Theorem 5.1: Recursive MKA Computes Gated Mixture Attention