[Feature] Support torch profiler across omni stages (#553)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: lishunyang <lishunyang12@163.com>
Co-authored-by: lishunyang <lishunyang12@163.com>

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -6,6 +6,7 @@ with the correct prompt format on Qwen2.5-Omni
 """
 
 import os
+import time
 from typing import NamedTuple
 
 import librosa
@@ -376,11 +377,17 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    if profiler_enabled:
+        omni_llm.start_profile(stages=[0])
     omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
+
+    total_requests = len(prompts)
+    processed_count = 0
     for stage_outputs in omni_generator:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
@@ -408,6 +415,18 @@ def main(args):
                 sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
                 print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
+        processed_count += len(stage_outputs.request_output)
+        if profiler_enabled and processed_count >= total_requests:
+            print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
+            # Stop the profiler while workers are still alive
+            omni_llm.stop_profile()
+
+            print("[Info] Waiting 30s for workers to write massive trace files to disk...")
+            time.sleep(30)
+            print("[Info] Trace export wait finished.")
+
+    omni_llm.close()
+
 
 def parse_args():
     parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")

examples/offline_inference/qwen3_omni/end2end.py

@@ -6,6 +6,7 @@ with the correct prompt format on Qwen3-Omni (thinker only).
 """
 
 import os
+import time
 from typing import NamedTuple
 
 import librosa
@@ -319,11 +320,17 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    if profiler_enabled:
+        omni_llm.start_profile(stages=[0])
     omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
 
+    total_requests = len(prompts)
+    processed_count = 0
+
     for stage_outputs in omni_generator:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
@@ -359,6 +366,16 @@ def main(args):
                 # Save audio file with explicit WAV format
                 sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
                 print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+        processed_count += len(stage_outputs.request_output)
+        if profiler_enabled and processed_count >= total_requests:
+            print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
+            # Stop the profiler while workers are still alive
+            omni_llm.stop_profile()
+
+            print("[Info] Waiting 30s for workers to write trace files to disk...")
+            time.sleep(30)
+            print("[Info] Trace export wait time finished.")
     omni_llm.close()
 
 

vllm_omni/entrypoints/async_omni.py

@@ -614,11 +614,39 @@ class AsyncOmni(OmniBase):
         """Generate outputs for a request from a pooling model."""
         raise NotImplementedError("encode() is not implemented for AsyncOmni")
 
-    async def start_profile(self) -> None:
-        raise NotImplementedError("start_profile() is not implemented for AsyncOmni")
+    async def start_profile(self, stages: list[int] | None = None) -> None:
+        """Start profiling for specified stages.
 
-    async def stop_profile(self) -> None:
-        raise NotImplementedError("stop_profile() is not implemented for AsyncOmni")
+        Async wrapper around the base implementation for API consistency.
+
+        Args:
+            stages: List of stage IDs to start profiling. If None, starts
+                profiling for all stages that have profiling enabled.
+
+        Example:
+            >>> await async_omni.start_profile()
+            >>> async for output in async_omni.generate(...):
+            ...     pass
+            >>> await async_omni.stop_profile()
+        """
+        super().start_profile(stages)
+
+    async def stop_profile(self, stages: list[int] | None = None) -> None:
+        """Stop profiling for specified stages.
+
+        Async wrapper around the base implementation for API consistency.
+
+        Args:
+            stages: List of stage IDs to stop profiling. If None, stops
+                profiling for all stages.
+
+        Example:
+            >>> await async_omni.start_profile()
+            >>> async for output in async_omni.generate(...):
+            ...     pass
+            >>> await async_omni.stop_profile()
+        """
+        super().stop_profile(stages)
 
     async def pause_generation(
         self,

vllm_omni/entrypoints/async_omni_llm.py

@@ -165,19 +165,21 @@ class AsyncOmniLLM(AsyncLLM):
         except RuntimeError:
             pass
 
-        if envs.VLLM_TORCH_PROFILER_DIR:
+        if envs.VLLM_TORCH_PROFILER_DIR and not envs.VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM:
             logger.info(
-                "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s",  # noqa: E501
+                "Torch profiler enabled. AsyncOmniLLM CPU traces will be collected under %s",
                 envs.VLLM_TORCH_PROFILER_DIR,
             )
-            worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
+            worker_name = f"{socket.gethostname()}_{os.getpid()}.async_omni_llm"
             self.profiler = torch.profiler.profile(
                 activities=[
                     torch.profiler.ProfilerActivity.CPU,
                 ],
                 with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
                 on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                    envs.VLLM_TORCH_PROFILER_DIR, worker_name=worker_name, use_gzip=True
+                    envs.VLLM_TORCH_PROFILER_DIR,
+                    worker_name=worker_name,
+                    use_gzip=envs.VLLM_TORCH_PROFILER_USE_GZIP,
                 ),
             )
         else:

vllm_omni/entrypoints/omni.py

@@ -30,7 +30,7 @@ from vllm_omni.distributed.ray_utils.utils import (
 )
 from vllm_omni.entrypoints.log_utils import OrchestratorMetrics
 from vllm_omni.entrypoints.omni_stage import OmniStage
-from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
+from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
@@ -334,6 +334,72 @@ class OmniBase:
 
         logger.error(f"[{self._name}] Stage initialization failed. Troubleshooting Steps:\n{formatted_suggestions}")
 
+    def start_profile(self, stages: list[int] | None = None) -> None:
+        """Start profiling for specified stages.
+
+        Sends start_profile command to stage workers. Profiling must be enabled
+        via VLLM_TORCH_PROFILER_DIR environment variable.
+
+        Args:
+            stages: List of stage IDs to start profiling. If None, starts
+                profiling for all stages that have profiling enabled.
+
+        Example:
+            >>> # Profile all stages
+            >>> omni.start_profile()
+            >>> outputs = omni.generate(prompts, sampling_params)
+            >>> omni.stop_profile()
+
+            >>> # Profile only stage 0 and 2
+            >>> omni.start_profile(stages=[0, 2])
+        """
+        if stages is None:
+            stages = list(range(len(self.stage_list)))
+
+        for stage_id in stages:
+            if stage_id < len(self.stage_list):
+                try:
+                    self.stage_list[stage_id].submit({"type": OmniStageTaskType.PROFILER_START})
+                    logger.info("[%s] Sent start_profile to stage-%s", self._name, stage_id)
+                except Exception as e:
+                    logger.warning(
+                        "[%s] Failed to send start_profile to stage-%s: %s",
+                        self._name,
+                        stage_id,
+                        e,
+                    )
+
+    def stop_profile(self, stages: list[int] | None = None) -> None:
+        """Stop profiling for specified stages.
+
+        Sends stop_profile command to stage workers to finalize and save traces.
+
+        Args:
+            stages: List of stage IDs to stop profiling. If None, stops
+                profiling for all stages.
+
+        Example:
+            >>> omni.start_profile()
+            >>> outputs = omni.generate(prompts, sampling_params)
+            >>> omni.stop_profile()
+            >>> # Traces saved to VLLM_TORCH_PROFILER_DIR/stage_X_<model_stage>/
+        """
+        if stages is None:
+            stages = list(range(len(self.stage_list)))
+
+        for stage_id in stages:
+            if stage_id < len(self.stage_list):
+                try:
+                    self.stage_list[stage_id].submit({"type": OmniStageTaskType.PROFILER_STOP})
+                    logger.info("[%s] Sent stop_profile to stage-%s", self._name, stage_id)
+                except Exception as e:
+                    logger.warning(
+                        "[%s] Failed to send stop_profile to stage-%s: %s",
+                        self._name,
+                        stage_id,
+                        e,
+                    )
+
     def close(self) -> None:
         """Close all stage processes and clean up resources."""
         if hasattr(self, "_weak_finalizer"):

vllm_omni/entrypoints/omni_stage.py

@@ -41,6 +41,7 @@ from vllm_omni.entrypoints.stage_utils import (
     SHUTDOWN_TASK,
     OmniStageTaskType,
     _to_dict,
+    is_profiler_task,
     maybe_dump_to_shm,
     set_stage_devices,
 )
@@ -606,6 +607,32 @@ def _stage_worker(
                 pass
     logger.debug("Engine initialized")
 
+    # Check if stage engine supports profiling (via vLLM's built-in profiler)
+    has_profiler = hasattr(stage_engine, "start_profile") and hasattr(stage_engine, "stop_profile")
+    if has_profiler:
+        logger.info(
+            "[Stage-%s] vLLM profiler support detected (model_stage=%s)",
+            stage_id,
+            engine_args.get("model_stage", None),
+        )
+
+    def handle_profiler_task(task_type: OmniStageTaskType) -> None:
+        """Handle profiler task."""
+        if task_type == OmniStageTaskType.PROFILER_START:
+            if has_profiler:
+                try:
+                    stage_engine.start_profile()
+                    logger.info("[Stage-%s] Profiler started via vLLM engine", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to start profiler: %s", stage_id, e)
+        elif task_type == OmniStageTaskType.PROFILER_STOP:
+            if has_profiler:
+                try:
+                    stage_engine.stop_profile()
+                    logger.info("[Stage-%s] Profiler stopped via vLLM engine", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to stop profiler: %s", stage_id, e)
+
     # Initialize OmniConnectors if configured
     connectors = {}
     if connectors_config:
@@ -636,6 +663,11 @@ def _stage_worker(
             logger.info("Received shutdown signal")
             break
 
+        # Handle profiler control commands
+        if is_profiler_task(task_type):
+            handle_profiler_task(task_type)
+            continue
+
         batch_tasks: list[dict[str, Any]] = [task]
         start_time = _time.time()
         if max_batch_size > 1:
@@ -645,6 +677,11 @@ def _stage_worker(
                     if extra == SHUTDOWN_TASK:
                         in_q.put(SHUTDOWN_TASK)
                         break
+                    # Handle profiler commands that arrive during batching
+                    extra_type = extra.get("type") if isinstance(extra, dict) else None
+                    if is_profiler_task(extra_type):
+                        handle_profiler_task(extra_type)
+                        continue  # Don't add to batch_tasks
                     batch_tasks.append(extra)
                     end_time = _time.time()
                     duration = end_time - start_time
@@ -1076,6 +1113,33 @@ async def _stage_worker_async(
     if stage_type != "diffusion":
         await stage_engine.reset_mm_cache()
     logger.debug("[Stage-%s] Engine initialized", stage_id)
+
+    # Check if stage engine supports profiling (via vLLM's built-in profiler)
+    has_profiler = hasattr(stage_engine, "start_profile") and hasattr(stage_engine, "stop_profile")
+    if has_profiler:
+        logger.info(
+            "[Stage-%s] vLLM profiler support detected for async stage (model_stage=%s)",
+            stage_id,
+            engine_args.get("model_stage", None),
+        )
+
+    async def handle_profiler_task_async(task_type: OmniStageTaskType) -> None:
+        """Handle profiler task asynchronously."""
+        if task_type == OmniStageTaskType.PROFILER_START:
+            if has_profiler:
+                try:
+                    await stage_engine.start_profile()
+                    logger.info("[Stage-%s] Profiler started via vLLM engine", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to start profiler: %s", stage_id, e)
+        elif task_type == OmniStageTaskType.PROFILER_STOP:
+            if has_profiler:
+                try:
+                    await stage_engine.stop_profile()
+                    logger.info("[Stage-%s] Profiler stopped via vLLM engine", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to stop profiler: %s", stage_id, e)
+
     # Signal readiness to orchestrator and send vllm_config back to main process
     try:
         # Send vllm_config back to main process so it can be accessed via
@@ -1183,8 +1247,11 @@ async def _stage_worker_async(
             elif task_type == OmniStageTaskType.ABORT:
                 rid = task["request_id"]
                 asyncio.create_task(stage_engine.abort(rid))
+            elif is_profiler_task(task_type):
+                await handle_profiler_task_async(task_type)
             else:
                 asyncio.create_task(generation_single_request(task))
+
         except queue.Empty:
             await asyncio.sleep(0.001)
         batch_request_outputs: list[Any] = []

vllm_omni/entrypoints/stage_utils.py

@@ -16,11 +16,17 @@ class OmniStageTaskType(enum.Enum):
     GENERATE = "generate"
     ABORT = "abort"
     SHUTDOWN = "shutdown"
+    PROFILER_START = "profiler_start"
+    PROFILER_STOP = "profiler_stop"
 
 
 SHUTDOWN_TASK = {"type": OmniStageTaskType.SHUTDOWN}
 
 
+def is_profiler_task(task_type: OmniStageTaskType) -> bool:
+    return task_type in (OmniStageTaskType.PROFILER_START, OmniStageTaskType.PROFILER_STOP)
+
+
 def set_stage_devices(
     stage_id: int,
     devices: str | int | None,