Optimizing Transformer Inference & Training

Learning Objectives
Notation
Core Intuition
The Inference Bottleneck: Memory-Bound vs Compute-Bound
Mixed Precision Training (FP16/BF16)
Operator Fusion and FlashAttention
Gradient Checkpointing
Tensor Parallelism
Speculative Decoding
Continuous Batching
Quantization for Inference
Common Pitfalls
Summary
Exercises

Learning Objectives

Classify transformer operations as memory-bound or compute-bound.
Derive the arithmetic intensity of attention and FFN layers.
Explain how mixed precision maintains accuracy with loss scaling.
Derive the memory savings from gradient checkpointing.
Explain speculative decoding and derive its expected speedup.

Notation

$B$ — batch size
$T$ — sequence length
$d$ — model dimension
$\text{AI}$ — arithmetic intensity (FLOPs/byte)
$\text{BW}$ — memory bandwidth (bytes/second)
$\text{FLOPS}$ — compute throughput (operations/second)

Core Intuition

Transformers are bottlenecked by different resources at different stages: training is typically compute-bound (massive matrix multiplications), while inference (especially autoregressive generation) is often memory-bandwidth-bound (loading model weights for each token). Optimization strategies must target the actual bottleneck.

FlashAttention Tiling

Block

Tile

Active tileProcessed

Explore: FlashAttention fuses QKᵀ, softmax, and ×V tile-by-tile in fast SRAM. Same exact output, but HBM usage scales O(N) instead of materializing O(N²).

The Inference Bottleneck: Memory-Bound vs Compute-Bound

Arithmetic intensity (AI): ratio of FLOPs to bytes transferred.

\text{AI} = \frac{\text{FLOPs}}{\text{Bytes loaded from memory}}. \tag{1}

A kernel is:

Compute-bound if $\text{AI} > \text{FLOPS/BW}$ (GPU compute is the bottleneck).
Memory-bound if $\text{AI} < \text{FLOPS/BW}$ (memory bandwidth is the bottleneck).

For autoregressive generation (batch size 1, generating one token at a time):

Each FFN layer loads $\sim 8d^2$ parameters (bytes in FP16) but performs $\sim 8d^2$ FLOPs on a single vector.
$\text{AI} = \frac{8d^2 \text{ FLOPs}}{8d^2 \times 2 \text{ bytes}} = 0.5$ FLOPs/byte.
GPU A100: FLOPS/BW $\approx 312\text{T}/2\text{TB/s} = 156$ . Since $0.5 \ll 156$ : severely memory-bound.

Implication: For inference, reducing memory footprint (quantization, pruning) helps more than adding compute.

Mixed Precision Training (FP16/BF16)

Strategy: Keep a master copy of weights in FP32; perform forward/backward passes in FP16/BF16.

Loss scaling: FP16 has limited dynamic range ( $\sim 10^{-8}$ to $65504$ ). Small gradients underflow to zero. Solution: multiply loss by a scale factor $S$ ; divide gradients by $S$ before weight update.

\mathbf{g}_{\text{FP16}} = S \cdot \nabla_\theta\mathcal{L}, \quad \mathbf{g}_{\text{FP32}} = \mathbf{g}_{\text{FP16}} / S. \tag{2}

BF16 advantage: Same dynamic range as FP32 (8 exponent bits) but reduced precision (7 mantissa bits vs 23). No loss scaling needed. Standard for LLM training.

Memory savings: Model weights: 2x reduction. Activations: 2x reduction. Optimizer states remain FP32.

Operator Fusion and FlashAttention

Problem: Each elementwise operation (add, multiply, softmax) requires a separate GPU kernel launch and HBM round-trip.

Solution: Fuse multiple operations into a single kernel:

Fuse bias + activation: one kernel instead of two.
Fuse LayerNorm into attention output.
FlashAttention: Fuses Q×K, scaling, masking, softmax, and ×V into a single tiled kernel.

FlashAttention memory: Standard attention stores full $T \times T$ matrix. FlashAttention never materializes it — only stores $O(T)$ intermediate values (row-wise max and sum for online softmax).

Speedup: 2–4x for training, enables 4–16x longer contexts.

Gradient Checkpointing

Problem: Storing all activations for backprop requires $O(L \cdot T \cdot d)$ memory.

Solution: Only store activations at checkpointed layers; recompute others during backward pass.

Memory: With checkpoints every $\sqrt{L}$ layers: memory reduces from $O(L)$ to $O(\sqrt{L})$ .

