← Back to DevBytes
Developer Tutorial August 18, 2026 12 min read

How to Profile LLM Inference Bottlenecks with PyTorch Profiler

Introduction to Profiling LLM Inference

Large Language Models (LLMs) have transformed how developers build applications, but their inference workloads are notoriously resource-intensive. Whether you are serving a 7B parameter model or a 70B one, understanding where time is spent during inference is critical for delivering responsive user experiences and controlling infrastructure costs. PyTorch Profiler is one of the most powerful tools available for diagnosing these performance issues at a granular level.

This tutorial walks you through what PyTorch Profiler is, why it matters for LLM workloads specifically, how to integrate it into your inference pipeline, and the best practices that will help you extract actionable insights without drowning in noise.

What Is PyTorch Profiler?

PyTorch Profiler is a built-in instrumentation tool that collects performance metrics during model execution. It captures CPU operations, CUDA kernel launches, memory allocations, and synchronization events, then presents them through a rich set of views including trace timelines, operator summaries, and memory profiles. The results can be exported to a Chrome Trace JSON file and visualized in chrome://tracing or in the PyTorch Profiler TensorBoard plugin.

For LLM inference, the profiler is particularly valuable because generation is an iterative process. Each token produced involves a forward pass through the transformer, and small inefficiencies compound across hundreds of decoding steps. A 5 millisecond overhead per token becomes a 5 second delay on a 1000-token response.

Why Profiling LLM Inference Matters

LLM inference has several characteristics that make profiling essential:

Without profiling, developers often guess at bottlenecks and waste effort optimizing the wrong layer. The profiler replaces guesswork with evidence.

Setting Up Your Environment

Before diving into profiling, make sure you have the necessary dependencies. You will need PyTorch with CUDA support, the profiler utilities, and optionally TensorBoard for visualization.

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install tensorboard torch-tb-profiler transformers accelerate

Verify that CUDA is available and that your GPU is recognized:

import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"Device: {torch.cuda.get_device_name(0)}")
print(f"PyTorch version: {torch.__version__}")

A Baseline LLM Inference Loop

Let us start with a simple inference loop using a Hugging Face model. This will be the code we instrument with the profiler. We use a small model here for demonstration, but the same approach scales to larger models.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-3.2-1B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
)
model.eval()

def generate_text(prompt, max_new_tokens=128):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# Warm-up run (important before profiling)
_ = generate_text("Hello, how are you?", max_new_tokens=10)

The warm-up run is essential. The first inference call triggers CUDA kernel compilation, memory allocation, and caching that would otherwise pollute your profiling data. Always discard the first run.

Wrapping Inference with PyTorch Profiler

Now we wrap the generation call with the profiler. The key configuration options determine what gets recorded and how much overhead is introduced.

from torch.profiler import profile, ProfilerActivity, schedule

def profiled_generate(prompt, max_new_tokens=128):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with profile(
        activities=[
            ProfilerActivity.CPU,
            ProfilerActivity.CUDA,
        ],
        schedule=schedule(
            wait=1,       # skip first step (warm-up within profile)
            warmup=1,     # discard this step's data
            active=3,     # record these steps
            repeat=1,     # do one cycle
        ),
        on_trace_ready=torch.profiler.tensorboard_trace_handler("./logs/llm_profile"),
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
    ) as prof:
        with torch.inference_mode():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )

    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

result = profiled_generate("Explain quantum computing in simple terms.", max_new_tokens=64)
print(result)

Let us break down the configuration:

Reading the Profiler Output

After running the profiled generation, you can inspect results programmatically or through TensorBoard. Let us look at the programmatic summaries first.

Operator-Level Summary

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

This prints a table of the top 20 operators ranked by total CUDA time. For LLM inference, you will typically see operations like aten::mm, aten::scaled_dot_product_attention, aten::index, and various element-wise kernels. Pay attention to:

Grouping by Shape

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

This is particularly useful for LLMs because the same operator runs with different shapes during prefill (processing the full prompt at once) versus decoding (processing one token at a time with a growing KV cache).

