[Interactive Drive] Annotate nvtx and Optimize

THIRD-PARTY-NOTICES

@@ -56,6 +56,7 @@ Optional integration: integrations/omnidreams
 
 mediapy                Apache-2.0   https://github.com/google/mediapy
 nvidia-cudnn-frontend  MIT          https://github.com/NVIDIA/cudnn-frontend
+nvtx                   Apache-2.0 WITH LLVM-exception https://github.com/NVIDIA/NVTX
 opencv-python-headless Apache-2.0   https://github.com/opencv/opencv-python
 grpcio, grpcio-tools   Apache-2.0   https://github.com/grpc/grpc
 shapely                BSD-3-Clause https://github.com/shapely/shapely

integrations/omnidreams/ludus-renderer/ludus_renderer/_cpp/bindings/torch_bindings_cuda.cpp

@@ -62,9 +62,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         pybind11::int_((int)CR::CudaRaster::RenderModeFlag_EnableDepthPeeling);
 
     // CUDA rendering ops
-    m.def("ludus_render_fwd_cuda", &ludus_render_fwd_cuda, "ludus f-theta CUDA rendering");
-    m.def("ludus_render_fwd_cuda_ts", &ludus_render_fwd_cuda_ts, "ludus f-theta CUDA rendering with timestamped cube pools");
-    m.def("ludus_render_fwd_cuda_timestamped", &ludus_render_fwd_cuda_timestamped, "ludus f-theta CUDA timestamped rendering from flat buffers");
+    // Rendering can wait behind world-model work on the same GPU. Keep those
+    // waits from blocking Python's presentation thread.
+    m.def("ludus_render_fwd_cuda", &ludus_render_fwd_cuda, "ludus f-theta CUDA rendering",
+          pybind11::call_guard<pybind11::gil_scoped_release>());
+    m.def("ludus_render_fwd_cuda_ts", &ludus_render_fwd_cuda_ts, "ludus f-theta CUDA rendering with timestamped cube pools",
+          pybind11::call_guard<pybind11::gil_scoped_release>());
+    m.def("ludus_render_fwd_cuda_timestamped", &ludus_render_fwd_cuda_timestamped, "ludus f-theta CUDA timestamped rendering from flat buffers",
+          pybind11::call_guard<pybind11::gil_scoped_release>());
 }
 
 //------------------------------------------------------------------------

integrations/omnidreams/omnidreams/conditioning/conditioning_wrapper.py

@@ -455,6 +455,7 @@ def start_generation(
             finalization_state={"autoregressive_index": 0},
         )
 
+    @nvtx.annotate()
     def continue_generation(
         self,
         state: OmnidreamsConditioningState,

integrations/omnidreams/omnidreams/interactive_drive/README.md

@@ -624,3 +624,44 @@ tests run with no checkpoint downloads.
 
 The hook does not auto-fix. Use `./scripts/check.sh --fix` to clean up lint and
 format issues explicitly.
+
+### To Profile
+
+```bash
+sudo sh -c 'echo 0 > /proc/sys/kernel/perf_event_paranoid'
+
+# unsafe, adds unpriv access to GPU perf counters
+cat > tmp.sh <<'EOF'
+#!/usr/bin/env bash
+set -euo pipefail
+# Unsafe: allows unprivileged access to NVIDIA GPU perf/profiling counters.
+# Run only on a trusted profiling machine.
+for cap in profiler-device profiler-context trace-device; do
+    minor="$(awk -v c="$cap" '$1 == c { print $2 }' /proc/driver/nvidia-caps/sys-minors)"
+    if [[ -z "${minor}" ]]; then
+        echo "Could not find sys-minor for ${cap}" >&2
+        exit 1
+    fi
+    sudo nvidia-modprobe -f "/proc/driver/nvidia/capabilities/${cap}"
+    sudo chmod a+r "/dev/nvidia-caps/nvidia-cap${minor}"
+    echo "DeviceFileModify: 0" | sudo tee "/proc/driver/nvidia/capabilities/${cap}" >/dev/null
+    echo "Enabled ${cap} via /dev/nvidia-caps/nvidia-cap${minor}"
+done
+EOF
+chmod +x tmp.sh
+./tmp.sh
+rm ./tmp.sh
+
+nsys profile --output=/tmp/flashdreams-nsys \
+ --force-overwrite=true \
+ --trace=cuda,nvtx,vulkan,python-gil \
+ --sample=process-tree --backtrace=dwarf --samples-per-backtrace=1 \
+ --python-sampling=true \
+ --python-sampling-frequency=1000 \
+ --python-backtrace=cuda \
+ --cudabacktrace=all \
+ --gpu-metrics-devices=all \
+ --gpu-metrics-frequency=10000 \
+ --stop-on-exit=true \
+ uv run --package flashdreams-omnidreams interactive-drive --manifest example_world_model_perf.yaml --auto-start --stop-after-chunks 48
+```

integrations/omnidreams/omnidreams/interactive_drive/app.py

@@ -16,6 +16,7 @@
     KeyboardState,
 )
 from omnidreams.interactive_drive.presenter import SlangPyPresenter
