LLM Systems Trade Memory, Compute and Communication Across Training and Inference

? afshine-amidi frames the systems question around two constraints: large language models are expensive to store and expensive to run. The decoder layers that make up an LLM consist largely of masked self-attention and feed-forward operations, and both rely heavily on matrix multiplication—work GPUs are designed to perform in parallel. A matrix product can be divided into many dot products, with different GPU cores handling different pieces. But a GPU’s computational capacity does not mean it can hold every model or intermediate result needed to use that capacity.

The lecture uses an NVIDIA H200 as a concrete reference point. Its specifications list 141 GB of GPU memory and 4.8 TB/s of memory bandwidth. Its compute rates vary substantially by number format: the slide gives 1,979 TFLOPS for BF16 Tensor Core operations and 3,958 TFLOPS for FP8. Those numbers describe different limits. Compute throughput measures operations per second; memory capacity determines how much data can be kept on the GPU; bandwidth measures how quickly that data can move.

The capacity constraint is easy to see with a rough calculation. A 100-billion-parameter model stored at two bytes per parameter—16-bit values—requires 200 GB just for its weights. That already exceeds the H200’s listed memory. Training needs more than the weights: it also has to retain activations, gradients and optimizer state. A model that fits for inference may therefore still be too large to train on one device.

The memory burden during training comes from different parts of the learning process. The forward pass computes a loss and produces activations that are needed to calculate gradients. The amount of activation state depends on model size, batch size and context length. The backward pass computes gradients used to update the weights. Optimizers such as Adam also maintain state for those updates. Thus, dividing the parameter count by the GPU’s memory capacity understates what training actually requires.

The second constraint is the amount of work involved. Training processes large volumes of tokens; serving processes user requests. Both raise the same practical questions: how can the necessary work fit on available hardware, and how can it run quickly? Those questions organize the lecture’s treatment of attention, multi-GPU training and inference.

Shervine Amidi begins with the memory traffic inside ordinary self-attention. For a sequence of length n, queries multiplied by keys produce an n × n matrix of attention scores. Each position in a decoder attends to itself and earlier positions. Even a sequence of 1,000 tokens produces a million score entries; longer contexts make the intermediate larger still.

The challenge is not only the number of calculations. Softmax couples every score in a row to every other score in that row: to calculate one normalized attention weight, the denominator must include the exponentials of all the row’s scores. A conventional implementation computes the score matrix, writes it to high-bandwidth memory (HBM), reads it back to apply softmax, writes the normalized matrix, then reads that matrix to multiply it by the values. The large intermediate makes this a costly pattern of memory reads and writes.

The GPU also has a smaller, faster memory, SRAM, close to its on-chip compute units. The lecture’s order-of-magnitude illustration puts HBM around 100 GB and SRAM around 100 MB. HBM has ample capacity for large intermediates but is relatively slow to read and write; SRAM is faster but too small to hold the full attention matrix. FlashAttention is motivated by this mismatch: use SRAM for small blocks of the computation and avoid repeatedly writing the full attention matrix to HBM.

The core idea can be seen in a single row of scores, s, and corresponding values, v. The desired result is the softmax-weighted sum of the values. Written as one fraction, it is a numerator—the sum of each value multiplied by its score’s exponential—divided by a denominator, the sum of those exponentials. Both sums can be accumulated as scores and values arrive. Instead of waiting to form and normalize a full row, the computation updates a running numerator and denominator.

There is a numerical issue: exponentials of large scores can become too large to represent reliably. The running sums can be rescaled without changing their ratio, provided numerator and denominator receive the same scaling. The algorithm tracks the largest score seen so far, m. When a new score arrives, it updates that maximum and rescales the sums already accumulated to match the new scale. In the lecture’s notation, the update uses the new maximum and a factor based on the difference between the old and new maxima. The new score’s contribution is exponentiated only after subtracting the updated maximum. This keeps the exponentials from growing without bound while preserving the exact result.

