Performance guide¶

Steering or patching over many forwards, or through generation? Read Performance: choosing a fast path first: the capture-free tl.when(site, action).bind(model) runs at about the cost of a plain forward hook.

TorchLens is fastest when it captures only the payloads you plan to inspect. A full tl.trace(model, x) records a complete operation graph and saves activations for the selected sites; tl.record(model, x, save=...) is the lighter path for tight loops where you only need selected records and can materialize a Trace later.

Backend note: these timings and sparse-recording examples are torch-oriented. tl.record() / sparse recording (tl.record()) is torch-only in the backend-v1 registry. Non-torch preview backends may have different capture costs; JAX, tinygrad, Paddle, and TensorFlow .tlspec saves materialize array payloads but loaded traces cannot replay-validate stripped runtime captures. Paddle preview capture has its own dygraph/eager replay and static-inventory audit costs; TensorFlow preview capture runs live eager op_callbacks plus self-consistency and per-op replay/perturbation accounting.

Decision tree¶

Fast activation pull¶

import torch
from torch import nn
import torchlens as tl
from torchlens.options import CaptureOptions


model = nn.Sequential(nn.Linear(4, 4), nn.ReLU(), nn.Linear(4, 2)).eval()
x = torch.randn(2, 4)

trace = tl.trace(
    model,
    x,
    save=tl.func("relu"),
    capture=CaptureOptions(save_source_context=False),
)
relu_site = trace.find_sites(tl.func("relu")).first()
relu_activation = relu_site.out

assert relu_activation.shape == (2, 4)

Use save=tl.func(...), save=tl.in_module(...), or composed predicates such as tl.func("conv2d") & tl.followed_by(tl.func("relu")) instead of saving every payload. The graph metadata remains available, but unsaved payloads are intentionally absent.

Sparse recording loop¶

import torch
from torch import nn
import torchlens as tl


model = nn.Sequential(nn.Linear(4, 4), nn.GELU(), nn.Linear(4, 2)).eval()
x = torch.randn(2, 4)

recording = tl.record(model, x, save=tl.func("gelu"))
trace = recording.to_trace()
gelu_site = trace.find_sites(tl.func("gelu")).first()

assert gelu_site.out.shape == (2, 4)

tl.record(save=...) is the only torch sparse-capture spelling; the old keep_op=/keep_module= alias kwargs are removed.

When a tl.record() forward raises, the default remains on_forward_error="raise". Opt into on_forward_error="attach_partial" to attach exc.partial_recording and re-raise, or on_forward_error="return_partial" to return a failed partial Recording (its output refuses with output_forward_failed). Failed partials set status="partial_error", failed=True, string-only error metadata, num_committed_ops, and best-effort last_event_* fields. user-op failures exclude the failing call; TL-side capture failures may include a skipped/partial current-call event. Failed partials cannot be converted with Recording.to_trace() or used with Recording.backward(). MemoryError is swallowed like any other forward exception under "return_partial"/"attach_partial", and materializing the partial allocates more memory inside an already-exhausted heap — under genuine memory pressure keep the default on_forward_error="raise". Full tl.trace(...) failures raise tl.errors.TraceCaptureError with the partial capture on exc.partial_trace.

Speed knobs¶

Escape diagnostics¶

The sys.modules crawler is deleted (replaced by the rescue re-run + the mechanical belt), so the default capture pays no crawl cost. The callable detector is a separate, opt-in diagnostic cost:

import torch
from torch import nn
import torchlens as tl
from torchlens.backends.torch.wrappers import unwrap_torch, wrap_torch


model = nn.ReLU()
x = torch.randn(4)
wrap_torch(escape_detector="shadow")
trace = tl.trace(model, x)
print(trace.health.num_escape_detector_events, trace.health.escape_detector_callback_ns)
unwrap_torch()

CaptureOptions(aten_ops=True) records the ATen rows behind each op and aten_ops="values" also keeps what each ATen call returned (AtenOp.out). Measured on one 4-thread CPU worker (torch 2.13, medians of 3): ResNet-18 at 224px, plain trace 0.53 s, rows 0.64 s, rows with values 0.61 s (107 rows, 44 MiB of ATen values against 31 MiB of Op.out); GPT-2 small at 1x32 tokens, plain trace 2.18 s, rows 2.9 s, rows with values 2.8 s (518 rows, 256 MiB against 96 MiB). The dispatch mode is the cost; the values add memory (roughly 1.4x to 2.7x the saved Op.out bytes, since every intermediate ATen result is kept), not time. With the switch off nothing is installed.

