[Core] Implement Diffusion Profiler Support (#709)

Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

examples/offline_inference/image_to_image/image_edit.py

@@ -347,6 +347,9 @@ def main():
     )
     print("Pipeline loaded")
 
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
     # Time profiling for generation
     print(f"\n{'=' * 60}")
     print("Generation Configuration:")
@@ -379,6 +382,10 @@ def main():
         "resolution": args.resolution,
     }
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
     outputs = omni.generate(**generate_kwargs)
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
@@ -386,6 +393,23 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
     if not outputs:
         raise ValueError("No output generated from omni.generate()")
     logger.info("Outputs: %s", outputs)

examples/offline_inference/image_to_video/image_to_video.py

@@ -19,6 +19,7 @@ Usage:
 """
 
 import argparse
+import os
 from pathlib import Path
 
 import numpy as np
@@ -95,6 +96,9 @@ def main():
     vae_use_slicing = is_npu()
     vae_use_tiling = is_npu()
 
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
     omni = Omni(
         model=args.model,
         vae_use_slicing=vae_use_slicing,
@@ -103,6 +107,10 @@ def main():
         flow_shift=args.flow_shift,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
         args.prompt,
@@ -178,6 +186,23 @@ def main():
     export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
 
 if __name__ == "__main__":
     main()

examples/offline_inference/text_to_image/text_to_image.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
+import os
 import time
 from pathlib import Path
 
@@ -148,6 +149,9 @@ def main():
         tensor_parallel_size=args.tensor_parallel_size,
     )
 
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
     omni = Omni(
         model=args.model,
         vae_use_slicing=vae_use_slicing,
@@ -158,6 +162,10 @@ def main():
         enforce_eager=args.enforce_eager,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
     # Time profiling for generation
     print(f"\n{'=' * 60}")
     print("Generation Configuration:")
@@ -189,6 +197,23 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
     # Extract images from OmniRequestOutput
     # omni.generate() returns list[OmniRequestOutput], extract images from the first output
     if not outputs or len(outputs) == 0:

examples/offline_inference/text_to_video/text_to_video.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
+import os
 from pathlib import Path
 
 import numpy as np
@@ -46,6 +47,9 @@ def main():
     vae_use_slicing = is_npu()
     vae_use_tiling = is_npu()
 
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
     omni = Omni(
         model=args.model,
         vae_use_slicing=vae_use_slicing,
@@ -54,6 +58,10 @@ def main():
         flow_shift=args.flow_shift,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
     frames = omni.generate(
         args.prompt,
         negative_prompt=args.negative_prompt,
@@ -126,6 +134,23 @@ def main():
     export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
 
 if __name__ == "__main__":
     main()

vllm_omni/diffusion/diffusion_engine.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import multiprocessing as mp
+import os
 import time
 import weakref
 from collections.abc import Callable, Iterable
@@ -290,6 +291,125 @@ class DiffusionEngine:
     def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
         return scheduler.add_req(requests)
 
+    def start_profile(self, trace_filename: str | None = None) -> None:
+        """
+        Start torch profiling on all diffusion workers.
+
+        Creates a directory (if needed) and sets up a base filename template
+        for per-rank profiler traces (typically saved as <template>_rank<N>.json).
+
+        Args:
+            trace_filename: Optional base filename (without extension or rank suffix).
+                            If None, generates one using current timestamp.
+        """
+        if trace_filename is None:
+            trace_filename = f"stage_0_diffusion_{int(time.time())}_rank"
+
+        trace_dir = os.environ.get("VLLM_TORCH_PROFILER_DIR", "./profiles")
+
+        # Expand ~ and ~user, then make absolute (robust against cwd changes)
+        trace_dir = os.path.expanduser(trace_dir)
+        trace_dir = os.path.abspath(trace_dir)
+
+        try:
+            os.makedirs(trace_dir, exist_ok=True)
+        except OSError as exc:
+            logger.error(f"Failed to create profiler directory {trace_dir}: {exc}")
+            raise
+
+        # Build final template path (without rank or extension — torch.profiler appends those)
+        full_template = os.path.join(trace_dir, trace_filename)
+
+        expected_pattern = f"{full_template}*.json"
+        logger.info(f"Starting diffusion profiling → {expected_pattern}")
+
+        # Also log the absolute directory once (useful in multi-node or containers)
+        logger.debug(f"Profiler output directory: {trace_dir}")
+
+        # Propagate to all workers
+        try:
+            self.collective_rpc(method="start_profile", args=(full_template,))
+        except Exception as e:
+            logger.error("Failed to start profiling on workers", exc_info=True)
+            raise RuntimeError(f"Could not start profiler: {e}") from e
+
+    def stop_profile(self) -> dict:
+        """
+        Stop profiling on all workers and collect the final trace/table paths.
+
+        The worker (torch_profiler.py) now handles trace export, compression to .gz,
+        and deletion of the original .json file. This method only collects and
+        reports the paths returned by the workers.
+
+        Returns:
+            dict with keys:
+            - "traces": list of final trace file paths (usually .json.gz)
+            - "tables": list of table strings (one per rank)
+        """
+        logger.info("Stopping diffusion profiling and collecting results...")
+
+        try:
+            # Give worker enough time — export + compression + table can be slow
+            results = self.collective_rpc(method="stop_profile", timeout=60000)
+        except Exception:
+            logger.error("Failed to stop profiling on workers", exc_info=True)
+            return {"traces": [], "tables": []}
+
+        output_files = {"traces": [], "tables": []}
+        successful_traces = 0
+
+        if not results:
+            logger.warning("No profiling results returned from any rank")
+            return output_files
+
+        for rank, res in enumerate(results):
+            if not isinstance(res, dict):
+                logger.warning(f"Rank {rank}: invalid result format (got {type(res)})")
+                continue
+
+            # 1. Trace file — should be .json.gz if compression succeeded
+            trace_path = res.get("trace")
+            if trace_path:
+                # We trust the worker — it created/compressed the file
+                logger.info(f"[Rank {rank}] Final trace: {trace_path}")
+                output_files["traces"].append(trace_path)
+                successful_traces += 1
+
+                # Optional: warn if path looks suspicious (e.g. still .json)
+                if not trace_path.endswith((".json.gz", ".json")):
+                    logger.warning(f"Rank {rank}: unusual trace path extension: {trace_path}")
+
+            # 2. Table file — plain text
+            table = res.get("table")
+            if table:
+                output_files["tables"].append(table)
+
+        # Final summary logging
+        num_ranks = len(results)
+        if successful_traces > 0:
+            final_paths_str = ", ".join(output_files["traces"][:3])
+            if len(output_files["traces"]) > 3:
+                final_paths_str += f" ... (+{len(output_files['traces']) - 3} more)"
+
+            logger.info(
+                f"Profiling stopped. Collected {successful_traces} trace file(s) "
+                f"from {num_ranks} rank(s). "
+                f"Final trace paths: {final_paths_str}"
+            )
+        elif output_files["traces"]:
+            logger.info(
+                f"Profiling stopped but no traces were successfully collected. "
+                f"Reported paths: {', '.join(output_files['traces'][:3])}"
+                f"{' ...' if len(output_files['traces']) > 3 else ''}"
+            )
+        else:
+            logger.info("Profiling stopped — no trace files were collected from any rank.")
+
+        if output_files["tables"]:
+            logger.debug(f"Collected {len(output_files['tables'])} profiling table(s)")
+
+        return output_files
+
     def _dummy_run(self):
         """A dummy run to warm up the model."""
         prompt = "dummy run"

vllm_omni/diffusion/profiler/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .torch_profiler import TorchProfiler
+
+# Default profiler – can be changed later via config
+CurrentProfiler = TorchProfiler
+
+__all__ = ["CurrentProfiler", "TorchProfiler"]

vllm_omni/diffusion/profiler/base.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ProfilerBase(ABC):
+    """
+    Abstract base class for all diffusion profilers.
+    Defines the common interface used by GPUWorker and DiffusionEngine.
+    """
+
+    @abstractmethod
+    def start(self, trace_path_template: str) -> str:
+        """
+        Start profiling.
+
+        Args:
+            trace_path_template: Base path (without rank or extension).
+                                 e.g. "/tmp/profiles/sdxl_run"
+
+        Returns:
+            Full path of the trace file this rank will write.
+        """
+        pass
+
+    @abstractmethod
+    def stop(self) -> str | None:
+        """
+        Stop profiling and finalize/output the trace.
+
+        Returns:
+            Path to the saved trace file, or None if not active.
+        """
+        pass
+
+    @abstractmethod
+    def get_step_context(self):
+        """
+        Returns a context manager that advances one profiling step.
+        Should be a no-op (nullcontext) when profiler is not active.
+        """
+        pass
+
+    @abstractmethod
+    def is_active(self) -> bool:
+        """Return True if profiling is currently running."""
+        pass
+
+    @classmethod
+    def _get_rank(cls) -> int:
+        import os
+
+        return int(os.getenv("RANK", "0"))

vllm_omni/diffusion/profiler/torch_profiler.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import subprocess
+from contextlib import nullcontext
+
+import torch
+from torch.profiler import ProfilerActivity, profile
+from vllm.logger import init_logger
+
+from .base import ProfilerBase
+
+logger = init_logger(__name__)
+
+
+class TorchProfiler(ProfilerBase):
+    """
+    Torch-based profiler configured for End-to-End continuous recording.
+    Uses 'on_trace_ready' to handle Trace export.
+    Compression is offloaded to a background subprocess to avoid blocking the worker loop.
+    """
+
+    _profiler: profile | None = None
+    _trace_template: str = ""
+
+    @classmethod
+    def start(cls, trace_path_template: str) -> str:
+        """
+        Start the profiler with the given trace path template.
+        """
+        # 1. Cleanup any existing profiler
+        if cls._profiler is not None:
+            logger.warning("[Rank %s] Stopping existing Torch profiler", cls._get_rank())
+            cls._profiler.stop()
+            cls._profiler = None
+
+        rank = cls._get_rank()
+
+        # 2. Make path absolute
+        trace_path_template = os.path.abspath(trace_path_template)
+        cls._trace_template = trace_path_template
+
+        # Expected paths
+        json_file = f"{trace_path_template}_rank{rank}.json"
+
+        os.makedirs(os.path.dirname(json_file), exist_ok=True)
+
+        logger.info(f"[Rank {rank}] Starting End-to-End Torch profiler")
+
+        # 3. Define the on_trace_ready handler
+        def trace_handler(p):
+            nonlocal json_file
+
+            # A. Export JSON Trace
+            try:
+                p.export_chrome_trace(json_file)
+                logger.info(f"[Rank {rank}] Trace exported to {json_file}")
+
+                try:
+                    subprocess.Popen(["gzip", "-f", json_file])
+                    logger.info(f"[Rank {rank}] Triggered background compression for {json_file}")
+                    # Update variable to point to the eventual file
+                    json_file = f"{json_file}.gz"
+                except Exception as compress_err:
+                    logger.warning(f"[Rank {rank}] Background gzip failed to start: {compress_err}")
+
+            except Exception as e:
+                logger.warning(f"[Rank {rank}] Failed to export trace: {e}")
+
+        # 4. Initialize profiler with long active period
+        cls._profiler = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            schedule=torch.profiler.schedule(
+                wait=0,
+                warmup=0,
+                active=100000,  # long capture window
+            ),
+            on_trace_ready=trace_handler,
+            record_shapes=True,
+            profile_memory=True,
+            with_stack=True,
+            with_flops=True,
+        )
+
+        # 5. Start profiling
+        cls._profiler.start()
+
+        # Return the expected final path
+        return f"{trace_path_template}_rank{rank}.json.gz"
+
+    @classmethod
+    def stop(cls) -> dict | None:
+        if cls._profiler is None:
+            return None
+
+        rank = cls._get_rank()
+
+        # Determine expected paths
+        base_path = f"{cls._trace_template}_rank{rank}"
+        gz_path = f"{base_path}.json.gz"
+
+        try:
+            # This triggers trace_handler synchronously
+            # Since we removed table generation and backgrounded compression, this returns fast.
+            cls._profiler.stop()
+        except Exception as e:
+            logger.warning(f"[Rank {rank}] Profiler stop failed: {e}")
+
+        cls._profiler = None
+
+        # We return the .gz path assuming background compression will succeed.
+        return {"trace": gz_path, "table": None}
+
+    @classmethod
+    def step(cls):
+        if cls._profiler is not None:
+            cls._profiler.step()
+
+    @classmethod
+    def is_active(cls) -> bool:
+        return cls._profiler is not None
+
+    @classmethod
+    def get_step_context(cls):
+        return nullcontext()

vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py

@@ -15,6 +15,7 @@ from collections.abc import Iterable
 from contextlib import nullcontext
 
 import torch
+from torch.profiler import record_function
 from vllm.config import LoadConfig
 from vllm.logger import init_logger
 from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes
@@ -160,6 +161,7 @@ class GPUDiffusionModelRunner:
             self.cache_backend.refresh(self.pipeline, req.num_inference_steps)
 
         with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config):
-            output = self.pipeline.forward(req)
+            with record_function("pipeline_forward"):
+                output = self.pipeline.forward(req)
 
         return output

vllm_omni/diffusion/worker/gpu_diffusion_worker.py

@@ -28,6 +28,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
+from vllm_omni.diffusion.profiler import CurrentProfiler
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
 
@@ -115,6 +116,16 @@ class GPUDiffusionWorker:
         """Generate output for the given requests."""
         return self.execute_model(requests, self.od_config)
 
+    @classmethod
+    def start_profile(cls, trace_path_template: str) -> str:
+        """Start profiling for this GPU worker."""
+        return CurrentProfiler.start(trace_path_template)
+
+    @classmethod
+    def stop_profile(cls) -> dict | None:
+        """Stop profiling and return the result dictionary."""
+        return CurrentProfiler.stop()
+
     def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
         """Execute a forward pass by delegating to the model runner."""
         assert self.model_runner is not None, "Model runner not initialized"

vllm_omni/entrypoints/omni.py

@@ -374,36 +374,71 @@ class OmniBase:
                         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>/
+    def stop_profile(self, stages: list[int] | None = None) -> dict:
+        """
+        Synchronously stop profiling for specified stages and collect
+        the file paths for traces and tables.
         """
         if stages is None:
             stages = list(range(len(self.stage_list)))
 
+        all_results = {"traces": [], "tables": []}
+
         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:
+                stage = self.stage_list[stage_id]
+
+                # Check if the stage object has our new bridge method
+                if hasattr(stage, "stop_profile"):
+                    logger.info("[%s] Requesting profile data collection from stage-%s", self._name, stage_id)
+
+                    # This is the blocking call that triggers the RPC chain
+                    stage_data = stage.stop_profile()
+
+                    if isinstance(stage_data, dict):
+                        # FIX: Handle both single key and list key formats
+                        traces = stage_data.get("trace") or stage_data.get("traces")
+                        tables = stage_data.get("table") or stage_data.get("tables")
+
+                        # Debug logging
+                        logger.debug(f"[{self._name}] Stage-{stage_id} returned: {stage_data.keys()}")
+                        if traces:
+                            logger.debug(f"[{self._name}] Stage-{stage_id} traces type: {type(traces)}")
+                        if tables:
+                            logger.debug(f"[{self._name}] Stage-{stage_id} tables type: {type(tables)}")
+
+                        # Handle single strings
+                        if traces:
+                            if isinstance(traces, str):
+                                all_results["traces"].append(traces)
+                            elif isinstance(traces, list):
+                                all_results["traces"].extend(traces)
+
+                        # Handle single strings
+                        if tables:
+                            if isinstance(tables, str):
+                                all_results["tables"].append(tables)
+                            elif isinstance(tables, list):
+                                all_results["tables"].extend(tables)
+                        else:
+                            logger.warning(f"[{self._name}] Stage-{stage_id} returned no table data")
+                    else:
+                        logger.warning(f"[{self._name}] Stage-{stage_id} returned non-dict data: {type(stage_data)}")
+                else:
+                    # Fallback for non-diffusion stages
                     logger.warning(
-                        "[%s] Failed to send stop_profile to stage-%s: %s",
+                        "[%s] Stage-%s does not support synchronous stop_profile. Falling back to async.",
                         self._name,
                         stage_id,
-                        e,
                     )
+                    stage.submit({"type": OmniStageTaskType.PROFILER_STOP})
+
+        # Final debug output
+        logger.info(
+            f"[{self._name}] Collected {len(all_results['traces'])} trace(s) and {len(all_results['tables'])} table(s)"
+        )
+
+        return all_results
 
     def close(self) -> None:
         """Close all stage processes and clean up resources."""
@@ -518,7 +553,6 @@ class Omni(OmniBase):
                 return self._run_generation_with_generator(prompts, sampling_params_list)
             else:
                 outputs = list(self._run_generation(prompts, sampling_params_list))
-                self.close()
                 return outputs
         except Exception as e:
             logger.exception("[Orchestrator] Failed to run generation: %s", e)

vllm_omni/entrypoints/omni_diffusion.py

@@ -126,3 +126,26 @@ class OmniDiffusion:
             self.close()
         except Exception:
             pass
+
+    def start_profile(self, trace_filename: str | None = None) -> None:
+        """Start profiling for the diffusion model.
+
+        Args:
+            trace_filename: Optional base filename for trace files.
+                           If None, a timestamp-based name will be generated.
+        """
+        if hasattr(self, "engine") and self.engine:
+            self.engine.start_profile(trace_filename)
+        else:
+            raise RuntimeError("Diffusion engine not initialized")
+
+    def stop_profile(self) -> dict:
+        """Stop profiling and return profiling results.
+
+        Returns:
+            Dictionary containing paths to trace and table files.
+        """
+        if hasattr(self, "engine") and self.engine:
+            return self.engine.stop_profile()
+        else:
+            raise RuntimeError("Diffusion engine not initialized")

vllm_omni/entrypoints/omni_stage.py

@@ -13,6 +13,7 @@ import multiprocessing as mp
 import os
 import queue
 import sys
+import time
 import traceback
 from dataclasses import fields
 from typing import Any
@@ -211,6 +212,38 @@ class OmniStage:
         self._in_q = in_q
         self._out_q = out_q
 
+    def stop_profile(self) -> dict:
+        """Stop profiling by sending a signal to worker and waiting for response."""
+        if self._in_q is None or self._out_q is None:
+            logger.warning(f"[Stage-{self.stage_id}] Queues not initialized, cannot stop profile.")
+            return {}
+
+        logger.info(f"[Stage-{self.stage_id}] Sending PROFILER_STOP to worker...")
+        self.submit({"type": OmniStageTaskType.PROFILER_STOP})
+
+        # Wait for result from worker
+        try:
+            # Profiling stop might take time to flush files, give it 180s
+            response = self._out_q.get(timeout=60000)
+
+            if isinstance(response, dict):
+                if response.get("type") == "profiler_result":
+                    return response.get("data", {})
+                elif "error" in response:
+                    logger.error(f"[Stage-{self.stage_id}] Profiler error: {response['error']}")
+                    return {}
+
+            # If we got something else (e.g. late generation result), we might lose it here,
+            # but usually profiling stop is called when generation is done.
+            logger.warning(
+                f"[Stage-{self.stage_id}] Received unexpected message while waiting for profiler: {response}"
+            )
+            return {}
+
+        except queue.Empty:
+            logger.error(f"[Stage-{self.stage_id}] Timeout waiting for profiler results.")
+            return {}
+
     def init_stage_worker(
         self,
         model: str,
@@ -620,33 +653,6 @@ def _stage_worker(
             except (OSError, ValueError):
                 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:
@@ -667,6 +673,45 @@ def _stage_worker(
     max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
     logger.info(f"Max batch size: {max_batch_size}")
 
+    def handle_profiler_task_local(task_type: OmniStageTaskType) -> dict:
+        """Handle profiler task locally in the worker process."""
+        if task_type == OmniStageTaskType.PROFILER_START:
+            if stage_type == "diffusion":
+                try:
+                    profile_dir = _os.environ.get("VLLM_TORCH_PROFILER_DIR", "./profiles")
+                    _os.makedirs(profile_dir, exist_ok=True)
+                    trace_filename = f"stage_{stage_id}_diffusion_{int(_time.time())}"
+                    stage_engine.start_profile(trace_filename=trace_filename)
+                    logger.info("[Stage-%s] Diffusion Torch profiler started", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to start diffusion profiler: %s", stage_id, e)
+            else:
+                try:
+                    stage_engine.start_profile()
+                    logger.info("[Stage-%s] vLLM profiler started", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to start vLLM profiler: %s", stage_id, e)
+            return {}
+
+        elif task_type == OmniStageTaskType.PROFILER_STOP:
+            if stage_type == "diffusion":
+                try:
+                    # CRITICAL: Capture return value
+                    result_data = stage_engine.stop_profile()
+                    logger.info("[Stage-%s] Diffusion Torch profiler stopped", stage_id)
+                    return result_data
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to stop diffusion profiler: %s", stage_id, e)
+                    return {}
+            else:
+                try:
+                    stage_engine.stop_profile()
+                    logger.info("[Stage-%s] vLLM profiler stopped", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to stop vLLM profiler: %s", stage_id, e)
+                return {}
+        return {}
+
     # Batch processing loop
     while True:
         task = in_q.get()
@@ -679,7 +724,10 @@ def _stage_worker(
 
         # Handle profiler control commands
         if is_profiler_task(task_type):
-            handle_profiler_task(task_type)
+            profiler_data = handle_profiler_task_local(task_type)
+            # If it was a STOP command, we must reply to the Orchestrator
+            if task_type == OmniStageTaskType.PROFILER_STOP:
+                out_q.put({"type": "profiler_result", "data": profiler_data})
             continue
 
         batch_tasks: list[dict[str, Any]] = [task]
@@ -694,8 +742,10 @@ def _stage_worker(
                     # 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
+                        p_data = handle_profiler_task_local(extra_type)
+                        if extra_type == OmniStageTaskType.PROFILER_STOP:
+                            out_q.put({"type": "profiler_result", "data": p_data})
+                        continue
                     batch_tasks.append(extra)
                     end_time = _time.time()
                     duration = end_time - start_time
@@ -1141,31 +1191,41 @@ async def _stage_worker_async(
         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."""