The rescaling matters because the earlier contributions were calculated using the previous maximum. If the running maximum changes, those contributions must be brought onto the new scale before they can be added to the new score’s contribution. The lecture’s update has two parts: a scaling factor for the running sum, and a new term whose exponent is shifted by the updated maximum. A contribution that becomes extremely small under the new scale can effectively be treated as zero; the more consequential numerical danger is an exponential too large to represent.

With these running quantities, a block of queries, keys and values can be loaded into SRAM, used to update the attention output, and then replaced by the next block. The full n × n score matrix need not be materialized in HBM. The method changes the order and location of the computation, not the attention result: the instructors emphasize that FlashAttention is exact, not an approximation.

Training adds another opportunity to reduce memory traffic. During the backward pass, gradients depend on quantities computed during the forward pass. A conventional approach stores intermediate attention matrices in HBM for later use. FlashAttention can instead recompute those quantities during the backward pass using the tiled operation. That means doing more arithmetic, but it can still reduce runtime by avoiding expensive transfers and storage.

The lecture’s comparison gives the tradeoff in one example. Standard attention is listed at 66.6 GFLOPs and 40.3 GB of HBM reads and writes, with a runtime of 41.7 ms. FlashAttention is listed at 75.2 GFLOPs and 4.4 GB of HBM reads and writes, with a runtime of 7.3 ms. The example shows how more computation can accompany less runtime when memory movement is the limiting cost.

The instructors distinguish two ideas worth retaining. First, tiled computation with incremental softmax avoids storing the full attention matrix and cuts HBM traffic. Second, during training, recomputation can be preferable to storing a large intermediate for later. Later FlashAttention versions, they note, add optimizations for newer GPU generations, but these principles explain the method’s basic motivation.

Even with more efficient attention, one GPU may not have enough memory for the model and the quantities needed during training. The memory pressure comes from several sources: parameters must be stored; the forward pass produces activations used to calculate loss and gradients; the backward pass produces gradients; and optimizers such as Adam maintain additional state for updating the weights. Activation requirements also depend on batch size and sequence length. Since attention processes tokens in relation to other tokens, longer sequences mean more token-specific intermediate values to manage.

Data parallelism divides a batch across devices. Each device performs the forward and backward passes for its portion of the data. In a simple implementation, every GPU has a full copy of the model parameters, gradients and optimizer state. This distributes the data computation but duplicates much of the model-related memory. The devices also need to synchronize gradient updates; data parallelism is therefore not communication-free even in this replicated form.

ZeRO, short for Zero Redundancy Optimizer, reduces duplication by sharding these quantities across GPUs. The lecture describes three stages: ZeRO-1 shards optimizer state; ZeRO-2 shards optimizer state and gradients; ZeRO-3 shards optimizer state, gradients and parameters. Sharding allows each device to hold less, but it is not free. If a device needs parameters stored elsewhere to perform its computation, it must communicate with other devices to retrieve them. The tradeoff is less memory per device in exchange for communication.

The instructors clarify why this remains data parallelism rather than model parallelism. In data parallelism, each device performs the model computation for its own sub-batch. With sharded parameters, a device may need to fetch missing parameters from its peers, but the computation for that sub-batch still takes place on the device handling it. Model parallelism, by contrast, splits the model’s computation itself across devices. As the instructors put it, if a model will not fit otherwise, the practical choice may be between training it more slowly and not being able to train it on that setup.

Tensor parallelism is one form of model parallelism. It divides operations within a layer across devices. For self-attention, different devices can process different attention heads and then combine results. The feed-forward layer can also be divided by splitting its matrices into submatrices. Because these methods operate on different parts of the computation, tensor parallelism can be combined with data parallelism rather than chosen instead of it.