The independent aten-level completeness witness is also opt-in and can run alone or alongside the callable detector. This is the only route to capture_verified=True. A plain tl.trace() never arms it: every default capture reads capture_verified=None, which every honesty surface renders as "not checked" -- not a clean bill. Escapes the wrappers themselves observe (worker-thread ops, torch.jit.script regions, direct torch.ops.aten.* calls, .data.copy_()) still ceiling the default capture to False, and tensor-to-scalar escapes (.item(), int(argmax)) raise the scalar_escape advisory; but host tensor sources the wrappers cannot see -- torch.from_numpy(...) / torch.as_tensor(np_array) / torch.from_dlpack(...) round-trips, an autograd.Function whose forward runs in NumPy, and in-forward storage writes through .numpy() / untyped_storage() aliases -- leave NO ceiling on the default capture (the op is visible to a graph reader as an internal-source op with parents=()). A default capture whose recorded ops include torch.from_numpy, torch.from_dlpack, torch.frombuffer, or torch.as_tensor on non-tensor data records one host_tensor_source row in trace.annotations["capture_advisories"] (count, sources and the first location when code context was kept; no warning, no verdict change), so the shared artifact says host data entered the graph unchecked. The other host sources above record nothing. Arm the witness (or run tl.validate) when that class of escape matters:

import torch
from torch import nn
import torchlens as tl
from torchlens.backends.torch.wrappers import unwrap_torch, wrap_torch


model = nn.ReLU()
x = torch.randn(4)
wrap_torch(
    escape_detector="shadow",
    completeness_witness=True,
)
trace = tl.trace(model, x)
print(trace.health.completeness_verified)
print(trace.health.num_unaccounted_witness_events)
unwrap_torch()

A value-free capture (structure_only=True, on meta or real tensors) is never verified: under the witness it still reads capture_verified=None, and completeness_witness_verified records the census result.

The witness attaches each aten dispatch to the live wrapper token, func_call_id, and leaf barcode; it does not correlate by clock time. A captured Python call may own several ordered aten operations, recorded in trace.health.completeness_decompositions. Dispatches with no owner, or owned by a wrapper whose func_call_id emits no logged op, produce a CaptureGapWarning, set trace._health_completeness_verified=False, and append structured rows to trace.completeness_diagnostics. Work under pause_logging() is outside the census. The documented torch.func/functorch transform boundary runs its interior under that paused scope, and fused or opaque kernel interiors below the dispatcher are not observable; the dispatched boundary operator itself is still accounted normally.

completeness_witness_verified=True makes only this dispatch-census claim: every owner-thread aten event observed during active logging was correlated to an emitted op or to an exact audited boundary. It does not prove that every user call route entered active logging. In particular, a transform callable built before wrapping can escape without a transform boundary op; TorchLens detects that route and leaves the census result True while setting capture_verified=False with capture_verification_reason="transform_call_route_unverified". escape_detector_verified=True makes the separate claim that the callable detector saw no observable raw-call escape. capture_verified=True is the combined result only when the enabled census/detector checks pass, no raw-transform escape was detected, and the owner-thread qualification remains valid. A bypassed torch.compile region reports the more specific capture_verification_reason="dynamo_region_not_logged" in preference to any of the above, since Dynamo's compile threads and unaccounted aten dispatches are symptoms of that one region.

The honest rollout comparison is legacy with no guard versus scoped with the requested guard, not crawl time in isolation. On Python 3.10–3.11, shadow mode uses sys.setprofile and can be expensive for Python-call-heavy models: a representative call-heavy 16-layer MLP measurement on Python 3.11 was +371% versus the unguarded capture. Treat shadow as an expensive, diagnostic-only soak tool, not a production capture setting. On Python 3.12+ it prefers local Python-start monitoring plus diagnostic CALL events. Shadow remains default-off.

The dispatcher witness is likewise default-off. Its diagnostic cost is environment- and workload- dependent; use an alternating on/off measurement on the target host rather than relying on a fixed cross-hardware percentage. On a Python 3.11.6 / PyTorch 2.8.0 CPU environment, a 16-pair Linear(64, 64)/ReLU MLP at batch size 16 measured +17.8% median overhead: 0.374 s/capture with the witness off and 0.440 s/capture with it on. Ten alternating batches of five full tl.trace calls per mode followed four two-capture warmup batches; the per-batch ranges were 0.341–0.379 s off and 0.415–0.453 s on. The witness-off route does not enter the witness-only pause_logging() contexts used to keep TorchLens bookkeeping dispatches out of the user-op census.

