What Is a GPU Doing With My Tensors?

Many LLM researchers can explain a forward pass from the model side. Hidden states flow through attention and MLPs, linear layers transform each token representation, and the final layer produces vocabulary logits. The GPU often remains just outside that picture. We know how much memory a card has, how many peak FLOPs it advertises, and how many GPUs a training run used. It is harder to say where the tensors actually live or how the computation physically unfolds.

That gap rarely matters when everything is fast enough. The framework dispatches model operations to the GPU and manages most of the memory on our behalf. Once performance becomes a problem, HBM, memory bandwidth, Tensor Cores, kernel fusion, NVLink, all-reduce, compute-bound, and communication-bound suddenly appear together. Each term explains one part of the behavior. What is usually missing is a map that connects them.

This post builds that map. We will start with where weights, activations, and the KV cache live inside one GPU; follow a tensor operation down to the hardware that executes it; examine how inference and training change the lifetimes of those tensors; and finally see what must move when a model spans multiple GPUs. HBM, Tensor Cores, kernel fusion, and all-reduce will eventually fit into the same four-part framework: capacity, bandwidth, compute, and communication.

Inside a GPU

At a useful level of abstraction, a GPU running an LLM has two major parts: a large collection of parallel compute resources and a memory system that continuously feeds them. The compute resources determine how much arithmetic can happen at once. The memory system determines whether the required values arrive in time. Much of GPU design is an attempt to keep those two sides in balance.

The model data first has to reach the GPU. Checkpoints normally live as files on an SSD. Loading moves their contents through CPU RAM and then into GPU memory. Once execution begins, weights, activations, and caches move among several locations inside the GPU. Figure 1 places both scales in one scene: warehouses outside the building handle long-term storage and transfer, while the kitchen inside the GPU keeps many parallel workstations supplied.

The largest memory routinely used on a GPU is high-bandwidth memory, or HBM. When someone describes a GPU as having 80 GB or 192 GB of memory, they are usually referring to the capacity of this layer. During inference, HBM holds model weights, temporary layer activations, and the KV cache that attention keeps for earlier tokens. Training also needs parameter gradients and optimizer state. A checkpoint primarily records weights, so its file size captures only part of the runtime footprint.

HBM is large; storage next to the compute units is much smaller. GPU compute is distributed across many Streaming Multiprocessors, abbreviated SMs. Each SM contains execution units, registers, and on-chip resources used by shared memory and the L1 cache. A larger L2 cache is shared across the SMs on the GPU. Inside a modern SM, Tensor Cores accelerate the matrix multiply-accumulate operations that dominate Transformer workloads.

A concrete program submitted to the GPU is called a kernel. A kernel may implement matrix multiplication, normalization, or an elementwise operation. It may also fuse several operations into one program. The GPU expands a kernel into a large amount of parallel work and assigns that work to available SMs. The SMs execute it while the memory hierarchy supplies the data.

This hierarchy has a crucial sense of direction. Storage becomes smaller and easier to access as it gets closer to the execution units. HBM can hold complete, large tensors. L2, shared memory, L1, and registers progressively narrow the view to the current working set, the small portion of data needed for the calculation at hand. During execution, the active pieces of a large tensor are brought closer to compute, and recently used data is kept nearby when it can be reused. The exact path depends on both the kernel and the hardware caches; data does not necessarily stop at every level in the diagram.

Inference and training create different footprints on the same hardware. Inference mainly retains the weights and a KV cache for each active request, while most intermediate activations live only briefly. A training forward pass must preserve some activations for the later backward pass. Backward produces gradients, and the optimizer may keep additional state for every parameter. A model that fits comfortably for inference can therefore exceed the same GPU configuration during training.

The kitchen analogy maps onto the hardware like this:

A GPU is constantly coordinating storage, movement, and computation. Model state travels from SSD and CPU RAM into HBM. Kernels then bring the active working set close to many SMs and perform the computation. If the required tensors exceed HBM capacity, the process runs out of memory. If execution spends most of its time waiting for data, the operation is memory-bound. If arithmetic dominates, it is compute-bound. Once a model spans multiple GPUs, a fourth question appears: how much data has to cross device boundaries, and how long does compute wait for that communication?

How Data Reaches Compute

The complete weights and intermediate results of a running model may occupy tens or hundreds of gigabytes. An execution unit needs only a tiny fraction of them at any instant. Keeping everything in large, distant memory would leave compute waiting. Keeping everything beside the execution units is physically impossible because the chip has too little space.

GPUs resolve this tension with several memory levels. Farther levels have more capacity; closer levels are smaller and better suited to frequent access. This arrangement is the memory hierarchy. The small subset required by the current operation is its working set. Using the same data several times before evicting it is reuse. The hierarchy gives large data somewhere to live while keeping the active pieces close to compute.

HBM: The Model’s Runtime Home

When PyTorch reports that a tensor is on cuda:0, the tensor usually occupies the first GPU’s device memory, the large memory directly accessible to that GPU. Data center accelerators typically implement device memory with high-bandwidth memory, or HBM. Model weights remain there for long periods. Activations created while evaluating a layer also appear there. Inference retains the keys and values computed for earlier tokens as a KV cache. Training adds parameter gradients and the history maintained by the optimizer. The smaller memories closer to compute cannot hold these objects for long.

