Profiling Inference Time with Python

Hello! Welcome to the third lesson in our module on the Request-to-Token Mental Model.

Introduction

In our previous lesson, we mapped the conceptual stages of inference—tokenization, prefill, decode, and detokenization—to the function calls and internal logic of the Hugging Face Transformers library. You learned that model.generate() is not a monolithic operation but an orchestration of a one-time prefill step and an iterative decode loop, with the KV cache being the critical component connecting them.

Today, we move from mapping concepts to measuring reality. Our goal is to measure the wall-clock time for each inference stage using Python profiling tools on a small model. This is a foundational skill in systems engineering; you cannot optimize what you cannot measure. By the end of this lesson, you will be able to:

• Use appropriate tools to accurately time GPU operations.
• Implement a manual generation loop to isolate and measure the prefill and decode stages.
• Use the PyTorch Profiler to gain deeper insights into the performance of a standard generate call.

This exercise will transform our abstract understanding of "prefill is compute-bound" and "decode is memory-bound" into concrete, quantifiable data.

Why We Must Measure Each Stage Separately

As a quick refresher, the performance characteristics of prefill and decode are fundamentally different.

• Prefill: Processes a long sequence of tokens in parallel. This involves large matrix multiplications and is highly parallelizable, making it compute-intensive.
• Decode: Processes one token at a time. This is a series of smaller operations that are latency-sensitive and often bottlenecked by the speed of memory access (reading the KV cache).

A simple timer around model.generate() would give us a single number, blurring these two distinct phases. To diagnose bottlenecks and optimize effectively, we must measure them independently. The document you reviewed earlier emphasizes this exact point.

The Right Tools for Timing GPU Operations

Measuring operations on a GPU requires care. Because the CPU issues commands to the GPU asynchronously, a standard Python timer like time.perf_counter() can be misleading. The CPU might finish queueing an operation long before the GPU has actually finished executing it.

To get accurate wall-clock time for GPU execution, we have two primary tools in PyTorch:

1. Synchronization: We can force the CPU to wait for the GPU to complete all pending work using torch.cuda.synchronize(). This allows us to use time.perf_counter() correctly.
2. CUDA Events: A more idiomatic and precise method is to use torch.cuda.Event. We can record events in the GPU's execution stream and then ask for the elapsed time between them, avoiding CPU-GPU synchronization issues during the measurement itself.

The following resource provides a clear, practical guide and code example for using torch.cuda.Event, which we will adopt for our hands-on work.

As the article notes, it's also crucial to perform "warm-up" runs. The first time a CUDA kernel is executed, it involves compilation overhead that doesn't reflect steady-state performance. We must discard timings from these initial runs.

Hands-On: A Manual Profiling Loop

To measure prefill and decode separately, we will manually implement the autoregressive generation loop that model.generate() abstracts away.

Here is the complete script. We will:

1. Set up the model and tokenizer.
2. Perform warm-up runs.
3. Time the tokenization step on the CPU.
4. Time the prefill step on the GPU.
5. Enter a loop to time each individual decode step.
6. Collect and analyze the timings.

import torch
import time
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM




# --- 1. SETUP ---
model_id = "gpt2" # Using a small model for speed
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cpu":
    print("Warning: Running on CPU. Timing results will not reflect GPU performance.")

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
model.eval() # Set model to evaluation mode

prompt = "The future of AI systems engineering is"
max_new_tokens = 20




# --- 2. WARM-UP RUNS ---
# Run a few generations to let CUDA kernels compile and warm up the caches
print("Performing warm-up runs...")
for _ in range(3):
    _ = model.generate(
        tokenizer(prompt, return_tensors="pt").to(device).input_ids,
        max_new_tokens=5,
        pad_token_id=tokenizer.eos_token_id
    )
print("Warm-up complete.")
print("-" * 30)





# --- 3. PROFILING ---




# --- A. Tokenization ---
t0 = time.perf_counter()
inputs = tokenizer(prompt, return_tensors="pt").to(device)
input_ids = inputs.input_ids
torch.cuda.synchronize() # Wait for data transfer to finish
t1 = time.perf_counter()
tokenization_time_ms = (t1 - t0) * 1000

generated_ids = input_ids.clone()
all_token_ids = [input_ids.squeeze().tolist()]




# Create CUDA events for timing
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)




# --- B. Prefill ---
with torch.no_grad():
    start_event.record()



    # First forward pass (prefill)
    outputs = model(input_ids=input_ids)
    past_key_values = outputs.past_key_values
    next_token_logits = outputs.logits[:, -1, :]
    next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(0)
    end_event.record()

torch.cuda.synchronize() # Wait for the event recording to finish
prefill_time_ms = start_event.elapsed_time(end_event)

generated_ids = torch.cat([generated_ids, next_token], dim=-1)
all_token_ids.append(next_token.squeeze().item())




