Chapter 3: Numerical Formats and Transformer Workloads

Transformer operations and tensor shapes

Map one declared decoder layer from mathematical operations to profiler-visible GPU work and tensor allocations

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

Matrix multiplication
Softmax
Tensor shapes

A model diagram can label one box ‘self-attention’ and another ‘feed-forward network.’ The machine does not execute those labels. The machine executes projections, matrix products, row reductions, exponentials, elementwise transforms, copies, and communication. Each operation receives tensors with concrete shapes, formats, strides, and lifetimes.

A PyTorch profiler records named operations such as `aten::linear`, `aten::softmax`, and `aten::add` while one decoder layer runs. Expanding the high-level diagram into these operations produces a workload graph: each node is a declared calculation over one or more tensors, and each arrow means that a consumer must wait for a producer's output. Tensor shapes and profiler durations explain why one layer contains both large matrix multiplications and short programs that mainly read and write full activation arrays.

The worked investigation annotates every tensor in one layer with its batch, sequence, hidden, and head axes; its bytes per element; its producer and consumer; and the interval during which its allocation remains needed. The annotated graph becomes an empirical prediction. A framework memory snapshot can confirm which tensors were physically allocated, while Nsight Systems can confirm the operation order and reveal any unexpected copy or layout conversion.

Choose symbolic dimensions

Let B be batch size, S sequence length, D hidden width, H the number of query heads, Hkv the number of key/value heads, Dh the head dimension, and F the feed-forward width. Usually D = H×Dh for the query representation. Keep these dimensions symbolic until the operator relationships are clear.

The input activation X has shape [B,S,D]. Framework code often flattens X to [B×S,D] for linear projections. The flattening is a metadata-only logical view when the underlying layout supports the requested strides. The product of B and S becomes the M dimension of several GEMMs, so batch and sequence jointly determine whether those GEMMs are large enough for efficient tiling.

Expand the query, key, and value projections

The query projection maps [B×S,D] to [B×S,H×Dh]. Key and value projections map to [B×S,Hkv×Dh]. Multi-head attention uses separate logical heads, but implementations commonly combine each set of head projections into one larger matrix multiplication and reshape the result.

When Hkv is smaller than H, several query heads share a key/value head. The grouped-query structure reduces K and V projection output and later KV-cache capacity. Grouped-query attention does not reduce the number of query heads or the Q projection in the same way.

Follow scaled dot-product attention

For one head, Q has query positions by Dh and K has key positions by Dh. The product QK^T creates one score for every query-key pair. Scaling by 1/sqrt(Dh) controls score magnitude, a causal or application mask removes invalid positions, and softmax normalizes each query row. Multiplying the probabilities by V produces one Dh-dimensional output per query.

The conceptual score tensor has shape [B,H,Sq,Sk]. During full self-attention, Sq and Sk can both equal S, which gives quadratic elements and arithmetic in sequence length. During autoregressive decode, Sq is often one per active sequence while Sk is the retained context length. The two shape patterns produce different computational regimes.

Distinguish conceptual tensors from materialized arrays

A mathematical expression can define a tensor that the implementation never writes to HBM. A naive attention implementation can materialize the complete score matrix, read the matrix for softmax, write probabilities, and read them for the value product. Repeated score and probability transfers can dominate long-sequence execution.

FlashAttention uses tiling and online softmax statistics to compute exact attention while reducing HBM reads and writes. The FlashAttention kernel processes blocks of Q, K, and V, keeps partial row maxima and normalization sums on chip, and avoids storing the complete score and probability matrices in HBM. The mathematical attention result remains the same within finite-precision effects; the physical dataflow changes substantially.

Expand the output and feed-forward projections

Attention head outputs are combined and projected from width D back to D. The residual path adds the layer input or an intermediate activation according to the architecture. Normalization can occur before or after the sublayer. Each placement changes dependencies and which values must be retained.

A feed-forward block maps each token independently through wider dimensions. A gated form often uses two projections from D to F, an elementwise activation and product, and one projection from F back to D. The three GEMMs can account for substantial model arithmetic and parameter bytes, while the activation and residual operations can remain dominated by full-tensor traffic.

Classify operations by reuse and arithmetic intensity

Projection GEMMs reuse weights across tokens and activations across output channels. Their efficiency depends on matrix shape, datatype, and tile compatibility. Larger B×S generally increases the M dimension and creates more work over which to amortize setup.

Normalization, bias, activation, and residual addition perform relatively few operations per element but can read and write the complete activation. Fusion can keep intermediate values in registers or shared memory and reduce HBM traffic. Fusion also increases kernel complexity, register pressure, and possibly recomputation, so its value depends on the complete resource model.

Annotate views, copies, dependencies, and lifetimes

A reshape changes only tensor metadata when the requested index mapping is compatible with existing strides. A transpose can also be represented as a view, but a later kernel may require a contiguous layout and trigger a physical copy. Mark the actual materialization rather than assuming every graph node moves data.

Draw the critical path from normalization through projections, attention, output projection, residual update, and feed-forward work. During training, mark which values backward requires. During inference, mark persistent K and V. Distributed parallelism adds collectives at selected projection boundaries. This annotated graph becomes the shared input for kernel, compiler, memory, and communication decisions.

Worked example: Annotating one decoder layer

