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

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

cs.LGMar 21, 2026

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

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.

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 $\downarrow$	Train Time (s) $\downarrow$	Decode (ms/tok) $\downarrow$
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 $\times$	4.05 $\times$	4.17 $\times$	4.25 $\times$	4.48 $\times$	4.87 $\times$	5.01 $\times$

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 $\times$	1.44 $\times$	1.52 $\times$	1.59 $\times$	1.68 $\times$	1.78 $\times$	1.86 $\times$

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

\leq

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 $\downarrow$	Train (tok/s) $\uparrow$	Decode (ms/tok) $\downarrow$
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 $\downarrow$	Train (tok/s) $\uparrow$	Decode (ms/tok) $\downarrow$
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 $\downarrow$	Train (s) $\downarrow$	Decode (ms/tok) $\downarrow$
Soft routing (learned)	3.22	2496.7	10.3
Fixed ( $\lambda=\frac{1}{3}$ )	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.

Explore similar work

Fast KV Compaction via Attention Matching
Feb 18, 2026Adam Zweiger, Xinghong Fu, Han Guo +1KV-Cache Management KV-Cache Compression
Massachusetts Institute of Technology.
MAC-Attention: a Match-Amend-Complete Scheme for Fast and Accurate Attention Computation
Mar 31, 2026Jinghan Yao, Sam Adé Jacobs, Walid Krichene +2LLM Inference Acceleration Efficient Attention
The Ohio State University · Microsoft, WA, USA · Anyscale, CA, USA
Latent-Condensed Transformer for Efficient Long Context Modeling
Apr 14, 2026Zeng You, Yaofo Chen, Qiuwu Chen +5Long-Context Language Modeling LLM Inference Acceleration
South China University of Technology · Pengcheng Laboratory · AIGCode +2

Feb 18, 2026J/K move · Enter open · S save

Fast KV Compaction via Attention Matching

Adam Zweiger, Xinghong Fu, Han Guo, Yoon Kim

Massachusetts Institute of Technology.

KV-Cache Management KV-Cache Compression Long-Context Language Modeling LLM Inference Acceleration

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Fast KV Compaction via Attention Matching

MAC-Attention: a Match-Amend-Complete Scheme for Fast and Accurate Attention Computation

Latent-Condensed Transformer for Efficient Long Context Modeling