Add experimental CUPTI Backend (#68)

* Add experimental CUPTI Backend

Signed-off-by: Alok Joshi <aljoshi@nvidia.com>

* Add CI job for testing CUPTI backend

Rename the existing NCU test job to test-ncu and add a test-cupti job
on both amd64 and arm64 GPU runners that runs the suite with
--tool-init=cupti to exercise the experimental CUPTI backend.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Signed-off-by: Alok Joshi <aljoshi@nvidia.com>

---------

Signed-off-by: Alok Joshi <aljoshi@nvidia.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: Alok-Joshi

.github/workflows/ci.yaml

@@ -34,8 +34,8 @@ jobs:
           pip install hatch
           hatch run docs:build
 
-  test:
-    name: test (${{ matrix.platform.runner }})
+  test-ncu:
+    name: test-ncu (${{ matrix.platform.runner }})
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
@@ -60,3 +60,21 @@ jobs:
         run: |
           pip install hatch
           hatch run test_cu13:run
+
+  test-cupti:
+    name: test-cupti (${{ matrix.platform.runner }})
+    runs-on: ${{ matrix.platform.runner }}
+    strategy:
+      matrix:
+        platform:
+          - runner: linux-amd64-gpu-h100-latest-1
+          - runner: linux-arm64-gpu-l4-latest-1
+    container:
+      image: nvcr.io/nvidia/pytorch:26.06-py3
+      options: "--cap-add=SYS_ADMIN --gpus all"
+    steps:
+      - uses: actions/checkout@v5
+      - name: "Running Tests (CUPTI backend)"
+        run: |
+          pip install hatch
+          hatch run test_cu13:run -- --tool-init=cupti

conftest.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+
+# How each --tool-init mode sets nsight up, as
+# (NSPY_NCU_INIT_AT_IMPORT, tool to activate in pytest_configure).
+_INIT_MODES = {
+    # nsight activates NCU itself while `import nsight` runs.
+    "ncu-at-import": ("1", None),
+    # Nothing is active up front, so nsight.analyze.kernel activates NCU on the
+    # first call to a decorated function.
+    "ncu-lazy": ("0", None),
+    # CUPTI is activated here instead; nsight must not claim NCU at import.
+    "cupti": ("0", "cupti"),
+}
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--tool-init",
+        default="ncu-at-import",
+        choices=sorted(_INIT_MODES),
+        help=(
+            "How the nsight tool is initialized for the suite. 'ncu-at-import' "
+            "(the default) lets nsight activate NCU while it is imported. "
+            "'ncu-lazy' leaves no tool active so the lazy NCU load in "
+            "nsight.analyze.kernel is exercised on the first decorated call. "
+            "'cupti' activates CUPTI before the suite runs. Only one tool can be "
+            "active per process, so this mode fully determines both "
+            "NSPY_NCU_INIT_AT_IMPORT and which tool is activated."
+        ),
+    )
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    init_at_import, tool_name = _INIT_MODES[config.getoption("--tool-init")]
+
+    # Set before nsight is imported: it reads this while importing, which is why
+    # the import below is here rather than at module scope. pytest_configure
+    # still runs before any test module is imported, so anything activated here
+    # is ahead of CUDA initialization.
+    os.environ["NSPY_NCU_INIT_AT_IMPORT"] = init_at_import
+
+    import nsight
+
+    if tool_name is not None:
+        nsight.activate(nsight.Tool(tool_name))

examples/test_examples.py

@@ -5,12 +5,18 @@
 
 import pytest
 
+import nsight
+
+_cupti_active = nsight.get_active_tool() == nsight.Tool.CUPTI
+_skip_cupti = pytest.mark.skipif(_cupti_active, reason="metric not supported in cupti")
+
 
 def test_00_minimal() -> None:
     minimal = importlib.import_module("examples.00_minimal")
     minimal.main()
 
 
+@_skip_cupti  # type: ignore[untyped-decorator]
 def test_01_compare_throughput() -> None:
     compare_throughput = importlib.import_module("examples.01_compare_throughput")
     compare_throughput.main()
@@ -47,21 +53,25 @@ def test_07_triton_minimal() -> None:
     triton_minimal.main()
 
 
+@_skip_cupti  # type: ignore[untyped-decorator]
 def test_08_multiple_metrics() -> None:
     multiple_metrics = importlib.import_module("examples.08_multiple_metrics")
     multiple_metrics.main()
 
 
+@_skip_cupti  # type: ignore[untyped-decorator]
 def test_09_advanced_metric_custom() -> None:
     advanced_custom = importlib.import_module("examples.09_advanced_metric_custom")
     advanced_custom.main()
 
 
