Profiling Memory Usage During LLM Inference

Profiling Memory Usage During LLM Inference

Large language models (LLMs) are notoriously memory-hungry. A 7B parameter model in FP16 requires roughly 14 GB just to hold its weights, and that is before you account for the KV cache, activations, and temporary buffers needed during generation. When inference slows down, crashes with an out-of-memory error, or behaves inconsistently across batch sizes, memory is almost always the culprit. Profiling memory usage during LLM inference is the disciplined practice of measuring exactly where bytes are allocated, when they are freed, and how they scale with sequence length, batch size, and decoding strategy. This tutorial walks through what memory profiling means in the context of LLM inference, why it matters, and how to instrument your stack with practical, runnable code.

What Memory Profiling Means for LLM Inference

During inference, GPU memory is consumed by several distinct categories. Understanding these categories is the first step toward profiling them effectively:

• Model weights: The static parameters of the transformer. Size scales linearly with parameter count and inversely with quantization precision.
• KV cache: Cached keys and values from attention layers. This grows linearly with sequence length, batch size, number of layers, and hidden dimension.
• Activations: Intermediate tensors produced during the forward pass, including attention scores and MLP intermediates.
• Temporary buffers: Short-lived allocations from kernels, such as workspace buffers for FlashAttention or cuBLAS GEMMs.
• Framework overhead: PyTorch's caching allocator, CUDA context, and any Python-side bookkeeping.

A good profiler separates these categories so you can answer targeted questions: Is the KV cache dominating? Are activations spiking during prefill? Is the caching allocator fragmenting? Without this breakdown, you are guessing.

Why Memory Profiling Matters

Memory profiling is not a luxury — it directly affects cost, latency, and reliability. The practical reasons to profile include:

• Capacity planning: Knowing the peak memory footprint lets you choose the right GPU (A10G vs. A100 vs. H100) without overpaying.
• Maximizing throughput: If you understand the KV cache footprint, you can tune batch size and max sequence length to fill the GPU without spilling.
• Diagnosing OOMs: Intermittent out-of-memory errors often come from fragmentation or unexpected activation spikes during long-context prefill.
• Validating optimizations: When you enable paged attention, quantization, or speculative decoding, profiling confirms the memory savings you expected.
• Reproducibility: Memory behavior can differ across PyTorch versions, CUDA drivers, and kernel implementations. Profiling gives you a baseline to detect regressions.

Tooling Landscape

Several tools are commonly used to profile GPU memory during LLM inference. Each operates at a different level of abstraction:

• torch.cuda.memory_allocated() / torch.cuda.max_memory_allocated() — Lightweight PyTorch APIs for tracking allocator state.
• torch.cuda.memory._record_memory_history() — Captures a detailed timeline of every allocation and free, exportable to a visualization.
• torch.profiler — The integrated profiler that records memory events alongside kernel execution and CPU activity.
• nvidia-smi / dcgmi — Coarse-grained system-level monitoring, useful for spotting leaks across processes.
• NVIDIA Nsight Systems (nsys) — A system-wide profiler that captures CUDA API calls, kernel launches, and memory transfers.
• Inference-engine built-ins: vLLM, TGI, and TensorRT-LLM expose metrics endpoints that report KV cache utilization and reserved memory.

For most developers, the PyTorch APIs and torch.profiler are the right starting point. Nsight Systems becomes essential when you need to attribute memory to specific kernels.

Profiling with PyTorch Memory APIs

Tracking Peak Allocation

The simplest useful measurement is peak allocated memory during a generation. This tells you the high-water mark you must provision for. The pattern is to reset the peak counter before inference and read it after:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda:0",
)
model.eval()

prompt = "Explain the theory of relativity in three paragraphs."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")

# Reset peak tracking right before inference
torch.cuda.reset_peak_memory_stats(device="cuda:0")
start_alloc = torch.cuda.memory_allocated(device="cuda:0")

with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
    )

end_alloc = torch.cuda.memory_allocated(device="cuda:0")
peak_alloc = torch.cuda.max_memory_allocated(device="cuda:0")

print(f"Model load allocated:    {start_alloc / 1e9:.2f} GB")
print(f"After generation:        {end_alloc / 1e9:.2f} GB")
print(f"Peak during generation:  {peak_alloc / e9:.2f} GB")
print(f"Delta (KV cache + acts): {(peak_alloc - start_alloc) / 1e9:.2f} GB")

The delta between peak and the post-load baseline is the most informative number: it isolates the memory consumed by the KV cache and activations during generation, separate from the static weight footprint.

Measuring the KV Cache Growth Curve

The KV cache grows token-by-token during autoregressive decoding. To see this curve, you can hook into the generation loop and sample memory after each step. The cleanest approach is to use a custom LogitsProcessor that records memory on every call:

import torch
from transformers import LogitsProcessor

class MemoryTracker(LogitsProcessor):
    def __init__(self, device):
        self.device = device
        self.samples = []

    def __call__(self, input_ids, scores):
        alloc = torch.cuda.memory_allocated(self.device)
        step = input_ids.shape[1]
        self.samples.append((step, alloc / 1e9))
        return scores

tracker = MemoryTracker("cuda:0")

with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        logits_processor=[tracker],
    )

for step, gb in tracker.samples[:10]:
    print(f"step={step:4d}  allocated={gb:.3f} GB")

# Compute the per-token KV cache slope
if len(tracker.samples) >= 2:
    first_step, first_gb = tracker.samples[0]
    last_step, last_gb = tracker.samples[-1]
    slope = (last_gb - first_gb) / (last_step - first_step)
    print(f"KV cache growth: ~{slope * 1024:.2f} MB per token")

The slope, in MB per token, is one of the most useful numbers in LLM capacity planning. Multiply it by your maximum sequence length and batch size to estimate the worst-case KV cache footprint.

Capturing a Detailed Memory Timeline

For deeper analysis, PyTorch can record every allocation and free as a timeline. This is invaluable for spotting fragmentation, identifying which module allocates the most, and correlating memory spikes with specific operations. The output is a JSON snapshot that can be visualized in Chrome's chrome://tracing tab or with Perfetto:

import torch

# Start recording with stack traces for every allocation
torch.cuda.memory._record_memory_history(
    max_entries=100000,
    device="cuda:0",
)

with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
    )

# Stop recording and dump the snapshot
torch.cuda.memory._dump_snapshot("llm_inference_memory.pickle")
torch.cuda.memory._record_memory_history(device="cuda:0", enabled=None)

print("Snapshot written. Load it into the PyTorch memory visualizer:")
print("https://pytorch.org/memory_viz")

Open the resulting pickle file in the PyTorch memory viz tool. You will see a stack chart of allocations over time, color-coded by category. Look for tall, persistent bars (those are your KV cache) and brief spikes (those are prefill activations or kernel workspaces). Fragmentation appears as a "Swiss cheese" pattern where free blocks are scattered between live allocations.

Using torch.profiler for Combined CPU/GPU/Memory Analysis

Memory does not exist in isolation — it is tied to kernel execution and CPU overhead. torch.profiler captures all three in a single trace. The following example records a generation pass and exports a Chrome trace plus a key averages table:

from torch.profiler import profile, ProfilerActivity, schedule

activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]

