vllm - 💡(How to fix) Fix [RFC]: Profiler support for CUDA graph capture tracing and roofline trace annotations [1 pull requests]

Motivation.

vLLM's PyTorch profiler integration today gives a good view of actual model execution under a real workload, but it has two gaps that make trace-driven performance analysis (especially attention roofline analysis) harder than it needs to be.

1. The CUDA graph capture phase is invisible to the profiler

When vLLM runs with CUDA graphs (cudagraph_mode=FULL, PIECEWISE, or FULL_AND_PIECEWISE), the profiler trace of the actual workload records graph replay: the GPU kernels execute from a pre-recorded graph, so the replay trace has GPU kernels but no CPU-side call stack and no input shapes for the ops inside the graph. The CPU operators, python call stack, and tensor shapes only exist during the graph capture phase at server startup, which the profiler never sees today.

This means that for any deployment using full/piecewise CUDA graphs (the common production configuration), a profiler trace cannot attribute GPU time to module / op / shape without additional, manual instrumentation. The information needed to enrich the replay trace already exists at capture time; we simply do not record it.

2. The default execution-step annotation lacks the statistics needed for roofline analysis — even though vLLM already computes them

Each execution step is annotated via torch.record_function, but the default annotation only encodes request and token counts:

execute_context_1(839)_generation_63(63)

For variable-length batched inference, the per-step attention compute and memory traffic depend on more than token counts. As derived in the attention roofline model, the only per-request quantities that appear in the FLOPs / bytes formulas are N_Q(i), N_KV(i), and their products, and these can be reduced to a handful of per-phase aggregates (computed separately for context and generation requests):

Because the default annotation does not carry Σ N_KV, Σ (N_Q·N_KV), or Σ (N_Q²), an analyst cannot reconstruct the attention roofline (arithmetic intensity, achieved vs. peak FLOPs/bandwidth) for a step directly from the trace. They must re-run with detailed logging or correlate against scheduler state that is not present in the trace file.

The detailed annotation emits that same reduction into each step's record_function event. The consequence is that all statistics needed for the roofline live in the single profiler trace file — kernel timings, shapes, and per-step compute/memory aggregates co-located. This is also complementary to the capture trace, which supplies the graph-mode shapes/call-stacks for those same kernels.

Proposed Change.

Add two opt-in, default-off ProfilerConfig flags (prototyped in PR #37524). Both are no-ops unless the torch profiler is enabled, so default behavior is unchanged.

Flag 1: capture_torch_profiler_dir (str, default "")

Absolute path to a directory in which to write a PyTorch profiler trace of the CUDA graph capture phase (rank 0 only). Requires --profiler-config.profiler torch. When set, the capture workflow in gpu_model_runner.py wraps graph capture in a torch.profiler context (and nullcontext() otherwise), producing one trace per captured graph (batch size × compilation mode):

torch_trace/
|-- capture_traces
|   |-- graph_capture_rank_0.<...>.pt.trace.json.gz   # per captured graph
|   `-- ...
|-- <workload replay trace>.pt.trace.json.gz
`-- profiler_out_0.txt

These capture traces retain the CPU-side call stack and input shapes for every op that ends up inside a replayed CUDA graph. A downstream tool can match a replay step (by batch size / compilation mode) to its capture trace and back-fill shapes and call stacks onto the replay tree.

Flag 2: detailed_trace_annotation (bool, default False)

When True, the torch.record_function annotation for each execution step is extended with per-phase roofline aggregates for context (prefill) and generation (decode) requests. The annotation format/labels are unchanged from the current prototype (this RFC does not propose changing the emitted string); variable names in the implementation were made more descriptive following review feedback.

Detailed format:

execute_886_context_1(sq824sk824sqsq678976sqsk678976)_generation_62(sq62sk93055sqsq62sqsk93055)

where, aggregated per phase:

• sq = Σ N_Q (scheduled/new tokens; Q read + output write traffic)
• sk = Σ N_KV (total KV length; K/V read traffic)
• sqsq = Σ (N_Q · N_Q) (causal self-attention correction term, dominant in pure prefill)
• sqsk = Σ (N_Q · N_KV) (QKᵀ / score·V compute term, dominant in decode / chunked prefill)
• leading bs = Σ N_Q across all requests (effective batch size)

Default (False) annotation is unchanged: execute_context_1(839)_generation_63(63).

These are written as user_annotation events, so roofline analysis is possible directly from the trace with no extra runtime logging.

Feedback Period.

No response

CC List.

@rasmith @njhill @WoosukKwon

Any Other Things.

How these features are used (offline analysis)

Both flags are additive metadata on the existing profiler trace — extra user_annotation events and an extra set of capture trace files — so they slot into any existing trace-based analysis flow without changing how the trace is collected or read. Any tool that already walks a PyTorch profiler trace can pick up the per-step roofline aggregates and the capture-phase shapes/call-stacks with no new instrumentation.

As a concrete example, TraceLens is an open-source PyTorch trace analysis toolkit. It builds a unified CPU/GPU call tree from a profiler trace, attaches per-op input shapes, and attributes GPU kernel time to analytic compute/memory performance ("perf") models per op category (GEMM, MoE, attention, norm, etc.) to place each op on a roofline. On top of that core it offers inference-specific workflows: trace splitting into per-step / per-phase (prefilldecode-mix, decode-only) windows, and a capture-merge step that augments graph-replay traces with the call-stacks and shapes recovered from graph-capture traces. The two vLLM flags map directly onto two TraceLens capabilities:

1. Attention roofline for variable-length inference. TraceLens parses the detailed annotation directly from the trace's user_annotation events and feeds the per-phase aggregates into its attention perf model:

• TraceLens/PerfModel/extensions/attention_perf_model_extensions.py consumes those fields in flops_func() and bytes_func(). The FLOPs implementation is the same equation as vLLM's perf.py (and the roofline derivation):

  # attention_perf_model_extensions.py (flops_func)
  ctx_flops_qk = H_Q * (2 * c_sqsk * d_h_qk) - H_Q * (1 * c_sqsq * d_h_qk)
  ctx_flops_pv = H_Q * (2 * c_sqsk * d_h_v)  - H_Q * (1 * c_sqsq * d_h_v)
  gen_flops_qk = H_Q * (2 * g_sqsk * d_h_qk)
  gen_flops_pv = H_Q * (2 * g_sqsk * d_h_v)

  i.e. Σ N_Q·N_KV drives the QKᵀ/score·V rectangle and Σ N_Q² is the causal prefill correction — exactly the sqsk / sqsq annotation terms. This yields per-step arithmetic intensity and achieved-vs-peak roofline placement, separating prefill-bound (compute) from decode-bound (memory) behavior, all from the single trace.

2. Analysis of the graph-captured region. Capture traces supply the shapes and call-stacks that graph-replay traces lack, so graph-mode kernels can be categorized and attributed to modules and shapes. At a high level, each captured graph is labeled with its batch size and compilation mode, and a runtime execution step is mapped to the matching captured graph using two signals already present in the runtime trace: the step's batch size and its number of graph-launch events (one *GraphLaunch for a FULL graph, several for a PIECEWISE graph). For example, a runtime decode step with batch size 32 and a single graph launch maps to the captured graph labeled {batch_size: 32, mode: FULL}; the analysis tool then back-fills that capture's call-stacks and shapes onto the step's replayed kernels.

Example TraceLens output. Running generate_perf_report_pytorch_inference with --capture_folder on the graph_full test trace produces a unified_perf_summary where each op carries its input shapes/dtypes (recovered from the capture trace), a modeled GFLOPs / data-moved / arithmetic-intensity, the underlying kernel(s), and its share of step time. Top rows (abbreviated) for a batch-size-32 FULL-graph decode step:

The InferenceAttention row's perf params are populated directly from the detailed annotation for that step (taken verbatim from the test's unified_perf_summary.csv):

{ 'B': 1, 'N_Q': 32, 'H_Q': 64, 'H_KV': 8, 'd_h_qk': 64, 'd_h_v': 64,
  'causal': False,
  'c_sq': 0, 'c_sk': 0, 'c_sqsq': 0, 'c_sqsk': 0,
  'g_sq': 32, 'g_sk': 55086, 'g_sqsq': 32, 'g_sqsk': 55086 }

This is a pure decode step (c_* = 0): 32 generation requests contribute g_sq = 32 query tokens against g_sk = 55086 cached KV tokens, and g_sqsk = 55086 drives the modeled attention FLOPs — values that exist in the trace only because detailed_trace_annotation emitted them. Without the flag, the same row would have no sk/sqsk and the attention roofline could not be computed for that step.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.