PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2603.20586v2cs.LGcs.AI

TL;DR

Long-context models face severe KV-cache memory and latency costs, while existing sharing and compression approaches can limit fidelity or flexibility. MKA routes queries across local, session, and long-term memories, and FastMKA fuses those sources before attention. The reported results describe competitive accuracy with substantially improved efficiency, including lower latency and faster training.

  • Problem

    Long-context attention is bottlenecked by the memory and access costs of large KV caches, while prior sharing and compression methods can trade off representation fidelity or memory-type flexibility.

  • Method

    MKA hierarchically organizes local, session, and long-term KV memories and dynamically routes each query, while FastMKA fuses routed memories before attention computation.

  • Results

    FastMKA preserves most of MKA’s accuracy benefits while significantly reducing training and inference latency across different sequence lengths and long-context benchmarks.

  • Takeaways & Limitations

    MKA provides a practical framework for efficient long-context attention with competitive perplexity and reduced compute cost.

  • Takeaways & Limitations

    The authors note that caching fused routed KV has semantic implications and evaluate it under a controlled comparison protocol discussed in the supplementary appendix.

Abstract

from arXiv · show

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.

1 Introduction

Long-context attention is bottlenecked by KV-cache memory and access costs, while existing sharing and compression methods trade away fidelity or flexibility. MKA addresses this with dynamically routed local, session, and long-term memories, and FastMKA further improves efficiency through broadcast routing.

  • Motivation: At 32K context, LLaMA-7B’s KV cache occupies 15.8GB and takes 11.3ms to access, exceeding 50% of inference latency.The cache and access costs constrain production throughput.
  • Motivation: MQA and GQA share KV features, while MLA compresses them through low-rank factorization, but these methods sacrifice fidelity or memory-type flexibility.The limitation motivates a mechanism that can distinguish heterogeneous memory sources.
  • MKA: MKA hierarchically organizes memory into local, session, and long-term levels and dynamically routes each query token across them.Lightweight routing gates modulate attention over heterogeneous memory types.
  • Implementation: MKA uses block-wise softmax, GPU-efficient kernel fusion, semantic chunking, and vectorized hashing to improve scalability and historical-content recall.The long-term memory is indexed for efficient retrieval.
  • FastMKA: FastMKA uses broadcast routing to improve training and inference efficiency while retaining MKA’s accuracy benefits.The paper presents it as suitable for high-throughput long-context inference.

2 Related Work

Prior long-context methods reduce KV or attention costs through sharing, compression, pruning, retrieval, or conditional access. Their limitations include irreversible information loss, static layouts, external indices, or insufficiently memory-specific routing, motivating MKA’s internal multi-level design.

  • Efficiency methods: FlashAttention improves memory throughput with tiled, IO-aware softmax, while MQA and GQA reduce cache size by sharing KV representations across heads or groups.These methods target attention computation or KV duplication rather than explicit heterogeneous memory routing.
  • Compression and retrieval: MLA factorizes KV memory into a smaller latent space, while DynamicKV, PyramidKV, InfiniGen, and Infinite Retrieval reduce or manage long-context cache usage through pruning or retrieval.The approaches differ in whether they compress, evict, dynamically manage, or retrieve cached information.
  • Limitations: Token-eviction methods can discard information irreversibly, and several prior approaches remain limited to static memory layouts.These limitations distinguish MKA’s routed memory organization.
  • Memory architectures: Memory Networks, Differentiable Neural Computers, Transformer-XL, and Compressive Transformers explore explicit or multi-level memory but face scalability or lossy-compression trade-offs.These systems provide historical precedents for multiple memory timescales.
  • Retrieval methods: RAG and RETRO retrieve external documents at sequence level, whereas MKA integrates local, session, and long-term memory internally with per-token routing.MKA’s routing is learned within the model rather than relying on an external index structure.
  • Routing and caching: Routing Transformers use static clustering, while TOVA and Quest focus on loading relevant KV pages rather than modeling multiple memory levels.These methods provide conditional or query-aware access without MKA’s three-tier memory structure.

