Optimizing Prefill: Prompt Caching Configuration

Hello! In our last lesson, we dove into quantization, a technique to shrink a model's memory footprint by reducing the precision of its weights. You learned how to serve these compact GPTQ and AWQ models using TGI and vLLM, a crucial skill for improving the cost-efficiency of an LLM service.

Today, we'll tackle another fundamental performance optimization: prompt caching. While quantization optimizes memory, prompt caching optimizes computation. You'll learn how LLM servers can avoid redundant work by reusing the results of previously processed prompts. For an LLM Ops engineer, this is a key lever for reducing latency and increasing throughput, especially in conversational applications.

This lesson will cover the "why" behind prompt caching by revisiting the inference lifecycle, and then show you how to enable and configure it in the vLLM serving framework.

1. The Two Phases of Inference: Prefill and Decode

To understand why prompt caching is so effective, we must first clearly distinguish between the two main stages of generating a response from an LLM:

1. Prefill (or Prompt Processing): In this first stage, the model processes the entire input prompt all at once. It's a highly parallel, compute-intensive operation. The primary output of this stage isn't a token, but the Key-Value (KV) cache—a large set of intermediate calculations (attention keys and values) for every token in the prompt. This KV cache is the model's "memory" of the prompt.
2. Decode (or Token Generation): In the second stage, the model generates the response one token at a time, autoregressively. For each new token, it reuses the entire KV cache from the prefill stage, appending the new token's own key and value to the cache. This phase is typically memory-bandwidth bound, as the main work is reading the large KV cache from GPU memory for each token.

The prefill stage is often the bottleneck for time-to-first-token (TTFT). For a long prompt, this can take a significant amount of time. Prompt caching is a strategy to eliminate this cost entirely if the prompt has been seen before.

2. How Prompt Caching Works

Prompt caching, also known as prefix caching, works by saving the KV cache generated during the prefill stage. When a new request arrives, the server checks if the prompt's prefix (or the entire prompt) matches one that's already in the cache.

• If a match is found, the server can skip the expensive prefill step and directly load the corresponding KV cache from memory.
• If there's no match, the server performs the prefill as usual and then stores the resulting KV cache for potential future use.

This is extremely powerful in common production scenarios:

• Multi-turn Chat: In a conversation, the history is a shared prefix. With each new turn, only the newest user message needs to be processed; the KV cache for the entire chat history can be reused.
• Retrieval-Augmented Generation (RAG): When you ask multiple questions about the same long document, the document's text is a shared prefix. It only needs to be processed once.
• Few-Shot Prompting: The examples provided in the prompt are a static prefix for every request.

Given your background in distributed systems, you can think of this as a specialized, application-level cache, similar to caching database query results or API responses, but applied to GPU computations.

3. Hands-On: Automatic Prefix Caching in vLLM

vLLM implements this feature as Automatic Prefix Caching (APC). It's simple to enable and works by storing the KV cache blocks of completed requests for a certain period.

Let's enable it. The process is straightforward and only requires adding a single flag to the vLLM server startup command.

1. Ensure any previous server is stopped (use Ctrl+C).

2. Launch the vLLM server with APC enabled. We'll use the same Mistral-7B-Instruct model from a previous lesson and add the --enable-prefix-caching flag.

   python -m vllm.entrypoints.openai.api_server \
       --model mistralai/Mistral-7B-Instruct-v0.2 \
       --enable-prefix-caching
   

3. Send a request with a long system prompt. Let's simulate a multi-turn chat scenario. From another terminal, send the first request:

   
   
   
   # First turn of the conversation
   curl http://localhost:8000/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "mistralai/Mistral-7B-Instruct-v0.2",
       "messages": [
         {
           "role": "system",
           "content": "You are a helpful assistant who provides concise, two-sentence answers. Your knowledge cutoff is March 2023."
         },
         {
           "role": "user",
           "content": "What is vLLM?"
         }
       ]
     }'
   

   The server will perform a full prefill for the entire prompt.

4. Send a second request with the same prefix. Now, let's simulate the next turn in the conversation. The prefix (system prompt + first user message + first assistant response) has already been processed.

   
   
   
   # Second turn of the conversation
   curl http://localhost:8000/v1/chat/completions \
     -H "Content-Type: application/json" \
     -d '{
       "model": "mistralai/Mistral-7B-Instruct-v0.2",
       "messages": [
         {
           "role": "system",
           "content": "You are a helpful assistant who provides concise, two-sentence answers. Your knowledge cutoff is March 2023."
         },
         {
           "role": "user",
           "content": "What is vLLM?"
         },
         {
           "role": "assistant",
           "content": "vLLM is a high-throughput and memory-efficient library for LLM inference and serving. It uses a novel memory management technique called PagedAttention to reduce memory waste and enable features like continuous batching."
         },
         {
           "role": "user",
           "content": "And what is PagedAttention?"
         }
       ]
     }'
   

   Although you can't see it directly, vLLM has reused the KV cache for the first three messages and only needed to perform the prefill computation for the final user message, "And what is PagedAttention?". This results in a much faster time-to-first-token for the second request.

4. Advanced Caching: RadixAttention in SGLang

While vLLM's APC is effective, other serving frameworks have developed even more sophisticated caching mechanisms. A prominent example is RadixAttention from the SGLang framework. This is relevant to your goal of gaining deep knowledge of the system layer.

RadixAttention stores the KV cache in a radix tree (or compressed trie), a data structure optimized for prefix matching. This allows it to automatically and efficiently find and share any common prefixes across all concurrent requests, not just reuse the cache of entire, completed requests.

The key architectural takeaway is that RadixAttention is designed for dynamic and automatic discovery of shared prefixes in complex, unpredictable workloads like multi-user chat applications. vLLM's APC is simpler and highly effective, but RadixAttention represents a more advanced approach to the same fundamental problem.

Conclusion

In this lesson, you've added another critical optimization technique to your LLM Ops toolkit. You now understand how to leverage prompt caching to significantly improve inference performance by avoiding redundant computation.

Here are the key takeaways:

• Inference is a two-stage process: A costly, parallel prefill stage to compute the KV cache, and a sequential decode stage to generate tokens.
• Prompt caching saves the KV cache from the prefill stage, allowing the server to skip this expensive computation for repeated prompts or prefixes.
• This technique is most effective for workloads with shared prefixes, such as multi-turn chat, RAG, and few-shot prompting.
• In vLLM, you can enable Automatic Prefix Caching with the --enable-prefix-caching flag to reuse the KV cache of previous requests.
• More advanced frameworks like SGLang use specialized data structures like radix trees (RadixAttention) for even more efficient and automatic prefix sharing.

In our next lesson, we'll focus on standardization. You will learn how to expose an OpenAI-compatible API endpoint from different serving frameworks, a crucial step for building a flexible and interoperable LLM infrastructure.