Elamur Research⌕
☰
home / papers / 2026 / MM-ShiftKV: Decode-Aware Prefill-Stage KV Selection for Multimodal Large Language Models

MM-ShiftKV: Decode-Aware Prefill-Stage KV Selection for Multimodal Large Language Models

Paper · 2026 · Jinsong Shu, Chenyang Wu, Zhongle Xie et al.

Summary

Multimodal LLMs turn high-resolution images into long token sequences, so the KV cache becomes a major memory and latency bottleneck. Prefill-stage KV selection methods prune this cache once, right after prefilling. They rank tokens using prefill statistics, which assumes that prompt-time queries resemble decoding-time queries. The authors tested this on Qwen2.5-VL-7B-Instruct with OCRBench and found that the assumption fails. During decoding, hidden states show markedly higher variance than during prefill, while mean shifts stay modest. As a result, prefill-based importance scores are under-scaled and unreliable, and under tight budgets they can drop key visual tokens.

MM-ShiftKV is a training-free fix that works only at prefill time. It fits mean and standard deviation over the current input's hidden states, inflates the variance, and samples synthetic query proxies from that distribution. Prompt keys are then scored by the attention mass these proxies give them, combined through group-wise voting and anchored by the last query. The retained cache stays fixed during decoding and remains compatible with FlashAttention.

Across OCR, grounding, and long-context VQA benchmarks, the method beats prior selectors under strict budgets while cutting latency and memory.

Key points

  • Core finding: a prefill–decode scale mismatch. Decode-stage hidden states are consistently more variable across layers. Selectors built on prefill statistics therefore retain less of the attention mass used at decode time. MM-ShiftKV achieves the highest coverage of the methods compared.
  • Ablation, built up step by step:
    • With variance expansion, the scores reach 59.6 for OCRBench and 48.6 CIDEr for TextCaps.
    • Adding group-wise voting raises them to 68.3 and 50.4.
  • Efficiency setup: per-head KV budget of 256, outputs of 100 tokens, inputs from 2K to 32K tokens.
  • Speed: per-token decoding latency stays nearly flat, giving up to a 1.9× speedup over FullKV at 32K input.
  • Memory: at 32K input, peak GPU memory falls from about 32.9 GB to 18.6 GB, roughly a 43% reduction. Code is released at github.com/zjuDBxAI/MM-ShiftKV.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

Cite

Cite as Jinsong Shu, Chenyang Wu, Zhongle Xie, Baokun Wang and Lidan Shou. 2026. MM-ShiftKV: Decode-Aware Prefill-Stage KV Selection for Multimodal Large Language Models. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.1447

@inproceedings{shu2026mm,
  title = {{MM-ShiftKV: Decode-Aware Prefill-Stage KV Selection for Multimodal Large Language Models}},
  author = {Jinsong Shu, Chenyang Wu, Zhongle Xie, Baokun Wang and Lidan Shou},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.1447},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.1447}
}

References

4 references, 2 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. 2023. GQA: Training generalized multi-query transformer models from multi-head checkpoints. Preprint, arXiv:2305.13245.
  2. All experiments in our work are conducted using open-source datasets and models. Our research is solely aimed at enabling efficient inference of multimodal large models (MLLMs), and does not involve any human subject, sensitive data, or commercial applications. find
  3. Kazi Hasan Ibn Arif, JinYi Yoon, Dimitrios S. Nikolopoulos, Hans Vandierendonck, Deepu John, and Bo Ji. 2025. HiRED: Attention-guided token dropping for efficient inference of high-resolution vision-language models. In Proceedings of the Thirty-Ninth AAAI Conference on Artificial Intelligence. find
  4. Shuai Bai, Keqin Chen, Xuejing Liu, Jialin Wang, Wenbin Ge, Sibo Song, Kai Dang, Peng Wang, Shijie Wang, Jun Tang, and 1 others. 2025. Qwen2.5-vl technical report. Preprint, arXiv:2502.13923. \#\# Limitations

Updated

This page in Markdown → https://md.research.elamur.ai/free/en/papers/shu-2026-mm-shiftkv-decode-aware-prefill/