Memory Profile

print(prof.key_averages().table(
    sort_by="self_cuda_memory_usage",
    row_limit=15,
))

This reveals which operators allocate the most GPU memory. In LLM inference, attention layers and the KV cache are usually the top consumers. If you see unexpected allocations, it may indicate that your KV cache is not being reused efficiently or that you are recomputing something that could be cached.

Exporting and Visualizing Traces

The TensorBoard trace handler we configured automatically writes trace files to the ./logs/llm_profile directory. To view them, launch TensorBoard:

tensorboard --logdir=./logs/llm_profile --port=6006

Open http://localhost:6006 in your browser and navigate to the PyTorch Profiler plugin. You will see several views:

You can also export a Chrome Trace JSON manually for quick inspection without TensorBoard:

prof.export_chrome_trace("./llm_trace.json")

Then open chrome://tracing in Chrome and load the file. The timeline view lets you zoom into individual decoding steps and see the exact sequence of CPU calls and GPU kernels.

Profiling the Prefill vs. Decode Phases Separately

LLM inference has two distinct phases with very different performance characteristics. The prefill phase processes the entire prompt in parallel and is compute-bound. The decode phase generates one token at a time and is memory-bandwidth-bound. Profiling them together can obscure the real bottleneck.

To profile them separately, you can run the model's forward pass manually instead of using model.generate:

def profile_prefill_and_decode(prompt, max_new_tokens=64):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to("cuda")

    # --- Profile prefill phase ---
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
    ) as prefill_prof:
        with torch.inference_mode():
            outputs = model(input_ids)
            past_key_values = outputs.past_key_values
            next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1, keepdim=True)

    print("=== PREFILL PHASE ===")
    print(prefill_prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

    # --- Profile decode phase ---
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
    ) as decode_prof:
        with torch.inference_mode():
            for _ in range(max_new_tokens):
                outputs = model(
                    next_token,
                    past_key_values=past_key_values,
                    use_cache=True,
                )
                past_key_values = outputs.past_key_values
                next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1, keepdim=True)

    print("=== DECODE PHASE ===")
    print(decode_prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))

    prefill_prof.export_chrome_trace("./prefill_trace.json")
    decode_prof.export_chrome_trace("./decode_trace.json")

profile_prefill_and_decode("Write a short poem about the ocean.", max_new_tokens=64)

This separation makes it immediately clear whether your bottleneck is in the initial prompt processing or in the token-by-token generation loop. In most production deployments, the decode phase dominates total latency, so that is usually where optimization effort should focus.

Identifying Common LLM Inference Bottlenecks

Once you have profiler output in hand, here are the patterns to look for and what they typically indicate:

1. Kernel Launch Overhead

If you see many small CUDA kernels with very short execution times but significant gaps between them on the trace timeline, you are likely CPU-bound on kernel launches. This is common in decode phase where each step involves many small operations. Solutions include using CUDA graphs, fused kernels, or compiled models with torch.compile.

2. Host-Device Synchronization

Look for cudaDeviceSynchronize or cudaStreamSynchronize events in the trace. These indicate the CPU is waiting for the GPU. In LLM inference, this often happens during sampling (when logits are moved to CPU for multinomial sampling) or when logging metrics. Minimize sync points by keeping operations on GPU and batching metric collection.

3. Memory Bandwidth Saturation

During decode, the model weights and KV cache must be loaded from GPU memory for every token. If CUDA time is dominated by memory-bound kernels (like aten::mm with low arithmetic intensity), you are memory-bandwidth-bound. Solutions include quantization (int8 or int4), weight tying, or using models with fewer parameters.

4. Attention Computation Scaling

Use the shape-grouped view to check how attention cost scales with sequence length. If you see quadratic growth, consider switching to FlashAttention or a variant that supports sliding window attention. The profiler will show aten::scaled_dot_product_attention and you can verify whether the FlashAttention kernel is being selected.

5. KV Cache Memory Pressure

