Stanford CS336 Language Modeling from Scratch | Spring 2025 | Inference

Key Concepts

• Inference: Generating responses from a trained model given prompts.
• Time to First Token (TTFT): Latency experienced by a user before the first token is generated.
• Latency: Speed at which tokens are generated after the first token.
• Throughput: Number of tokens generated per unit of time, useful for batch processing.
• KV Cache: Caching mechanism to store key and value pairs from previous layers to avoid redundant computations during inference.
• Arithmetic Intensity: Ratio of floating-point operations (flops) to bytes transferred, indicating whether a process is compute-bound or memory-bound.
• Group Query Attention (GQA): Attention mechanism that reduces the number of key-value heads to decrease KV cache size.
• Multi-Head Latent Attention (MLA): Projects key and value pairs into a lower-dimensional space to reduce KV cache size.
• Cross-Layer Attention (CLA): Shares key-value projections across different layers of the transformer.
• Local Attention: Attending only to the past K tokens, limiting the KV cache size.
• State Space Models (SSM): Alternative architectures to transformers, such as Mamba, that aim to model long-range dependencies more efficiently.
• Diffusion Models: Generate all tokens in parallel and iteratively refine the generation.
• Quantization: Reducing the precision of numerical representations (e.g., from FP16 to INT8) to decrease memory usage and increase speed.
• Model Pruning: Removing unimportant parts of a model (layers, heads, or hidden dimensions) to reduce its size and computational cost.
• Speculative Decoding/Sampling: Using a smaller, faster "draft model" to generate candidate tokens, then verifying them with a larger, more accurate "target model."
• Continuous Batching: Adding new requests to the current batch as they arrive, without waiting for the batch to be full.
• Selective Batching: Batching together tensors of different sizes by flattening them for MLP layers.
• Page Attention: Dividing the KV cache into a sequence of contiguous blocks and allocating them wherever free space is available.

Inference Workload and Implications

• Inference is crucial for various applications, including chatbots, code completion, batch data processing, model evaluation, and reinforcement learning.
• Efficiency is paramount in inference because it is repeated multiple times, unlike training, which is a one-time cost.
• OpenAI generates 100 billion words a day, and Cursor generates a billion lines of accepted code daily, highlighting the scale of inference.
• Good inference is measured by TTFT, latency, and throughput.
• Unlike training, inference in transformers is sequential, making it harder to utilize available compute and leading to memory limitations.
• Providers serving both closed and open-weight models invest heavily in optimizing inference.

Transformer Math Review

• The lecture reviews transformer math, including the computation graph, attention, and MLP layers.
• Key notations:
  • B: Number of sequences in a batch.
  • L: Number of layers.
  • T: Sequence length (number of tokens generated).
  • S: Sequence length (number of tokens in the prompt).
  • V: Vocabulary size.
  • D: Dimensionality of the model.
  • F: MLP hidden dimension (usually 4D).
  • H: Attention head dimension.
  • N: Number of query heads.
  • K: Number of key-value heads.
  • G: Number of groups in GQA.
• Flops required for a forward pass are 6 * B * T * (number of parameters) + O(T^2) for attention.

Arithmetic Intensity and Memory Limitations

• Arithmetic intensity is the ratio of flops to bytes transferred.
• For a matrix multiplication X (B x D) * W (D x F), the arithmetic intensity is B.
• Accelerator intensity is the ratio of accelerator flops per second to memory bandwidth.
• If computational intensity > accelerator intensity, the process is compute-limited. Otherwise, it's memory-limited.
• When B = 1 (matrix-vector product), the arithmetic intensity is 1, indicating a memory-limited scenario.
• In inference, generating tokens sequentially results in low arithmetic intensity, making it memory-limited.

Inference Process and KV Cache

• The naive inference approach involves feeding the prompt through the transformer, sampling the next token, attaching it to the prompt, and repeating.
• This is inefficient due to redundant work in encoding the prefix.
• The solution is to cache the key and value pairs in HBM.
• The KV cache stores, for every sequence, token, layer, and head, an H-dimensional vector.
• Inference has two stages:
  • Prefill: Encoding the prompt (like training, parallelizable, compute-limited).
  • Generation: Generating response tokens one by one (sequential, memory-limited).

