Add CUDA Graph annotation postprocessing

## Human Note

## Agent note

Turn captured `mark_kernels` metadata into GPU annotation spans automatically for both Kineto and
the CUPTI monitor backend. The postprocessor joins graph and node IDs to the live annotation
registry, groups contiguous nodes with the same label, and emits boxes that the existing TrackEvent
writer places on its dedicated GPU annotations track while preserving overlap lane assignment.

The monitor keeps its native direct-export path when no graph annotations exist. Annotated traces
use the JSON merge path so the shared postprocessors can run. On a representative 1.8K-event trace,
annotation synthesis measured about 0.60 ms median; complete native TrackEvent export measured about
43 ms median.

## Test Plan

```bash
~/.venvs/nightly/bin/python -m pytest test/test_profiler.py test/test_perfetto.py -q
uvx ruff check transformer_nuggets/utils/benchmark.py transformer_nuggets/utils/perfetto.py transformer_nuggets/utils/__init__.py test/test_profiler.py test/test_perfetto.py
uvx ruff format --check transformer_nuggets/utils/benchmark.py transformer_nuggets/utils/perfetto.py transformer_nuggets/utils/__init__.py test/test_profiler.py test/test_perfetto.py
prek
```

Author: drisspg

test/test_perfetto.py

@@ -4,6 +4,7 @@
 import pytest
 
 from transformer_nuggets.utils.perfetto import (
+    add_cuda_graph_annotation_boxes,
     default_trace_path,
     default_track_event_path,
     read_trace,
@@ -17,6 +18,87 @@ def _duration_events(trace):
     return [event for event in trace["traceEvents"] if event.get("ph") == "X"]
 
 
+def test_cuda_graph_annotations_become_contiguous_gpu_boxes():
+    graph_id = 2
+    annotations = {
+        (graph_id << 32) | 1: [{"name": "attention"}],
+        (graph_id << 32) | 2: [{"name": "attention"}],
+        (graph_id << 32) | 3: [{"name": "loss"}],
+    }
+    trace = {
+        "traceEvents": [
+            {
+                "ph": "X",
+                "cat": "kernel",
+                "name": "kernel_a",
+                "pid": 0,
+                "tid": 7,
+                "ts": 10,
+                "dur": 3,
+                "args": {"graph id": graph_id, "graph node id": 1},
+            },
+            {
+                "ph": "X",
+                "cat": "kernel",
+                "name": "kernel_b",
+                "pid": 0,
+                "tid": 7,
+                "ts": 13,
+                "dur": 5,
+                "args": {"graph id": graph_id, "graph node id": 2},
+            },
+            {
+                "ph": "X",
+                "cat": "kernel",
+                "name": "kernel_c",
+                "pid": 0,
+                "tid": 7,
+                "ts": 18,
+                "dur": 2,
+                "args": {"graph id": graph_id, "graph node id": 3},
+            },
+        ]
+    }
+
+    processed = add_cuda_graph_annotation_boxes(trace, annotations)
+    boxes = [
+        event for event in processed["traceEvents"] if event.get("cat") == "gpu_user_annotation"
+    ]
+
+    assert [(box["name"], box["ts"], box["dur"]) for box in boxes] == [
+        ("attention", 10.0, 8.0),
+        ("loss", 18.0, 2.0),
+    ]
+    assert all(box["tid"] == 7 for box in boxes)
+    assert len(trace["traceEvents"]) == 3
+
+
+def test_cuda_graph_annotation_boxes_accept_monitor_embedded_metadata():
+    trace = {
+        "traceEvents": [
+            {
+                "ph": "X",
+                "cat": "kernel",
+                "name": "kernel",
+                "pid": 0,
+                "tid": 7,
+                "ts": 10,
+                "dur": 3,
+                "args": {
+                    "graph id": 2,
+                    "graph node id": 1,
+                    "annotation": '[{"name": "attention"}]',
+                },
+            }
+        ]
+    }
+
+    processed = add_cuda_graph_annotation_boxes(trace)
+
+    assert processed["traceEvents"][-1]["name"] == "attention"
+    assert processed["traceEvents"][-1]["cat"] == "gpu_user_annotation"
+
+
 def test_split_overlapping_slices_creates_adjacent_lanes():
     trace = {
         "traceEvents": [
@@ -163,7 +245,7 @@ def test_default_track_event_path_uses_native_perfetto_suffix():
 
 
 def test_track_event_conversion_preserves_instants_counters_and_warns_on_unsupported():
-    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import TrackEvent, Trace
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import Trace, TrackEvent
 
     trace = {
         "traceEvents": [
@@ -187,7 +269,7 @@ def test_track_event_conversion_preserves_instants_counters_and_warns_on_unsuppo
 
 
 def test_track_event_conversion_puts_gpu_annotations_on_separate_track():
-    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import TrackEvent, Trace
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import Trace, TrackEvent
 
     trace = {
         "traceEvents": [
@@ -224,7 +306,7 @@ def test_track_event_conversion_puts_gpu_annotations_on_separate_track():
 
 
 def test_track_event_conversion_attaches_paired_flows_to_slices():
-    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import TrackEvent, Trace
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import Trace, TrackEvent
 
     trace = {
         "traceEvents": [
@@ -255,9 +337,9 @@ def test_track_event_conversion_attaches_paired_flows_to_slices():
 
 def test_track_event_conversion_splits_crossing_slices_and_keeps_nested_slices():
     from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import (
+        Trace,
         TrackDescriptor,
         TrackEvent,
-        Trace,
     )
 
     trace = {
@@ -302,7 +384,7 @@ def test_track_event_conversion_splits_crossing_slices_and_keeps_nested_slices()
 
 
 def test_track_event_conversion_keeps_back_to_back_slices_separate():
-    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import TrackEvent, Trace
+    from perfetto.protos.perfetto.trace.perfetto_trace_pb2 import Trace, TrackEvent
 
     trace = {
         "traceEvents": [

transformer_nuggets/utils/__init__.py

@@ -20,6 +20,7 @@
 from transformer_nuggets.utils.merge_traces import merge_traces
 from transformer_nuggets.utils.perfetto import (
     TraceFormat,
+    add_cuda_graph_annotation_boxes,
     default_trace_path,
     default_track_event_path,
     perfetto_trace_path,

transformer_nuggets/utils/benchmark.py

@@ -21,6 +21,7 @@
 from torch.utils import benchmark
 
 from transformer_nuggets.utils.perfetto import (
+    add_cuda_graph_annotation_boxes,
     perfetto_trace_path,
     read_trace,
     write_perfetto_trace,
@@ -551,6 +552,15 @@ def benchmark_cuda_function_in_microseconds_triton(func: Callable, *args, **kwar
         return do_bench(no_args) * 1e3
 
 
+def _cuda_graph_annotations() -> dict[int, list[object]]:
+    """Return captured CUDA Graph annotations when the prototype API is available."""
+    try:
+        from torch.cuda.graph_annotations import get_kernel_annotations
+    except ImportError:
+        return {}
+    return get_kernel_annotations()
+
+
 def _write_profiler_trace(
     prof: torch.profiler.profile,
     trace_path: Path,
@@ -561,8 +571,12 @@ def _write_profiler_trace(
     gzip_trace: bool = False,
 ) -> None:
     """Export a torch profiler trace through the canonical Perfetto writer."""
+    annotations = _cuda_graph_annotations()
     can_export_direct_json = (
-        trace_format == "chrome_json" and not split_overlaps and trace_path.suffix != ".gz"
+        trace_format == "chrome_json"
+        and not split_overlaps
+        and trace_path.suffix != ".gz"
+        and not annotations
     )
     if can_export_direct_json:
         prof.export_chrome_trace(str(trace_path))
@@ -573,7 +587,7 @@ def _write_profiler_trace(
         prof.export_chrome_trace(str(export_path))
         write_perfetto_trace(
             trace_path,
-            read_trace(export_path),
+            add_cuda_graph_annotation_boxes(read_trace(export_path), annotations),
             trace_format=trace_format,
             split_overlaps=split_overlaps,
             track_pattern=track_pattern,
@@ -605,8 +619,37 @@ def _write_cupti_monitor_trace(
     trace_path: Path,
     *,
     trace_format: Literal["chrome_json", "track_event"],
+    split_overlaps: bool = True,
+    track_pattern: str | None = "stream.*",
+    gzip_trace: bool = False,
 ) -> None:
-    """Export a monitor trace while hiding its unconditional ``.gz`` suffix."""
+    """Export a monitor trace and apply graph-aware postprocessing when needed."""
+    annotations = _cuda_graph_annotations()
+    if annotations:
+        export_path = trace_path.with_name(f"{trace_path.stem}.monitor-export.json")
+        monitor_path = Path(f"{export_path}.gz")
+        export_path.unlink(missing_ok=True)
+        monitor_path.unlink(missing_ok=True)
+        try:
+            prof.export_chrome_trace(str(export_path))
+            generated_path = export_path if export_path.exists() else monitor_path
+            if not generated_path.exists():
+                raise FileNotFoundError(
+                    f"CUPTI monitor export produced neither {export_path} nor {monitor_path}"
+                )
+            write_perfetto_trace(
+                trace_path,
+                add_cuda_graph_annotation_boxes(read_trace(generated_path), annotations),
+                trace_format=trace_format,
+                split_overlaps=split_overlaps,
+                track_pattern=track_pattern,
+                gzip_trace=gzip_trace,
+            )
+        finally:
+            export_path.unlink(missing_ok=True)
+            monitor_path.unlink(missing_ok=True)
+        return
+
     export_path = trace_path.with_name(f"{trace_path.name}.monitor-export")
     monitor_path = Path(f"{export_path}.gz")
     export_path.unlink(missing_ok=True)
@@ -1008,7 +1051,14 @@ def profiler(
 
     def trace_handler(prof) -> None:
         if use_cupti_monitor:
-            _write_cupti_monitor_trace(prof, path, trace_format=trace_format)
+            _write_cupti_monitor_trace(
+                prof,
+                path,
+                trace_format=trace_format,
+                split_overlaps=fix_overlapping_events,
+                track_pattern=overlap_track_pattern,
+                gzip_trace=gzip_trace,
+            )
             return
         _write_profiler_trace(
             prof,

transformer_nuggets/utils/perfetto.py

@@ -22,9 +22,9 @@
 import re
 import zlib
 from collections import defaultdict
+from collections.abc import Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
-from collections.abc import Iterator
 from pathlib import Path
 from re import Pattern
 from typing import Any, Literal
@@ -34,7 +34,6 @@
     write_track_event_trace,
 )
 
-
 TraceFormat = Literal["chrome_json", "track_event"]
 """Perfetto-compatible output format selector."""
 
@@ -67,6 +66,105 @@ def write_trace(path: str | Path, trace: dict[str, Any], *, indent: int | None =
         json.dump(trace, f, indent=indent)
 
 
+def _annotation_name(entries: Sequence[Any] | None) -> str | None:
+    name = None
+    for annotation in entries or ():
+        if isinstance(annotation, dict) and "name" in annotation:
+            name = str(annotation["name"])
+        elif isinstance(annotation, str):
+            name = annotation
+    return name
+
+
+def _graph_annotation_box(
+    name: str,
+    pid: Any,
+    tid: Any,
+    start: float,
+    end: float,
+) -> dict[str, Any]:
+    return {
+        "ph": "X",
+        "cat": "gpu_user_annotation",
+        "name": name,
+        "pid": pid,
+        "tid": tid,
+        "ts": start,
+        "dur": end - start,
+        "args": {"transformer_nuggets.graph_annotation": True},
+    }
+
+
+def add_cuda_graph_annotation_boxes(
+    trace: dict[str, Any],
+    annotations: Mapping[int, Sequence[Any]] | None = None,
+) -> dict[str, Any]:
+    """Add GPU annotation spans for contiguous CUDA Graph regions.
+
+    Graph annotation metadata is joined by ``(graph id, graph node id)`` when an
+    annotation registry is supplied. CUPTI monitor JSON traces can also carry the
+    annotation list directly in each GPU event's ``annotation`` argument.
+    """
+    events = list(trace.get("traceEvents", []))
+    if any(event.get("args", {}).get("transformer_nuggets.graph_annotation") for event in events):
+        return trace.copy()
+
+    annotated_work: dict[tuple[Any, Any], list[tuple[dict[str, Any], str]]] = defaultdict(list)
+    for event in events:
+        if event.get("ph") != "X" or event.get("cat") not in {
+            "kernel",
+            "gpu_memcpy",
+            "gpu_memset",
+        }:
+            continue
+        args = event.get("args", {})
+        graph_id = args.get("graph id")
+        graph_node_id = args.get("graph node id")
+        if graph_id is None or not graph_node_id:
+            continue
+
+        entries = None
+        if annotations is not None:
+            entries = annotations.get((int(graph_id) << 32) | int(graph_node_id))
+        if entries is None and isinstance(args.get("annotation"), str):
+            try:
+                embedded = json.loads(args["annotation"])
+            except json.JSONDecodeError:
+                embedded = None
+            if isinstance(embedded, list):
+                entries = embedded
+        name = _annotation_name(entries)
+        if name is not None:
+            annotated_work[(event.get("pid"), event.get("tid"))].append((event, name))
+
+    annotation_boxes: list[dict[str, Any]] = []
+    for (pid, tid), stream_events in annotated_work.items():
+        region_name = None
+        region_start = 0.0
+        region_end = 0.0
+        for event, name in sorted(stream_events, key=lambda item: item[0]["ts"]):
+            start = float(event["ts"])
+            end = start + float(event.get("dur", 0.0))
+            if name != region_name:
+                if region_name is not None:
+                    annotation_boxes.append(
+                        _graph_annotation_box(region_name, pid, tid, region_start, region_end)
+                    )
+                region_name = name
+                region_start = start
+                region_end = end
+            else:
+                region_end = max(region_end, end)
+        if region_name is not None:
+            annotation_boxes.append(
+                _graph_annotation_box(region_name, pid, tid, region_start, region_end)
+            )
+
+    processed = trace.copy()
+    processed["traceEvents"] = [*events, *annotation_boxes]
+    return processed
+
+
 def default_trace_path(file_path: str | Path, *, gzip_by_default: bool = True) -> Path:
     """Return a Chrome JSON trace path, treating suffix-less paths as stems.