Inference Optimization

KV Cache Optimization — The Memory Bottleneck

The KV cache stores key-value tensors for all previous tokens during generation. For large models, it often exceeds the model's own memory footprint. Optimization is essential for serving.

PagedAttention — Non-contiguous memory allocation eliminates fragmentation
RadixAttention — Prefix sharing across requests reduces redundant computation
Cache Quantization — Store KV cache in lower precision for 2-4x memory reduction

The KV cache is the hidden cost of autoregressive generation.

KV Cache Optimization for LLMs

During autoregressive generation, each new token requires attention over all previous tokens. The KV cache stores pre-computed key and value tensors, avoiding redundant computation. However, for long sequences and large models, the KV cache becomes the primary memory bottleneck.

KV Cache Memory Requirements

PagedAttention

The Problem: Memory Fragmentation

The Solution: PagedAttention

Architecture Diagram

Logical KV Cache:  [Page 0] [Page 1] [Page 2] [Page 3]
                         |         |         |         |
Physical Memory:  [Block 5] [Block 2] [Block 8] [Block 1]
                         |         |         |         |
                   GPU Memory: Non-contiguous allocation

Implementation

class PagedKVCache:
    def __init__(self, page_size=16, num_pages=1024, num_layers=32, num_heads=32, head_dim=128):
        self.page_size = page_size
        self.num_pages = num_pages
        self.page_table = {}  # request_id -> list of physical pages
        self.free_pages = list(range(num_pages))
        
        # Pre-allocate all pages
        self.k_cache = torch.zeros(num_pages, page_size, num_heads, head_dim, device="cuda")
        self.v_cache = torch.zeros(num_pages, page_size, num_heads, head_dim, device="cuda")
    
    def allocate_page(self, request_id):
        if not self.free_pages:
            raise MemoryError("No free pages available")
        page = self.free_pages.pop()
        if request_id not in self.page_table:
            self.page_table[request_id] = []
        self.page_table[request_id].append(page)
        return page
    
    def free_request(self, request_id):
        if request_id in self.page_table:
            self.free_pages.extend(self.page_table[request_id])
            del self.page_table[request_id]
    
    def get_kv(self, request_id, position):
        page_idx = position // self.page_size
        offset = position % self.page_size
        physical_page = self.page_table[request_id][page_idx]
        return self.k_cache[physical_page, offset], self.v_cache[physical_page, offset]
    
    def store_kv(self, request_id, position, k, v):
        page_idx = position // self.page_size
        offset = position % self.page_size
        
        if page_idx >= len(self.page_table.get(request_id, [])):
            self.allocate_page(request_id)
        
        physical_page = self.page_table[request_id][page_idx]
        self.k_cache[physical_page, offset] = k
        self.v_cache[physical_page, offset] = v

RadixAttention

Prefix Caching Benefits

Scenario	Prefix Size	Requests	Memory Saved	Latency Reduction
System prompt sharing	2000 tokens	1000	99%	40% (fewer prefill steps)
Document QA with context	8000 tokens	100	95%	60%
Few-shot prompting	500 tokens	500	80%	20%
Code completion	1000 tokens	2000	90%	30%

KV Cache Quantization

def quantize_kv_cache(k_cache, v_cache, bits=8):
    """Quantize KV cache to lower precision."""
    if bits == 8:
        k_quantized = k_cache.to(torch.int8)
        v_quantized = v_cache.to(torch.int8)
        k_scale = k_cache.abs().max() / 127
        v_scale = v_cache.abs().max() / 127
        return k_quantized, v_quantized, k_scale, v_scale
    elif bits == 4:
        # Group quantization with group size 128
        group_size = 128
        k_groups = k_cache.reshape(-1, group_size)
        k_max = k_groups.abs().max(dim=1, keepdim=True).values
        k_scale = k_max / 7
        k_quantized = (k_groups / k_scale).round().to(torch.int8).reshape(k_cache.shape)
        return k_quantized, k_scale

Grouped-Query Attention (GQA)

Practice Exercises

Memory Analysis: Calculate the KV cache memory for a 70B model serving 100 concurrent requests with 2048 token sequences. How much GPU memory is needed?
PagedAttention Design: Design a page allocation policy that minimizes fragmentation for a workload with 30% short requests (<256 tokens) and 70% long requests (>1024 tokens).
Prefix Caching: If you serve a chatbot with a 1000-token system prompt and 1000 concurrent users, how much KV cache memory is saved by prefix caching?
Quantization Tradeoff: Compare INT8 vs INT4 KV cache quantization in terms of memory savings, latency overhead, and quality degradation on a long-context benchmark.

Key Takeaways

What to Learn Next

-> Flash Attention and Memory Efficiency IO-aware attention algorithms that reduce memory.

-> Continuous Batching for LLMs Dynamic batching for maximum GPU utilization.

-> Speculative Decoding Generating multiple tokens per step.

-> LLM Inference Optimization Broader inference optimization strategies.

-> Long Context and Context Window Handling very long sequences efficiently.

-> Building Production LLM Applications End-to-end production systems.

KV Cache Optimization