efficient inference kv cache

**Efficient Inference (KV Cache, Speculative Decoding, Continuous Batching)** is **the set of systems-level optimizations that reduce the latency, throughput, and cost of serving large language model predictions in production** — transforming LLM deployment from a prohibitively expensive endeavor into a scalable service capable of handling millions of concurrent requests. **The Inference Bottleneck** LLM inference is fundamentally memory-bandwidth-bound during autoregressive decoding: each generated token requires reading the entire model weights from GPU memory, but performs very little computation per byte loaded. For a 70B parameter model in FP16, generating one token reads ~140 GB of weights but performs only ~140 GFLOPS—far below the GPU's compute capacity. The arithmetic intensity (FLOPS/byte) is approximately 1, while modern GPUs offer 100-1000x more compute than memory bandwidth. This makes serving costs proportional to memory bandwidth rather than compute throughput. **KV Cache Mechanism and Optimization** - **Cache purpose**: During autoregressive generation, each new token's attention computation requires key and value vectors from all previous tokens; the KV cache stores these to avoid redundant recomputation - **Memory consumption**: KV cache size = 2 × num_layers × num_heads × head_dim × seq_len × batch_size × dtype_bytes; for LLaMA-70B with 4K context, this is ~2.5 GB per request - **PagedAttention (vLLM)**: Manages KV cache as virtual memory pages, eliminating fragmentation and enabling 2-4x more concurrent requests; pages allocated on-demand and freed when sequences complete - **KV cache compression**: Quantizing KV cache to INT8 or INT4 halves or quarters memory with minimal quality impact; KIVI and Gear achieve 2-bit KV quantization - **Multi-Query/Grouped-Query Attention**: Reduces KV cache size by sharing key-value heads across query heads (8x reduction for MQA, 4x for GQA) - **Sliding window eviction**: Discard oldest KV entries beyond a window size; StreamingLLM maintains initial attention sink tokens plus recent window for infinite-length generation **Speculative Decoding** - **Core idea**: Use a small draft model to generate k candidate tokens quickly, then verify all k tokens in parallel with the large target model in a single forward pass - **Acceptance criterion**: Each draft token is accepted if the target model would have generated it with at least as high probability; rejected tokens are resampled from the corrected distribution - **Speedup**: 2-3x faster inference with zero quality degradation—the output distribution is mathematically identical to the target model alone - **Draft model selection**: The draft model must be significantly faster (7B drafting for 70B target) while sharing vocabulary and producing reasonable approximations - **Self-speculative decoding**: Uses early exit from the target model's own layers as the draft, avoiding the need for a separate draft model - **Medusa**: Adds multiple prediction heads to the target model that predict future tokens in parallel, achieving speculative decoding without a separate draft model **Continuous Batching** - **Problem with static batching**: Naive batching waits until all sequences in a batch finish before starting new requests, wasting GPU cycles on padding for shorter sequences - **Iteration-level scheduling**: Continuous batching (Orca, vLLM) inserts new requests into the batch as soon as existing sequences complete, maximizing GPU utilization - **Preemption**: Lower-priority or longer requests can be preempted (KV cache swapped to CPU) to serve higher-priority incoming requests - **Throughput gains**: Continuous batching achieves 10-20x higher throughput than static batching for variable-length workloads - **Prefill-decode disaggregation**: Separate GPU pools for compute-intensive prefill (processing the prompt) and memory-bound decode (generating tokens), optimizing each phase independently **Model Parallelism for Serving** - **Tensor parallelism**: Split weight matrices across GPUs within a node; all-reduce synchronization per layer adds latency but enables serving models larger than single-GPU memory - **Pipeline parallelism**: Distribute layers across GPUs; micro-batching hides pipeline bubbles; suitable for multi-node serving - **Expert parallelism for MoE**: Route tokens to experts on different GPUs; all-to-all communication overhead managed by high-bandwidth interconnects - **Quantization**: GPTQ, AWQ, and GGUF quantize weights to 4-bit with minimal accuracy loss, halving GPU memory requirements and doubling throughput **Serving Frameworks and Infrastructure** - **vLLM**: PagedAttention-based serving engine with continuous batching, tensor parallelism, and prefix caching; standard for open-source LLM serving - **TensorRT-LLM (NVIDIA)**: Optimized inference engine with INT4/INT8 quantization, in-flight batching, and custom CUDA kernels for maximum GPU utilization - **SGLang**: Compiler-based approach with RadixAttention for automatic KV cache sharing across requests with common prefixes - **Prefix caching**: Reuse KV cache for shared prompt prefixes across requests (system prompts, few-shot examples), reducing first-token latency by 5-10x for repeated prefixes **Efficient inference optimization has reduced LLM serving costs by 10-100x compared to naive implementations, with innovations in memory management, speculative execution, and batching strategies making it economically viable to serve frontier models to billions of users at interactive latencies.** --- **AI Accelerator Architecture — Compute, Memory, and Interconnect.** Modern AI chips are purpose-built for matrix multiplication: a systolic array or tensor core computes thousands of multiply-accumulate (MAC) operations per cycle, fed by a memory hierarchy (registers → SRAM → HBM) connected through a network-on-chip (NoC) that determines whether the compute units starve or stay busy. The single metric that captures this interaction is the roofline model: peak performance (TFLOPS) vs memory bandwidth (TB/s), where the arithmetic intensity of the workload (FLOPs/byte) determines which resource limits throughput. **Tensor Cores — The Matrix Multiply Unit.** NVIDIA tensor cores perform 4$\times$4 matrix multiply-accumulate (D = A$\times$B + C) in a single clock cycle at mixed precision (FP16 inputs, FP32 accumulate). The H100 has 528 tensor cores across 132 SMs, delivering 989 TFLOPS at FP16 or 1,979 TFLOPS at FP8 — a 3$\times$ generational improvement over A100 (312 TFLOPS FP16). Programming tensor cores requires structuring data in tile-friendly layouts (16$\times$16 or 32$\times$8 fragments) via CUDA WMMA or MMA PTX instructions. Utilization typically reaches 60–80% in production training (compute-bound GEMM) but drops to 10–30% during inference decode (memory-bound, limited by KV cache reads). AMD CDNA3 Matrix Cores and Google TPU v5 MXUs provide equivalent functionality at comparable TFLOPS/W. **KV Cache and Inference Efficiency.** During autoregressive LLM inference, each generated token requires reading the full key-value cache of all prior tokens — creating a memory-bandwidth bottleneck where arithmetic intensity drops to 1–5 FLOPs/byte (far left of the roofline). A 70B-parameter model at sequence length 4096 stores 40 GB of KV cache in HBM; generating each token reads 40 GB at 3.35 TB/s = 12 ms latency per token — regardless of compute capacity. Solutions: PagedAttention (vLLM) eliminates KV cache fragmentation; multi-query attention (MQA/GQA) reduces KV size by 8$\times$; speculative decoding verifies 4–8 draft tokens per forward pass, increasing effective throughput 2–4$\times$; continuous batching (Orca) amortizes KV reads across multiple sequences in flight. **Network-on-Chip (NoC) for AI Accelerators.** The NoC connects hundreds of compute tiles (tensor cores, memory controllers, I/O ports) through a mesh, ring, or hierarchical topology — and its bisection bandwidth determines the maximum data rate for all-reduce operations during distributed training. An H100 has a 12$\times$11 crossbar connecting 132 SMs, 6 HBM3 stacks, and 18 NVLink ports. The total internal bandwidth exceeds 30 TB/s. For multi-chip training, NVLink 4.0 provides 900 GB/s chip-to-chip (18 links $\times$ 50 GB/s each) while PCIe 5.0 adds 128 GB/s for host communication. The NoC design determines whether the GPU can keep all tensor cores fed during a 2048-GPU training run where each iteration requires an all-reduce of 1–10 GB of gradients across the fabric. **Mixture of Experts (MoE) — Hardware Implications.** MoE models (GPT-4, Mixtral, Switch Transformer) activate only 2–8 experts per token out of 64–256 total, reducing compute by 10–30$\times$ relative to a dense model of equivalent capacity — but at the cost of massive memory footprint (every expert's weights must reside in HBM) and irregular memory access patterns that stress the NoC and memory controller. A Mixtral 8$\times$7B model has 46.7B total parameters but only 12.9B active per token; the challenge is that expert routing is data-dependent and unpredictable, causing load imbalance across GPU SMs and across nodes in distributed inference. Hardware solutions include expert parallelism (each GPU holds a subset of experts), capacity factors limiting expert overload, and all-to-all communication patterns that require high bisection bandwidth.