Developing a prompt caching technique with retrieval-augmented generation for efficient LLM inference

Retrieval-Augmented Generation (RAG) improves the factual correctness of Large Language Models (LLMs) by using external sources of information. However, when adding fetched documents to user requests, the system creates large prompts resulting in increased latency in the prefill stage, ultimately harming the user experience. To resolve this issue, this report presents a single system which consolidates RAG with a modular Prompt Caching technology. Unlike traditional Key Value (KV) caching which concentrates on decoding, our strategy incorporates attention reuse with the prompt the processing phase. We view retrieved text portions as modular and, with the use of a Prompt Markup Language (PML) schema, attempt to patch pre-calculated attention states to new queries in order to avoid repetitive processing for documents that are frequently accessed. We used this architecture with the Falcon-7B model and a BM25 retrieval system to evaluate it on the LongBench benchmark. We used the Prompt Cache + RAG system and M. Y. demonstrated that longitudinal processing with Falcon-7B and BM25 yielded a greater temporal efficiency on the LongBench benchmark. Our experiments indicate that the Prompt Cache + RAG system greatly reduces Time-To-First-Token (TTFT) (12 to 66 times) for both CPUs and GPUs when compared to the traditional RAG system. Most importantly, the generation quality remains the same resulting in a statistical equivalence in accuracy (F1 and ROUGE)on cached and non-cached executions. These findings confirm that the combining retrieval with modular cached attention provides a scalable approach for low-latency, knowledge intensive LLM inference.