+import nvtx
 from omnidreams.interactive_drive.runtime.loop import (
     LoopConfig,
     PresenterBackend,
@@ -167,6 +168,7 @@ def first_chunk_produced(self) -> bool:
         """``True`` once the model has produced its first generated chunk."""
         return self._pipeline.first_chunk_produced.is_set()
 
+    @nvtx.annotate()
     def load_scene(
         self, scene_path: object, variant: str, prompt_override: str | None
     ) -> bool:
@@ -303,6 +305,7 @@ def _preload_worker(self, pending: list[tuple[object, str, str | None]]) -> None
                 f"{Path(str(scene_path)).name} variant={variant!r}",
             )
 
+    @nvtx.annotate()
     def _resolve_scene_assets(
         self, scene_path: object, variant: str, prompt_override: str | None
     ) -> tuple[SceneBundle, MapBounds | None, GroundSnapper | None]:
@@ -417,6 +420,7 @@ def _loading_base_frame(self) -> np.ndarray:
             self._loading_base_rgb = np.zeros((height, width, 3), dtype=np.uint8)
         return self._loading_base_rgb
 
+    @nvtx.annotate()
     def run_scene(self) -> None:
         """Drive the current scene until the presenter closes or switches.
 
@@ -502,6 +506,7 @@ def run_scene(self) -> None:
             # rather than after the rebuild completes.
             self._present_loading_once(loading_status)
 
+    @nvtx.annotate()
     def _present_loading_once(self, loading_status: Callable[[], str]) -> None:
         """Render a single loading-overlay frame immediately (used on reset)."""
         if self._scene is None:
@@ -535,6 +540,7 @@ def _resetting_status_message(self) -> str:
         """Phase text shown while a reset / respawn re-primes the rollout."""
         return "Resetting..."
 
+    @nvtx.annotate()
     def run(self) -> None:
         """Single-scene convenience: load the configured scene, run, tear down.
 

integrations/omnidreams/omnidreams/interactive_drive/backends/raster.py

@@ -5,6 +5,7 @@
 
 from omnidreams.interactive_drive.backends.base import RenderBackend
 from omnidreams.interactive_drive.config import BevConfig, ChunkConfig, RasterConfig
+import nvtx
 from omnidreams.interactive_drive.rasterizer import LudusConditionRasterizer
 from omnidreams.interactive_drive.types import FrameChunk, SceneBundle, TrajectoryChunk
 
@@ -24,16 +25,20 @@ def warmup_model(self) -> None:
         # No model to load; per-scene work happens in load_scene.
         return
 
+    @nvtx.annotate()
     def load_scene(self, scene: SceneBundle) -> None:
         self._scene = scene
         self._rasterizer.load_scene(scene)
 
+    @nvtx.annotate()
     def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
         return self._render_chunk(trajectory)
 
+    @nvtx.annotate()
     def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
         return self._render_chunk(trajectory)
 
+    @nvtx.annotate()
     def _render_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
         raster_chunk = self._rasterizer.render_chunk(
             rig_poses_world=trajectory.rig_poses_world,

integrations/omnidreams/omnidreams/interactive_drive/backends/world_model.py

@@ -17,6 +17,7 @@
     RasterConfig,
     WorldModelProfileConfig,
 )
+import nvtx
 from omnidreams.interactive_drive.rasterizer import LudusConditionRasterizer
 from omnidreams.interactive_drive.types import (
     FrameChunk,
@@ -66,6 +67,7 @@ def optimizes_on_first_chunk(self) -> bool:
         # which can take minutes on the first launch.
         return True
 
+    @nvtx.annotate()
     def warmup_model(self) -> None:
         if self._manifest.resolution_wh != self._raster.resolution_wh:
             raise ValueError(
@@ -91,6 +93,7 @@ def warmup_model(self) -> None:
             f"[world-model] model warmup session_ms={(time.perf_counter() - start) * 1000.0:.1f}",
         )
 
+    @nvtx.annotate()
     def load_scene(self, scene: SceneBundle) -> None:
         self._scene = scene
         self._next_chunk_count = 0
@@ -115,6 +118,7 @@ def load_scene(self, scene: SceneBundle) -> None:
             f"total_ms={(prepare_end - load_start) * 1000.0:.1f}",
         )
 
+    @nvtx.annotate()
     def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
         scene = self._require_scene()
         chunk_start = time.perf_counter()
