Key Moments
Stanford CME295 Transformers & LLMs | Autumn 2026 | Lecture 5 - LLM Systems
Stanford OnlineKey Moments
Training massive LLMs requires multiple GPUs, but techniques like FlashAttention optimize single-GPU performance by reducing slow memory reads/writes. However, distributing models across GPUs for speed and size introduces communication overheads, making efficiency a complex balancing act.
Key Insights
Flash Attention achieves exact self-attention calculations by optimizing reads and writes between High Bandwidth Memory (HBM) and on-chip SRAM, avoiding the need to store the full N x N attention matrix.
To train LLMs exceeding single-GPU memory capacity, techniques like data parallelism (splitting data across GPUs) and model parallelism (splitting model computations, e.g., tensor, pipeline, expert) are crucial.
Zero Redundancy Optimizer (ZeRO) is a data parallelism technique that shards model parameters, gradients, and optimizer states across GPUs to reduce memory footprint per device.
Inference, especially for generative tasks, involves a 'prefill' phase for initial input processing and a slower 'autoregressive decoding' phase for generating subsequent tokens, leading to different performance bottlenecks.
PagedAttention manages KV cache memory more efficiently by using fixed-size blocks and metadata, reducing fragmentation and improving throughput compared to naive allocation.
Speculative decoding uses a smaller 'draft' model to propose multiple future tokens, which are then verified by the larger target model, speeding up generation by leveraging parallel input processing.
The need for efficient LLM systems
Large Language Models (LLMs) are characterized by their immense size, often in the billions or even trillions of parameters, and are trained on vast datasets, sometimes tens of trillions of tokens. This scale necessitates specialized hardware and infrastructure for both training and inference. The core computational operations in LLMs, particularly within the Transformer's decoder block (masked multi-head attention and feed-forward networks), are matrix multiplications. GPUs excel at these parallel computations due to their numerous cores and specialized hardware. However, a significant challenge is the limited memory capacity of even high-end GPUs (e.g., 141 GB for the NVIDIA H200). Storing a 100 billion parameter model with 16-bit precision alone requires 200 GB, exceeding single-GPU memory. Training further exacerbates this with the need to store activations, gradients, and optimizer states. Therefore, the primary goals are to make LLMs fit within available hardware and to execute them quickly.
Flash Attention: Optimizing attention on a single GPU
Flash Attention is a technique designed to optimize the self-attention mechanism, a computationally intensive part of Transformers. The bottleneck in standard attention lies in repeatedly reading from and writing to the High Bandwidth Memory (HBM), which is slower than on-chip SRAM. Specifically, calculating the softmax over a large attention matrix (N x N, where N can be tens of thousands) requires loading the entire matrix into HBM, computing the softmax, writing it back, and then multiplying by the values. Flash Attention circumvents this by performing the attention computation iteratively. It leverages the faster SRAM by loading small blocks of query (Q), key (K), and value (V) matrices. The core idea is to compute the attention output iteratively without materializing the full N x N attention matrix. This is achieved by calculating intermediate numerator and denominator terms for the softmax, which can be updated incrementally. To prevent numerical instability from large exponents, scaling factors are applied, effectively re-normalizing the computation. This method reduces HBM reads and writes, making attention computation significantly faster and more memory-efficient without approximations. Notably, during the backward pass for training, Flash Attention recomputes these intermediate quantities instead of storing them in HBM, further accelerating the process and increasing FLOP utilization.
Distributed training: Parallelism strategies for multi-GPU setups
When a single GPU's memory is insufficient, distributed training across multiple GPUs becomes necessary. Several parallelism strategies are employed. Data Parallelism (DP) involves replicating the model across GPUs and partitioning the data, processing different subsets concurrently. Techniques like the Zero Redundancy Optimizer (ZeRO) optimize DP by sharding parameters, gradients, and optimizer states across GPUs, reducing memory redundancy. Model Parallelism splits the model's computations themselves. Tensor Parallelism divides operations like matrix multiplications across GPUs, often by splitting matrices along rows or columns. Pipeline Parallelism partitions the model's layers sequentially across GPUs, with data flowing through them in a pipeline fashion, though this can introduce 'bubbles' of idle time. Context Parallelism handles long sequences by distributing query computations across devices and communicating keys and values. Expert Parallelism, relevant for Mixture-of-Experts (MoE) models, routes tokens to specific 'expert' subnetworks residing on different GPUs, activating only a subset of parameters per token.
Inference optimization: Speeding up token generation
Inference presents different challenges, particularly in generative tasks where models produce output token by token. The process typically involves an initial 'prefill' phase, where the input prompt is processed in parallel, followed by an autoregressive decoding phase, generating one token at a time. The autoregressive nature makes decoding significantly slower than prefill. Key metrics for inference include Time To First Token (TTFT), Time Per Output Token (TPOT), and Time To Last Token (TTLT). Efficient inference requires balancing memory bandwidth and computational throughput. The 'Roofline model' illustrates this trade-off, showing performance limits imposed by either memory or computation. For LLM inference, which is often memory-bound due to loading model weights and the KV cache, optimizing memory access is critical.
Efficient KV cache management and speculative decoding
The Key-Value (KV) cache stores intermediate attention computations for previously processed tokens, significantly speeding up autoregressive decoding. PagedAttention optimizes KV cache management by dividing the cache into fixed-size blocks, similar to virtual memory paging. This reduces fragmentation and improves memory utilization compared to naive block allocation. Spectulative decoding further accelerates generation by using a smaller, faster 'draft' model to predict multiple future tokens. These candidate tokens are then verified by the main 'target' model. If the target model's probability for a token aligns with or exceeds the draft model's prediction, it's accepted. Otherwise, it's rejected, and the process continues, effectively allowing the model to 'look ahead' and generate tokens much faster by parallelizing the verification step.
Model compression and advanced attention mechanisms
Techniques like Grouped-Query Attention (GQA) reduce the memory footprint of the attention mechanism by decreasing the number of key and value projection heads (G) relative to query heads (H), offering a trade-off between memory savings and model expressiveness. Multi-query attention (MQA) is an extreme case where G=1. Another approach is Multi-vector Latent Attention (MLA), which compresses multiple query, key, and value vectors into a single, smaller latent vector per token, further reducing memory requirements. These methods aim to make models more efficient without significantly sacrificing performance, often involving careful hyperparameter tuning to find the optimal balance.
Mentioned in This Episode
●Products
●Software & Apps
●Organizations
●Concepts
Common Questions
GPUs are critical because LLMs heavily rely on matrix multiplication, a task at which GPUs excel due to their large number of parallel processing cores and specialized hardware, allowing for efficient computation of billions of parameters.
Topics
Mentioned in this video
A parameter-efficient fine-tuning technique that decomposes matrix multiplication, allowing for smaller, trainable representations to be used.
A technique where a smaller model (draft model) is trained to mimic the predictions of a larger, more powerful model (teacher model), helping to align their token distributions for speculative decoding.
An advanced inference technique that accelerates token generation by using a smaller 'draft' model to predict multiple subsequent tokens, which are then verified in parallel by the larger 'target' model.
An approach where the main LLM itself is used as both the draft and target model, training K heads to predict future tokens simultaneously, leveraging the first head as the verifier.
The foundational neural network architecture upon which Large Language Models (LLMs) are based, specifically utilizing its decoder part.
Models that scale the Transformer decoder, possessing billions or trillions of parameters and trained on vast amounts of data (trillions of tokens).
An optimization technique for self-attention on a single GPU that reduces HBM read/write operations by performing computations iteratively and combining operations, leading to significant speedups without approximations.
A theoretical model that plots computational throughput against computational intensity, showing the maximum achievable performance limited by either memory bandwidth or computational capacity.
A technique in LLM inference where previously computed keys and values for attention are stored and reused, significantly reducing computational cost for subsequent tokens.
A memory management technique that allocates KV cache memory in fixed-size blocks (pages) to prevent fragmentation and efficiently manage dynamic output lengths, improving GPU utilization.
A technique that transforms multiple query, key, and value vectors into a single, smaller latent vector by sharing projection matrices, dramatically reducing memory footprint per token.
An attention mechanism that groups multiple queries to share a single key and value projection, reducing the number of key/value matrices and thus memory consumption, while maintaining performance.
Hardware component specializing in matrix multiplication, crucial for training and operating LLMs due to their parallel processing capabilities and numerous cores.
An LLM model mentioned as utilizing GQA with G=2 and H=24 for memory efficiency.
A popular optimization algorithm used in training neural networks, which requires tracking several quantities that consume significant memory.
More from Stanford Online
View all 167 summaries
72 minStanford AA274A Principles of Robotic Autonomy | Autumn 2019 | Decision Making Under Uncertainty
70 minStanford Robotics Seminar ENGR319 | Autumn 2026 | Skydio: Autonomy for a Family of Flying Robots
62 minStanford Robotics Seminar ENGR319 | Autumn 2026 | Bridging Rigorous Methods with Machine Learning
77 minStanford AA274A Principles of Robotic Autonomy | Autumn 2019 | Multi-Sensor Fusion
Ask anything from this episode.
Save it, chat with it, and connect it to Claude or ChatGPT. Get cited answers from the actual content — and build your own knowledge base of every podcast and video you care about.
Get Started Free