Elamur Research⌕
☰
home / papers / 2023 / Prompt Cache: Modular Attention Reuse for Low-Latency Inference

Prompt Cache: Modular Attention Reuse for Low-Latency Inference

Paper · 2023 · In Gim, Guojun Chen, Seung-seob Lee et al.

Summary

LLM prompts often repeat the same material: system instructions, templates, and context documents drawn from a shared pool. Standard KV Cache avoids recomputation only within a single request, so these shared pieces are re-encoded every time they appear. The authors propose Prompt Cache. It precomputes key-value attention states for recurring segments once, keeps them in CPU or GPU memory, and splices them into new prompts, so only the uncached text needs fresh attention computation. Reusable segments are declared as prompt modules in a schema written in a Prompt Markup Language (PML). Each module is assigned its own position IDs. The authors also find empirically that LLMs still work when position IDs have gaps, so modules can be selected and combined flexibly. A prototype built on the HuggingFace transformers library was tested with Llama2, Falcon and MPT on LongBench tasks. Time-to-first-token fell by up to an order of magnitude on GPU and by tens of times on CPU. Accuracy showed no significant loss, and model weights were not changed. This is most valuable for long, document-heavy prompts, where the prefill stage dominates latency.

Key points

  • When modules are kept in GPU memory and inference runs on the GPU, TTFT improves by between 1.5× and 10×. CPU inference sees improvements of 20× to 70×. The abstract summarizes this as about 8× (GPU) and 60× (CPU).
  • Evaluation used the LongBench suite, covering multi-document question answering and recommendation tasks, with three open model families: Llama2, Falcon, MPT.
  • Two storage options were compared:
    • CPU memory scales to terabytes but adds host-to-device copy cost.
    • GPU memory avoids copying but has limited capacity.
  • The paper also analyzes the memory footprint of cached states. It demonstrates personalization, code generation and parameterized prompts with negligible quality degradation.
  • Code and evaluation data are released at github.com/yale-sys/prompt-cache. The approach claims compatibility with any Transformer that supports KV Cache.

Interesting ideas

  • I1. Explicit prompt structure. Making prompt structure explicit through a markup schema, instead of detecting repeated text automatically, makes cache lookup easy. This pattern could fit any system that assembles prompts from a fixed library of documents or instructions, such as RAG over a stable corpus.
  • I2. Fixed position slots per module. Each reusable segment gets its own reserved range of position IDs, so its cached states stay valid wherever the module is used. This is a simple way to work around the position dependence of attention states.
  • I3. Tolerance of position gaps. The observation that models handle non-contiguous position IDs suggests that a large schema can include optional modules that are dropped per request without recomputing everything. That is useful for configurable templates, for example a different set of reference documents per user.
  • I4. Cost argument. Prefill attention cost grows quadratically with prompt length, while the cache grows only linearly. Caching therefore pays off most for long contexts, which is a useful rule of thumb when deciding where caching is worth the memory.
  • I5. Tiered storage. Treating CPU DRAM and GPU memory as storage tiers for cached states is a design choice that could carry over to serving setups with modest GPUs. Large document pools could live in host memory while frequently used modules stay on the GPU.

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

 Institutions

🏢GoogleMountain View · United States492 works · 440,462 citations🎓Yale UniversityNew Haven · United States101 works · 13,302 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 works

Cite

Cite as In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal and Lin Zhong. 2023. Prompt Cache: Modular Attention Reuse for Low-Latency Inference. MLSys 2024.

The paper itself, not this page. BibTeX from the library's record of the work. arXiv 2311.04934

@inproceedings{gim2023prompt,
  title = {{Prompt Cache: Modular Attention Reuse for Low-Latency Inference}},
  author = {In Gim, Guojun Chen, Seung-seob Lee, Nikhil Sarda, Anurag Khandelwal and Lin Zhong},
  booktitle = {MLSys 2024},
  year = {2023},
  eprint = {2311.04934},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2311.04934}
}

