Memory-Bound Inference

Table of Contents

• 1. Why Inference is Memory-Bound
  • 1.1 Autoregressive Decoding as a Stream of Matrix-Vector Products
  • 1.2 Formal Derivation of Arithmetic Intensity for a Linear Layer
  • 1.3 Contrast with Training and Prefill
• 2. The Bandwidth-Scaling Gap
  • 2.1 Ridge Point Trend Across GPU Generations
  • 2.2 Consequences for Inference Workloads
• 3. KV-Cache Bandwidth Math
  • 3.1 KV-Cache Size
  • 3.2 Per-Step Bandwidth Cost
  • 3.3 Combined Arithmetic Intensity with KV Cache
• 4. Diagnosing the Regime: Practical Heuristics
  • 4.1 Theoretical Classification
  • 4.2 Empirical Signals Without Profiling
  • 4.3 Profiling Tools
  • 4.4 Operation-Level Rules of Thumb
• 5. Implications for Hardware Choice
• References

1. Why Inference is Memory-Bound 🧠

1.1 Autoregressive Decoding as a Stream of Matrix-Vector Products

Autoregressive decoding generates tokens one at a time: at each step \(t\), the model takes a single new token embedding \(x_t \in \mathbb{R}^d\) and produces the next hidden state by passing \(x_t\) through all transformer layers. Crucially, at each layer, every weight matrix is multiplied against a single vector — this is a matrix-vector product (GEMV), not a matrix-matrix product (GEMM).

The asymmetry is decisive. A GEMM of two \(d \times d\) matrices requires \(2d^3\) FLOPs and transfers \(O(d^2)\) bytes, yielding arithmetic intensity \(I \sim d/3\), which scales to hundreds of FLOP/byte for large \(d\). A GEMV of a \(d \times d\) matrix against a vector requires \(2d^2\) FLOPs and transfers \(2d^2 \cdot 2\) bytes (loading the matrix in BF16), yielding:

\[I_{\text{GEMV}} = \frac{2d^2}{2 \cdot 2d^2} = \frac{1}{2} \approx 1 \text{ FLOP/byte}\]

This is the core observation: at batch size 1, every weight matrix is loaded from HBM exactly once per token, and only one FLOP of compute is performed per byte transferred. The GPU’s Tensor Cores sit idle; the bottleneck is the HBM bus.

1.2 Formal Derivation of Arithmetic Intensity for a Linear Layer

Definition (Linear Layer Arithmetic Intensity). Consider a single linear layer \(y = Wx\) with \(W \in \mathbb{R}^{M \times K}\) and a batch of \(B\) tokens, so \(x \in \mathbb{R}^{K \times B}\) and \(y \in \mathbb{R}^{M \times B}\). In BF16 (2 bytes per element), the data-movement and compute costs are:

\[F(B) = 2MKB \quad \text{(FLOPs: one mul + one add per inner product entry)}\]

\[\mathcal{B}(B) = 2 \underbrace{MK}_{\text{weights}} \cdot 2 + 2 \underbrace{KB}_{\text{input}} \cdot 2 + 2 \underbrace{MB}_{\text{output}} \cdot 2 \quad \text{(bytes, BF16)}\]

\[= 4MK + 4KB + 4MB\]

The arithmetic intensity as a function of batch size is:

\[I(B) = \frac{F(B)}{\mathcal{B}(B)} = \frac{2MKB}{4MK + 4KB + 4MB} = \frac{MKB}{2MK + 2KB + 2MB}\]

For large \(M, K \gg B\) (the typical regime: dimension \(d \sim 4096\)–\(8192\), batch \(B \leq 32\)), the weight term \(2MK\) dominates the denominator:

\[I(B) \approx \frac{MKB}{2MK} = \frac{B}{2} \quad \text{FLOP/byte}\]

Consequence. The layer is memory-bound whenever \(I(B) < I^*\), i.e., whenever \(B \lesssim I^* / 2\). For the H100 SXM (\(I^* = 591\) FLOP/byte), this means batch sizes up to \(B \approx 300\) keep the layer memory-bound. For the A100 (\(I^* = 156\)), the crossover is around \(B \approx 78\).

1.3 Contrast with Training and Prefill

During training and prefill (processing a full input prompt of length \(S\)), the input is a matrix \(x \in \mathbb{R}^{K \times S}\) rather than a vector. The GEMM \(Wx\) with \(W \in \mathbb{R}^{M \times K}\) has arithmetic intensity:

