MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning
Authors: Dong Liu, Yanxuan Yu, Ben Lengerich, Ying Nian Wu
Organizations: University of California, Los Angeles, Los Angeles, California, USA · Columbia University, New York, New York, USA · University of Wisconsin-Madison, Madison, Wisconsin, USA
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.
Figures & tables
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.
Method
PPL ↓
Train Time (s) ↓
Decode (ms/tok) ↓
MHA (baseline)
3.31
6234.7
21.4
GQA
3.28
5012.4
18.6
MLA
3.22
4456.9
12.8
FastMKA (ours)
3.26
1248.3
8.4
Table 1. Comparison of attention mechanisms on Qwen2.5-7B (sequence length 16K). FastMKA achieves the best trade-off between accuracy and efficiency.
Method
4K
8K
16K
32K
64K
128K
256K
MHA
347
289
231
184
142
98
68
GQA
424
382
341
296
243
178
124
MLA
468
428
387
342
287
212
148
FastMKA (ours)
1847
1732
1614
1453
1287
1032
742
Speedup vs MLA
3.94 ×
4.05 ×
4.17 ×
4.25 ×
4.48 ×
4.87 ×
5.01 ×
Table 2. Training throughput (tokens/second) on Qwen2.5-7B with batch size 2, bf16, FlashAttention-2. Note non-linear scaling beyond 64K due to memory bandwidth and kernel launch overhead.
Method
4K
8K
16K
32K
64K
128K
256K
MHA
14.2
16.8
21.4
28.6
39.2
58.3
87.6
GQA
12.4
14.8
18.6
24.3
33.4
49.8
75.2
MLA
8.7
10.2
12.8
16.4
22.8
32.7
48.9
FastMKA (ours)
6.2
7.1
8.4
10.3
13.6
18.4
26.3
Speedup vs MLA
1.40 ×
1.44 ×
1.52 ×
1.59 ×
1.68 ×
1.78 ×
1.86 ×
Table 3. Decode latency (ms/token, batch=1) on Qwen2.5-7B with bf16. Latency scales non-linearly beyond 64K due to KV cache paging and memory bandwidth saturation.
Context
Method
Prefill (s)
Decode (ms/tok)
Prefill Attn %
Decode Attn %
32K
MHA
0.58
28.6
68.4%
71.2%
GQA
0.49
24.3
64.7%
67.8%
MLA
0.35
16.4
61.3%
64.2%
FastMKA
0.21
10.3
55.8%
58.6%
64K
MHA
1.12
39.2
70.2%
73.1%
GQA
0.94
33.4
66.8%
69.7%
Table 4. Prefill vs Decode latency breakdown on Qwen2.5-7B (batch=1, bf16). Prefill: total time to process input tokens. Decode: latency per generated token. Measurements use contiguous KV layout for ≤ 32K and paged layout for longer contexts.
Method
KV Cache (GB)
KV Read (GB)
HBM BW (GB/s)
Utilization
MHA
18.7
18.7
1240
78.2%
GQA
12.4
12.4
1156
72.9%
MLA
8.9
8.9
1087
68.5%
FastMKA (ours)
6.2
6.2
1324
83.5%
Reduction vs MHA
66.8%
66.8%
+6.8%
+5.3%
Table 5. Memory footprint and HBM bandwidth on Qwen2.5-7B. KV cache memory (GB) and effective bandwidth (GB/s) measured via Nsight Compute at 128K context, batch=1, bf16. FastMKA’s higher HBM bandwidth utilization despite smaller KV cache is due to more contiguous memory access patterns (route-fusion creates a single fused KV tensor) and better kernel saturation (fewer but larger kernels), as validated by Nsight Compute roofline analysis.
Model
Method
PPL ↓
Train (tok/s) ↑
Decode (ms/tok) ↓
Qwen2.5-7B
GQA (baseline)
3.24
296
24.3
MLA
3.18
342
16.4
FastMKA (ours)
3.22
1453
10.3
Qwen2.5-14B
GQA (baseline)
3.12
158
32.7
MLA
3.06
184
21.8
FastMKA (ours)
3.10
642
13.6
Table 6. Comparison across different model architectures on WikiText-2 (sequence length 32K, batch=2 for training, batch=1 for inference). Qwen2.5-14B experiments use 4 A800 GPUs with tensor parallelism (TP=4); all other models use a single A800 GPU. FastMKA maintains efficiency advantages across architectures.
Method
QA
Summarization
Code
Avg
MHA
42.3
58.7
51.2
50.7
GQA
44.8
61.3
53.4
53.2
MLA
46.2
63.8
55.1
55.0
FastMKA (ours)
45.7
63.2
54.6
54.5
Table 7. Performance on LongBench benchmark (Qwen2.5-7B, sequence length 128K). We use the official LongBench evaluation protocol: base models (not instruction-tuned) with max_new_tokens=128, evaluated using the official LongBench evaluation scripts. Results show accuracy across different task categories (QA, Summarization, Code).
Method
4K
8K
16K
32K
64K
128K
MHA
98.4
97.2
94.8
87.3
72.6
45.2
GQA
98.6
97.8
96.2
91.4
81.7
58.3
MLA
98.8
98.2
97.1
94.6
88.4
74.8
FastMKA (ours)
98.7
98.1
96.9
94.2
87.8
73.4
Table 8. Passkey retrieval accuracy on RULER benchmark (Qwen2.5-7B). We use the official RULER passkey retrieval task: a random number is inserted at a random position in a long context, and the model must retrieve it. Evaluation uses base models (not instruction-tuned) with max_new_tokens=1, following the official RULER evaluation protocol. Higher is better. Demonstrates L3 memory’s effectiveness for long-distance retrieval. FastMKA uses L3 retrieval with R=8 (top-8 retrieved chunks per query).
Variant
PPL ↓
Train (tok/s) ↑
Decode (ms/tok) ↓
FastMKA (L1+L2+L3)
3.22
1453
10.3
L1+L2 only
3.31
1562
9.2
L1+L3 only
3.28
1508
9.6
L1 only (=MHA)
3.51
184
28.6
Table 9. Tier ablation study on Qwen2.5-7B (sequence length 32K, batch=2). Removing memory tiers degrades performance, validating the hierarchical design.
Routing Method
PPL ↓
Train (s) ↓
Decode (ms/tok) ↓
Soft routing (learned)
3.22
2496.7
10.3
Fixed ( λ=31 )
3.34
2534.2
10.5
Hard top-2 routing
3.29
2512.8
10.4
Hard top-1 routing
3.41
2498.3
10.3
Table 10. Routing mechanism ablation (Qwen2.5-7B, sequence length 32K). Learned soft routing outperforms fixed and hard routing.
Model
Method
FLOPs (B)
KV Read (GB)
Kernel Launches
Attn Paths
Qwen2.5-7B
GQA (baseline)
2.89
18.7
1
1
MLA
3.12
15.2
1
1
MKA (full, ours)
6.87
14.8
9
3
FastMKA (ours)
2.34
6.2
3
1
Qwen2.5-14B
GQA (baseline)
5.76
37.4
1
1
MLA
6.24
30.4
1
1
Table 11. Per-token computational cost breakdown across different model architectures (128K context, batch=1, bf16). FLOPs and memory access measured via Nsight Compute. FastMKA achieves speedup through algorithmic reduction of attention paths rather than implementation differences.
Scaling language models to long contexts is often bottlenecked by the size of the key-value (KV) cache. In deployed settings, long contexts are typically managed through compaction in token space via summarization. However, summarization can be highly lossy, substantially harming downstream performance. Recent work on Cartridges has shown that it is possible to train highly compact KV caches in latent space that closely match full-context performance, but at the cost of slow and expensive end-to-end optimization. This work describes an approach for fast context compaction in latent space through Attention Matching, which constructs compact keys and values to reproduce attention outputs and preserve attention mass at a per-KV-head level. We show that this formulation naturally decomposes into simple subproblems, some of which admit efficient closed-form solutions. Within this framework, we develop a family of methods that significantly push the Pareto frontier of compaction time versus quality, achieving up to 50x compaction in seconds on some datasets with little quality loss.
Long-context decoding in LLMs is IO-bound: each token re-reads an ever-growing KV cache. Prior accelerations cut bytes via compression, which lowers fidelity, or selection/eviction, which restricts what remains accessible, and both can degrade delayed recall and long-form generation. We introduce MAC-Attention, a fidelity- and access-preserving alternative that accelerates decoding by reusing prior attention computations for semantically similar recent queries. It starts with a match stage that performs pre-RoPE L2 matching over a short local window; an amend stage rectifies the reused attention by recomputing a small band near the match boundary; and a complete stage fuses the rectified results with fresh attention computed on the KV tail through a numerically stable merge. On a match hit, the compute and bandwidth complexity is constant regardless of context length. The method is model-agnostic and composes with IO-aware kernels, paged-KV managers, and MQA/GQA. Across LongBench v2 (120K), RULER (120K), and LongGenBench (16K continuous generation), compared to the latest FlashInfer library, MAC-Attention reduces KV accesses by up to 99%, cuts token generation latency by over 60% at 128K, and achieves over 14.3x attention-phase speedups, up to 2.6x end-to-end, while maintaining full-attention quality. By reusing computation, MAC-Attention delivers long-context inference that is both fast and faithful. Code is available here: https://github.com/YJHMITWEB/MAC-Attention.git
Jinghan Yao, Sam Adé Jacobs, Walid Krichene +2
The Ohio State University · Microsoft, WA, USA · Anyscale, CA, USA
Large language models (LLMs) face significant challenges in processing long contexts due to the linear growth of the key-value (KV) cache and quadratic complexity of self-attention. Existing approaches address these bottlenecks separately: Multi-head Latent Attention (MLA) reduces the KV cache by projecting tokens into a low-dimensional latent space, while sparse attention reduces computation. However, sparse methods cannot operate natively on MLA's compressed latent structure, missing opportunities for joint optimization. In this paper, we propose Latent-Condensed Attention (LCA), which directly condenses context within MLA's latent space, where the representation is disentangled into semantic latent vectors and positional keys. LCA separately aggregates semantic vectors via query-aware pooling and preserves positional keys via anchor selection. This approach jointly reduces both computational cost and KV cache without adding parameters. Beyond MLA, LCA's design is architecture-agnostic and readily extends to other attention mechanisms such as GQA. Theoretically, we prove a length-independent error bound. Experiments show LCA achieves up to 2.5× prefilling speedup and 90% KV cache reduction at 128K context while maintaining competitive performance.
Zeng You, Yaofo Chen, Qiuwu Chen +5
South China University of Technology · Pengcheng Laboratory · AIGCode +2