Two HBM numbers describe different constraints. Memory capacity is the number of bytes that can coexist. Memory bandwidth is the maximum number of bytes per second that can travel between HBM and the GPU chip. Insufficient capacity produces an out-of-memory error. Insufficient bandwidth leaves the execution units waiting even though the data technically fits. NVIDIA specifies the H100 SXM with 80 GB of HBM and 3.35 TB/s of peak memory bandwidth (NVIDIA, 2026). Even at that peak rate, streaming through 80 GB once would take roughly 24 ms.

HBM gets its bandwidth from its physical design. Manufacturers vertically stack memory dies and connect several stacks to the accelerator through a very wide interface (SK Group, 2024). Think of that interface as a road. A single lane can only become so fast. HBM lays down many lanes in parallel so that far more data can depart at once. This improves aggregate throughput; an individual access still takes time.

Intermediate results retained by a training forward pass for backward are called saved activations. Temporary space requested by kernels or communication libraries is a workspace. Combining them with persistent model state gives two useful, approximate inventories:

\[\text{training HBM} \approx \text{weights} + \text{saved activations} + \text{gradients} + \text{optimizer states} + \text{temporary workspace}.\] \[\text{inference HBM} \approx \text{weights} + \text{KV caches of active requests} + \text{temporary workspace}.\]

Training must fit weights, saved activations, gradients, and optimizer state at the same time. Forward, backward, and parameter updates also read and write this data repeatedly. Distributing model state across more GPUs brings more HBM into the system. A B200 provides 180 GB of HBM and up to 8 TB/s of bandwidth; an eight-GPU server therefore has 1.44 TB. The 72 Blackwell GPUs in a GB200 NVL72 provide 13.4 TB in total (NVIDIA, 2026; NVIDIA, 2026). Large training clusters concentrate an enormous demand for HBM.

Inference keeps fewer kinds of persistent state, though that state may remain active for much longer. Every deployed copy of the full model, or model replica, keeps another set of weights resident in HBM. Each active request also retains a KV cache that grows with context length (Luo & Shen, 2026). Generating every new token reads the weights and the existing cache again. Long contexts, agent workloads, and long-reasoning models such as o1 extend that process, causing the same weights to be read for more rounds and request state to remain live for longer (Gandhi et al., 2026; Micron Technology, 2026). Training and inference both demand HBM capacity and bandwidth, with different tensor lifetimes and access patterns.

This is also why SK hynix now appears so often in conversations that ostensibly begin with GPUs. HBM production involves memory-die fabrication, stacking, interconnection, packaging, and testing. Expansion is constrained by both wafer supply and advanced packaging. By December 2025, Micron had completed price and volume agreements for its entire 2026 HBM supply. SK hynix later reported demand beyond its supply capability and continued to raise investment (Micron Technology, 2025; SK hynix, 2026).

AI’s appetite for HBM has driven a striking rise in SK hynix’s share price and earnings. From August 10, 2021 to August 10, 2026, the stock rose from ₩112,500 to ₩1.42 million, nearly thirteenfold. The financial results tell the same story: annual revenue grew from ₩32.8 trillion in 2023 to ₩97.1 trillion in 2025, while the company moved from a ₩7.7 trillion operating loss to a ₩47.2 trillion operating profit. Citing Counterpoint Research, SK hynix reported 62% of HBM shipments in the second quarter of 2025 and 57% of HBM revenue in the third quarter (SK hynix, 2026).

HBM is valuable because several pressures converge: each accelerator needs more capacity and bandwidth, deployments keep expanding, and supply cannot scale instantly. Back inside one GPU, reaching HBM is only the beginning. Once data enters the chip, compute would still prefer to avoid returning to HBM for every reuse. The smaller L2 cache takes over the next leg of the journey.

L2 Cache: Avoiding Another HBM Trip

HBM is fast by memory standards, yet fetching from it remains expensive compared with using data already on the chip. Computation often reuses the same value within a short interval. Repeating the HBM transfer each time would spend scarce bandwidth on duplicate movement. A GPU therefore includes a much smaller L2 cache that automatically retains recently accessed data. If that data is needed again soon, the GPU may use the nearby copy.

L2 is shared by every SM on the GPU. If data fetched by one SM remains in L2, a later access from another SM may avoid another HBM transfer. When the requested data is already present, the access is a cache hit. Hardware generally decides what remains cached. PyTorch has no operation resembling tensor.to("l2").

L2 is tiny beside HBM. An H100 has 80 GB of HBM and only 50 MB of L2 cache (Andersch et al., 2022). It cannot hold a model. It can retain only a small active portion, such as tiles from a large tensor that a kernel is currently processing or an intermediate result that will be consumed again soon. Every avoided HBM read leaves more bandwidth for genuinely new data. That is the practical value of reuse.

Shared Memory and L1: One SM’s Workspace

L2 serves the entire GPU. Closer to execution, each SM needs a local working area. A kernel asks many threads to handle small pieces of the work. Threads that need to cooperate are grouped into a thread block, and threads in the same block can share data near their SM.