Pipeline parallelism divides the model across layers. One device handles an earlier part of the network, another handles the next part, and so on. Activations pass forward through the devices; the backward pass sends gradient information in the other direction. A difficulty is idle time, often called a “bubble”: while one device waits for its turn to process a forward or backward step, it may do no useful work. Scheduling methods aim to reduce that idle time, though the lecture does not go into them. The diagram shown in the lecture makes the cost visible as gaps in device schedules between forward and backward work.

Context parallelism distributes work along the sequence dimension. Each device handles a group of queries, while keys and values are moved between devices so that each query can attend to the relevant tokens. In the two-device illustration, each device computes its local query, key and value, and performs the portion of attention available from its local information. The devices then exchange key-value information and continue. The same idea extends to groups of tokens in longer sequences: queries are assigned to devices, while key and value blocks circulate so that each group can attend across the sequence.

Expert parallelism applies to models that route tokens to different feed-forward experts. Rather than use every expert for every token, the routing mechanism selects an expert for each token. The experts can be placed on different devices; tokens are dispatched to the device holding the chosen expert, processed there and then returned. The lecture describes this as another dimension of model parallelism, alongside tensor, pipeline and context parallelism. Its example shows an all-to-all dispatch to experts and a corresponding combine step after expert computation.

These approaches can be mixed, but parallelizing more is not automatically better. The instructors identify low utilization, GPU expense and communication costs as constraints. A method that reduces memory use may require transfers; a pipeline may leave devices idle; and GPUs cost money even when they are not doing useful work. The practical goal is not to maximize the number of splits in isolation. It is to fit parameters, activations, gradients and optimizer state while keeping computation and communication efficient.

In the lecture’s summary, distributing those quantities across GPUs addresses fit. Processing different batches in parallel and splitting model computations across devices are ways to increase speed. The distinction is useful: data parallelism divides examples, while model-parallel methods divide computation. Their combination is possible, but its value depends on whether the memory savings and additional parallel work justify the communication and utilization costs.

Shervine Amidi contrasts two requests: a short prompt followed by a very long output, and a very long prompt followed by a short output. The longer prompt can be processed in parallel during prefill. Output generation is autoregressive: each new token depends on the preceding tokens, so the model generates them one at a time. For that reason, a request with a short input and many output tokens can take longer than one with a much longer input and few output tokens.

The request’s total time can be separated into queue time, prefill time and output generation time. Queue time reflects other users and serving or batching strategies. Prefill processes the prompt and computes the key-value information needed to begin generation. Output time accumulates token by token. The lecture names three operational metrics: time to first token (TTFT), time per output token (TPOT) and time to last token (TTLT). TTFT reflects queueing and prefill as well as the start of generation; TPOT captures the pace at which the user sees subsequent tokens. The instructors emphasize the first two as particularly useful infrastructure measures. TTLT is also reported, but it depends on how much the model produces—for example, a response involving more reasoning takes longer to finish.

The lecture treats prefill as relatively fast compared with autoregressive generation, while acknowledging that the comparison depends on context. A long prompt still requires substantial processing. But prompt tokens can be provided together, and in repeated interactions some prompt state may already be cached. By contrast, new output tokens depend on previous outputs, which makes generation sequential.

The distinction also helps explain why input and output tokens can have different costs. The pricing screenshots discussed in the lecture show output priced at roughly five times input for the examples displayed, though the instructors stress that five is not a universal rule. Output generation is costly in part because it is sequential, while prompt tokens can be processed together. The displayed price contrast is an observation about the examples shown, not a fixed relationship that the system must follow.

To reason about the hardware limit, the lecture introduces arithmetic intensity: the amount of computation performed per byte moved. Alongside computational throughput and memory bandwidth, it helps identify whether a workload is limited by compute or by memory movement. The roofline model illustrates the relationship. Below its ridge point, performance is memory-bound; beyond it, it is compute-bound. The instructors place autoregressive decoding in the memory-bound regime: the system loads model weights but uses them to predict one token at a time. Prefill, by contrast, can be compute-bound when many prompt tokens are processed in parallel and memory use is handled efficiently.