# --- C. Decode ---
decode_times_ms = []
for _ in range(max_new_tokens - 1):
    with torch.no_grad():
        start_event.record()



        # Subsequent forward passes (decode)
        outputs = model(input_ids=next_token, past_key_values=past_key_values)
        past_key_values = outputs.past_key_values
        next_token_logits = outputs.logits[:, -1, :]
        next_token = torch.argmax(next_token_logits, dim=-1).unsqueeze(0)
        end_event.record()
    
    torch.cuda.synchronize()
    decode_times_ms.append(start_event.elapsed_time(end_event))
    
    generated_ids = torch.cat([generated_ids, next_token], dim=-1)
    all_token_ids.append(next_token.squeeze().item())





# --- 4. ANALYSIS & OUTPUT ---
avg_decode_time_ms = np.mean(decode_times_ms)
total_generation_time_ms = prefill_time_ms + sum(decode_times_ms)
final_text = tokenizer.decode(generated_ids.squeeze(), skip_special_tokens=True)

print("--- Profiling Results ---")
print(f"Tokenization: {tokenization_time_ms:.2f} ms")
print(f"Prefill Time: {prefill_time_ms:.2f} ms")
print(f"Average Decode Time per Token: {avg_decode_time_ms:.2f} ms")
print(f"Total Generation Time (Prefill + Decode): {total_generation_time_ms:.2f} ms")
print("-" * 30)
print(f"Generated Text: '{final_text}'")

When you run this code, you should observe that the Prefill Time is significantly larger than the Average Decode Time per Token. This is the empirical evidence for the different performance profiles we've been discussing. The prefill processes multiple tokens at once, resulting in a single, large upfront cost. The decode phase processes one token at a time, resulting in a smaller, repeated cost.

From Manual Timing to the PyTorch Profiler

Manually instrumenting code is excellent for understanding the high-level stages, but it's cumbersome for deep-diving into which specific CUDA kernels are slow. For that, we use the built-in PyTorch Profiler. It traces all PyTorch operations on both the CPU and GPU, providing a detailed breakdown of execution time and memory usage.

A key feature of the profiler is its ability to export a trace file that can be visualized in a Chrome browser (chrome://tracing), giving you a timeline of every operation. This is an indispensable tool for serious performance engineering.

Let's watch a brief tutorial on how to use the profiler and analyze its output.

Now, let's apply this to our model.generate() call. We no longer need the manual loop; we can simply wrap the standard Hugging Face method and let the profiler do the work.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM




# --- SETUP (same as before) ---
model_id = "gpt2"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
model.eval()

prompt = "The future of AI systems engineering is"
inputs = tokenizer(prompt, return_tensors="pt").to(device)




# --- PROFILING WITH TORCH PROFILER ---
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'),
    record_shapes=True,
    with_stack=True,
    profile_memory=True
) as prof:
    for _ in range(5): # The schedule will only record steps 3, 4, 5
        with torch.no_grad():
            _ = model.generate(**inputs, max_new_tokens=20, pad_token_id=tokenizer.eos_token_id)
        prof.step() # Signal the profiler to move to the next step in the schedule




# --- ANALYSIS & OUTPUT ---
# The profiler automatically saves the trace to the './log' directory.
# You can also print a summary table.
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

After running this code, a .json trace file will be created in the log directory. Open a Chrome browser, navigate to chrome://tracing, and load this file. You will see a detailed timeline. By zooming in, you can identify the large initial block of computation corresponding to prefill and the subsequent series of smaller, repeating blocks corresponding to the decode steps. This visual representation powerfully confirms our manual measurements.

Conclusion

In this lesson, you have moved from theory to practice by measuring the wall-clock time of key inference stages.

Key Takeaways:

• Accurate GPU timing requires handling asynchronicity, making torch.cuda.Event the preferred tool over simple CPU timers.
• By manually implementing the generation loop, we can isolate and time the prefill and decode stages, empirically confirming that prefill has a high initial cost while decode has a lower, repeated cost per token.
• The PyTorch Profiler automates this process, providing detailed traces of CPU and GPU activity that can be visualized to identify bottlenecks at the kernel level.
• Our measured Prefill Time is the main contributor to Time-To-First-Token (TTFT), a key user-facing metric.
• Our measured Average Decode Time per Token is the Inter-Token Latency (ITL) or Time Per Output Token (TPOT), which determines the streaming speed of the response.

Preview of the Next Lesson:

We have successfully measured the duration of the prefill and decode stages. In our next lesson, "Analyze profiling data to characterize the distinct compute and memory profiles of the prefill vs. autoregressive decode phases," we will dive deeper into the why. We'll analyze the profiling data we just learned to generate, connecting the timing differences to the underlying computational workloads (e.g., matrix multiplication sizes) and memory access patterns that define these two critical phases of LLM inference.