Skip to content

viable/strict/1786707154

@ngimel ngimel tagged this 14 Aug 06:22
Adds a second way for mark_kernels to discover which nodes a scope contains, on
top of the shared subscriber-callback registry: a RESOURCE / GRAPHNODE_CREATED
handler records each node against whatever scope is open when CUPTI announces
its creation. Selected via torch.cuda.graph(annotation_config={"backend": ...}),
which defaults to "auto" -- CUPTI when the monitor already holds a subscription,
and otherwise today's dependent-edge walk, unchanged.

Annotation options go in a dict rather than as separate arguments: enable_annotations
has already shipped as a flat argument, but the options still to come (backward
projection, Python launch stacks) would each widen the signature again. Keys and
values are both validated, so a typo raises rather than silently leaving the
default in place.

Two measured limitations of the walk motivate the new backend. First, a scope
entered while the *current* stream is not yet capturing records nothing at all:
the walk snapshots that stream's capture state on entry, finds none, and no-ops,
even though the work inside the scope is captured. Measured 0 nodes annotated
against 3 for the same capture. Second, the walk rescans a nested scope's nodes
once per enclosing scope, so its cost is quadratic in nesting depth where
CUPTI's is linear: at depth 16 over ~2550 nodes, annotation overhead is 73.6ms
against 33.5ms, isolated from the 18.2ms the capture itself costs.

Worth being explicit about what this does NOT fix, since it is a natural thing
to assume: the walk does not miss work merely because another stream issued it.
It follows dependency edges across streams, so long as the current stream was
capturing when the scope opened.

This also folds the capture-graph-id read into one place. #191845 added
maybe_stamp_capture_root at capture_begin for its conditional-body checks, while
maybe_stamp_capture_graph_id read the same id again at capture_end for the
remap. The template graph keeps one id for its whole life, so the capture_end
read is gone and the capture_begin one stamps the graph object as well. That
drops a cudaGraphGetId per capture and leaves one definition of "this capture's
graph id", shared by the conditional-body checks, the CUPTI node filter, and
remap_to_exec_graph.

Review in dependency order: _graph_annotations.py first (the ambient scope
stack, the backend switch, and the unified stamp, all of which the handler
reads), then the new _graph_node_callbacks.py, then the config parsing and
arm/disarm wiring in graphs.py, and finally has_live_subscription() in
monitor.py -- a probe that deliberately does not construct the monitor, since
constructing it would take a CUPTI subscription and lock kineto out for the rest
of the process.

Three choices worth challenging. The single-threaded-autograd requirement gates
only the CUPTI backend rather than enable_annotations as a whole: attributing a
node needs the scope open on the thread that created it, so multithreaded
autograd would mis-attribute nodes its engine threads create -- but making that
unconditional would break every existing annotated capture, because
multithreaded autograd is on by default. So "cupti" raises, "auto" falls back to
the walk, and "edge_walk" is untouched. (capture_error_mode is not the right
gate: it scopes capture's safety checks, not which threads contribute nodes.)
Registration happens before capture_begin and arming after, so that failing to
obtain CUPTI cannot raise with a capture already underway. And nodes belonging
to a child-graph or conditional body are dropped rather than recorded: the
exec-graph ids that traces report for body work are announced by no CUPTI
callback at any point, so recording them would produce keys matching nothing.
The handler counts those drops and disarm() warns once with the total, which is
the edge walk's conditional-body warning without its per-scope cost -- and it
reports what was actually lost, so it covers a scope containing a body as well
as one inside it. The count is warned about outside the callback because a
warnings filter can promote warnings to errors, which must not reach CUPTI's C
dispatch.

One ordering subtlety to check while reviewing: _set_annotations_enabled has to
run before capture_begin, because maybe_stamp_capture_root is gated on it, while
the backend is only final after arming, which needs the capture live. Hence the
separate _set_annotation_backend rather than one call carrying both -- with them
combined the stamp silently no-ops and the conditional-body detection from

Test Plan:

New tests, including the two that carry the justification above --
test_scope_entered_before_stream_joins_capture (0 annotated under the walk, 3
under CUPTI) and test_parity_with_edge_walk. The forward-compat driver is
required for the graph annotation APIs:

```bash
LD_LIBRARY_PATH=/usr/local/cuda-13.2/compat:$LD_LIBRARY_PATH \
  python test/test_cuda_graph_utils.py TestCuptiAnnotationBackend -v
```

11 passed.

Regression coverage, since mark_kernels gained a branch and the capture-id stamp
moved:

```bash
LD_LIBRARY_PATH=/usr/local/cuda-13.2/compat:$LD_LIBRARY_PATH \
  python test/test_cuda_graph_utils.py
LD_LIBRARY_PATH=/usr/local/cuda-13.2/compat:$LD_LIBRARY_PATH \
  python test/test_cuda.py -k graph
```

65 passed, and 141 run with one failure. That failure,
test_graph_capture_error_releases_reserved_segments, is pre-existing and not
from this change: it passes in isolation, and clean origin/main reproduces the
identical 141-run/1-failure result.

test_scope_inside_conditional_body_records_nothing passing is the check that the
capture_begin stamp actually fires -- it would fail if the stamp no-opped.

Authored with assistance from Claude Code.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/191999
Approved by: https://github.com/ezyang
ghstack dependencies: #191944
Assets 2
Loading