3 Motivation: Beyond MLA and MHA

MLA compresses KV representations but does not explicitly represent heterogeneous memory sources or selective memory-slot reuse. MKA extends this direction with three memory levels and query-dependent routing.

  • Design gap: MLA uses low-rank projections and shared K/V structures but lacks explicit heterogeneous memory sources and selective memory-slot reuse.This defines the design gap addressed by MKA.
  • Memory hierarchy: MKA uses L1 for current-window tokens, L2 for session summaries or gated history, and L3 for indexed long-term retrieval.The three levels correspond to distinct temporal and retrieval roles.
  • Routing: A routing gate λ_l ∈ R^3 is learned dynamically for each query token to select among the three memory sources.Routing is query-dependent rather than fixed across tokens.

4 Methodology

MKA computes attention over hierarchically organized memories using per-token routing, causal summaries, block-wise normalization, and optional long-term retrieval. FastMKA fuses memory representations before one KV projection and one attention computation to reduce runtime and bandwidth costs.

  • MKA design: MKA is presented as a hierarchical attention design with symbolic algorithms, tiled execution, and recursive attention computation for memory efficiency.The methodology covers both the conceptual hierarchy and hardware-oriented execution.
  • Hierarchical routing: The architecture uses causal L1 local memory, L2 session summaries, and L3 long-term memory with routing computed per token and per layer.The routing weights are produced by an MLP and normalized across three levels.
  • FastMKA: FastMKA performs token-wise fusion of local, session, and optional long-term memories before one key-value projection and one attention computation.It caches routed fused KV rather than raw token KV.
  • FastMKA implementation: FastMKA is specified as a route-fused algorithm with projection matrices, a routing MLP, and cached concatenation of previous and current routed keys and values.The fused representation is carried forward through the cache update.
  • Operating modes: Without retrieval, FastMKA becomes a two-tier L1/L2 baseline; causal L2 prevents future-token leakage, while optional L3 enables long-distance access.This defines the retrieval-disabled and retrieval-enabled operating modes.
  • Block-MKA: Block-MKA uses L1 on-chip SRAM, L2 HBM, and L3 vectorized hash-based DRAM for hierarchical computation and chunk-based recall.Vectorized hashing supports reuse of past attention patterns and reduces redundant calculations.
  • Block computation: Standard attention forms full S ∈ R^(N×N), whereas Block-MKA partitions sequences into blocks and performs online normalization within each block.The block procedure uses Q, K, and V partitions with on-chip and HBM memory management.
  • Long-term retrieval: L3 chunk recall uses vectorized hashing to retrieve similar historical attention patterns with amortized subquadratic complexity O(BTd + BRd), where R ≪ T.The bounded recall count limits the number of retrieved chunks per block.

5 Theoretical Formulation: Recursive MKA with Online Softmax

Recursive MKA combines gated, unnormalized attention scores across hierarchical memory levels and normalizes them globally. Its online, max-shifted recursion supports stable, subquadratic computation for local, session, and retrieved long-term memory.

  • 5 Theoretical Formulation: Recursive MKA with Online Softmax: MKA computes a gated mixture of exponentiated scores from multiple memory levels, followed by one global normalization rather than separate per-level softmaxes.This formulation avoids explicitly storing full attention maps.
  • 5.1 Recursive Reformulation: Online computation accumulates unnormalized scores and values before applying a single normalization step.This avoids storing attention weights explicitly.
  • 5.1 Recursive Reformulation: The recursive formulation is derived by expanding accumulated weighted scores and values across memory levels, yielding the same gated-mixture attention expression.The proof sums contributions element-wise over the sequence dimension.
  • 5.2 Numerical Stability via Max-Shift: Hierarchical max-shift mirrors FlashAttention’s scan update and stabilizes recursive computation in low-precision or long-context regimes.The update begins with μ(0) = −∞.
  • 5.3 Local vs. Global MKA Modes: Local-MKA uses local and session memory with windowed, block-parallel O(n) computation, whereas Global-MKA adds hash-retrieved long-term memory with amortized sublinear recursive scanning.The two modes differ in whether they include L3 memory.
  • 5.4 Runtime Bounds and Complexity: L1 costs O(B^2d), L2 costs O(BTd), and L3 costs O(BRd), where R ≪ T; tiled local computation and chunk recall reduce memory access.N is total sequence length and B is block size.
  • 5.4 Runtime Bounds and Complexity: The resulting runtime is subquadratic in N rather than O(N^2d) full attention.The analysis attributes the advantage to tiled L1 computation and chunk-based L3 recall.
  • 5.4 Runtime Bounds and Complexity: Across 4K–256K sequences, FastMKA achieves 3.9–5.0× higher training throughput than MLA and 1.4–1.9× lower decode latency.These measurements empirically validate the predicted subquadratic scaling.

