[Profiler] Add Nsight Systems support for serving (#1098)

Signed-off-by: Jinheng Li <ahengljh@gmail.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Canlin Guo <961750412@qq.com>

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -20,6 +20,7 @@ import zmq
 from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
+from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.utils.mem_utils import GiB_bytes
 from vllm.v1.worker.workspace import init_workspace_manager
@@ -83,15 +84,7 @@ class DiffusionWorker:
             od_config=self.od_config,
             device=self.device,
         )
-        # Initialize profiler if configured
-        self.profiler: OmniTorchProfilerWrapper | None = None
-        profiler_config = self.od_config.profiler_config
-        if profiler_config and profiler_config.profiler == "torch":
-            self.profiler = create_omni_profiler(
-                profiler_config=profiler_config,
-                worker_name=f"diffusion_worker_{self.rank}",
-                local_rank=self.local_rank,
-            )
+        self.profiler: WorkerProfiler | None = self._create_profiler()
         if not skip_load_model:
             self.load_model(load_format=self.od_config.diffusion_load_format)
             self.init_lora_manager()
@@ -122,6 +115,7 @@ class DiffusionWorker:
         vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size
         vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size
         vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel
+        vllm_config.profiler_config = self.od_config.profiler_config
         self.vllm_config = vllm_config
 
         # Initialize distributed environment
@@ -147,6 +141,24 @@ class DiffusionWorker:
             )
             init_workspace_manager(self.device)
 
+    def _create_profiler(self) -> WorkerProfiler | None:
+        profiler_config = self.od_config.profiler_config
+        profiler_type = getattr(profiler_config, "profiler", None)
+        if profiler_type == "torch":
+            return create_omni_profiler(
+                profiler_config=profiler_config,
+                worker_name=f"diffusion_rank{self.rank}",
+                local_rank=self.local_rank,
+            )
+        if profiler_type == "cuda":
+            return CudaProfilerWrapper(profiler_config)
+        if profiler_type is not None:
+            logger.warning("Unknown profiler backend %r on diffusion worker %s", profiler_type, self.rank)
+        return None
+
+    def _get_profiler(self) -> WorkerProfiler | None:
+        return getattr(self, "profiler", None)
+
     def load_model(self, load_format: str = "default", custom_pipeline_name: str | None = None) -> None:
         """Load the diffusion model using DiffusionModelRunner."""
         with (
@@ -192,27 +204,21 @@ class DiffusionWorker:
 
         Args:
             is_start: True to start profiling, False to stop.
-            profile_prefix: Optional prefix for trace filename (vLLM compat).
-
-        Note:
-            Matches vLLM's worker.profile() signature for consistency.
-            Traces are saved automatically via on_trace_ready callback.
+            profile_prefix: Optional prefix for trace filename.
         """
-        if self.profiler is None:
-            logger.warning("Profiler not initialized, skipping profile(%s)", is_start)
+        profiler = self._get_profiler()
+        if profiler is None:
             return
 
         if is_start:
-            from vllm_omni.profiler import OmniTorchProfilerWrapper
-
-            if isinstance(self.profiler, OmniTorchProfilerWrapper):
+            if isinstance(profiler, OmniTorchProfilerWrapper):
                 import time
 
-                filename = profile_prefix or f"diffusion_{int(time.time())}"
-                self.profiler.set_trace_filename(filename)
-            self.profiler.start()
+                filename = profile_prefix or f"diffusion_rank{self.rank}_{int(time.time())}"
+                profiler.set_trace_filename(filename)
+            profiler.start()
         else:
-            self.profiler.stop()
+            profiler.stop()
 
     def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput:
         """Execute a forward pass by delegating to the model runner."""
@@ -224,7 +230,13 @@ class DiffusionWorker:
                 if req.sampling_params.lora_request is not None:
                     raise
                 logger.warning("LoRA activation skipped: %s", exc)
-        return self.model_runner.execute_model(req)
+        profiler = self._get_profiler()
+        ctx = profiler.annotate_context_manager("diffusion_forward") if profiler else nullcontext()
+        with ctx:
+            output = self.model_runner.execute_model(req)
+        if profiler:
+            profiler.step()
+        return output
 
     def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput:
         """Execute one diffusion step by delegating to the model runner."""
@@ -236,8 +248,13 @@ class DiffusionWorker:
 
         if any(new_req.req.sampling_params.lora_request is not None for new_req in scheduler_output.scheduled_new_reqs):
             raise ValueError("Step mode does not support LoRA yet.")
-
-        return self.model_runner.execute_stepwise(scheduler_output)
+        profiler = self._get_profiler()
+        ctx = profiler.annotate_context_manager("diffusion_step") if profiler else nullcontext()
+        with ctx:
+            output = self.model_runner.execute_stepwise(scheduler_output)
+        if profiler:
+            profiler.step()
+        return output
 
     def load_weights(self, weights) -> set[str]:
         """Load weights by delegating to the model runner."""