Cost: ~33% additional compute (recomputing one forward pass segment).

\text{Memory} = O(\sqrt{L} \cdot T \cdot d), \quad \text{Compute} = \frac{4}{3} \times \text{standard}. \tag{3}

Tensor Parallelism

Split large matrix multiplications across GPUs.

For FFN layer $\mathbf{Y} = \text{GELU}(\mathbf{XA})\mathbf{B}$ :

Split $\mathbf{A}$ column-wise across $N$ GPUs: each GPU computes $\text{GELU}(\mathbf{XA}_i)$ .
Split $\mathbf{B}$ row-wise: each GPU computes partial output.
All-reduce to combine results.

For attention: Split heads across GPUs (each GPU handles $H/N$ heads).

Communication: 2 all-reduce operations per layer (one for attention, one for FFN). Each all-reduce transfers $O(Td)$ bytes.

Speculative Decoding

Idea: Use a small "draft" model to generate $K$ candidate tokens quickly, then verify all $K$ in parallel with the large model.

Algorithm:

Draft model generates $K$ tokens: $\hat{x}_1, \ldots, \hat{x}_K$ .
Large model scores all $K$ tokens in one forward pass (parallel, not sequential).
Accept tokens that match the large model's distribution; reject from the first mismatch.

Expected speedup: If draft model matches large model with probability $\alpha$ per token:

\text{Speedup} = \frac{K(1-\alpha^{K+1})}{(1-\alpha)(K+1-\alpha)} \cdot \frac{c_{\text{large}}}{c_{\text{large}} + Kc_{\text{draft}}}, \tag{4}

where $c$ is the per-token cost. Typical speedup: 2–3x for well-matched draft models.

Continuous Batching

Problem: Static batching wastes compute — shorter sequences finish early but GPUs wait for the longest sequence.

Solution: Insert new requests as soon as others complete.

Iteration-level scheduling: At each decoding step, decide independently for each sequence whether to continue or swap in a new request. Maximizes GPU utilization.

PagedAttention (vLLM): Store KV-cache in non-contiguous memory pages, like virtual memory. Avoids memory fragmentation; enables sharing cache between requests with common prefixes.

Quantization for Inference

Reduce weight precision from FP16 to INT8/INT4:

w_{\text{int}} = \text{round}\left(\frac{w - z}{s}\right), \quad w \approx s \cdot w_{\text{int}} + z, \tag{5}

where $s$ (scale) and $z$ (zero-point) are calibration parameters.

INT8 (W8A8): 2x memory reduction, ~1% accuracy loss. Works for most models.

INT4 (W4A16): 4x memory reduction; weights in INT4, activations in FP16. GPTQ, AWQ achieve near-lossless quality.

Impact on inference: For memory-bound generation, 4x weight compression → ~4x speedup (limited by bandwidth).

Common Pitfalls

Pitfall 1. Optimizing compute when memory-bound. For single-request inference, reducing FLOPs (e.g., pruning 20% of parameters) barely helps if you're still loading the same weights.

Pitfall 2. Applying FP16 without loss scaling. Small gradients underflow silently, causing divergence after thousands of steps.

Pitfall 3. Using static batching for serving. Wastes 30–50% of GPU compute compared to continuous batching.

Summary

Autoregressive inference is memory-bandwidth-bound; training is compute-bound.
Mixed precision halves memory; BF16 avoids loss scaling complexity.
FlashAttention fuses attention into a single IO-efficient kernel.
Gradient checkpointing trades $O(\sqrt{L})$ memory for 33% more compute.
Speculative decoding achieves 2–3x inference speedup with a draft model.
Quantization (INT4/INT8) provides near-linear speedup for memory-bound workloads.

Exercises

Exercise 1. Compute the arithmetic intensity of a matrix multiply $\mathbf{A} \in \mathbb{R}^{M \times K}$ by $\mathbf{B} \in \mathbb{R}^{K \times N}$ in FP16.

Exercise 2. For a model with $L=32$ layers and activation memory $m$ per layer, compute the memory with gradient checkpointing every 4 layers.

Exercise 3. Derive the expected number of accepted tokens in speculative decoding with acceptance rate $\alpha$ and draft length $K$ .

Exercise 4. For an A100 GPU (312 TFLOPS FP16, 2 TB/s bandwidth), determine whether a batch-1 FFN layer with $d=4096$ is compute-bound or memory-bound.

Exercise 5. Compute the memory savings from INT4 quantization for a 70B parameter model (original in FP16).