GPU memory profiling

Yesterday I looked at routing across regions. Today I am going inward, to the GPU itself. When I first started serving models, I would just load the weights and see how many concurrent requests I could handle before hitting an out-of-memory error. That is trial and error, not engineering. Proper memory profiling lets you predict capacity, choose the right quantization level, and size your KV cache before a single request hits the server.

The four tenants of GPU memory

During LLM inference, GPU HBM is shared between four categories of allocations:

• Model weights: the parameters of the model. For a 7B parameter model in FP16, this is 7 billion * 2 bytes = 14 GB. In INT8, it is 7 GB. In INT4 (GPTQ or AWQ), roughly 3.5 GB.
• KV cache: the key and value tensors for all active sequences. This is the big variable. It grows with sequence length, batch size, and number of layers and attention heads.
• Activations: temporary tensors created during the forward pass. In inference (unlike training), we do not store activations for backpropagation, so these are relatively small and short-lived.
• CUDA context and fragmentation: the CUDA runtime itself consumes memory (typically 300-800 MB depending on the driver version and loaded libraries). PyTorch's memory allocator also introduces fragmentation overhead.

Let me walk through how to measure each one.

Measuring model weight memory

The simplest measurement. Load the model and check before any inference happens:

import torch
from transformers import AutoModelForCausalLM

# Before loading
torch.cuda.reset_peak_memory_stats()
baseline = torch.cuda.memory_allocated()

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    torch_dtype=torch.float16,
    device_map="cuda",
)

weights_mem = torch.cuda.memory_allocated() - baseline
print(f"Model weights: {weights_mem / 1e9:.2f} GB")
# Llama-3.1-8B in FP16: ~16.07 GB

You can also compute this analytically. Llama-3.1-8B has approximately 8.03 billion parameters. At 2 bytes per parameter in FP16, that is 16.06 GB. The measured value will be slightly higher due to buffer allocations and embedding tables.

Profiling KV cache growth

The KV cache is where inference memory gets interesting. For each layer, each attention head stores a key tensor and a value tensor for every token in every active sequence. The formula is:

KV cache per token = 2 * num_layers * num_kv_heads * head_dim * dtype_bytes

# For Llama-3.1-8B:
#   num_layers = 32
#   num_kv_heads = 8 (GQA)
#   head_dim = 128
#   dtype = FP16 (2 bytes)

per_token = 2 * 32 * 8 * 128 * 2  # = 131,072 bytes = 128 KB per token

For a batch of 32 sequences, each 2048 tokens long, that is 32 * 2048 * 128 KB = 8 GB of KV cache. This is why the KV cache, not the model weights, is usually the binding constraint on how many requests you can serve concurrently.

To measure KV cache growth empirically, I run a controlled experiment:

import torch

results = []
for seq_len in [128, 256, 512, 1024, 2048, 4096]:
    torch.cuda.reset_peak_memory_stats()
    mem_before = torch.cuda.memory_allocated()

    # Generate tokens to fill the KV cache
    input_ids = torch.randint(0, 32000, (1, seq_len), device="cuda")
    with torch.no_grad():
        output = model(input_ids, use_cache=True)

    mem_after = torch.cuda.memory_allocated()
    kv_mem = mem_after - mem_before
    results.append((seq_len, kv_mem / 1e6))
    print(f"seq_len={seq_len:5d}  KV cache: {kv_mem/1e6:.1f} MB")

    # Clear the cache
    del output
    torch.cuda.empty_cache()

The growth should be linear with sequence length. If it is not, something else (activation buffers, CUDA allocator overhead) is contaminating the measurement.

Activation memory during inference

Activation memory during inference is much smaller than during training because we do not retain the computation graph for backpropagation. However, it is not zero. The forward pass creates intermediate tensors for each layer: the attention scores matrix (which is batch_size * num_heads * seq_len * seq_len for prefill), layer norm intermediates, and MLP activations.

The peak activation memory occurs during prefill, when the full sequence is processed at once. During decode, each step processes only one new token, so activations are tiny.

# Measure peak activation memory during prefill
torch.cuda.reset_peak_memory_stats()
mem_before = torch.cuda.max_memory_allocated()

input_ids = torch.randint(0, 32000, (1, 4096), device="cuda")
with torch.no_grad():
    _ = model(input_ids, use_cache=True)

peak = torch.cuda.max_memory_allocated()
activation_peak = peak - mem_before - weights_mem
print(f"Peak activation memory (prefill, seq=4096): {activation_peak/1e6:.1f} MB")

For Llama-3.1-8B at sequence length 4096, expect peak activation memory in the range of 500 MB to 1.5 GB depending on the attention implementation (FlashAttention is much more memory-efficient than standard attention).

Using torch.cuda.memory_summary

PyTorch provides a detailed memory summary that breaks down allocated, reserved, and fragmented memory:

print(torch.cuda.memory_summary(abbreviated=True))

The key fields to look at:

• Allocated: memory actively used by tensors. This is what you control.
• Reserved: memory claimed by PyTorch's caching allocator but not currently used by any tensor. This looks like "wasted" memory but it is actually a performance optimization, since re-allocating from the CUDA driver is expensive.
• Active: subset of allocated that is currently referenced. If active is much less than allocated, you have tensors that could be freed.

Profiling with torch.profiler

For a timeline view of memory allocations, torch.profiler can record every CUDA allocation and deallocation:

from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CUDA],
    profile_memory=True,
    record_shapes=True,
    with_stack=True,
) as prof:
    input_ids = torch.randint(0, 32000, (1, 512), device="cuda")
    with torch.no_grad():
        _ = model(input_ids, use_cache=True)

# Export to Chrome trace for visual inspection
prof.export_chrome_trace("memory_trace.json")

# Or print the top memory-consuming operations
print(prof.key_averages().table(
    sort_by="self_cuda_memory_usage", row_limit=20
))

Open the trace in chrome://tracing or Perfetto and you get a timeline showing exactly when each tensor was allocated, how large it was, and which Python function created it. This is invaluable for finding unexpected memory spikes, like a poorly-placed .contiguous() call that doubles memory usage temporarily.

Practical memory budget for serving

When I set up a vLLM server, I use the --gpu-memory-utilization flag (default 0.9) which tells vLLM how much of the GPU's HBM it is allowed to use. vLLM then computes how many KV cache blocks fit in the remaining space after loading model weights.

The calculation goes like this for an 80 GB H100 serving Llama-3.1-8B in FP16:

• Total usable: 80 GB * 0.9 = 72 GB
• CUDA context: ~0.7 GB
• Model weights: ~16 GB
• Available for KV cache: 72 - 16 - 0.7 = ~55.3 GB
• KV cache per token: 128 KB
• Total tokens: 55.3 GB / 128 KB = ~432,000 tokens
• At 2048 tokens per sequence: ~211 concurrent sequences

That is the theoretical max. In practice, fragmentation and activation memory reduce it by 10-15%. The point is that you can predict this number before deploying, which matters for capacity planning.

Memory profiling is not optional in production inference. Every byte of HBM you waste on fragmentation or oversized buffers is a request you cannot serve. Measure first, optimize second, and always check your math against nvidia-smi before you trust it.

Tomorrow I will explore what happens to throughput when you change the weight precision, comparing FP16 vs INT8 vs INT4 on real workloads.