Phase timing buckets¶

trace._phase_timings groups wall-clock timings by stable snake_case bucket names. Capture buckets include ctx_build_*, dispatch_*, clone_save_*, backward_grad_event_* and object_construction_op. Postprocess buckets use the postprocess pipeline's step names (add_output_layers through release_param_refs). Graphviz rendering records render_graphviz_forward, render_graphviz_backward, or render_graphviz_combined when those render entrypoints run.

Chunked forward capture¶

import torch
from torch import nn
import torchlens as tl


model = nn.Sequential(nn.Linear(1024, 1024), nn.ReLU()).eval()
x = torch.randn(128, 1024)

trace = tl.trace(model, x, save=tl.func("relu"), capture=tl.options.CaptureOptions(chunk_size=16))

assert trace.find_sites(tl.func("relu")).first().out.shape[0] == 128

chunk_size= is forward-only sugar over a first trace(...) followed by trace.append(...) for later chunks. It splits selected positional tensor leaves along dimension 0, executes one sub-batch at a time, and returns one accumulated in-memory Trace.

The v1 limits are intentionally narrow: torch backend only, positional inputs only, no backward_ready, no save_grads, no live hooks=, no public intervene=, and no storage=tl.to_disk(...)/storage=. Loaded or live chunked traces also reject backward() because they do not retain one full-batch autograd graph.

Auto mode splits only when there is exactly one ndim > 0 tensor leaf under standard Python containers (list, tuple, dict, or namedtuple). If there are several candidates, pass explicit paths:

import torch
from torch import nn
import torchlens as tl


class MaskedModel(nn.Module):
    def forward(self, tokens, attention_mask):
        return (tokens * attention_mask).sum(dim=-1)


model = MaskedModel().eval()
tokens = torch.randn(64, 16)
attention_mask = torch.ones(64, 16)

trace = tl.trace(model, (tokens, attention_mask), chunk_paths=["0", "1"], capture=tl.options.CaptureOptions(chunk_size=8))

Unlisted leaves are passed unchanged to every chunk, which is useful for shared masks or bias tables. The memory contract is forward-pass peak only: final saved activations are concatenated and retained according to save=, and preprocessing still sees the full batch if CaptureOptions(input_transform=...) materializes it.

For activation extraction without a Trace, use tl.extract_dataset(...) (the former tl.batched_extract alias is removed); that path returns tensors or .pt files rather than accumulated graph metadata. chunk_size= covers the remaining "dataloader wrapper" case for stacked multi-pass trace capture. Disk mode (path=) writes atomic shards plus a self-describing manifest.json (site identity, stimulus provenance, axis semantics, dtypes); an interrupted long extraction continues from its last completed shard with resume=True, and torchlens.load(path) reads the artifact back with its metadata (both names are experimental).

Windowed and disk-backed capture¶

from pathlib import Path

import torch
from torch import nn
import torchlens as tl


model = nn.Sequential(nn.Conv2d(1, 2, 3), nn.ReLU()).eval()
x = torch.randn(1, 1, 5, 5)
path = Path(DOCS_TMPDIR) / "windowed.tlspec"

predicate = tl.func("conv2d") & tl.followed_by(tl.func("relu"))
trace = tl.trace(
    model,
    x,
    save=predicate,
    storage=tl.to_disk(path), capture=tl.options.CaptureOptions(lookback=4, lookback_payload_policy="detached_raw"),
)

assert trace.find_sites(tl.func("conv2d")).first().out.shape == (1, 2, 3, 3)

Use disk-backed storage for selected payloads that are too large or numerous to keep in memory. Trace repr, notebook HTML, and JSON explanation do not materialize lazy disk payloads for their NaN/Inf summary; those refs are reported as unexamined until you call op.payload.materialize_out(). This exemption applies to predicate-selected disk-only payloads. Exhaustive saving (save="all", the default) keeps RAM copies until postprocess, so combine a narrower save= with disk streaming to reduce peak retained memory. Note that save= itself takes a predicate, selector, or CaptureOptions — never the string "all". Portable .tlspec/ bundles store manifest data plus tensor sidecars when the backend supports materialized payloads; executable Python callables are not portable. Backend-aware manifest schema v2 adds backend, backend_runtime, nullable torch-specific fields, and payload_policy. JAX, tinygrad, Paddle, and TensorFlow preview bundles materialize array payloads; loaded traces still report replay validation as unavailable because portable save strips runtime replay captures.

Intervention cost¶