@@ -128,23 +132,24 @@ def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
             display_frames = raster_chunk.frames
         else:
             raster_end = time.perf_counter()
-            condition_frames = [
-                frame.copy() for frame in self._debug_first_chunk_condition_frames
-            ]
-            display_frames = tuple(
-                PresentedFrame(
-                    timestamp_us=int(timestamp_us),
-                    rgb_host_uint8=frame.copy(),
-                    depth_host_f32=None,
-                    rgb_native=None,
-                    depth_native=None,
-                )
-                for timestamp_us, frame in zip(
-                    trajectory.timestamps_us,
-                    self._debug_first_chunk_condition_frames,
-                    strict=True,
+            with nvtx.annotate("backend.world_model.debug_condition_frames", color="yellow"):
+                condition_frames = [
+                    frame.copy() for frame in self._debug_first_chunk_condition_frames
+                ]
+                display_frames = tuple(
+                    PresentedFrame(
+                        timestamp_us=int(timestamp_us),
+                        rgb_host_uint8=frame.copy(),
+                        depth_host_f32=None,
+                        rgb_native=None,
+                        depth_native=None,
+                    )
+                    for timestamp_us, frame in zip(
+                        trajectory.timestamps_us,
+                        self._debug_first_chunk_condition_frames,
+                        strict=True,
+                    )
                 )
-            )
             logger.info(
                 "[world-model] first_chunk using official hdmap override "
                 f"dir={self._manifest.debug_condition_frame_dir}",
@@ -182,6 +187,7 @@ def render_first_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
             ),
         )
 
+    @nvtx.annotate()
     def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
         self._require_scene()
         chunk_start = time.perf_counter()
@@ -225,10 +231,12 @@ def render_next_chunk(self, trajectory: TrajectoryChunk) -> FrameChunk:
             ),
         )
 
+    @nvtx.annotate()
     def reset(self) -> None:
         self._session.reset()
         self._next_chunk_count = 0
 
+    @nvtx.annotate()
     def reset_scene_conditioning(self) -> None:
         self._session.reset(clear_precomputed_embeddings=True)
         self._next_chunk_count = 0
@@ -265,6 +273,7 @@ def _load_debug_condition_frames(
                 frames.append(np.array(rgb, dtype=np.uint8))
         return tuple(frames)
 
+    @nvtx.annotate()
     def _merge_frames(
         self,
         raster_frames: Sequence[PresentedFrame],

integrations/omnidreams/omnidreams/interactive_drive/cuda_host_prefetch.py

@@ -7,6 +7,7 @@
 from typing import Any
 
 import numpy as np
+import nvtx
 
 _STREAMS_LOCK = threading.Lock()
 _HOST_COPY_STREAMS: dict[int, Any] = {}
@@ -22,6 +23,7 @@ def __init__(self, tensor: Any, *, source_event: Any | None = None) -> None:
         self._done_event: Any | None = None
         self._started = False
 
+    @nvtx.annotate()
     def start(self) -> bool:
         if self._started:
             return self._host_tensor is not None
@@ -48,10 +50,11 @@ def start(self) -> bool:
                 if self._source_event is not None:
                     copy_stream.wait_event(self._source_event)
                 with torch.cuda.stream(copy_stream):
-                    host_tensor.copy_(tensor, non_blocking=True)
-                    tensor.record_stream(copy_stream)
-                    done_event = torch.cuda.Event()
-                    done_event.record(copy_stream)
+                    with nvtx.annotate("cuda_host_prefetch.copy_to_host", color="yellow"):
+                        host_tensor.copy_(tensor, non_blocking=True)
+                        tensor.record_stream(copy_stream)
+                        done_event = torch.cuda.Event()
+                        done_event.record(copy_stream)
         except Exception:
             self._host_tensor = None
             self._done_event = None
@@ -61,6 +64,7 @@ def start(self) -> bool:
         self._done_event = done_event
         return True
 
+    @nvtx.annotate()
     def to_numpy(self) -> np.ndarray:
         host_tensor = self._host_tensor
         if host_tensor is None:

integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py

@@ -5,6 +5,7 @@
 import time
 
 from omnidreams.interactive_drive.input.backend import InputBackend, SampledInput
+import nvtx
 from omnidreams.interactive_drive.types import (
     ControlSnapshot,
     DriverCommand,
@@ -149,6 +150,7 @@ class KeyboardInputBackend(InputBackend):
     def __init__(self, keyboard: KeyboardState) -> None:
         self._keyboard = keyboard
 
+    @nvtx.annotate()
     def sample(self) -> SampledInput:
         sample_time = time.perf_counter()
         return SampledInput(command=self._keyboard.command(), sample_time=sample_time)