Performance Profile

After tuning a model, AITune’s performance profile answers four runtime-attribution questions:

• How much runtime is inside AITune-tuned modules?
• How much runtime is inside other model modules that AITune did not tune?
• How much runtime is outside any module (pipeline glue, scheduler math, framework dispatch)?
• What do the underlying PyTorch operators and CUDA kernels look like?

The report combines PyTorch Profiler data with region annotations emitted from forward hooks on instrumented modules and from monkey-patched wrappers on their user-defined entry-point methods. Two kinds of regions are tracked: AITune-managed (AOT-wrapped) modules, and untuned modules discovered by walking the profiled object’s submodule tree (skipping AITune-managed and JIT-managed subtrees). The API returns an in-memory profile with structured data and a Markdown renderer. Users choose whether and where to persist JSON or Markdown; the raw PyTorch Profiler Chrome trace is exported only when you pass an explicit trace_file.

The report is factual: it summarizes measured runtime, not recommendations.

Quick Start

The typical flow is to profile before tuning (to see where time goes today and pick high-value tuning targets) and after tuning (to verify what actually moved). Both reports use the same call.

1
import json
2
from pathlib import Path
3
4
import torch
5
from diffusers import DiffusionPipeline
6
7
import aitune.torch as ait
8
from aitune.torch.backend import TorchEagerBackend
9
10
# 1. Load the pipeline you want to profile.
11
pipe = DiffusionPipeline.from_pretrained(
12
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
13
    torch_dtype=torch.float16,
14
).to("cuda")
15
16
inputs = [{"prompt": "A beautiful landscape with mountains and a lake"}]
17
18
19
def infer(prompt):
20
    return pipe(prompt, width=512, height=512, num_inference_steps=5)
21
22
23
# 2. (Optional) Baseline report before tuning. Discovered `nn.Module` entries
24
#    surface as `untuned_module:*` regions; anything else (plain `list` /
25
#    `dict` containers, root methods, pipeline glue) lands in residual. Useful
26
#    for seeing where time goes today and picking what to tune.
27
baseline = ait.profile(
28
    obj=pipe,
29
    input_data=inputs[0],
30
    inference_function=infer,
31
)
32
Path("baseline_profile.json").write_text(json.dumps(baseline.data, indent=2) + "\n")
33
Path("baseline_profile.md").write_text(baseline.markdown())
34
35
# 3. Inspect, wrap, and tune the modules that dominate runtime.
36
#    `TorchEagerBackend` is the no-compile baseline used here for a runnable
37
#    example; swap in `TensorRTBackend` / `TorchInductorBackend` / etc. for
38
#    real acceleration.
39
modules_info = ait.inspect(pipe, inputs, inference_function=infer)
40
modules = modules_info.get_modules(min_execution_percentage=0.05)
41
strategy = ait.OneBackendStrategy(TorchEagerBackend()).enable_find_max_batch_size(False)
42
pipe = ait.wrap(pipe, modules, strategy=strategy)
43
ait.tune(infer, inputs, batch_sizes=[1])
44
45
# 4. Profile the tuned pipeline. Tuned modules now appear as `aot_module:*`
46
#    regions; the rest stay as `untuned_module:*`. Comparing against the
47
#    baseline shows what tuning moved.
48
tuned = ait.profile(
49
    obj=pipe,
50
    input_data=inputs[0],
51
    inference_function=infer,
52
    warmup_runs=3,
53
    measured_runs=10,
54
    trace_file="tuned_trace.json",
55
)
56
57
Path("tuned_profile.json").write_text(json.dumps(tuned.data, indent=2) + "\n")
58
Path("tuned_profile.md").write_text(tuned.markdown())
59
print(tuned.trace_file.name)  # tuned_trace.json (full absolute path on .trace_file)

Each returned profile’s .data field (e.g. baseline.data, tuned.data) is the JSON-ready dict if you want to query in-process or persist it yourself.

Step 2 is optional but recommended for an unfamiliar pipeline: without any AOT-managed modules, the report attributes discovered module entries via untuned-module discovery (including method-style entry points like vae.decode) and leaves non-module / framework / pipeline-glue code in residual, so you can see the shape of inference before deciding what to tune. After tuning, the same call produces a report where the tuned modules are visible as aot_module:* regions.

Parameters

Warmup and measured invocations run under torch.no_grad(), matching AITune inspection and tuning.

What’s in the report

The JSON has stable top-level keys:

1
{
2
  "schema_version": 1,
3
  "created_at": "...",
4
  "aitune_version": "...",
5
  "config": { "warmup_runs": ..., "measured_runs": ..., "uses_inference_function": ... },
6
  "target": { "type": "..." },
7
  "input": { "args_count": ..., "kwargs": [...] },
8
  "runs": [...],
9
  "profiler": { "activities": [...], "key_averages": {...} },
10
  "regions": [...],
11
  "warnings": [...]
12
}

Per-run timing and attribution

Each entry in runs covers one measured iteration:

1
{
2
  "run_index": 0,
3
  "timing": { "wall_time_us": ..., "cpu_time_us": ..., "device_time_us": ... },
4
  "regions": [
5
    { "region_id": "aot_module:unet", "calls": 6,
6
      "cpu_time_us": ..., "cpu_time_fraction": ...,
7
      "device_time_us": ..., "device_time_fraction": ... }
8
  ],
9
  "residual": {
10
    "cpu_time_us": ..., "cpu_time_fraction": ...,
11
    "device_time_us": ..., "device_time_fraction": ...
12
  }
13
}

• timing — wall-clock measurement (Python-side time.perf_counter_ns) plus PyTorch Profiler’s CPU and device timings for the run.
• regions — per-region aggregates for every region observed during the run (both AOT-managed and untuned). calls aggregates multiple invocations within one measured run (e.g., a UNet called 6 times across denoising steps appears as one row with calls: 6).
• residual — time spent outside any region — neither AOT-managed nor untuned — per time domain. Negative residuals from float precision or async timing are clamped to zero.

Fractions are run-relative. Per-domain math: sum(regions[].cpu_time_fraction) + residual.cpu_time_fraction ≈ 1.0.

Region metadata

regions at the top level describes each region observed across the run. Two kinds are emitted:

1
[
2
  { "id": "aot_module:unet", "name": "unet", "kind": "aot_managed_module",
3
    "module_type": "diffusers.models.unets.unet_2d_condition.UNet2DConditionModel",
4
    "wrapper_state": "tuned" },
5
  { "id": "untuned_module:text_encoder", "name": "text_encoder", "kind": "untuned_module",
6
    "module_type": "transformers.models.clip.modeling_clip.CLIPTextModel" }
7
]

• aot_managed_module — an AITune-wrapped module. Carries wrapper_state reflecting the live wrapper FSM (tuned, recording, passthrough). If a region was observed in the profile but the wrapper is no longer in MODULE_REGISTRY (e.g., it was unwrapped between profiling and report assembly), the state is "unknown" and an UNREGISTERED_AOT_REGION warning is emitted.
• untuned_module — an nn.Module reachable from obj that is not AITune-managed and not JIT-patched. Discovered recursively, with subtrees containing managed descendants descended-into rather than attributed as a whole (so siblings of a tuned module still get their own region). wrapper_state is omitted — the field is AOT-specific.

Method-style entry points

Some pipelines invoke an nn.Module via a custom method rather than via __call__ — e.g., Diffusers’ pipe.vae.decode(z), or HuggingFace model.generate(...) when the model is held as an attribute of the profiled object rather than being the profiled object itself (see the root-method caveat below). Forward hooks alone never see those calls. Discovery additionally wraps user-defined methods declared on the module’s class hierarchy (anything beyond what nn.Module itself provides, excluding dunders, forward, and non-function descriptors). Each wrapped method produces a region whose path is <module_path>.<method_name>:

1
{ "id": "untuned_module:vae.decode", "name": "vae.decode", "kind": "untuned_module",
2
  "module_type": "diffusers.models.autoencoders.autoencoder_kl_ltx2.AutoencoderKLLTX2Video" }

Method-region module_type is the type of the underlying module (the parent path). Method wrappers preserve the original signature (inspect.signature round-trips) so framework code that introspects methods sees the real shape with the wrapper active. Methods that aren’t actually invoked during measurement produce no region rows.

This applies to untuned targets only. AOT-managed modules receive forward hooks but not method wrappers — wrapping a tuned module’s methods would overlap with its own forward span when the method internally calls __call__.

PyTorch Profiler drilldown

profiler.key_averages contains two bounded sorted views (cpu_time_total and device_time_total) of the profiler’s own per-operator aggregates. Each row aggregates across all measured_runs iterations combined:

1
{
2
  "key": "aten::scaled_dot_product_attention",
3
  "device_type": "CPU",
4
  "count": 723,
5
  "self_cpu_time_total_us": ..., "cpu_time_total_us": ...,
6
  "self_device_time_total_us": ..., "device_time_total_us": ...
7
}

Use this to find the dominant operators inside (or outside) tuned regions. The view is bounded to 20 rows per sort key to keep the JSON compact; when trace_file is provided, the full event stream lives in that Chrome trace.

The Markdown view

profile.markdown() renders a view over profile.data for readability. It contains:

• Overview — created timestamp, AITune version, target, configured warmup/measured counts, input positional-arg count and keyword-arg names.
• Runs — per-run wall, CPU, and device timing.
• Regions — region metadata table (name, kind, wrapper state, module type). For untuned regions, the State column is -.
• Per-Run Attribution — region + residual table. The _(residual)_ row appears after each run’s regions so you can visually verify the column sums to ~100%.
• Profiler Key Averages — top operators by CPU total and device total.
• Warnings — structured warnings (see below).

The Markdown is a pure view; rerunning profile.markdown() against the same data always produces identical output.

Interpreting residual

Residual is the time NOT covered by any region — neither AOT-managed nor untuned. Discovery walks the submodule tree of obj reached via direct nn.Module attributes and transparent nn.ModuleList / nn.ModuleDict containers; plain Python list and dict attributes that happen to hold modules are not entered. In typical pipelines this covers everything, and the residual is almost entirely code that runs outside any nn.Module.forward:

• Pipeline-level Python (e.g., a diffusion __call__ running scheduler.step(), classifier-free guidance combine, latent scaling, image post-processing).
• Standalone tensor ops between module calls (torch.cat, .chunk, arithmetic on logits/latents).
• Host-side CUDA driver work — cudaLaunchKernel, cudaMemcpyAsync, cudaStreamSynchronize (the last typically triggered by .item() calls inside scheduler math).
• Framework dispatch (PyTorch’s overhead between operator calls).

If a pipeline stores modules in a plain list or dict rather than a ModuleList / ModuleDict, those modules will not be discovered and their compute will land in residual. Convert to a ModuleList / ModuleDict (the PyTorch-idiomatic choice anyway) to get them attributed.

On a Stable Diffusion v1.5 pipeline with UNet + VAE decoder AOT-wrapped (and text_encoder, safety_checker, vae.post_quant_conv automatically captured as untuned regions), you typically see:

• CPU residual ~28% — scheduler math, CFG combine, latent prep, image post-processing, host-side CUDA driver overhead.
• Device residual ~0.2% — essentially all GPU work lives inside a region.

A large CPU residual on a pipeline-heavy workload is expected (Python orchestration is real). A large device residual means there’s significant GPU work happening outside any captured nn.Module.forward — that’s unusual and worth investigating in the Chrome trace.

Residual is computed only for CPU and device. Per-region wall time isn’t available from PyTorch Profiler, so residual doesn’t include a wall figure.

Configuration tips

Warmup

The default warmup_runs=3 works for most real-sized models. AITune always runs one additional warmup iteration inside the profiler to absorb buffer allocation, so a value of 0 is permitted and still produces a clean first measured run — but pure model warmup (CUDA kernel autotune, cuDNN heuristics, allocator state) benefits from a few iterations before measurement.

For sub-millisecond workloads where micro-warmup effects on the CPU side are visible (branch predictor, cache state), increase to warmup_runs=5 or higher. On real inference models (tens to hundreds of milliseconds), warmup_runs=3 is plenty.

Measured runs

measured_runs=10 gives stable averages on most workloads. Increase if you see run-to-run variance you want to characterize statistically. The cost is roughly linear: 10 measured runs of a 100 ms inference takes ~1 second to profile plus negligible analysis overhead.

Caveats

• Single-input profiling. V1 profiles one representative input scenario, repeated. Multi-scenario aggregation (different batch sizes, different prompts) is future work.
• CUDA-side device time discrepancy. In profiler.key_averages, you may see two entries per AITune region — one with device_type: "CPU" and one with device_type: "CUDA". The CPU-side row’s device_time_total is the sum of kernel device times launched within the region. The CUDA-side row’s device_time_total is the GPU-timeline span of the region, including idle gaps between kernels. Our per-region device_time_us in runs[] uses the CPU-side view (sum of kernels), which is the right value for attribution.
• Trace size. Chrome traces can be tens of megabytes for typical models. Compress before sharing if needed.
• AOT-wrapped modules that are not yet tuned. AITune wrappers in recording or passthrough state still emit their AOT region annotation. You can therefore run aitune.torch.profile(...) before tuning to scope candidate modules; the region appears with wrapper_state set to "recording" or "passthrough" rather than "tuned". (Distinct from the untuned_module kind, which refers to modules AITune never wrapped at all.)
• Methods on a bare nn.Module object are not instrumented. When obj is itself an nn.Module, discovery starts at obj.named_children() — so methods on the root (e.g. model.generate) are not wrapped, and their pipeline-glue cost (loops, sampling, KV cache management) lands in residual. Per-child compute is captured correctly. To attribute root-level methods, wrap the model in a small pipeline-like container holding it as an attribute and pass that container as obj (Container(model=my_model) → obj=container). The cleaner fix needs hierarchical attribution and is tracked as follow-up work.
• Pre-bound method references bypass the wrapper. Method wrapping monkey-patches at the instance attribute level. A reference captured before aitune.torch.profile(...) enters (saved = pipe.vae.decode) won’t go through the wrapper if invoked later. Pipeline code that resolves the attribute at call site (self.vae.decode(...)) is captured normally — this is the common case.
• Direct calls into descendants of a clean untuned parent are missed. Untuned-module discovery stops at the topmost clean module (no AITune-managed descendants) and installs hooks only there. If pipeline code skips the parent and calls a deeper submodule directly (self.block.inner(x) instead of self.block(x)), the parent hook never fires and the deeper submodule was never instrumented — the call’s compute lands in residual. Per-event nested-untuned suppression (the inverse of the existing AOT case) would let us instrument at every level safely; tracked as follow-up.

Warnings

Structured warnings appear in the warnings array with stable codes:

Relationship to other observability tools

Runtime attribution is complementary to AITune’s other observability features:

They measure different things on different time models. For a complete diagnostic view, enable all three with their respective env vars.

The Chrome trace at profile.trace_file, when requested, opens directly in Perfetto or chrome://tracing for kernel-level inspection — useful when the JSON summary points to a region of interest and you want to drill into individual operators.