+        """Handle profiler task asynchronously for both LLM and diffusion stages."""
         if task_type == OmniStageTaskType.PROFILER_START:
-            if has_profiler:
+            if stage_type == "diffusion":
+                try:
+                    # Sync call is safe here — diffusion profiling is lightweight
+                    profile_dir = os.environ.get("VLLM_TORCH_PROFILER_DIR", "./profiles")
+                    os.makedirs(profile_dir, exist_ok=True)
+                    trace_filename = f"stage_{stage_id}_diffusion_{int(time.time())}"
+                    stage_engine.start_profile(trace_filename=trace_filename)
+                    logger.info("[Stage-%s] Diffusion Torch profiler started", stage_id)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to start diffusion profiler: %s", stage_id, e)
+            else:
                 try:
                     await stage_engine.start_profile()
-                    logger.info("[Stage-%s] Profiler started via vLLM engine", stage_id)
+                    logger.info("[Stage-%s] vLLM profiler started", stage_id)
                 except Exception as e:
-                    logger.warning("[Stage-%s] Failed to start profiler: %s", stage_id, e)
+                    logger.warning("[Stage-%s] Failed to start vLLM profiler: %s", stage_id, e)
+
         elif task_type == OmniStageTaskType.PROFILER_STOP:
-            if has_profiler:
+            if stage_type == "diffusion":
+                try:
+                    trace_files = stage_engine.stop_profile()
+                    logger.info("[Stage-%s] Diffusion Torch profiler stopped", stage_id)
+                    if trace_files:
+                        logger.info("Diffusion trace files: %s", trace_files)
+                except Exception as e:
+                    logger.warning("[Stage-%s] Failed to stop diffusion profiler: %s", stage_id, e)
+            else:
                 try:
                     await stage_engine.stop_profile()
-                    logger.info("[Stage-%s] Profiler stopped via vLLM engine", stage_id)
+                    logger.info("[Stage-%s] vLLM profiler stopped", stage_id)
                 except Exception as e:
-                    logger.warning("[Stage-%s] Failed to stop profiler: %s", stage_id, e)
+                    logger.warning("[Stage-%s] Failed to stop vLLM profiler: %s", stage_id, e)
 
     # Signal readiness to orchestrator and send vllm_config back to main process
     try: