torch.cuda.graph_annotations.get_kernel_annotations

torch.cuda.graph_annotations.get_kernel_annotations() → Mapping[int, list]

Return the live registry of recorded kernel annotations.

Keys are opaque integers matching the graph node id field that CUPTI-based profilers attach to kernel events; values are the lists of annotation dicts recorded for that node. The registry accumulates across captures and is global to the process.

The returned mapping is a live view: it is updated in place when a graph is instantiated (annotation keys are rekeyed to the executable graph’s ids), so a reference obtained early stays current. Keys are valid for joining against a profiler trace once the corresponding graphs have been instantiated. Treat the mapping as read-only; snapshot it with dict(...) if isolation is needed.

Warning

This API is in prototype and may change in future releases.

Example:

>>> annotations = torch.cuda.graph_annotations.get_kernel_annotations()
>>> with open("annotations.pkl", "wb") as f:
...     pickle.dump(dict(annotations), f)

Return type:

Mapping[int, list[Any]]