The roofline framing connects the hardware specifications to the two different phases. Computational throughput is the number of operations the device can perform per second; memory bandwidth is the rate at which data moves from HBM toward computation. Arithmetic intensity compares the work done with the bytes moved. If a workload does little computation per byte, memory movement can set the pace even when the GPU has substantial compute capacity. At decode time, the weights are used to produce one next token, leaving the workload relatively memory-bound. During prefill, many tokens are processed together, so more computation can be performed against the data moved.

This distinction guides the inference optimizations that follow. Some reduce the amount of state that must be stored for each request; others try to make sequential generation produce more accepted tokens per expensive target-model step.

During decoding, a new token needs to attend to earlier tokens. Recomputing the keys and values for the whole history at every step would repeat work, so inference systems retain those past keys and values in a KV cache and reuse them as generation proceeds. The cache grows as new tokens are generated, supplying the keys and values that later tokens need to attend to.

The lecture also distinguishes this decode-time cache from prompt caching across turns. A provider may reuse previously computed prompt state when a user continues an interaction, and the displayed examples show cached input priced below new input. That pricing illustration concerns provider-side reuse of previously computed state. It is separate from the per-token KV cache used while generating a response, whose cost the instructor says is included in output pricing.

Caching creates a memory-management problem. A serving system may reserve space for the maximum output length allowed for each request. But actual outputs are often shorter than that maximum. The reserved remainder cannot readily be assigned to another request under this naive arrangement, leaving memory that is technically allocated but unusable. Across many requests, these gaps create fragmentation: there can be free space in the aggregate but not in a form the system can use for another allocation.

PagedAttention addresses the allocation problem by dividing the cache into blocks and maintaining a lookup table that maps each request’s keys and values to their locations. A request need not occupy one large, contiguous reservation. Its state can be placed in blocks wherever the lookup structure directs the system, limiting fragmentation to unused space within blocks. Block size is a tradeoff: smaller blocks reduce wasted space but require more metadata and management; larger blocks reduce that overhead but can recreate greater fragmentation. The lecture gives 16 as a typical block size.

The cache can also be made smaller by reducing how many key and value vectors are stored. In standard multi-head attention, each query head has its own key and value heads. Grouped-query attention (GQA) shares key and value heads among groups of query heads. If the number of key/value heads, G, equals the number of query heads, h, this is ordinary multi-head attention. If G is one, it becomes multi-query attention. The lecture presents the latter as an extreme that can leave the model with too little key-value representation, while intermediate group sizes offer a tradeoff between memory savings and expressiveness.

For the example of 24 query heads and two key/value heads, the lecture estimates more than 90% fewer key and value vectors than in the one-per-query-head arrangement. The reduction is not costless: fewer distinct key/value heads can make the model less expressive. The instructors describe selecting G as an empirical choice to be tested through grid searches or ablation studies. In practice, the aim is not simply to minimize the number of heads, but to choose a grouping that improves memory use while retaining adequate model performance.

Multi-head latent attention (MLA) takes a different approach to reducing cache size. Instead of storing many full-sized key and value vectors, it stores a smaller compressed representation and reconstructs the vectors when needed. The lecture describes applying a decomposition into successive projections: first compute a smaller representation, then use learned projections to reproduce the key and value representations. Those expanded key and value vectors are not what the cache stores. The cache retains the compact shared representation instead.

As presented in the lecture, MLA shares a projection across keys, values and heads, so a single per-token representation can stand in for the collection of larger vectors. This reduces the amount of per-token state held in the cache, at the cost of reconstructing the needed representations for attention. The emphasis is on what is stored and how it is recovered, rather than on a particular low-rank interpretation.

Together, PagedAttention, GQA and MLA address distinct sources of cache pressure: fragmented allocation, the number of vectors, and their size. They are not interchangeable operations, but they all make it possible to accommodate more inference state within limited GPU memory.