An SM has both an L1 cache and shared memory. L1 resembles L2 in that hardware automatically keeps recently used data, though its scope is closer to one SM. Shared memory is an explicitly managed workspace that all threads in a block can access. On modern NVIDIA GPUs, L1 and shared memory may draw from some of the same physical on-chip resources, while exposing different programming models (NVIDIA, 2026).

Suppose a kernel processes a tensor far larger than shared memory. The full tensor stays in HBM, and the kernel divides the work into smaller pieces so that each block handles one at a time. A piece selected for local computation is a tile. Once a tile reaches the SM, threads in the block can repeatedly read it from shared memory, compute their portions, and then move on to the next tile. A linear layer reuses input and weight tiles. Attention, normalization, and other operations choose working sets according to their own dependencies. Tile size and retention policy determine how often data can be reused and how many HBM accesses are required.

Shared memory lives for the duration of the current block. When the block finishes, the space becomes available to later work. Its capacity also limits parallelism. If one block consumes a large amount of shared memory, fewer blocks can reside on the same SM at once. Keeping data near compute reduces movement while consuming scarce on-chip space, so kernel design must trade one resource against the other.

Registers: Values in Use

Registers sit at the execution end of this path. Every thread has private registers that hold the few values used by its current instructions, such as input numbers, array indices, or an unfinished accumulation. The compiler assigns registers when it turns kernel code into GPU instructions. Programmers do not usually manage them as they would tensors.

Consider a small segment of a vector dot product. The full vectors may live in HBM, the current tiles may pass through L2 or shared memory, and a partial sum maintained by one thread can remain in a register. The thread repeatedly updates that sum as it performs multiplication and addition. Only after the local computation finishes does the result move back to larger memory. Matrix multiplication, normalization, and elementwise operations follow the same general division of labor.

Register capacity is also finite. As each thread requires more registers, an SM can often keep fewer threads resident. If the compiler cannot find enough registers, some temporary values fall into slower, more distant memory, a condition called register spilling. On-chip storage saves data movement and simultaneously constrains how much work the GPU can keep in flight.

Three questions summarize how an operation uses the memory hierarchy: how many bytes must coexist, how many bytes must cross the HBM boundary, and how many times nearby data can be reused. The full runtime state lives primarily in HBM. L2 automatically retains data that may be reused across the GPU. L1 and shared memory serve the tiles currently processed by one SM. Registers hold the values currently used by individual threads. A real access path may skip levels, and a cache may evict data before the next use.

How a GPU Executes Tensor Operations

Once weights and activations reach the GPU, a line of linear, softmax, or elementwise addition still has to become a large amount of concurrent work on the chip. At the model level, an operation specifies the mathematical result a tensor should have. A kernel is the program the GPU actually runs. It specifies how to divide the work, where to read data, which instructions to execute, and where to write the result. One operation may invoke several kernels, while a compiler may fuse several adjacent operations into one kernel. The framework running on the CPU initiates a kernel launch, and the GPU executes the selected program.

It helps to keep three different hierarchies separate:

All the work created by one launch is a grid. The grid is divided into independently schedulable thread blocks, and each block contains many threads. A thread represents the smallest logical piece of work.

The GPU assigns a complete block to a Streaming Multiprocessor, or SM, with enough available resources. A resident block stays on one SM for its lifetime. One SM can host several blocks at the same time.

Once a block reaches an SM, consecutive groups of 32 threads are assembled into warps. A warp is the unit whose instructions the SM actually advances. Each position in the warp is a lane corresponding to one logical thread. The SM scheduler selects a ready warp and sends its instruction to hardware that performs memory access or arithmetic, such as a Tensor Core. These pieces of hardware are collectively called execution units. The warp therefore connects logical threads to physical hardware. Figure 3 places these relationships in one scheduling snapshot.

A Linear Layer on the GPU

Take a linear layer, one of the most common operations in an LLM. Fold batch and sequence positions into $N$ rows. Let the input be $X\in\mathbb{R}^{N\times d_{\text{in}}}$ and the weight be $W\in\mathbb{R}^{d_{\text{out}}\times d_{\text{in}}}$. The output is

\[Y = XW^\top, \qquad Y\in\mathbb{R}^{N\times d_{\text{out}}}.\]

The $6\times6$ example below shows why this work can be divided. Each matrix is split into a $3\times3$ grid of tiles. The animation begins with the upper-left $2\times2$ output tile in $Y$. That tile depends on one row of tiles from $X$ and one column from $W^\top$. Moving along the inner dimension, three input-tile pairs multiply to produce three $2\times2$ partial products. Those products accumulate into the same output tile. Once it is complete, the process continues through the other eight tiles in $Y$.

A real linear layer may have feature dimensions in the thousands or tens of thousands. The same process expands into many output tiles, each accumulating many partial results. Whether an operation can be tiled depends on its mathematical dependencies. The kernel identifies pieces that can progress independently and maps them onto blocks, warps, and threads. This is one of the central ways a GPU exposes massive parallelism.

