Add CUPTI monitor profiler support

## Human Note

## Agent note

Add an opt-in CUPTI monitor configuration to the shared profiler wrapper, including PM metric
sampling, environment counters, CUDA Graph dependency recording, event-node attribution, and
metric discovery. Constructing the profiler now prepares the backend immediately, which lets graph
recorders arm before capture while preserving existing context-manager usage.

Monitor traces use PyTorch's native Perfetto exporter and normalize its unconditional `.gz` suffix,
so callers receive the requested `.pftrace` or JSON path instead of having to discover a differently
named artifact. The optional dependency keeps cupti-python out of the default installation.

## 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/__init__.py test/test_profiler.py
uvx ruff format --check transformer_nuggets/utils/benchmark.py transformer_nuggets/utils/__init__.py test/test_profiler.py
prek
```

Author: drisspg

pyproject.toml

@@ -42,6 +42,10 @@ dev = [
     "docstring-parser",
 ]
 
+cupti-monitor = [
+    "cupti-python>=13.3",
+]
+
 llama = [
     "sentencepiece==0.1.99",
     "datasets==2.15.0",

test/test_profiler.py

@@ -0,0 +1,136 @@
+import gzip
+from pathlib import Path
+
+import pytest
+
+from transformer_nuggets.utils import benchmark
+from transformer_nuggets.utils.benchmark import (
+    DEFAULT_CUPTI_MONITOR_PM_METRICS,
+    CuptiMonitorConfig,
+    _write_cupti_monitor_trace,
+)
+
+
+def test_cupti_monitor_defaults_enable_rich_trace_data():
+    config = CuptiMonitorConfig()
+
+    assert config.environment_counters
+    assert config.pm_metrics == DEFAULT_CUPTI_MONITOR_PM_METRICS
+    assert config.graph_dependencies
+    assert config.event_node_ids
+
+
+def test_cupti_monitor_config_builds_pytorch_options():
+    config = CuptiMonitorConfig(
+        environment_counters=True,
+        pm_metrics=("sm__cycles_active.avg.pct_of_peak_sustained_elapsed",),
+        graph_dependencies=True,
+        event_node_ids=True,
+        cuda_sync_events=True,
+        pftrace_compression_level=4,
+    )
+
+    assert config.custom_profiler_config() == {
+        "backend": "cupti_monitor",
+        "enable_environment_counters": True,
+        "enable_pm_sampling": True,
+        "pm_metrics": ["sm__cycles_active.avg.pct_of_peak_sustained_elapsed"],
+        "enable_graph_dependencies": True,
+        "enable_event_node_ids": True,
+        "enable_cuda_sync_events": True,
+        "pftrace_compression_level": 4,
+    }
+
+
+@pytest.mark.parametrize("level", [-1, 10])
+def test_cupti_monitor_config_rejects_invalid_compression(level):
+    with pytest.raises(ValueError, match="between 0 and 9"):
+        CuptiMonitorConfig(pftrace_compression_level=level)
+
+
+class GzipExportProfiler:
+    def __init__(self, payload: bytes) -> None:
+        self.payload = payload
+
+    def export_chrome_trace(self, path: str) -> None:
+        with gzip.open(f"{path}.gz", "wb") as output:
+            output.write(self.payload)
+
+
+def test_cupti_monitor_pftrace_export_hides_gzip_suffix(tmp_path):
+    path = tmp_path / "trace.pftrace"
+    profiler = GzipExportProfiler(b"compressed native trace")
+
+    _write_cupti_monitor_trace(profiler, path, trace_format="track_event")
+
+    assert path.exists()
+    assert not Path(f"{path}.gz").exists()
+    with gzip.open(path, "rb") as trace:
+        assert trace.read() == b"compressed native trace"
+
+
+def test_cupti_monitor_json_export_hides_gzip_suffix(tmp_path):
+    path = tmp_path / "trace.json"
+    profiler = GzipExportProfiler(b'{"traceEvents": []}')
+
+    _write_cupti_monitor_trace(profiler, path, trace_format="chrome_json")
+
+    assert path.read_bytes() == b'{"traceEvents": []}'
+    assert not Path(f"{path}.gz").exists()
+
+
+def test_cupti_monitor_profiler_is_constructed_before_context_entry(monkeypatch, tmp_path):
+    experimental_config = object()
+    constructed = []
+
+    class FakeProfiler:
+        def start(self):
+            pass
+
+        def stop(self):
+            pass
+
+    def fake_profile(**kwargs):
+        constructed.append(kwargs)
+        return FakeProfiler()
+
+    monkeypatch.setattr(
+        benchmark,
+        "_cupti_monitor_experimental_config",
+        lambda config: experimental_config,
+    )
+    monkeypatch.setattr(benchmark.torch.profiler, "profile", fake_profile)
+
+    context = benchmark.profiler(
+        tmp_path / "trace",
+        backend="cupti_monitor",
+    )
+
+    assert constructed[0]["experimental_config"] is experimental_config
+    with context:
+        pass
+
+
+def test_cupti_monitor_config_selects_monitor_backend(monkeypatch, tmp_path):
+    constructed = []
+
+    class FakeProfiler:
+        def start(self):
+            pass
+
+        def stop(self):
+            pass
+
+    monkeypatch.setattr(
+        benchmark,
+        "_cupti_monitor_experimental_config",
+        lambda config: constructed.append(config) or object(),
+    )
+    monkeypatch.setattr(benchmark.torch.profiler, "profile", lambda **kwargs: FakeProfiler())
+
+    benchmark.profiler(
+        tmp_path / "trace",
+        cupti_monitor_config=CuptiMonitorConfig(environment_counters=False),
+    )
+
+    assert not constructed[0].environment_counters

transformer_nuggets/utils/__init__.py

@@ -1,22 +1,23 @@
 from transformer_nuggets.utils.benchmark import (
-    benchmark_torch_function_in_microseconds,
+    DEFAULT_CUPTI_MONITOR_PM_METRICS,
+    CudaBenchmarkStats,
+    CuptiMonitorConfig,
+    ProfileConfig,
+    attach_oom_observer,
     benchmark_cuda_function_in_microseconds,
     benchmark_cuda_function_in_microseconds_triton,
     benchmark_cuda_function_stats,
+    benchmark_torch_function_in_microseconds,
+    cuda_memory_usage,
     locked_clocks,
     max_memory_usage,
-    cuda_memory_usage,
     profile_function,
-    ProfileConfig,
-    CudaBenchmarkStats,
-    save_memory_snapshot,
     profiler,
-    attach_oom_observer,
+    save_memory_snapshot,
+    supported_cupti_monitor_metrics,
 )
-from transformer_nuggets.utils.tracing import LoggingMode, NanInfDetect
-from transformer_nuggets.utils.triton import print_sass
-from transformer_nuggets.utils.merge_traces import merge_traces
 from transformer_nuggets.utils.memory_viz import generate_memory_comparison_html
+from transformer_nuggets.utils.merge_traces import merge_traces
 from transformer_nuggets.utils.perfetto import (
     TraceFormat,
     default_trace_path,
@@ -28,5 +29,7 @@
     write_trace,
     write_track_event_trace,
 )
+from transformer_nuggets.utils.tracing import LoggingMode, NanInfDetect
 from transformer_nuggets.utils.track_event import chrome_trace_to_track_event_trace
+from transformer_nuggets.utils.triton import print_sass
 # from transformer_nuggets.utils.model_extraction import extract_attention_data