References

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

  1. Aminabadi, R. Y., Rajbhandari, S., Awan, A. A., Li, C., Li, D., Zheng, E., Ruwase, O., Smith, S., Zhang, M., Rasley, J., and He, Y. Deepspeed- inference: Enabling efficient inference of transformer models at unprecedented scale. In Wolf, F., Shende, S., Culhane, C., Alam, S. R., and Jagode, H. (eds.… find
  2. IEEE, 2022. doi: 10.1109/SC41404.2022 .00051. URL https://doi.org/10.1109/SC41 404.2022.00051. — [doi:10.1109/sc41404.2022](https://doi.org/10.1109/sc41404.2022) ↗
  3. Bai, Y., Lv, X., Zhang, J., Lyu, H., Tang, J., Huang, Z., Du, Z., Liu, X., Zeng, A., Hou, L., Dong, Y., Tang, J., and Li, J. Longbench: A bilingual, multitask benchmark for long context understanding. CoRR, abs/2308.14508, 2023. doi: 10.48550/ARXIV.2308.14508. URL https://doi.org/10.48550/arXiv.230… ↗
  4. Beltagy, I., Peters, M. E., and Cohan, A. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL https://arxiv.org/abs/2004.0 5150. ↗
  5. Beurer-Kellner, L., Fischer, M., and Vechev, M. Prompting is programming: A query language for large language models. Proceedings of the ACM on Programming Languages, 7(PLDI):1946–1969, 2023. find
  6. Cui, J., Li, Z., Yan, Y., Chen, B., and Yuan, L. Chatlaw: Open-source legal large language model with integrated external knowledge bases, 2023. find
  7. Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344–16359, 2022. find
  8. Driess, D., Xia, F., Sajjadi, M. S. M., Lynch, C., Chowdhery, A., Ichter, B., Wahid, A., Tompson, J., Vuong, Q., Yu, T., Huang, W., Chebotar, Y., Sermanet, P., Duckworth, D., Levine, S., Vanhoucke, V., Hausman, K., Toussaint, M., Greff, K., Zeng, A., Mordatch, I., and Florence, P. Palm-e: An embodie… find
  9. Dufter, P., Schmitt, M., and Schütze, H. Position information in transformers: An overview. Computational Linguistics, 48(3):733–763, 2022.
  10. Fabbri, A. R., Li, I., She, T., Li, S., and Radev, D. Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 1074–1084, 2019.
  11. Feng, Y., Jeon, H., Blagojevic, F., Guyot, C., Li, Q., and Li, D. Attmemo : Accelerating transformers with memoization on big memory systems, 2023. find
  12. GitHub. Github copilot · your ai pair programmer, 2023. URL https://github.com/features/copi lot. ↗
  13. Guidence. A guidance language for controlling large language models. https://github.com/guidanc e-ai/guidance, 2023. ↗
  14. Guo, D., Xu, C., Duan, N., Yin, J., and McAuley, J. Longcoder: A long-range pre-trained language model for code completion. arXiv preprint arXiv:2306.14893, 2023. — [arXiv:2306.14893](https://arxiv.org/abs/2306.14893) find
  15. Ho, X., Nguyen, A.-K. D., Sugawara, S., and Aizawa, A. Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. In Proceedings of the 28th International Conference on Computational Linguistics, pp. 6609–6625, 2020.
  16. Huang, L., Cao, S., Parulian, N., Ji, H., and Wang, L. Efficient attentions for long document summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1419–1436, 2021.
  17. Huang, W., Abbeel, P., Pathak, D., and Mordatch, I. Language models as zero-shot planners: Extracting actionable knowledge for embodied agents. arXiv preprint arXiv:2201.07207, 2022.
  18. Joshi, M., Choi, E., Weld, D. S., and Zettlemoyer, L. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1601–1611, 2017.
  19. Keles, F. D., Wijewardena, P. M., and Hegde, C. On the computational complexity of self-attention, 2022. find
  20. Kočiskỳ, T., Schwarz, J., Blunsom, P., Dyer, C., Hermann, K. M., Melis, G., and Grefenstette, E. The narrativeqa reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317–328, 2018.
  21. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. arXiv preprint arXiv:2309.06180, 2023.
  22. Lew, Z., Walther, J. B., Pang, A., and Shin, W. Interactivity in online chat: Conversational contingency and response latency in computer-mediated communication. Journal of Computer-Mediated Communication, 23(4):201–221, 2018. find
  23. Liu, T., Xu, C., and McAuley, J. Repobench: Benchmarking repository-level code auto-completion systems. arXiv preprint arXiv:2306.03091, 2023a.
  24. Liu, Y., Li, H., Du, K., Yao, J., Cheng, Y., Huang, Y., Lu, S., Maire, M., Hoffmann, H., Holtzman, A., Ananthanarayanan, G., and Jiang, J. Cachegen: Fast context loading for language model applications, 2023b. find
  25. MosaicML. Introducing mpt-7b: A new standard for open-source, commercially usable llms, 2023. URL www.mo saicml.com/blog/mpt-7b. Accessed: 2023-05-05. find
  26. Nay, J. J., Karamardian, D., Lawsky, S. B., Tao, W., Bhat, M., Jain, R., Lee, A. T., Choi, J. H., and Kasai, J. Large language models as tax attorneys: A case study in legal capabilities emergence, 2023.
  27. Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-perfor… — PyTorch: An Imperative Style, High-Performance Deep Learning Library
  28. Penedo, G., Malartic, Q., Hesslow, D., Cojocaru, R., Cappelli, A., Alobeidli, H., Pannier, B., Almazrouei, E., and Launay, J. The refinedweb dataset for falcon llm: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116, 2023. find
  29. Pope, R., Douglas, S., Chowdhery, A., Devlin, J., Bradbury, J., Levskaya, A., Heek, J., Xiao, K., Agrawal, S., and Dean, J. Efficiently scaling transformer inference, 2022. find
  30. Press, O., Smith, N. A., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapolation. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. find
  31. OpenReview.net, 2022. URL https://openreview.net/forum?id=R8sQPpGCv0. ↗
  32. Qin, Y., Hu, S., Lin, Y., Chen, W., Ding, N., Cui, G., Zeng, Z., Huang, Y., Xiao, C., Han, C., Fung, Y. R., Su, Y., Wang, H., Qian, C., Tian, R., Zhu, K., Liang, S., Shen, X., Xu, B., Zhang, Z., Ye, Y., Li, B., Tang, Z., Yi, J., Zhu, Y., Dai, Z., Yan, L., Cong, X., Lu, Y., Zhao, W., Huang, Y., Yan,… — Tool Learning with Foundation Models
  33. Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language understanding by generative pre-training. 2018. find
  34. Rasmy, L., Xiang, Y., Xie, Z., Tao, C., and Zhi, D. Medbert: pretrained contextualized embeddings on large-scale structured electronic health records for disease prediction. npj Digit. Medicine, 4, 2021. doi: 10.1038/S41746-021 -00455-Y. URL https://doi.org/10.1038/s4 1746-021-00455-y. — [doi:10.1038/s41746-021](https://doi.org/10.1038/s41746-021) find
  35. Rozière, B., Gehring, J., Gloeckle, F., Sootla, S., Gat, I., Tan, X. E., Adi, Y., Liu, J., Remez, T., Rapin, J., Kozhevnikov, A., Evtimov, I., Bitton, J., Bhatt, M., Canton-Ferrer, C., Grattafiori, A., Xiong, W., Défossez, A., Copet, J., Azhar, F., Touvron, H., Martin, L., Usunier, N., Scialom, T.… — Code Llama: Open Foundation Models for Code
  36. Scao, T. L., Fan, A., Akiki, C., Pavlick, E., Ilic, S., Hesslow, D., Castagné, R., Luccioni, A. S., Yvon, F., Gallé, M., Tow, J., Rush, A. M., Biderman, S., Webson, A., Ammanamanchi, P. S., Wang, T., Sagot, B., Muennighoff, N., del Moral, A. V., Ruwase, O., Bawden, R., Bekman, S., McMillan-Major,… find
  37. Shen, J. T., Yamashita, M., Prihar, E., Heffernan, N. T., Wu, X., and Lee, D. Mathbert: A pre-trained language model for general NLP tasks in mathematics education. CoRR, abs/2106.07340, 2021. URL https://arxiv.org/abs/2106.07340. — [arXiv:2106.07340](https://arxiv.org/abs/2106.07340) ↗
  38. Sheng, Y., Zheng, L., Yuan, B., Li, Z., Ryabinin, M., Chen, B., Liang, P., Ré, C., Stoica, I., and Zhang, C. Flexgen: High-throughput generative inference of large language models with a single GPU. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Internat…
  39. PMLR, 2023. URL https://procee dings.mlr.press/v202/sheng23a.html. find
  40. Steinberg, E., Jung, K., Fries, J. A., Corbin, C. K., Pfohl, S. R., and Shah, N. H. Language models are an effective representation learning technique for electronic health record data. J. Biomed. Informatics, 113:103637, 2021. doi: 10.1016/J.JBI.2020.103637. URL https://doi. org/10.1016/j.jbi.2020.… — [doi:10.1016/j.jbi.2020.103637](https://doi.org/10.1016/j.jbi.2020.103637) find
  41. Su, J., Lu, Y., Pan, S., Wen, B., and Liu, Y. Roformer: Enhanced transformer with rotary position embedding. CoRR, abs/2104.09864, 2021. URL https://arxiv.org/abs/2104.09864. ↗
  42. Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey. ACM Comput. Surv., 55(6):109:1– 109:28, 2023. doi: 10.1145/3530811. URL https://doi.org/10.1145/3530811. — [doi:10.1145/3530811](https://doi.org/10.1145/3530811) ↗
  43. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D., Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Esiobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V., Goyal, N., Hartshorn, A., Hossei… — Llama 2: Open Foundation and Fine-Tuned Chat Models
  44. Trivedi, H., Balasubramanian, N., Khot, T., and Sabharwal, A. Musique: Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539–554, 2022. — ♫ MuSiQue: Multihop Questions via Single-hop Question Composition
  45. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need, 2023. find
  46. White, J., Fu, Q., Hays, S., Sandborn, M., Olea, C., Gilbert, H., Elnashar, A., Spencer-Smith, J., and Schmidt, D. C. A prompt pattern catalog to enhance prompt engineering with chatgpt. CoRR, abs/2302.11382, 2023. doi: 10.485 50/ARXIV.2302.11382. URL https://doi.org/10.48550/arXiv.2302.11382. ↗
  47. Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. L., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. Huggingface’s transformers: State-of-th… — Transformers: State-of-the-Art Natural Language Processing
  48. Wu, L., Zheng, Z., Qiu, Z., Wang, H., Gu, H., Shen, T., Qin, C., Zhu, C., Zhu, H., Liu, Q., Xiong, H., and Chen, E. A survey on large language models for recommendation. CoRR, abs/2305.19860, 2023. doi: 10.48550/ARXIV.2 305.19860. URL https://doi.org/10.48550/arXiv.2305.19860. — [doi:10.48550/arxiv.2](https://doi.org/10.48550/arxiv.2) ↗
  49. Yang, Z., Qi, P., Zhang, S., Bengio, Y., Cohen, W., Salakhutdinov, R., and Manning, C. D. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2369–2380, 2018.
  50. Zhang, Z., Sheng, Y., Zhou, T., Chen, T., Zheng, L., Cai, R., Song, Z., Tian, Y., Ré, C., Barrett, C., Wang, Z., and Chen, B. H 2 o: Heavy-hitter oracle for efficient generative inference of large language models, 2023. — H$_2$O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
  51. Zhong, M., Yin, D., Yu, T., Zaidi, A., Mutuma, M., Jha, R., Hassan, A., Celikyilmaz, A., Liu, Y., Qiu, X., et al. Qmsum: A new benchmark for query-based multi-domain meeting summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Lingui…

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/gim-2023-prompt-cache-modular-attention-reuse/