The framework chooses a matrix-multiplication kernel. Its launch creates many blocks, often with each block responsible for a tile of $Y$, and the GPU schedules those blocks onto SMs with available resources. Each SM advances the threads in warp-sized groups and issues suitable matrix instructions to Tensor Cores. Grids and blocks describe how the work is partitioned; SMs and Tensor Cores provide the hardware that carries it out.

Meanwhile, the complete $X$, $W$, and $Y$ remain in HBM. To compute one output tile, the kernel fetches the corresponding input and weight tiles in stages. Recent accesses may hit L2. Tiles reused within the current block enter shared memory. Smaller fragments are fed to Tensor Cores, while unfinished partial sums remain in registers whenever possible. The output tile returns to HBM only after the kernel has traversed $d_{\text{in}}$. The grid partitions the output, the SMs and Tensor Cores execute the tiles, and the memory hierarchy supplies inputs while preserving partial sums.

This arrangement extracts more computation from every transfer. One tile of $X$ can combine with several tiles of $W$, while one tile of $W$ can serve multiple rows of $X$. Partial sums remain near the execution units until the result is complete. High-performance kernels choose finer tile shapes according to the tensor shape and dtype, and they overlap data movement with computation when possible (NVIDIA, 2026).

One Kernel, Millions of Threads

Elementwise addition gives the hierarchy a concrete scale. Suppose two vectors of length 1,048,576 satisfy $z_i=x_i+y_i$. Every output position is independent, so a kernel can create 1,048,576 logical threads, each using its identifier to find one $i$. With 256 threads per block, the launch contains 4,096 blocks, and every block contains eight warps. The GPU schedules those blocks dynamically according to the resources currently available on each SM. All 4,096 blocks do not need to reside on the chip at once.

Hiding Memory Latency

After a warp issues a memory request, the data takes time to return. A warp scheduler can advance another ready warp during that wait, keeping the execution units occupied. Covering memory latency with other work is called latency hiding, and it is a major reason GPUs need a large pool of parallel tasks.

Two effects can reduce that parallelism. With a data-dependent branch, lanes in the same warp may have to execute different paths in separate groups. This is branch divergence. Registers are allocated per thread and shared memory per block. If one block consumes too many of those resources, fewer warps can remain resident on the SM. The ratio of active warps to the hardware limit is called occupancy. It indicates whether the scheduler has enough work available to switch between, though a higher value does not automatically imply a faster kernel (NVIDIA, 2026).

Once the scheduler selects a warp, the instruction goes to an appropriate execution unit. Load/store units handle memory access. General arithmetic units perform operations such as addition, multiplication, and address calculation. Tensor Cores handle suitable matrix multiply-accumulate instructions. Elementwise operations, normalization, attention, and matrix multiplication all use this scheduling framework, with different execution units, reuse patterns, and HBM traffic.

Kernel boundaries also change data movement. Suppose a model performs an elementwise addition followed by an activation. Two separate kernels will often write the intermediate tensor to HBM and read it back. If a compiler fuses them into one kernel, the intermediate values may remain in registers, saving one launch and a round trip through HBM. This is kernel fusion (Turnansky, 2026).

The work decomposition determines how an operation runs. Tensor lifetime determines when its memory overlaps with everything else and creates the HBM peak.

Runtime State in HBM

A byte’s position in the memory hierarchy determines its access cost. A tensor’s lifetime determines when it begins occupying HBM and when that space can be reused. Peak HBM usage emerges from the overlap among these lifetimes.

Loading a Checkpoint

A checkpoint is a collection of persistent files stored on SSD. It usually contains model parameters. It may also include fixed runtime data, scales associated with quantized weights, and, for a training checkpoint, optimizer state and the current step. During loading, the framework creates runtime tensors with specific shapes, dtypes, and layouts. The raw size of a tensor follows the simplest possible formula:

\[\text{tensor bytes}=\text{number of elements}\times\text{bytes per element}.\]

A seven-billion-parameter model whose weights all reside in HBM as BF16 uses about 14 GB for the weights alone. Alignment and reserved blocks in the memory allocator, metadata, temporary buffers, and the KV cache raise the actual footprint. Checkpoint files may also be sharded or compressed, and runtime dtypes may differ from stored dtypes. File size is therefore an imperfect proxy for device-memory use.

Runtime tensors have very different lifetimes. Weights usually remain for the entire process. The space behind a temporary activation can be reused after its final consumer finishes. A KV cache follows one inference request. Optimizer state lasts for the entire training run. Training also retains selected activations until the corresponding backward computation completes.

Another easy-to-miss category is the workspace. Kernels and libraries may request temporary storage for reductions, sorting, matrix computation, or communication. These bytes are absent from the model’s mathematical definition and still contribute to peak memory. A framework’s memory allocator requests and reuses device memory and may retain blocks after their tensors have been freed. As a result, the process memory shown by nvidia-smi does not always match the sum of currently live tensors.

Tensor Lifetimes

Arranging runtime state by lifetime makes the sources of the memory peak easier to see:

This explains the large gap between fitting inference and fitting training. Continue with the rough seven-billion-parameter example. BF16 weights occupy about 14 GB. BF16 gradients add another 14 GB. Adam’s two FP32 moment arrays add about 56 GB, bringing those three categories to roughly 84 GB. Some training configurations also retain an FP32 parameter copy for updates, commonly called master weights, which adds another 28 GB. Saved activations, workspaces, and allocator overhead are still absent from this estimate. The exact total depends on runtime dtypes, low-precision representations, and whether model state is distributed across GPUs.