Consider one decoder layer with B = 8, S = 2,048, D = 4,096, H = 32 query heads, Hkv = 8 key/value heads, Dh = 128, and F = 11,008. Activations and weights are stored in BF16.

Flatten tokens so the projection input is [16,384, 4,096]. Q output width is H×Dh = 4,096, so Q has shape [16,384, 4,096] before a head view. K and V output width is Hkv×Dh = 1,024, so each has shape [16,384, 1,024].
Calculate Q projection work as 2×16,384×4,096×4,096, about 550 GFLOP. Each K or V projection is one quarter of that output width and requires about 137 GFLOP. The operation counts ignore bias and format conversion.
The conceptual full score tensor has B×H×S×S = 8×32×2,048², or about 1.074 billion elements. In BF16, one materialized score tensor alone would occupy about 2.15 GB. A second probability tensor of the same shape would add similar storage and traffic.
Represent exact tiled attention differently in the physical graph. Show blocks of K and V loaded for a Q block, on-chip partial score processing, and retained row maximum and normalization sum. Do not add the complete score tensor to the peak allocation when the implementation never materializes it.
Calculate the base activation size B×S×D×2 bytes = about 134 MB. Each separate normalization, residual, or activation pass can read and write a substantial fraction of this tensor. Mark adjacent elementwise operations that could share one pass.
Draw dependencies and lifetimes through QKV projections, attention, output projection, residual, feed-forward projections, and the next residual. Add saved values required by backward or persistent KV values required by inference according to the workload being modeled.

Conclusion: The annotated graph connects model structure to operation shapes, approximate arithmetic, possible HBM traffic, and memory lifetimes. Later profiling can test this model instead of beginning with an unstructured list of kernels.

Common errors

Counting every mathematical tensor as an HBM allocation. Views, fused intermediates, tiled algorithms, and recomputation change physical storage.
Treating all projections as identical GEMMs. Q, K, V, output, and feed-forward matrices have different dimensions and can select different kernels.
Calling every transpose free because it can be represented as a view. The consuming kernel may require a contiguous physical order and force a copy.
Ignoring the backward or decode graph. Saved activations, gradients, and persistent KV state can change both dominant traffic and peak memory.

Reference summary

For the recurring 7-billion-parameter decoder, one layer receives an activation tensor whose hidden axis contains 4,096 values per token. Linear projections become matrix multiplications. Attention also performs masking, row reductions, exponentiation, and normalization. The feed-forward path adds wider projections and an activation. Residual addition and normalization perform less arithmetic but can read and write the complete activation tensor. PyTorch profiler operation names and GPU timelines expose the resulting sequence.

attention(Q,K,V) = softmax(QKᵀ / √d + mask)V

Projection shapes depend on batch size, sequence length, hidden width, head count, and feed-forward width. The five dimensions determine whether a matrix multiplication exposes enough parallel tile work to use the intended GPU matrix instructions efficiently.
The conceptual attention-score tensor has query-length by key-length entries for every batch item and head. An IO-aware implementation can compute exact attention without storing the complete score tensor in HBM.
A reshape can be a metadata-only view when the required strides already exist. A transpose that changes physical order can require a complete copy.
Fusion can remove intermediate HBM reads and writes, but it can also increase register use and reduce occupancy. Evaluate the complete fused kernel rather than assuming that fewer launches are always faster.
Training retains selected intermediate values for backward computation. Autoregressive inference instead retains keys and values across decoding steps.

Knowledge check

For a residual addition followed by normalization, which intermediate tensor would an unfused path write to HBM and read back, and which profiler byte counters would test whether fusion removed that round trip?

Show the answer guide

In an unfused residual-add followed by normalization, the residual sum tensor `r = x + f(x)` is the intermediate. The addition kernel writes the complete `r` tensor to HBM, and the normalization kernel later reads `r` from HBM.
For a tensor containing N elements of b bytes each, the removable round trip is ideally N·b bytes written plus N·b bytes read, or 2N·b bytes, before accounting for cache effects and write allocation.
In Nsight Compute, compare the device-memory or DRAM Bytes Written and DRAM Bytes Read values in the Memory Workload Analysis report; on supported versions the corresponding raw metrics include `dram__bytes_write.sum` and `dram__bytes_read.sum`.
The fused path should also remove one materialized tensor and usually one launch, while preserving residual and normalization semantics. Requested and transferred bytes can differ, so retain the profiler's named boundary and metric definitions.
Ordinary unprofiled timings, generated launches, and correctness results must accompany the byte counters because lower DRAM traffic alone does not prove an end-to-end improvement.

Practice

Calculate and inspect a residual round trip

Use a BF16 tensor of shape `[2, 1024, 4096]`. Calculate the bytes in `r = x + f(x)` and the predicted HBM write-plus-read removed by fusing residual addition with normalization. Run or inspect an eager and compiled/fused implementation on one declared GPU when available, using the same inputs and normalization definition.

Evidence to keep: Keep the tensor-byte and 2× round-trip calculations, an output comparison, the launch list, the Nsight Compute DRAM read/write report or a clearly labeled unavailable-hardware prediction, and a conclusion that separates calculated from measured bytes.

The operator graph defines which tensors exist and when they are consumed. The next section turns those lifetimes into a peak-memory model for training.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.