+@_skip_cupti  # type: ignore[untyped-decorator]
 def test_10_combine_kernel_metrics() -> None:
     combine_metrics = importlib.import_module("examples.10_combine_kernel_metrics")
     combine_metrics.main()
 
 
+@_skip_cupti  # type: ignore[untyped-decorator]
 def test_11_output_csv() -> None:
     output_csv = importlib.import_module("examples.11_output_csv")
     output_csv.main()

nsight/analyze.py

@@ -17,6 +17,7 @@
 import nsight.collection as collection
 import nsight.visualization as visualization
 from nsight.info_collector import CollectionScope, InfoCollector
+from nsight.tools_manager import Tool, activate, get_active_tool
 from nsight.utils import VerbosityLevel
 
 
@@ -106,7 +107,23 @@ def kernel(
     | Callable[[Callable[..., None]], Callable[..., collection.core.ProfileResults]]
 ):
     """
-    A decorator that collects profiling data using NVIDIA Nsight Compute.
+    A decorator that collects profiling data from GPU kernels.
+
+    **Tool selection**
+
+    By default, ``@nsight.analyze.kernel`` profiles using NVIDIA Nsight Compute
+    (NCU). NCU's injection library is loaded automatically on the first call to
+    a decorated function, but it **must** be loaded before CUDA is initialized.
+    This means the first decorated function call must occur before any CUDA-initializing code.
+    To guarantee early loading and avoid this constraint, call  ``nsight.activate(nsight.Tool.NCU)``
+    explicitly at the very top of your script, before any CUDA imports.
+
+    CUPTI is also supported *(experimental)* for lightweight collection of
+    ``"gpu__time_duration.sum"`` using kernel timestamps, without launching
+    ``ncu``. To enable it, call ``nsight.activate(nsight.Tool.CUPTI)`` before
+    profiling. CUPTI does not have the CUDA initialization ordering constraint.
+    See :class:`~nsight.collection.cupti.CUPTICollector` for supported metrics
+    and limitations.
 
     Can be used with or without parentheses:
         - ``@nsight.analyze.kernel`` (no parentheses)
@@ -159,7 +176,7 @@ def wrapped_function(*args, configs=None, **kwargs) -> ProfileResults
         derive_metric:
             A function to transform the collected metrics.
             This can be used to compute derived metrics like TFLOPs that cannot
-            be captured by ncu directly. The function takes the metric values and
+            be captured directly. The function takes the metric values and
             the arguments of the profile-decorated function and returns the new
             metrics. Return value can be either:
 
@@ -182,7 +199,7 @@ def wrapped_function(*args, configs=None, **kwargs) -> ProfileResults
             Combine with ``derive_metric`` to compute speedup (reciprocal of normalized value).
         metrics: The metrics to collect. By default, kernel runtimes in nanoseconds
             are collected. Default: ``["gpu__time_duration.sum"]``. To see the available
-            metrics on your system, use the command: ``ncu --query-metrics``.
+            metrics on your system when using ncu, use the command: ``ncu --query-metrics``.
         ignore_kernel_list:
             List of kernel names to ignore. If you call a library within an annotated range context, you might not have precise control over which and how many kernels are being launched.
             If some of these kernels should be ignored in the profile, their names can be provided in this parameter. Default: ``None``
@@ -364,6 +381,22 @@ def my_benchmark(n, dtype):
     # Strip whitespace
     metrics = [m.strip() for m in metrics]
 
+    # These are decorator parameters, so reject illegal values now rather than
+    # when the decorated function is first called. Whether the active tool
+    # supports a legal value is a separate check, made by the collector.
+    if clock_control not in ("base", "none"):
+        raise ValueError(
+            f"Invalid clock_control {clock_control!r}. Expected 'base' or 'none'."
+        )
+    if cache_control not in ("all", "none"):
+        raise ValueError(
+            f"Invalid cache_control {cache_control!r}. Expected 'all' or 'none'."
+        )
+    if replay_mode not in ("kernel", "range"):
+        raise ValueError(
+            f"Invalid replay_mode {replay_mode!r}. Expected 'kernel' or 'range'."
+        )
+
     # Validate thermal parameters if both are provided
     if thermal_wait is not None and thermal_cont is not None:
         if thermal_cont <= thermal_wait:
@@ -461,15 +494,26 @@ def _canonical_scope(scope: Any, label: str) -> str:
             output_csv=output_csv,
             info_collectors=normalized_collectors,
         )
-        ncu = collection.ncu.NCUCollector(
-            metrics=metrics,
-            ignore_kernel_list=ignore_kernel_list,
-            combine_kernel_metrics=combine_kernel_metrics,
-            clock_control=clock_control,
-            cache_control=cache_control,
-            replay_mode=replay_mode,
-        )
-        return collection.core.NsightProfiler(settings, ncu)
+
+        def get_collector() -> collection.core.NsightCollector:
+            if get_active_tool() == Tool.CUPTI:
+                # clock_control, cache_control and replay_mode have no
+                # meaning for CUPTI, so the collector does not take them.
+                return collection.cupti.CUPTICollector(
+                    metrics=metrics,
+                    ignore_kernel_list=ignore_kernel_list,
+                    combine_kernel_metrics=combine_kernel_metrics,
+                )
+            return collection.ncu.NCUCollector(
+                metrics=metrics,
+                ignore_kernel_list=ignore_kernel_list,
+                combine_kernel_metrics=combine_kernel_metrics,
+                clock_control=clock_control,
+                cache_control=cache_control,
+                replay_mode=replay_mode,
+            )
+
+        return collection.core.NsightProfiler(settings, get_collector)
 
     # Support both @kernel and @kernel() syntax
     if _func is None:

nsight/collection/__init__.py

@@ -4,7 +4,8 @@
 import functools
 
 import nsight.collection.core as core
+import nsight.collection.cupti as cupti
 import nsight.collection.ncu as ncu
 import nsight.utils as utils
 
-__all__ = ["ncu", "core"]
+__all__ = ["ncu", "cupti", "core"]

nsight/collection/core.py

@@ -633,7 +633,7 @@ class ProfileSettings:
     """
     A function to transform the collected metrics.
     This can be used to compute derived metrics like TFLOPs that cannot