\[I_{\text{prefill}}(S) \approx \frac{S}{2} \quad \text{FLOP/byte (same formula, with } S \text{ replacing } B\text{)}\]

A prompt of length \(S = 2048\) yields \(I \approx 1024\) FLOP/byte — above the ridge point of any current GPU. Prefill is compute-bound; decode is memory-bound. This is why optimizing decode throughput requires maximally batching requests, while prefill optimization focuses on kernel efficiency (FlashAttention, etc.).

2. The Bandwidth-Scaling Gap 📈

2.1 Ridge Point Trend Across GPU Generations

The ridge point \(I^* = \Pi / \beta\) has risen monotonically across every data-center GPU generation because compute throughput (TFLOPS) has grown faster than HBM bandwidth (TB/s). Quantitatively:

• Compute (FP16 TFLOPS): roughly 3–4× per generation (2 years).
• HBM bandwidth (TB/s): roughly 1.5–1.8× per generation.

The ratio \(I^* = \Pi / \beta\) therefore grows by a factor of roughly \(3/1.6 \approx 2\) per generation.

The table below tracks this trend from Pascal (the first HBM data-center GPU) through Blackwell:

2.2 Consequences for Inference Workloads

The rising \(I^*\) means the fraction of inference workloads that are memory-bound increases with each GPU generation. A workload that was compute-bound on V100 may be memory-bound on H100.

Concretely: a batch of \(B = 64\) tokens has \(I(64) \approx 32\) FLOP/byte, which is: - Compute-bound on P100 (\(I^* = 29\)): barely above ridge point. - Memory-bound on V100 (\(I^* = 139\)): comfortably below. - Deeply memory-bound on H100 (\(I^* = 591\)): 18× below ridge point.

Surprisingly, GPUs are becoming arithmetically richer relative to their bandwidth — a property that accelerates training (which exploits TFLOPS at large batch) but actively hurts inference efficiency (which is already bandwidth-starved and benefits only from bandwidth improvements). For inference at low batch size, the figure of merit is GB/s, not TFLOPS.

3. KV-Cache Bandwidth Math 🗄️

3.1 KV-Cache Size

The KV cache stores the key and value tensors for all past tokens so that attention at position \(t\) does not recompute them. For a transformer with: - \(L\) layers - \(h\) attention heads - Head dimension \(d_h\) (so \(d_{\text{model}} = h \cdot d_h\)) - Current context length \(S\) tokens - Stored in BF16 (2 bytes per element)

the total KV-cache memory is:

\[\text{KV}_{\text{bytes}}(S) = 2 \cdot 2 \cdot L \cdot h \cdot d_h \cdot S = 4 L h d_h S\]

The leading \(2\) accounts for both keys and values; the second \(2\) is bytes per BF16 element.

Example. For a LLaMA-3-70B model (\(L = 80\), \(h = 64\), \(d_h = 128\), context \(S = 4096\)):

\[\text{KV}_{\text{bytes}}(4096) = 4 \times 80 \times 64 \times 128 \times 4096 = 4 \times 80 \times 33{,}554{,}432 \approx 10.7 \text{ GB}\]

At \(S = 32768\) tokens the KV cache alone is \(\approx 85.9\) GB — exceeding the H100’s total 80 GB VRAM.

3.2 Per-Step Bandwidth Cost

At each decode step, the full KV cache must be loaded from HBM to compute attention for the current query. The bandwidth cost per decode step for attention alone is:

\[\mathcal{B}_{\text{kv}}(S) = 4 L h d_h S \quad \text{bytes}\]

This grows linearly with context length. At long contexts, KV-cache bandwidth can match or exceed weight-loading bandwidth.

Example. For the same 70B model with full-head attention (\(h=64\)) at \(S = 4096\): - Weight bytes: \(70 \times 10^9 \times 2 = 140\) GB - KV-cache bytes per step: \(10.7\) GB (computed above) - KV cache is \(\approx 7.6\%\) of weight bandwidth at this context length

At \(S = 32768\): - KV-cache bytes: \(\approx 85.9\) GB — now \(\approx 61\%\) of weight bandwidth

3.3 Combined Arithmetic Intensity with KV Cache

Let \(P\) be the total parameter count (so weight bytes \(= 2P\) in BF16). At batch size \(B\) and context length \(S\), the combined arithmetic intensity for a decode step is:

\[I(B, S) = \frac{2P \cdot B}{2P \cdot 2 + \mathcal{B}_{\text{kv}}(S)} = \frac{2PB}{4P + 4LhdS}\]

where we write \(4P\) for the weight bytes (\(2P\) parameters \(\times\) 2 bytes, loaded \(\times\) 2 for the factor-of-2 I/O convention) and have simplified \(4Lhd_hS\) as \(4LhdS\).

For small \(S\) (short context), the weight term dominates: \(I(B, S) \approx B/2\) as before. For large \(S\), the KV-cache term dominates:

\[I(B, S) \xrightarrow{S \to \infty} \frac{2PB}{4LhdS} = \frac{PB}{2LhdS}\]

This decreases with \(S\): long-context decoding becomes even more memory-bandwidth-limited than short-context decoding, even as the compute cost (FLOPs for attention QK products) grows with \(S\).

4. Diagnosing the Regime: Practical Heuristics 🔬

4.1 Theoretical Classification

Given the roofline model from Roofline Model, classifying a workload is straightforward in principle:

1. Compute the arithmetic intensity \(I = F / B\) for the kernel or end-to-end step.
2. Compare to the GPU’s ridge point \(I^*\).

If \(I < I^*\): memory-bound. Attainable throughput is \(P \approx \beta \cdot I\); Tensor Cores are underutilized.

If \(I > I^*\): compute-bound. Attainable throughput is \(P \approx \Pi\); bandwidth is underutilized.

The quick rule: autoregressive decode with \(B < 100\) on any GPU from V100 onward is almost certainly memory-bound. For H100, the threshold rises to \(B \approx 300\)–\(600\) (Section 1).

4.2 Empirical Signals Without Profiling

When computing \(I\) analytically is cumbersome, three empirical signals reveal the regime:

📐 Signal 1: Batch-size scaling. Double the batch size \(B \to 2B\) and measure throughput change: - If throughput roughly doubles: the layer was compute-bound (FLOPs doubled, bandwidth unchanged). This is rare at low batch. - If throughput barely increases: memory-bound. The bandwidth was already saturated; doubling \(B\) doubled FLOPs but the bottleneck is bytes/second, not FLOPs/second. - In practice: a linear throughput increase with \(B\) until a knee, then saturation, is the signature of memory-to-compute crossover.

📐 Signal 2: nvidia-smi bandwidth vs. SM utilization. The command nvidia-smi dmon -s u (polling interval 100 ms) reports two key columns: sm (SM compute utilization, %) and mem (HBM bandwidth utilization, %). The diagnostic: - High mem, low sm → memory-bound. The HBM bus is the bottleneck. - High sm, low mem → compute-bound. Tensor Cores are the bottleneck. - Both high → well-balanced kernel near the ridge point (rare in practice). - Both low → kernel is too short (latency-bound by kernel launch overhead or small problem size).

📐 Signal 3: Manual bandwidth saturation check. Estimate whether you are saturating bandwidth as follows. Let \(T\) be the observed token generation time (seconds/token) and \(P_{\text{bytes}}\) be the model size in bytes. Then: \[\text{Effective memory bandwidth used} \approx \frac{P_{\text{bytes}}}{T}\] If this value is close to the GPU’s peak \(\beta\) (within 20–30%), the kernel is memory-bandwidth-saturated. If it is far below \(\beta\), the bottleneck may be compute, kernel launch latency, or CPU-GPU synchronization.

4.3 Profiling Tools 🛠️

NVIDIA Nsight Compute (ncu). This is the authoritative kernel-level profiler. Run with:

ncu --metrics sm__throughput.avg.pct_of_peak_sustained_active,\
l1tex__t_bytes_pipe_lsu_mem_global_op_ld.sum.pct_of_peak_sustained_elapsed \
python inference_script.py

The key output section is GPU Speed of Light (or Memory Throughput and Compute (SM) Throughput in newer versions). These metrics express the fraction of peak sustained throughput for memory and compute respectively:

• Memory Throughput near 100%, Compute Throughput low: memory-bound. The kernel is streaming data as fast as possible; Tensor Cores are idle.
• Compute Throughput near 100%, Memory Throughput low: compute-bound. Tensor Cores are saturated; bandwidth is available.
• Both near 100%: roofline-optimal (uncommon).

