This ropo contains code implemnentation of our paper MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning

The main idea include:

• MKA (3-path hierarchical memory attention)
• FastMKA (route-fused variant for speed)
• CUDA extensions for fused routing + online softmax
• Reproducible training/evaluation scripts

Our code repo include follwing:

• mka/layers/: PyTorch modules (MKAFullAttention, FastMKAAttention)
• mka/hf/: HuggingFace monkey patch support (Qwen/Llama style self_attn)
• mka/cuda/: CUDA extensions (fastmka_attn, optional fused_route_mka)
• mka/config/: optional YAML fields for memory_hierarchy etc.
• mka/utils/repro.py: global RNG seeding for reproducible runs
• scripts/: train/eval/benchmark entry points
• configs/: experiment configs

To get started, run following scripts

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
python scripts/train_wikitext2.py --config configs/qwen7b_fastmka.yaml

1. FastMKA forward path in PyTorch:

   • L1 local memory (X)
   • L2 causal session summary (prefix EMA)
   • Optional L3 retrieved memory
   • Learned routing gate (softmax(MLP(Q)))
   • Route-fusion before single KV projection

2. MKA full path in PyTorch:

   • Per-level attention over L1/L2/L3
   • Soft mixture over outputs

3. CUDA design:

   • Tiled QK score calculation
   • Online max/denominator (m, z) update
   • Fused route application before attention
   • Causal masking support

4. HuggingFace direct patch path:

   • mka/hf/attention.py: HFFastMKAAttention wrapper
   • mka/hf/patch.py: monkey patch over decoder self_attn
   • scripts/train_hf_patch.py: train loop with patched attention
   • configs/hf_qwen_fastmka.yaml, configs/hf_llama_fastmka.yaml
   • scripts/launch_dp_torchrun.sh: DP launch (torchrun)
   • scripts/launch_tp_dp_accelerate.sh: TP+DP launch path (accelerate + HF TP)

cd mka/cuda
python build.py build_ext --inplace
cd ../..

If build fails with CUDA_HOME environment variable is not set, export your CUDA path first, e.g. export CUDA_HOME=/usr/local/cuda.

2.1 For Single GPU

python scripts/train_hf_patch.py --config configs/hf_qwen_fastmka.yaml

2.2 For Multi-GPU DP

bash scripts/launch_dp_torchrun.sh configs/hf_qwen_fastmka.yaml 4

2.3 For Multi-GPU TP+DP

1. Set tp_size in config (>1).
2. Launch:

bash scripts/launch_tp_dp_accelerate.sh configs/hf_qwen_fastmka.yaml 4

Notes:

• TP relies on HuggingFace native tp_plan="auto" support for the model/version.
• Dependencies use lower bounds in requirements.txt (adjust torch for your CUDA wheel).
• Training throughput (train_throughput_tok_s) is measured after warmup_steps (see YAML) and includes forward + backward + optimizer. Inference prefill/decode (forward-only) is reported by scripts/bench_inference_metrics.py.
• FastMKA CUDA kernel is used automatically when:
  • extension fastmka_cuda is available,
  • tensor is CUDA,
  • head_dim <= 256,
  • no extra additive attention mask is required.

• YAML: seed, warmup_steps (exclude cold-start from timed throughput), optional deterministic: true (slower, stricter cudnn).
• train_hf_patch.py: logs train_mean_loss, train_total_elapsed_s, train_throughput_tok_s (post-warmup), peak_gpu_memory_gb, optional --eval-ppl for validation PPL.
• bench_inference_metrics.py: prefill_tok_s, decode_tok_s, per-phase peak GPU memory, kv_cache_bytes_* from past_key_values. HBM bandwidth is not available from PyTorch alone; use Nsight / nvidia-smi dmon on the host.

Block-MKA (§4.2) maps memory to compute tiers: L1 on-chip SRAM (tiled attention, online softmax with running max and partition sum); L2 HBM (activations, Q/K/V, fused KV cache); L3 DRAM (vectorized hash, chunk recall). FastMKA (Algorithm 2) route-fuses L1/L2/(L3) into one hidden representation, then one KV projection and one causal attention — the dominant data path is fused activations → KV on HBM → attention (see detials in our paper Tables 4–6).

YAML memory_hierarchy records these tiers for reproducibility (mka/config/memory_hierarchy.py). Older keys (hbm_enabled, dram_staging, ssd_tier_path) still parse as aliases.

Scripts vs metrics

• Training: train_hf_patch.py — tokens/s includes backward + optimizer
• Inference-style: bench_inference_metrics.py — prefill vs decode, KV bytes, per-phase GPU peak memory.
• HBM bandwidth: use Nsight Compute / vendor tools, not PyTorch alone.

References for implementation quality

• FlashAttention-2 / SDPA for L1+L2 fused attention paths.
• Paged KV / disk offload patterns (vLLM, FlashInfer-class) when extending L3 or spill.
• ZeRO-Infinity–style paths for ssd_spill_path experiments.

For evaluation of LongBench and RULER. Use:

• scripts/run_longbench.sh for LongBench workflow
• scripts/run_ruler.sh for RULER workflow

@inproceedings{mka2026,
   author = {Liu, Dong and Yu, Yanxuan and Lengerich, Ben and Wu, Ying Nian},
   title = {MKA: Memory-Keyed Attention for Efficient Long-Context Reasoning},
   year = {2026},
   isbn = {9798400725685},
   publisher = {Association for Computing Machinery},
   address = {New York, NY, USA},
   url = {https://doi.org/10.1145/3801487.3801812},
   doi = {10.1145/3801487.3801812},
   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 enhanced efficiency. Experiments on different sequence lengths show that FastMKA achieves a favorable accuracy-efficiency trade-off: comparable perplexity to MLA while achieving up to 5 \texttimes{} faster training throughput and 1.8 \texttimes{} lower evaluation latency. These results highlight MKA as a practical and extensible framework for efficient long-context attention.},
   booktitle = {Proceedings of the 23rd ACM International Conference on Computing Frontiers},
   pages = {280–289},
   numpages = {10},
   location = {Catania, Italy},
   series = {CF '26}
}