with profile(
    activities=activities,
    schedule=schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=lambda p: p.export_chrome_trace("llm_trace.json"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    with torch.inference_mode():
        for _ in range(5):
            _ = model.generate(
                **inputs,
                max_new_tokens=64,
                do_sample=False,
            )
            prof.step()

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

Sorting by self_cuda_memory_usage reveals which operations allocate the most GPU memory. For LLMs, you will typically see attention kernels, layer norm, and the final LM head near the top. The Chrome trace lets you hover over any kernel to see its memory allocations and frees, which is essential for understanding why a particular decoding step spikes.

Profiling vLLM and Other Inference Engines

Production inference engines like vLLM manage their own memory pools, so PyTorch's allocator APIs report only part of the picture. vLLM exposes a Prometheus metrics endpoint that reports KV cache utilization directly. To enable it, launch the server with metrics enabled:

vllm serve meta-llama/Llama-2-7b-hf \
  --port 8000 \
  --gpu-memory-utilization 0.9 \
  --max-model-len 4096

Then scrape the metrics endpoint:

import requests

metrics = requests.get("http://localhost:8000/metrics").text

# Filter for the memory-related metrics
for line in metrics.splitlines():
    if any(k in line for k in [
        "vllm:gpu_cache_usage_perc",
        "vllm:num_preemption",
        "vllm:request_queue_time",
        "vllm:gpu_prefix_cache_hit_rate",
    ]):
        print(line)

The vllm:gpu_cache_usage_perc metric tells you what fraction of the KV cache blocks are in use. If it consistently approaches 1.0, you are at capacity and requests will be preempted (visible in vllm:num_preemption). Preemptions are a strong signal that you should either lower --max-model-len, reduce the batch size, or move to a larger GPU.

Profiling with NVIDIA Nsight Systems

When you need kernel-level attribution, Nsight Systems is the gold standard. It captures every CUDA allocation, kernel launch, and synchronization event. Run your inference script under nsys:

nsys profile \
  --trace=cuda,nvtx,osrt \
  --output=llm_inference \
  --force-overwrite=true \
  python inference_benchmark.py

Open the resulting llm_inference.nsys-rep in the Nsight Systems GUI. The memory timeline view shows allocations color-coded by stream, and you can correlate spikes with the kernel that caused them. This is the tool to reach for when PyTorch's profiler cannot explain a memory regression — for example, when a new attention kernel silently allocates a larger workspace buffer.

Best Practices

Always Warm Up Before Profiling

The first inference call triggers lazy initialization: cuBLAS handle creation, kernel autotuning, and CUDA graph capture. Profile these separately if you care about cold-start, but for steady-state analysis, run a few warmup generations before recording:

# Warmup
with torch.inference_mode():
    for _ in range(3):
        _ = model.generate(**inputs, max_new_tokens=8, do_sample=False)

torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats("cuda:0")

# Now profile the real run
with torch.inference_mode():
    _ = model.generate(**inputs, max_new_tokens=256, do_sample=False)

Profile with Realistic Workloads

Memory behavior depends heavily on sequence length distribution, batch composition, and sampling parameters. A profile run with a single 32-token prompt tells you almost nothing about production behavior. Always profile with representative inputs — ideally replayed from production logs. If you serve mixed short and long requests, profile both extremes and the median.

Separate Prefill from Decode

Prefill (processing the prompt) and decode (generating tokens) have radically different memory profiles. Prefill allocates large activation tensors proportional to prompt length squared in naive attention, while decode allocates mostly KV cache. Profile them separately to understand which phase is the bottleneck:

torch.cuda.reset_peak_memory_stats("cuda:0")
with torch.inference_mode():
    # Prefill only: forward pass on the prompt
    _ = model(**inputs)
prefill_peak = torch.cuda.max_memory_allocated("cuda:0")

torch.cuda.reset_peak_memory_stats("cuda:0")
with torch.inference_mode():
    # Decode only: generate one token at a time from cached state
    past = model(**inputs, use_cache=True)
    for _ in range(64):
        past = model(past_key_values=past.past_key_values,
                     inputs_embeds=model.get_input_embeddings()(
                         torch.tensor([[past.logits[0, -1].argmax()]],
                                      device="cuda:0")))
decode_peak = torch.cuda.max_memory_allocated("cuda:0")

print(f"Prefill peak: {prefill_peak / 1e9:.2f} GB")
print(f"Decode peak:  {decode_peak / 1e9:.2f} GB")

Watch for Fragmentation

PyTorch's caching allocator can hold large reserved memory even when allocated memory is low, because freed blocks may not be coalescible. Compare memory_allocated against memory_reserved:

allocated = torch.cuda.memory_allocated("cuda:0")
reserved = torch.cuda.memory_reserved("cuda:0")
fragmentation = 1.0 - (allocated / reserved) if reserved > 0 else 0.0
print(f"Allocated:   {allocated / 1e9:.2f} GB")
print(f"Reserved:    {reserved / 1e9:.2f} GB")
print(f"Fragmentation: {fragmentation:.1%}")

If fragmentation is consistently above 20%, consider calling torch.cuda.empty_cache() between batches, or switching to a memory pool manager like vLLM's PagedAttention, which avoids fragmentation by design.

Automate Profiling in CI

Memory regressions are easy to introduce and hard to catch manually. Add a profiling step to your CI pipeline that runs a fixed workload and asserts the peak memory stays below a budget. A simple guard looks like this:

import sys

BUDGET_GB = 16.0  # Adjust to your GPU's capacity

torch.cuda.reset_peak_memory_stats("cuda:0")
with torch.inference_mode():
    _ = model.generate(**inputs, max_new_tokens=256, do_sample=False)

peak_gb = torch.cuda.max_memory_allocated("cuda:0") / 1e9
print(f"Peak memory: {peak_gb:.2f} GB (budget: {BUDGET_GB} GB)")

if peak_gb > BUDGET_GB:
    print(f"FAIL: peak memory {peak_gb:.2f} GB exceeds budget {BUDGET_GB} GB")
    sys.exit(1)
print("PASS")

This catches regressions from model changes, kernel updates, or framework upgrades before they reach production.

Conclusion

Profiling memory usage during LLM inference is the bridge between "it works on my GPU" and "it serves production traffic reliably." By combining PyTorch's lightweight allocator APIs for quick measurements, torch.profiler and memory snapshots for detailed attribution, engine-level metrics for production visibility, and Nsight Systems for kernel-level debugging, you build a complete picture of where bytes go and why. The discipline pays off in concrete ways: tighter batch sizing, fewer OOM crashes, validated optimizations, and confident capacity planning. Start with peak allocation tracking, graduate to timeline capture when you need to explain a specific behavior, and automate the whole process in CI so regressions never slip through. Memory is the scarcest resource in LLM serving, and profiling is how you make every byte count.