Flops and Memory IO for Transformer Inference

• The lecture breaks down the flops and memory IO for MLP and attention layers during prefill and generation.
• For MLP layers, the arithmetic intensity is B * T.
  • Prefill: BT can be large enough to be compute-limited.
  • Generation: T = 1, so B (number of concurrent requests) needs to be large for efficiency.
• For attention layers, the arithmetic intensity is S * T / (S + T).
  • Prefill: T = S, so intensity is order S (good).
  • Generation: T = 1, so intensity is approximately 1 (bad).
• Attention is always memory-limited because each sequence has its own unique KV cache.

Throughput and Latency Analysis

• Assumptions: Communication and compute can be perfectly overlapped.
• Example: Llama 2 13B on H100.
• Total memory = (batch size * cache per sequence) + parameter size.
• Latency = memory / memory bandwidth.
• Throughput = (B * tokens) / latency.
• Smaller batch sizes yield better latency, while larger batch sizes yield better throughput.
• There are limits to how much throughput can be increased due to memory constraints.
• Model parallelism for inference is simple: launch M copies of the model, increasing throughput by M.

Techniques to Reduce KV Cache Size

• Group Query Attention (GQA): Reduces the number of key-value heads.
  • Llama 3 uses GQA.
  • Improves latency and throughput with negligible accuracy loss.
• Multi-Head Latent Attention (MLA): Projects key and value pairs into a lower-dimensional space.
  • Used in DeepSeek v2.
  • Reduces KV cache size.
• Cross-Layer Attention (CLA): Shares key-value projections across layers.
  • Improves the Pareto frontier of accuracy and KV cache size.
• Local Attention: Attends only to the past K tokens.
  • KV cache size remains constant.
  • Can be interleaved with full global attention layers.

Radical Architectural Changes for Faster Inference

• State Space Models (SSM):
  • Alternative to transformers, drawing from signal processing and control theory.
  • Examples: S4, Hyena, H3, Mamba.
  • Mamba has been scaled up to 52B parameters.
  • Linear attention is a revival of an older idea, scaling up to 456 billion parameters.
• Diffusion Models:
  • Generate all tokens in parallel and iteratively refine the generation.
  • Session Labs has produced interesting models.
  • Potentially much faster than transformers in terms of tokens per second.

Quantization and Model Pruning

• Quantization: Reduces the precision of numbers.
  • FP32 (training), BF16 (default for inference), FP8, INT8, INT4.
  • Post-training quantization is common.
  • Element-wise INT8 handles outliers separately in full 16-bit precision.
  • Activation-aware quantization quantizes weights based on activations.
• Model Pruning: Removes parts of an expensive model.
  • Identify important layers/heads/hidden dimensions.
  • Remove unimportant ones.
  • Distill the original model into the pruned model.

Speculative Decoding/Sampling

• Exploits the symmetry between checking (prefill) and generation.
• Uses a cheap draft model to generate candidate tokens.
• Evaluates those tokens with a target model.
• Accepts tokens if they look good.
• Guaranteed to get an exact sample from the target model.
• Medusa and Eagle are ways to improve the draft model.

Serving Live Traffic

• Live traffic requests come at different times, finish at different times, and have different lengths.
• Continuous Batching: Add new requests to the current batch as they arrive.
• Selective Batching: Batch together tensors of different sizes by flattening them for MLP layers.
• Page Attention: Divides the KV cache into a sequence of contiguous blocks and allocates them wherever free space is available.
  • Copy-on-write for sharing prefixes.

Synthesis/Conclusion

Inference is a critical aspect of deploying language models, with unique challenges compared to training. The lecture covers a wide range of techniques to improve inference efficiency, from architectural modifications like GQA and Mamba to systems-level optimizations like page attention. The key takeaway is that inference is not just about optimizing a particular model but about delivering good accuracy within a given resource budget, often requiring radical changes in architecture and a deep understanding of memory limitations.