Peak memory comes from overlapping lifetimes. During inference, weights persist, KV caches grow with active requests, and temporary activations appear and release storage across layers. Training adds saved activations, gradients, and optimizer state that coexist within the same step. Even within inference, prompt processing and token-by-token generation produce two distinct performance profiles.

The GPU During Inference

When a serving process starts, the weights settle into HBM and are shared by incoming requests. A decoder-only request then has two phases: prefill, which processes the known prompt, and decode, which repeatedly produces new tokens. Both phases traverse the same layers and use the same weights. Their performance profiles differ because they expose different numbers of token rows at once and preserve different amounts of historical state.

Prefill: Processing the Prompt

Suppose the prompt has length $T$. Before the first layer begins, all $T$ token embeddings already exist. Causal attention allows position $t$ to read only positions no later than $t$, but every prompt token is already known. The Q, K, and V projections for all positions in one layer can therefore run together. A causal mask removes forbidden future pairs. There is no need to advance through the prompt one token at a time. Dependencies remain between layers, so the second layer waits for the first layer’s complete hidden state.

As prefill traverses each decoder layer, it computes K and V for every prompt position and writes them into that layer’s KV cache. The final layer produces logits for all positions, with the last row giving the distribution for the first generated token.

Prefill exposes many input rows. Once a linear-layer weight tile reaches the compute units, it can serve many prompt positions, giving the matrix operation substantial data reuse. Attention must also compare many query-key pairs. For a prompt of length $T$, each attention head logically contains a $T\times T$ score matrix. Both arithmetic and the intermediates produced by a direct implementation grow quickly with prompt length.

Decode: One Token at a Time

After prefill selects the first generated token, that token becomes the input to the next model call. At a given layer, it contributes only one new row of hidden state. The layer computes a new Q, K, and V for that row. The query reads K/V from every visible earlier position, and the new K/V is appended to the cache for later tokens. After the row passes through the full layer stack, the newest logits select another token.

The next round must wait for that selection because its input embedding depends on the token ID just produced. Decode is therefore serial along generation length. Each step still launches many parallel kernels; the serial dependency sits between adjacent generated tokens.

For one request, a decode input matrix usually contains a single row. Every layer still reads a large set of weights, though each weight tile now serves very few rows. Attention also reads an expanding KV cache. The execution units may finish the arithmetic quickly and then wait for HBM to supply the next data. This is why small-batch decode is often limited by memory bandwidth. Weight traffic may dominate at short contexts. As the context grows, reads of historical K/V become increasingly important.

Serving systems place several active requests into the same decode iteration. With $B$ requests in one iteration, each linear layer can apply one set of weights to $B$ input rows, improving weight reuse and aggregate token throughput. Prompts and generation lengths vary, so completed requests should leave promptly and new requests should enter. Rebuilding the batch after each generated token is continuous batching. Larger batches usually increase throughput, consume more KV-cache capacity, and may make individual requests wait longer for scheduling.

Dynamic KV-cache growth introduces an allocation problem as well. Reserving the maximum context for every request wastes space. Requiring each cache to remain physically contiguous can break free memory into unusable fragments, a condition called fragmentation. PagedAttention maps the logical KV positions of a request onto fixed-size pages that may occupy noncontiguous locations in HBM. New pages are allocated as the request grows and returned when it finishes. Pages for a shared prefix can also be reused when the system permits it. PagedAttention improves allocation and reuse of the KV cache. The K/V contents required per token by a given attention architecture remain the same (Kwon et al., 2023).

FlashAttention: Why Softmax Matters

The linear layer above can be tiled from the start. Each output tile has an accumulator. Incoming input tiles produce partial products, and the completed accumulator eventually returns to HBM. The two matrix multiplications in attention can also be tiled. Softmax, sitting between them, is what breaks the straightforward pipeline. For one attention head, scaled dot-product attention is

\[S=\frac{QK^\top}{\sqrt d}, \qquad P=\operatorname{softmax}(S), \qquad O=PV.\]

The first matrix multiplication uses Q and K to produce the score matrix $S$. Softmax converts each row of $S$ into probabilities, which the second matrix multiplication uses to take a weighted sum of V. Every probability in a row depends on every score in that row:

\[p_i=\frac{e^{s_i}}{\sum_j e^{s_j}}.\]

Suppose a kernel has seen only the first four keys. It can calculate four scores, but it cannot finalize their probabilities. A later tile may contain a larger score, and every later score contributes to the denominator. Either event changes all earlier probabilities. A score tile cannot immediately flow into the final output as a partial product does in a linear layer.

A literal implementation separates the formula into three operations. First, $QK^\top$ produces the complete $S$ and writes it to HBM. Softmax reads $S$, produces the complete $P$, and writes $P$ back to HBM. The final matrix multiplication reads $P$ and V to obtain O. Softmax has created a materialization boundary between two matrix multiplications. Although $S$ and $P$ are temporary, both occupy HBM and both travel across the HBM boundary.