The central inference bottleneck is that ordinary decoding produces one token at a time. Speculative decoding attempts to make that process more efficient by using a faster draft model to propose several tokens, then asking the target model to verify the proposed sequence together. Verification can process the sequence in parallel, like prefill, rather than generating every candidate sequentially with the target model.

The target’s probability distribution, not the draft’s, must govern the final output. For a proposed token, let D be the draft model’s probability and T the target model’s probability for that token. The proposal is accepted with probability equal to the smaller of 1 and T/D. When the target assigns at least as much probability to the token as the draft does, the proposal is accepted. When the target assigns less, it can still be accepted, but only with probability T/D. This acceptance rule is part of a procedure over the proposed sequence; it should not be read as an assurance that every draft token is accepted or as an independent guarantee that a whole sequence is valid.

If a proposal is rejected, the algorithm samples the next token from a residual distribution proportional to the positive part of the difference between the target and draft distributions: [T − D]+, normalized across tokens. It then exits that speculative sequence and continues from the resampled token. Proposals after the rejection are not used. The residual step matters because it accounts for probability that the target assigns to a token beyond what the draft proposal distribution accounts for.

The lecture works through why the procedure preserves the target distribution. For a given token, its probability of appearing has two contributions: it can be accepted as a draft proposal, or it can be produced after rejection through the residual distribution. The accepted contribution is the minimum of the draft and target probabilities. The residual contributes the positive difference when the target probability exceeds the draft’s. Adding the two gives the target probability in either case. If the draft probability is lower, the residual fills the gap; if it is higher, the accepted contribution is already the target probability. Thus, the accepted-proposal and residual paths together recover the target distribution.

This guarantee concerns the output distribution, not the amount of wasted work. When a rejection occurs, the draft tokens generated after the rejected point do not contribute to the final output. The approach therefore works best when the draft is noticeably faster than the target and its next-token distribution is well aligned with the target’s. Better alignment means fewer rejections and less discarded draft work. Knowledge distillation is one way to train a smaller draft model to match a larger teacher’s predictions.

A further benefit is that verifying proposed tokens also yields the target model’s next-token prediction after them. The lecture describes this as getting the next token “for free” relative to the verification pass: the target computation over the proposed sequence produces an additional next-token result. This does not mean that generating proposals has no cost; the point is that the verification pass evaluates multiple positions in parallel and provides a next prediction as part of that work.

Multi-token prediction (MTP) uses a different arrangement: the same LLM serves as both draft and verifier. Instead of training one output prediction head, the model is trained with k heads for future tokens. Those heads can propose multiple tokens, while the first head serves as the verifier. The common aim of speculative decoding and MTP is to produce or assess several tokens in a batch, reducing reliance on expensive one-token-at-a-time generation.

Across the methods, the lecture returns to two operational questions. To make models fit, avoid materializing the full attention matrix, distribute parameters and training state across GPUs, organize cached vectors into blocks, and store fewer or smaller key-value representations. To make systems fast, reduce HBM traffic, use parallel computation where it helps, and find ways to generate more than one token’s worth of progress per target-model step.

Neither objective stands alone. A memory-saving technique can add communication; a faster method can require more computation; and aggressive parallelism can leave expensive hardware underused. FlashAttention illustrates that extra arithmetic can still accompany lower runtime when it substantially reduces memory movement. ZeRO illustrates that reducing per-device memory can require communication to retrieve sharded state. Speculative decoding illustrates that parallel verification can reduce serial generation work, but rejected proposals waste computation.

The choice of optimization therefore depends on which resource constrains the workload. Training can be limited by the memory needed for parameters, activations, gradients and optimizer state, as well as by the time required to process tokens. Inference has its own constraints: prefill and decoding stress hardware differently, and the KV cache grows with the state retained for active requests. There is no single technique in the lecture that removes all of these limits. Each changes where work happens, what must be stored, or how much work is done before the system can produce its next result.