import torch
from torch import nn
import torchlens as tl


model = nn.Sequential(nn.Linear(4, 4), nn.ReLU(), nn.Linear(4, 2)).eval()
x = torch.randn(2, 4)

patched = tl.trace(
    model,
    x,
    save=tl.func("relu"),
    intervene=tl.when(tl.func("relu"), tl.zero_ablate()),
)

assert torch.count_nonzero(patched.find_sites(tl.func("relu")).first().out) == 0

Use intervene= when the edited value must affect downstream execution. For post-hoc experiments on an existing trace, prefer trace.fork() and do(site, edit), which stages the edit and pushes it downstream in one call (which engine fits which question: intervention explainers).

Long sequences and long contexts¶

Saved activations grow with sequence length at every op, and attention probabilities grow with its square: one layer's (batch, heads, q, kv) matrix for a 4096-token sequence with 32 heads in bfloat16 is 32 x 4096 x 4096 x 2 bytes, 1 GiB, so a 32-layer model's set is 32 GiB. Plan what to keep before the first long capture:

• Name the sites with save= (tl.module("model.layers.20"), tl.func("linear") & tl.in_module("model.layers.20")) instead of the whole-graph default.
• CaptureOptions(save_budget=...) (default "auto") caps saved bytes; an exceeded budget raises SaveBudgetExceededError rather than exhausting memory, and output_device="cpu" keeps saved activations in host memory.
• For a logit-lens sweep over many positions, torchlens._semantic.logit_lens_predictions projects one layer at a time and keeps only top-k summaries.
• For many forwards or generation, a bound spec (and its generate) saves nothing at all; editing the cached past of a long prompt is the KV-cache recipe.
• For a dataset of long sequences, tl.extract_dataset(..., path=...) writes the named layers to disk batch by batch.

Counting the backward that ran¶

tl.report.backward_status(trace) reports the compute of the backward pass that actually ran, counted call by call. When a capture requested a backward (backward_ready=True, save_grads= or default_grad=), every ATen call that backward dispatches is recorded on its BackwardPass.aten_ops, with its shapes, its dtypes and the GradFn call it ran under. Each call is then counted with the same convention as torch.utils.flop_counter: one multiply-add counts 2 FLOPs, and only matmul and convolution calls are counted. Any other call is listed by name in uncounted_ops; it is never counted as zero and never estimated. Frozen weights, partial requires_grad and second-order passes are counted exactly as they ran. A create_graph=True gradient and the backward run through it are separate passes in executed_flops_per_pass.

import torch
from torch import nn
import torchlens as tl

model = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 4))
model[0].requires_grad_(False)  # frozen first layer, and x needs no gradient
x = torch.randn(2, 16)

trace = tl.trace(model, x, capture=tl.options.CaptureOptions(backward_ready=True))
trace.backward(trace[trace.output_layers[0]].out.sum())

status = tl.report.backward_status(trace)
# Only the last layer's weight gradient runs (2 * 2 * 32 * 4 FLOPs): nothing upstream of
# it needs a gradient, so no input-gradient matmul is dispatched.
assert status.executed_flops == 2 * 2 * 32 * 4
print(status.status_line)
print(status.uncounted_op_counts)  # the bias-gradient sums and other calls, named, not zeroed

from torch.nn.attention import SDPBackend, sdpa_kernel


class Attend(nn.Module):
    def __init__(self):
        super().__init__()
        self.qkv = nn.Linear(8, 24)

    def forward(self, x):
        q, k, v = self.qkv(x).view(1, 6, 3, 2, 4).permute(2, 0, 3, 1, 4)
        return nn.functional.scaled_dot_product_attention(q, k, v)


def executed(model, x):
    trace = tl.trace(model, x, capture=tl.options.CaptureOptions(backward_ready=True))
    trace.backward(trace[trace.output_layers[0]].out.sum())
    return tl.report.backward_status(trace).executed_flops


attend, x = Attend(), torch.randn(1, 6, 8)
fused = executed(attend, x)  # CPU picks its fused attention kernel
with sdpa_kernel(SDPBackend.MATH):
    assert executed(attend, x) == fused  # the decomposed math counts the same

Calls TorchLens makes inside its own GradFn posthooks (gradient logging and the backward edits they apply) are not recorded; the copies its tensor hooks make when saving gradients are recorded as zero-arithmetic calls and never change the count. The record exists only in the session that ran the backward: a saved and reloaded trace reports the executed backward as unknown. A capture that requests no backward never installs the recorder and pays nothing for it. The training-cost estimate tl.report.backward_estimate(trace) is a separate figure, always labeled hypothetical, and it never appears where an executed cost belongs.