These intermediates become large very quickly. With sequence length $T$, each attention head logically has $T\times T$ scores. At $T=8192$, one matrix contains roughly 67 million elements. Even at two bytes per element, it occupies 128 MiB. Together, $S$ and $P$ reach 256 MiB, while the final O has shape only $T\times d_v$. FlashAttention addresses a precise question: can execution cross the softmax while keeping those two large intermediate matrices out of HBM?

The answer begins with a simple property of softmax. Subtracting the same constant $c$ from every score in a row leaves the probabilities unchanged. The numerator and denominator acquire the same factor $e^{-c}$, which cancels:

\[\frac{e^{s_i-c}}{\sum_j e^{s_j-c}} = \frac{e^{-c}e^{s_i}}{e^{-c}\sum_j e^{s_j}} = \frac{e^{s_i}}{\sum_j e^{s_j}}.\]

In practice, $c$ is usually chosen as $\max_j s_j$. The largest exponential is then $e^{\max_j s_j-\max_j s_j}=e^0=1$, and every other exponential is at most 1. This prevents large exponentials from overflowing. For any group of scores processed so far, a kernel needs to retain only three quantities:

\[m=\max_i s_i, \qquad \ell=\sum_i e^{s_i-m}, \qquad u=\sum_i e^{s_i-m}v_i.\]

Here, $m$ is the largest score seen so far, $\ell$ is the running softmax denominator measured relative to that maximum, and $u$ is the corresponding unnormalized weighted sum of values. After processing the complete row, the attention output is exactly $u/\ell$. These three quantities contain everything needed to continue the calculation, so the individual scores can be discarded.

A small example makes the compression concrete. Consider one query attending to eight keys, split into two tiles. The first tile has scores $[2,1,-1,0]$ and corresponding values $[1,2,3,4]$. To keep the arithmetic readable, let each value have only one dimension. The maximum in the first tile is 2, giving

\[\begin{aligned} \ell_1&=e^{2-2}+e^{1-2}+e^{-1-2}+e^{0-2}\approx1.55,\\ u_1&=e^{2-2}\cdot1+e^{1-2}\cdot2+e^{-1-2}\cdot3+e^{0-2}\cdot4\approx2.43. \end{aligned}\]

The four scores in this tile can now disappear. The triple $(m_1,\ell_1,u_1)=(2,1.55,2.43)$ captures their entire contribution to the eventual output.

The second tile has scores $[3,-2,1,2]$ and values $[5,6,7,8]$. The same calculation compresses it to $(m_2,\ell_2,u_2)=(3,1.51,8.93)$. The combined maximum is now 3, so the two sums retained from the first tile must move from a reference point of 2 to a reference point of 3. Multiplying them by $e^{2-3}$ performs that conversion:

\[\begin{aligned} \ell&=e^{2-3}\ell_1+\ell_2\approx2.08,\\ u&=e^{2-3}u_1+u_2\approx9.82,\\ O&=u/\ell\approx4.72. \end{aligned}\]

Keeping all eight scores, applying softmax once, and taking the weighted sum produces the same $4.72$. The reason is visible in the state itself: $\ell$ retains the total exponential mass, $u$ retains how that same mass weights V, and $m$ lets summaries from different tiles be expressed against a common reference and added. In real attention, each value is a vector, so $u$ becomes a vector as well; all of its dimensions share the same $m$ and $\ell$. Updating softmax one tile at a time in this way is called online softmax (Dao et al., 2022).

This property restores a continuous tiled pipeline. A kernel keeps a tile of Q near the execution units in shared memory or registers, then streams through the corresponding K/V tiles from HBM. For each K/V tile, Tensor Cores compute a score tile, online softmax updates $(m,\ell,u)$, and the temporary scores and probabilities are released after use. The next K/V tile repeats the process. After traversing the sequence, the kernel computes $u/\ell$ and writes the final output tile to HBM. The full $S$ and $P$ are never materialized there.

FlashAttention preserves dense-attention semantics and still evaluates the same set of query-key pairs. Its gain comes from execution order. Temporary scores and probabilities remain close to compute and disappear as soon as they have contributed to the output. During prefill, many query rows are active, so avoiding the $T\times T$ intermediates substantially reduces temporary capacity and HBM traffic. Training forward receives the same benefit. During backward, the kernel can recompute the score and probability tiles from Q, K, and V, trading some additional arithmetic for fewer saved intermediates.

Decode can use the same tiled attention and online softmax, though its shape shifts the bottleneck. Each request usually contributes one new query per iteration, producing one score row whose length equals the context length. The full $T\times T$ prefill intermediate no longer exists. Reading the growing KV cache becomes the dominant concern. Flash-Decoding divides the K/V sequence into chunks, lets separate blocks compute their own $(m,\ell,u)$ summaries, and combines those summaries with a reduction. This exposes more parallel work along the context dimension. The required K/V bytes still have to be read (Dao et al., 2023).