6 Experiments

The experiments compare MKA variants with established attention mechanisms across model architectures, context lengths, hardware settings, and long-context evaluation protocols. Measurements include quality, throughput, latency, and memory-related inference behavior.

  • Models: Evaluation covers Qwen2.5-7B/14B, Llama 3.1-8B, and DeepSeek-V3, spanning GQA, standard long-context, and MLA-based architectures.DeepSeek-V3 enables direct comparison with MLA compression.
  • Dataset: WikiText-2 is the primary fine-tuning and evaluation dataset, supplemented by long-context tasks with sequences up to 128K tokens.The dataset contains 36,718 training, 3,760 validation, and 4,358 test sentences.
  • Hardware: Experiments run on NVIDIA A800 80GB GPUs, using one GPU for shorter 7B sequences and 4–8 GPUs for larger models or sequences.DeepSeek-V3 experiments focus on inference latency and KV-cache efficiency.
  • Training Details: Models are fine-tuned for one epoch at 4K–256K sequence lengths, with batch sizes adjusted by model size and context length.Reported throughput and latency include routing-MLP and memory-fusion overhead.
  • Inference Details: Inference separately measures prefill and decode using batch sizes 1, 4, 8, and 16, with contiguous KV storage through 32K and paged storage at longer contexts.Latency, throughput, and memory bandwidth are measured after warmup and repeated inference iterations.
  • Metric: Language-model quality is measured with cross-entropy loss L and perplexity defined as PPL = e^L.

6.2 Main Results

The main-results tables compare FastMKA with attention baselines across quality, throughput, latency, memory bandwidth, architectures, and long-context benchmarks. The supplied captions define each evaluation setting, while the reported results emphasize efficiency and benchmark coverage.

  • 6.2 Main Results: Table 1 compares attention mechanisms on Qwen2.5-7B at 16K, with FastMKA identified as achieving the best accuracy–efficiency trade-off.
  • 6.2 Main Results: Table 2 reports Qwen2.5-7B training throughput in tokens/second across sequence lengths using batch size 2, bf16, and FlashAttention-2.The caption notes non-linear scaling beyond 64K from memory bandwidth and kernel-launch overhead.
  • 6.2 Main Results: Table 3 reports Qwen2.5-7B decode latency in ms/token at batch size 1, with non-linear scaling beyond 64K attributed to KV-cache paging and bandwidth saturation.
  • 6.2 Main Results: Table 4 separates prefill total time from per-token decode latency and compares contiguous KV layout through 32K with paged layout for longer contexts.
  • 6.2 Main Results: Table 5 measures KV-cache memory and effective HBM bandwidth at 128K using Nsight Compute, linking FastMKA’s utilization to fused memory access and kernel saturation.
  • 6.2 Main Results: Table 6 compares methods across architectures on WikiText-2 at 32K, using batch 2 for training and batch 1 for inference.Qwen2.5-14B uses four A800 GPUs with TP=4; other models use one A800 GPU.
  • 6.2 Main Results: Table 7 evaluates LongBench accuracy across QA, summarization, and code categories at 128K using official base-model evaluation protocols.
  • 6.2 Main Results: Table 8 evaluates RULER passkey retrieval accuracy with a random number inserted into long contexts, using top-8 L3 chunks per query for FastMKA.Higher accuracy is better under the official base-model protocol.