Nsight Compute adds overhead (typically 10–100× slowdown per kernel) and should be used on a representative single batch, not in a production loop.

PyTorch Profiler. For end-to-end attribution (identifying which ops dominate), use:

import torch
from torch.profiler import profile, ProfilerActivity, schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=2, active=5),
    with_stack=True,
) as prof:
    for _ in range(8):
        output = model.generate(inputs, max_new_tokens=1)
        prof.step()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

Look at the cuda_time_total column — linear layer kernels (e.g., ampere_sgemm or cutlass_gemm_*) dominate for weight-loading bottlenecks. A high ratio of memory-copy time to kernel execution time is another signal of bandwidth-bound behavior.

Simple Python micro-benchmark. For quick order-of-magnitude checks:

import torch, time

# Warm-up
for _ in range(5):
    y = model(x)
torch.cuda.synchronize()

t0 = time.perf_counter()
for _ in range(N):
    y = model(x)
    torch.cuda.synchronize()
elapsed = (time.perf_counter() - t0) / N  # seconds per step

tflops_observed = flops_per_step / elapsed / 1e12
bw_observed     = bytes_per_step / elapsed / 1e12  # TB/s

print(f"Observed: {tflops_observed:.1f} TFLOPS, {bw_observed:.2f} TB/s")
print(f"Peak:     {peak_tflops:.1f} TFLOPS, {peak_bw:.2f} TB/s")
print(f"Compute util: {tflops_observed/peak_tflops*100:.1f}%")
print(f"Bandwidth util: {bw_observed/peak_bw*100:.1f}%")

If bandwidth utilization is high (>50%) and compute utilization is low (<5%), the workload is memory-bound.

4.4 Operation-Level Rules of Thumb 📋

The following table summarizes typical arithmetic intensities for operations encountered in transformer inference. Values assume BF16, large-\(d\) regime, short sequences unless noted.

5. Implications for Hardware Choice 💡

The memory-bound nature of low-batch inference has concrete implications for GPU selection, deployment strategy, and cost.

📐 At low batch size (\(B \leq 32\)): bandwidth is the only figure of merit.

An H100 doing inference at \(B = 1\) operates at \(< 0.1\%\) of its peak TFLOPS. It is running as a 3.35 TB/s memory bandwidth device. An H200, with 4.8 TB/s HBM3e, delivers \(\approx 43\%\) more tokens/second at \(B = 1\) — not because of its identical GH100 compute, but purely from the bandwidth upgrade. A B200 at 8.0 TB/s delivers \(\approx 2.4\times\) the H100’s throughput per GPU for single-request latency.

TFLOPS advantages between these GPUs are irrelevant for \(B = 1\) decode.

📐 At high batch size (\(B \geq I^*\)): compute matters.

Continuous batching systems like vLLM aggregate requests across users, effectively building batch size from incoming request streams. At sufficient load, \(B\) can approach or exceed the ridge point, transitioning the workload toward compute-bound. In this regime, higher TFLOPS (H100 vs. A100, for instance) yield proportional throughput gains.

The implication: for a latency-sensitive deployment (e.g., interactive chat with p95 TTFT < 500 ms), deploy fewer H200s or B200s and prioritize bandwidth. For a throughput-optimized deployment (offline batch inference, serving many users simultaneously), increase batch size until compute-bound and prioritize TFLOPS.

📐 Blackwell and the convergence of regimes.

The B200 targets both regimes simultaneously: 8.0 TB/s HBM3e addresses latency-critical memory-bound workloads, while 4,500 FP16 TFLOPS (9,000 with sparsity) addresses throughput-critical compute-bound workloads. Its ridge point of \(I^* \approx 563\) FLOP/byte (dense FP16) is similar to the H100’s, meaning the same batch-size crossover thresholds apply. The B200 advantage for inference is primarily bandwidth, not a structural change in the compute-vs-memory balance.

📐 TPU note. Google TPUs optimize for high-throughput, compute-bound training workloads: they have high TFLOPS but HBM bandwidth that is competitive with, not superior to, NVIDIA GPUs. TPU inference typically requires large batches or speculative decoding (which increases effective FLOPs per step) to achieve high utilization. Single-request TPU inference is similarly memory-bound.

References