What is counted, and where FlopCounterMode differs¶

Forward, each op's flops_forward follows its rule, with one multiply-add counted as 2 FLOPs:

• A linear or convolution weight computed inside forward (weight norm, spectral norm, a weight passed in as an input) is read from the call's arguments, as a parameter weight is.
• The fused recurrent kernels behind nn.RNN, nn.GRU, nn.LSTM and their cells count one multiply-add per weight-matrix element per output row, plus their bias adds; the gate nonlinearities inside the kernel are not counted. A sequence kernel's whole call is counted on its sequence output and its final states read 0. lstm_cell returns two outputs of one shape that cannot be told apart, so half the call is counted on each.
• torch.linalg.multi_dot counts its chain in the order it runs, the cheapest one.
• einsum equations with ... are counted; repeated labels within one operand and equations over three or more operands stay unknown.
• mv, dot, vdot and spmm count as matrix products; a sparse spmm operand is counted as dense, from its shape.

Backward, tl.report.backward_status(trace) also counts two fused kernels as the math they replace: the CPU fused attention backward (_scaled_dot_product_flash_attention_for_cpu_backward: the four matrix products of the attention gradient over the full query x key square, which is what the MATH attention backend dispatches) and the oneDNN LSTM layer backward (mkldnn_rnn_layer_backward: dX, dW_ih, dH and dW_hh).

The count follows torch.utils.flop_counter.FlopCounterMode's convention and differs from it in two named ways. FlopCounterMode counts a grouped or depthwise convolution's weight gradient groups times over; TorchLens counts it once. And FlopCounterMode has no rule for mv, addmv, dot, vdot, the CPU fused attention kernels or the oneDNN LSTM layer kernels, so it reports 0 for that work; TorchLens counts it.

When another tool is faster¶

Use raw PyTorch when you only need the final output. Use torch.profiler when the question is kernel timing rather than activation provenance. Use TransformerLens when your workflow is entirely inside its supported transformer families and you already want its named hook points. Use Captum or Inseq when you want mature attribution algorithms out of the box rather than the lower-level activation and gradient substrate.

Measured numbers (canonical bench host)¶

Measured on the canonical bench host (Intel i9-9900X) at SHA 712a4e5 on 2026-06-16, from a full non-smoke run (baseline_status: "canonical"). The complete 196-row table (CPU + CUDA, every model and row) lives in docs/_perf_numbers.md and the raw baseline in benchmarks/perf_baselines/linux-cpu.json.

Headline: with sparse recording and an early halt at ~25% depth, capture runs faster than the raw forward pass itself — fastlog_halt_25 is 0.84x raw forward on ResNet-18 (CPU) and 0.83x on GPT-2 (HookedTransformer, CPU). You only pay for the layers you actually reach.

Representative CPU rows ("vs raw" = multiple of the raw forward pass):

(Full exhaustive capture (tl_trace) costs ~14x the forward and amortizes on large models; the fastlog_* rows show selective / early-exit capture, where halting at 25% depth drops below the raw-forward cost. TinyNet ratios are dominated by fixed per-capture overhead — see the full table.)

Re-capturing baselines (canonical host)¶

# Full suite on the canonical bench host (quiet machine):
python -m benchmarks.perf_suite --rerun --baseline-status canonical \
  --out-json benchmarks/perf_baselines/<host>-<device>.json \
  --out-md benchmarks/perf_results_<date>.md
# Sanity self-compare:
python -m benchmarks.perf_gate --baseline benchmarks/perf_baselines/<host>-<device>.json \
  --current benchmarks/perf_baselines/<host>-<device>.json   # expect "passed": true
# Regenerate the docs table:
python -m benchmarks.generate_perf_numbers benchmarks/perf_baselines/<host>-<device>.json \
  --out docs/_perf_numbers_provisional.md

On the canonical host, drop the -provisional filename suffix; --baseline-status canonical requires a full non-smoke, non-addendum run and emits generated speed headlines.

The gate judges TorchLens-owned rows on process-CPU statistics (cpu_median_ms/cpu_iqr_ms); the wall-clock fallback for pre-CPU-metric payloads is not authoritative — a TorchLens row judged on wall clock fails the gate by default (the committed 2026-06-16 linux-cpu.json baseline predates the CPU metrics, so comparisons against it need either a rebaseline on the canonical host or the explicit, disclosed --allow-wall-clock-only legacy opt-out).