Later versions refine how this pipeline maps onto newer hardware. FlashAttention-2 changes the division of work across blocks and warps to improve occupancy and reduce communication through shared memory (Dao, 2024). FlashAttention-3 targets Hopper and overlaps data movement, matrix multiplication, and softmax more tightly (Shah et al., 2024). FlashAttention-4 redesigns the execution pipeline for Blackwell (Zadouri et al., 2026). The durable mental model across these versions is simple: complete the path from score tile to output contribution as close to the execution units as possible.

The GPU During Training

An inference activation can release its storage after its final use. Training introduces another dependency: backward will need some of the information produced during forward. A training step has three consecutive phases. Forward computes the result and preserves the required records, backward computes gradients, and the optimizer step updates parameters.

Why Forward Saves Activations

Mathematically, a training forward pass looks like any other forward pass. Token embeddings move through attention, MLPs, normalization, and residual paths before the model produces logits and a loss. At the execution level, it uses the same kernels, tiles, warps, and memory hierarchy described above.

The difference lies in intermediate lifetimes. Backward needs selected forward inputs or outputs. The weight gradient of a linear layer depends on the input activation seen during forward. The backward rule for an activation function depends on its earlier input or output. Normalization needs its associated statistics. A framework’s automatic differentiation system, or autograd, retains enough information until backward reaches the corresponding operation. Deeper models, longer sequences, and larger batches generally increase the saved-activation footprint.

Attention backward also uses FlashAttention’s tiled execution. Forward saves compact normalization statistics from online softmax. Backward recreates the score and probability tiles as needed, then computes gradients for Q, K, and V. This exchanges some recomputation for fewer saved intermediates and less HBM traffic.

The framework does not have to preserve every forward intermediate in its original form. Fused kernels and autograd can retain smaller sufficient records, while other buffers release their storage after the final use. Peak activation memory depends on the architecture, sequence length, batch size, dtype, and implementation.

Backward and the Optimizer

Once the loss is available, backward proceeds from the final layer toward the input. Each operation receives an output gradient from later computation and produces two kinds of results: an input gradient that continues toward earlier layers, and gradients accumulated onto its trainable parameters. At the end of backward, every parameter has a gradient for the current batch.

The optimizer then reads each parameter, its gradient, and any persistent optimizer state, and writes an updated parameter. Adam’s moments remain live across training steps. The optimizer step therefore reads and writes model-sized arrays and can carry a substantial HBM-bandwidth cost.

Compared with one inference forward pass, training adds pressure along two dimensions. For capacity, forward records, gradients, and optimizer state coexist in HBM. For compute and bandwidth, backward and the parameter update add more matrix operations and model-sized data movement. Multi-GPU training must also combine gradients for shared parameters before the update, making those gradients a major communication payload.

Reducing Training Memory

When each value occupies too many bytes, mixed precision runs suitable operations in BF16, FP16, or FP8 while preserving higher precision where greater range or more stable accumulation is needed. It can shrink some weights, activations, and gradients and unlock low-precision Tensor Core instructions. BF16 training often retains higher-precision accumulations or optimizer state, so total HBM usage does not simply fall by half (PyTorch Foundation, 2026).

When activations from an entire batch cannot coexist comfortably, microbatching and gradient accumulation divide that batch into smaller pieces and run them sequentially. Each microbatch releases its activations after forward and backward. Parameter gradients continue accumulating, and the optimizer steps after all microbatches finish. Peak activation memory falls, while individual matrix operations become smaller and the number of kernel launches increases.

When a deep model preserves too many forward records, activation checkpointing stores only selected layer boundaries. As backward reaches a missing interval, the system reruns that part of forward to reconstruct the intermediates and then computes the gradients. The technique converts saved-activation capacity into extra computation. It does not directly reduce weights, gradients, or optimizer state (PyTorch Foundation, 2026).

Mixed precision shrinks individual values. Microbatching reduces how much batch state coexists. Activation checkpointing reduces the intermediates that must wait in HBM for backward. If memory pressure still exceeds one GPU, state or computation can be distributed across several GPUs. That adds capacity and compute, along with a new requirement to move data between devices.

Beyond One GPU

Adding a GPU provides more HBM and more execution units. It also introduces a device boundary. One GPU cannot treat another GPU’s HBM as freely accessible local memory. Once model state or computation is partitioned, some tensors must travel over an interconnect, and the receiving GPU may have to wait for them.

GPUs in one server may communicate through NVLink, NVSwitch, or PCIe. Communication across servers uses network fabrics such as InfiniBand or Ethernet. Bandwidth now refers to two distinct boundaries: HBM bandwidth supplies data within one GPU, while interconnect bandwidth moves data between GPUs. Choosing a parallel strategy begins with identifying what no longer fits or runs fast enough.

The major forms of parallelism first differ by the dimension they divide:

Real systems often combine several of these strategies because they divide different dimensions. The right combination depends on which state does not fit and which computation does not scale.

When the full model fits, replication is the most direct way to scale. An inference service can assign different requests to different model replicas. Each request completes decode within its replica and usually avoids exchanging hidden states at every layer. Data parallelism (DP) follows the same pattern in training. Every GPU holds a model copy, reads a different local batch, runs forward and backward independently, and then synchronizes parameter gradients (PyTorch Foundation, 2026).