6.3 Experimental Results Analysis

FastMKA is reported to preserve competitive quality while improving efficiency across context lengths, model scales, architectures, and long-context tasks. Ablations attribute quality differences to routing and hierarchical memory choices.

  • 6.3 Experimental Results Analysis: FastMKA reaches 3.26 PPL versus MLA’s 3.22 at 16K while training 3.6× faster and decoding 1.5× faster on Qwen2.5-7B.The paper characterizes this as a 1.2% accuracy exchange for substantial compute savings.
  • 6.3 Experimental Results Analysis: At 256K, FastMKA reaches a 5.0× training speedup and 1.86× decode speedup, with the pattern extending across 7B and 14B models and Llama and DeepSeek architectures.The paper links the speedup to route fusion, which uses three rather than nine kernel launches.
  • 6.4 Long-Context Benchmark Evaluation: LongBench and RULER evaluations report competitive long-context accuracy while substantially reducing latency.
  • 6.3 Experimental Results Analysis: Removing memory tiers degrades performance, supporting the hierarchical design in the tier ablation.

6.5 Ablation Studies

The ablations show that hierarchical memory and learned soft routing improve both perplexity and efficiency, while route fusion reduces attention overhead across prefill and decode.

  • Memory-tier ablations: 3.51 PPL and 7.9× slower training throughput with L1-only memory confirm the necessity of hierarchical memory.Removing L3 raises perplexity to 3.31, while removing L2 yields 3.28 PPL.
  • Routing-strategy ablations: 3.22 PPL from learned soft routing outperforms fixed uniform routing at 3.34 PPL by 3.7%.Hard top-k routing performs intermediately, while routing overhead adds only 1–2% latency.
  • Latency breakdown: 1.63× lower prefill latency and 1.78× lower decode latency distinguish FastMKA from MLA at 128K context.FastMKA measures 0.87s versus 1.42s for prefill and 18.4ms versus 32.7ms per token for decode.
  • Computational mechanism: FastMKA reduces attention computations from three memory-level paths to one while maintaining comparable accuracy through learned routing.The reduction is attributed to route fusion rather than implementation differences.

6.7 Discussion

MKA routes attention dynamically across hierarchical memory, while FastMKA fuses those memories before attention to reduce latency and training cost. The resulting design preserves most accuracy benefits and supports high-throughput long-context inference.

  • MKA design: MKA dynamically routes across local, session, and long-term memory without increasing the model’s parameter count.The hierarchy is intended to preserve representation fidelity and cache reusability in memory-intensive inference.
  • FastMKA design: FastMKA fuses routed memory levels before attention, using one KV projection and one attention computation.It caches the fused routed KV rather than raw token KV, reducing memory bandwidth.
  • Empirical trade-off: FastMKA significantly reduces training and inference latency while preserving most of MKA’s accuracy benefits across sequence lengths and long-context benchmarks.The claim is supported by the reported experimental validation across different sequence lengths and benchmark settings.
  • Practical implications: FastMKA’s lightweight single-projection design is positioned for high-throughput inference and edge deployment with a balance among accuracy, latency, and memory usage.The discussion presents it as a drop-in replacement for traditional attention in long-context scenarios.

7 Conclusion

The paper introduces MKA as hierarchical query routing across multiple memory levels for efficient long-context modeling, then extends it with FastMKA to fuse memory before attention. Together, they target competitive perplexity with lower compute cost, latency, and training expense.

  • MKA: MKA routes queries across multiple memory levels to enable efficient long-context modeling.The mechanism is presented as a hierarchical attention framework for memory-aware transformer design.
  • FastMKA: FastMKA performs memory fusion before attention computation to reduce overhead while retaining MKA’s architectural benefits.The broadcast-routed variant is designed for scalable inference with long-sequence inputs.
Loading 2603.20586v2…