-    be captured by ncu directly. The function takes the metric values and
+    be captured directly. The function takes the metric values and
     the arguments of the profile-decorated function and returns the new
     metrics. See the examples for concrete use cases.
     """
@@ -807,11 +807,14 @@ class NsightProfiler:
     a collector to gather raw profiling data and processes it according to
     the provided settings.
 
-    Attributes:
-        settings: Configuration settings for profiling,
-            including normalization, and other options.
-        collector: The collector responsible for gathering
-            raw profiling data.
+    Args:
+        settings: Configuration settings for profiling (runs, verbosity,
+            derive_metric, output prefix, etc.).
+        get_collector: Factory returning the collector to profile with. It is
+            called lazily, the first time a decorated function runs, and the
+            result is reused for the lifetime of this profiler, so the
+            collector matches whichever tool is active by then rather than at
+            decoration time.
 
     Methods:
         __call__(func):
@@ -820,14 +823,24 @@ class NsightProfiler:
             Returns the processed data.
     """
 
-    def __init__(self, settings: ProfileSettings, collector: NsightCollector):
+    def __init__(
+        self,
+        settings: ProfileSettings,
+        get_collector: Callable[[], NsightCollector],
+    ) -> None:
         self.settings = settings
-        self.collector = collector
+        self._get_collector_fn = get_collector
+        self._collector: NsightCollector | None = None
+
+    def _get_collector(self) -> NsightCollector:
+        if self._collector is None:
+            self._collector = self._get_collector_fn()
+        return self._collector
 
     def __call__(
         self, func: Callable[..., None]
     ) -> Callable[..., ProfileResults | None]:
-        func._nspy_ncu_run_id = 0  # type: ignore[attr-defined]
+        func._nspy_run_id = 0  # type: ignore[attr-defined]
 
         # Fail fast (before starting the profiling session) on collector names
         # that would silently clobber profiler columns or function-arg columns.
@@ -843,7 +856,7 @@ def wrapper(
             **kwargs: Any,
         ) -> ProfileResults | None:
 
-            tag = f"{func.__name__}-{func._nspy_ncu_run_id}"  # type: ignore[attr-defined]
+            tag = f"{func.__name__}-{func._nspy_run_id}"  # type: ignore[attr-defined]
 
             configs = _sanitize_configs(
                 func,
@@ -853,7 +866,7 @@ def wrapper(
                 **kwargs,
             )
 
-            raw_df = self.collector.collect(func, configs, self.settings)
+            raw_df = self._get_collector().collect(func, configs, self.settings)
 
             # Check if the function has a return type
             if raw_df is not None:
@@ -891,7 +904,7 @@ def wrapper(
                             f"[NSIGHT-PYTHON] Refer to {processed_csv_path} for the processed profiling data"
                         )
 
-                func._nspy_ncu_run_id += 1  # type: ignore[attr-defined]
+                func._nspy_run_id += 1  # type: ignore[attr-defined]
 
                 return ProfileResults(results=processed)