The replicas in training produce different gradients. If two GPUs obtain $g_0$ and $g_1$ for the same parameter, a sum all-reduce leaves both GPUs with $g_0+g_1$. Dividing by the replica count gives the average when that is the desired gradient convention.

If persistent model state no longer fits, replication duplicates the very cost that needs to shrink. ZeRO and Fully Sharded Data Parallel (FSDP) partition parameters, gradients, and optimizer state across a data-parallel group. Before a layer runs, GPUs gather the required parameter shards. After backward produces gradients, those gradients are aggregated and sharded again. HBM across the group can now hold a larger training state, at the cost of parameter and gradient communication around each layer (Rajbhandari et al., 2020; PyTorch Foundation, 2026).

If the pressure comes from one layer, tensor parallelism (TP) divides large tensors within that layer. For a linear layer, different GPUs may store different sections of the weight matrix and compute separate output features, or they may produce partial sums for the same output. Later computation may need those pieces recombined through cross-GPU communication. Attention heads and MLP dimensions provide natural partitioning axes in a Transformer (Shoeybi et al., 2019). This communication often occurs in every Transformer block, so GPUs in one TP group usually require a fast link.

If the full layer stack must be distributed by depth, pipeline parallelism (PP) assigns contiguous layers to different stages. Forward sends activations across stage boundaries, and backward sends activation gradients in the opposite direction. Training usually divides the batch into microbatches so that different stages can process different points in the pipeline concurrently. Some stages still sit idle while the pipeline fills and drains. This idle region is the pipeline bubble (Narayanan et al., 2021).

If long sequences create the pressure, context parallelism (CP) assigns different ranges of token positions to different GPUs, distributing activations that grow with sequence length. Linear layers and normalization can operate on local sequence chunks. Attention has dependencies across tokens, so local queries still need K/V from the full sequence. GPUs exchange K/V chunks around attention and reconstruct the corresponding gradients during backward (NVIDIA, 2026). CP reduces the activation footprint on each GPU and adds communication that grows with context length.

Mixture-of-Experts models introduce another partitioning axis. Expert parallelism (EP) places different experts on different GPUs. After the router selects experts for each token, every GPU may need to send different tokens to different destinations. This exchange is an all-to-all. Expert outputs then return to the original token order. The payload depends on routing decisions and load balance (Fedus et al., 2022).

All of these strategies lead back to one question: after a tensor is partitioned, does the next operation still have the data it needs locally? An all-reduce aggregates contributions and gives every participant the result. An all-gather reconstructs shards held by other participants. A reduce-scatter aggregates values while leaving the result partitioned. An all-to-all lets every participant send different data to every destination. These are collective communication operations, coordinated data rearrangements across a group of devices (NVIDIA, 2026). Gradients in DP, parameter shards in FSDP, activations in TP, and routed tokens in EP are different payloads carried by these same communication patterns.

When Communication Dominates

Communication pays at least two costs. Launch and synchronization take time, and the payload is limited by link bandwidth. Many small messages repeatedly pay the first cost. Large payloads are more likely to hit the second. Collectives also require participants to become ready at compatible times. If one GPU computes slowly or one network path becomes congested, the others may wait.

Implementations try to overlap communication with computation. Data-parallel backward moves from the final layer toward the first. As soon as one layer’s gradient is ready, its synchronization can begin while the GPU computes gradients for earlier layers. Communication in tensor and context parallelism sits closer to computation inside each layer and is often harder to hide completely. Pipeline and expert parallelism also depend on microbatch scheduling and routing balance.

If adding GPUs barely improves step time or makes it worse, three checks are useful: whether each GPU now receives too little computation, whether communication payload or frequency grows with the parallel degree, and whether collectives overlap with useful work. Once the bottleneck sits at the device boundary, adding theoretical FLOPs does not resolve it.

Four Ledgers for GPU Performance

Most GPU-performance questions can be reduced to four ledgers:

One optimization often changes several ledgers at once. Quantization shrinks weights or the KV cache and reduces both capacity and HBM traffic, provided that the hardware and kernels support the chosen low-precision representation. Batching lets one set of weights serve more token rows, improving reuse and throughput while increasing concurrent request state. Kernel fusion keeps intermediates in registers or shared memory, reducing launches and HBM round trips. FlashAttention reorders attention so that large intermediates avoid HBM while the logical result stays the same.

During training, activation checkpointing reduces saved activations and adds recomputation in backward. Data, tensor, pipeline, context, and expert parallelism distribute computation and state across more GPUs, each creating its own communication payload. Sharding can solve a single-GPU capacity problem and move the bottleneck to the interconnect. The name of an optimization does not reveal its benefit by itself. Trace which tensors became smaller, which bytes stopped moving, which operations disappeared or were recomputed, and which devices now have to wait for one another.

The next time you see .to("cuda"), remember that it marks only the beginning of the path. Weights move from the host into HBM. Kernels expand operations into grids, blocks, warps, and threads. The memory hierarchy continually feeds the current working set. Inference maintains growing request state, while training preserves the information needed for backward and parameter updates. When a model spans several GPUs, some tensors continue across the interconnect. Most GPU-performance discussions can be placed back onto this map and charged to one of four ledgers: capacity, bandwidth, compute, or communication.