In the memory view, watch for steadily growing allocations during decode. If memory usage climbs faster than expected, you may have a KV cache implementation that is reallocating rather than pre-allocating. Pre-allocating the cache with the maximum sequence length avoids fragmentation and allocation overhead.

Using torch.compile with Profiling

torch.compile can dramatically reduce kernel launch overhead and fuse operations. You should profile both before and after applying it to measure the actual improvement.

# Compile the model
compiled_model = torch.compile(model, mode="reduce-overhead")

def profile_compiled_generate(prompt, max_new_tokens=64):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    # Warm up compiled model (first calls trigger compilation)
    for _ in range(3):
        with torch.inference_mode():
            _ = compiled_model(**inputs)

    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
    ) as prof:
        with torch.inference_mode():
            output_ids = compiled_model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )

    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=15))
    prof.export_chrome_trace("./compiled_trace.json")
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

result = profile_compiled_generate("List three benefits of exercise.", max_new_tokens=64)

Compare the compiled trace to the eager trace. You should see fewer, larger kernels and reduced gaps between them. The reduce-overhead mode uses CUDA graphs under the hood, which is especially effective for the decode phase where kernel launch overhead dominates.

Best Practices for Profiling LLM Inference

Putting It All Together: A Reusable Profiling Utility

Here is a reusable utility function that encapsulates the best practices discussed above. You can drop this into your project and call it whenever you need to profile a generation.

import torch
from torch.profiler import profile, ProfilerActivity, schedule
from contextlib import contextmanager
import os
from datetime import datetime

@contextmanager
def llm_profiler(
    trace_dir="./profiler_logs",
    trace_name=None,
    record_shapes=True,
    profile_memory=True,
    with_stack=False,
    wait=1,
    warmup=1,
    active=5,
):
    os.makedirs(trace_dir, exist_ok=True)
    if trace_name is None:
        trace_name = f"llm_trace_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    trace_path = os.path.join(trace_dir, trace_name)

    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        schedule=schedule(wait=wait, warmup=warmup, active=active, repeat=1),
        on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_path),
        record_shapes=record_shapes,
        profile_memory=profile_memory,
        with_stack=with_stack,
    ) as prof:
        yield prof

    # Print summary after the context exits
    print("\n=== Top Operators by CUDA Time ===")
    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

    if profile_memory:
        print("\n=== Top Operators by Memory Usage ===")
        print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10))

    # Also export a Chrome trace
    chrome_path = os.path.join(trace_dir, f"{trace_name}.json")
    prof.export_chrome_trace(chrome_path)
    print(f"\nChrome trace saved to: {chrome_path}")
    print(f"TensorBoard logs saved to: {trace_path}")


# Usage example
def benchmark_generation(prompt, max_new_tokens=128):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    # Warm up
    with torch.inference_mode():
        _ = model.generate(**inputs, max_new_tokens=5, pad_token_id=tokenizer.eos_token_id)

    # Profile
    with llm_profiler(trace_name="baseline_generation"):
        with torch.inference_mode():
            output_ids = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=tokenizer.eos_token_id,
            )

    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

result = benchmark_generation("What are the main causes of climate change?", max_new_tokens=100)
print(result)

This utility handles directory creation, naming, warm-up separation, and output formatting. It prints both the time and memory summaries and exports traces for visualization. You can easily toggle with_stack on when you need deeper debugging.

Conclusion

Profiling LLM inference with PyTorch Profiler transforms optimization from guesswork into a data-driven process. By instrumenting your generation loop, separating prefill and decode phases, and learning to read the operator, memory, and trace views, you can pinpoint exactly where latency originates and direct your optimization effort where it matters most. Whether you are evaluating torch.compile, testing a quantization scheme, or tuning your batching strategy, the profiler gives you the before-and-after evidence you need to make confident decisions. Start with a warm-up run, profile a realistic workload, examine the trace timeline for gaps and sync points, and iterate. The bottlenecks you find will rarely be where you expected them, and that is precisely why profiling is indispensable for anyone serious about efficient LLM inference.

— Ad —

Google AdSense will appear here after approval

← Back to all articles