[Refactor] Unify torch profiler for omni and diffusion models (#2099)

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

docs/contributing/profiling.md

@@ -4,48 +4,54 @@
 
 vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
 
-### 1. Set the Output Directory
-Before running any script, set this environment variable. The system detects this and automatically saves traces here.
+### 1. Configure Profiling in the Stage YAML
 
-```bash
-export VLLM_TORCH_PROFILER_DIR=./profiles
+Enable profiling by adding `profiler_config` under `engine_args` for the stage(s) you want to profile in your stage config YAML:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    engine_args:
+      # ... other engine args ...
+      profiler_config:
+        profiler: torch
+        torch_profiler_dir: ./perf
 ```
 
+| Field | Description |
+|---|---|
+| `profiler` | Profiler backend to use. Currently supports `torch`. |
+| `torch_profiler_dir` | Directory where trace files are saved. Created automatically if it doesn't exist. |
+
+> **Tip:** Only enable `profiler_config` on stages you actually need to profile. Stages without it will not start a profiler, keeping overhead minimal.
+
 ### 2. Profiling Omni-Modality Models
 
-It is best to limit profiling to one iteration to keep trace files manageable.
-
-```bash
-export VLLM_PROFILER_MAX_ITERS=1
-```
-
 **Selective Stage Profiling**
-The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
+
+It is highly recommended to profile specific stages to prevent producing overly large trace files:
+
 ```python
 # Profile all stages
-omni.start_profile()
+omni_llm.start_profile()
 
 # Only profile Stage 1
-omni.start_profile(stages=[1])
+omni_llm.start_profile(stages=[1])
+
+# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
+omni_llm.start_profile(stages=[0, 2])
 ```
 
-```python
-# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
-omni.start_profile(stages=[0, 2])
-```
+> **Important:** Always pass the same `stages` list to both `start_profile()` and `stop_profile()`. If you omit `stages` from `stop_profile()`, it defaults to stopping all stages — including ones that were never started — which will produce errors.
 
 **Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni.entrypoints.omni import Omni
+profiler_stages = [0]  # Only profile the stages you need
 
-omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
-
-profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-
-# 1. Start profiling if enabled
-if profiler_enabled:
-    omni.start_profile(stages=[0])
+# 1. Start profiling
+omni.start_profile(stages=profiler_stages)
 
 # Initialize generator
 omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
@@ -59,42 +65,57 @@ for stage_outputs in omni_generator:
     # ... [Output processing logic for text/audio would go here] ...
 
     # Update count to track when to stop profiling
-    processed_count += 1
+    processed_count += len(stage_outputs.request_output)
 
     # 2. Check if all requests are done to stop the profiler safely
     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 active
-        omni.stop_profile()
+        # Pass the same stages list used in start_profile()
+        omni_llm.stop_profile(stages=profiler_stages)
 
         # Wait for traces to flush to disk
         print("[Info] Waiting 30s for workers to write trace files to disk...")
         time.sleep(30)
         print("[Info] Trace export wait time finished.")
 
-omni.close()
+omni_llm.close()
 ```
 
 
+**CLI Usage** (using `end2end.py`):
+```bash
+# Profile only Stage 0 (Thinker)
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler --profiler-stages 0
+
+# Profile Stage 0 and Stage 2
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler --profiler-stages 0 2
+
+# Profile all stages (omit --profiler-stages)
+python end2end.py --output-wav output_audio \
+    --query-type text --enable-profiler
+```
+
 **Examples**:
 
 1. **Qwen2.5-Omni**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
 
 2. **Qwen3-Omni**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
 
-
 ### 3. Profiling diffusion models
 
-Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
+Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding. Standalone diffusion scripts use `--profiler-dir` to enable profiling.
 
 **CLI Usage:**
-```python
-
+```bash
 python image_to_video.py \
     --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
     --image qwen-bear.png \
     --prompt "A cat playing with yarn, smooth motion" \
+    --profiler-dir \
     \
     # Minimize Spatial Dimensions (Optional but helpful):
     #    Drastically reduces memory usage so the profiler doesn't
@@ -124,25 +145,72 @@ python image_to_video.py \
     --flow-shift 12.0 \
     --fps 16 \
     --output i2v_output.mp4
-
 ```
 
+> **Note:** For diffusion stages within a multi-stage omni pipeline, use `profiler_config` in the stage YAML instead (see Section 1).
+
 **Examples**:
 
 1. **Qwen image edit**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
 
 2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
 
-### 4. Analyzing Omni Traces
+### 4. Profiling Online Serving
 
-Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
+When `profiler_config` is set in the stage YAML, the server automatically exposes `/start_profile` and `/stop_profile` HTTP endpoints.
+
+**1. Start the server** with a stage YAML that has `profiler_config` enabled:
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B \
+    --omni \
+    --stage-configs-path qwen2_5_omni.yaml \
+    --port 8091
+```
+
+Or for one stage diffusion models:
+
+```bash
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091 --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
+```
+
+**2. Start profiling** by sending a POST request:
+```bash
+# Profile all stages that have profiler_config set
+curl -X POST http://localhost:8091/start_profile
+
+# Profile specific stages only
+curl -X POST http://localhost:8091/start_profile \
+    -H "Content-Type: application/json" \
+    -d '{"stages": [0]}'
+```
+
+**3. Send your inference requests** as normal while the profiler is running.
+
+**4. Stop profiling** and collect traces:
+```bash
+# Stop all stages
+curl -X POST http://localhost:8091/stop_profile
+
+# Stop specific stages (must match the stages you started)
+curl -X POST http://localhost:8091/stop_profile \
+    -H "Content-Type: application/json" \
+    -d '{"stages": [0]}'
+```
+
+Trace files are written to the `torch_profiler_dir` specified in your stage YAML.
+
+> **Important:** Always stop the same stages you started. Stopping a stage that was never started will produce errors.
+
+### 5. Analyzing Traces
+
+Output files are saved to the `torch_profiler_dir` specified in your stage YAML config.
 
 **Output**
-**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
+**Chrome Trace** (`.json.gz`): Visual timeline of kernels and stages. Open in Perfetto UI.
 
 **Viewing Tools:**
 
-- [Perfetto](https://ui.perfetto.dev/)(recommended)
-- ```chrome://tracing```(Chrome only)
+- [Perfetto](https://ui.perfetto.dev/) (recommended)
+- `chrome://tracing` (Chrome only)
 
 **Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)

examples/offline_inference/qwen3_omni/end2end.py

@@ -353,9 +353,9 @@ def main(args):
         for i, prompt in enumerate(prompts):
             prompt["modalities"] = output_modalities
 
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.enable_profiler
     if profiler_enabled:
-        omni.start_profile(stages=[0])
+        omni.start_profile(stages=args.profiler_stages)
     omni_generator = omni.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
@@ -405,7 +405,7 @@ def main(args):
         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.stop_profile()
+            omni.stop_profile(stages=args.profiler_stages)
 
             print("[Info] Waiting 30s for workers to write trace files to disk...")
             time.sleep(30)
@@ -532,6 +532,19 @@ def parse_args():
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--enable-profiler",
+        action="store_true",
+        default=False,
+        help="Enables profiling when set.",
+    )
+    parser.add_argument(
+        "--profiler-stages",
+        type=int,
+        nargs="*",
+        default=None,
+        help="List of stage IDs to profile. If not set, profiles all stages.",
+    )
 
     return parser.parse_args()
 

tests/conftest.py

@@ -1907,6 +1907,33 @@ class OmniRunner:
         )
         return self.generate(omni_inputs, sampling_params_list)
 
+    def start_profile(
+        self,
+        profile_prefix: str | None = None,
+        stages: list[int] | None = None,
+    ) -> list[Any]:
+        """Start profiling specified stages.
+
+        Args:
+            profile_prefix: Optional prefix for the trace file names.
+            stages: List of stage IDs to profile. If None, profiles all stages.
+
+        Returns:
+            List of results from each stage.
+        """
+        return self.omni.start_profile(profile_prefix=profile_prefix, stages=stages)
+
+    def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
+        """Stop profiling specified stages.
+
+        Args:
+            stages: List of stage IDs to profile. If None, stops all stages.
+
+        Returns:
+            List of results from each stage.
+        """
+        return self.omni.stop_profile(stages=stages)
+
     def _cleanup_process(self):
         try:
             keywords = ["enginecore"]
@@ -2020,6 +2047,18 @@ class OmniRunnerHandler:
         assert_omni_response(response, request_config, run_level="L2")
         return response
 
+    def start_profile(
+        self,
+        profile_prefix: str | None = None,
+        stages: list[int] | None = None,
+    ) -> list[Any]:
+        """Start profiling specified stages."""
+        return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages)
+
+    def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
+        """Stop profiling specified stages."""
+        return self.runner.stop_profile(stages=stages)
+
 
 @pytest.fixture
 def omni_runner_handler(omni_runner):

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -274,20 +274,19 @@ class TestSerialOperations:
         assert result.error == "result_for_Y"
 
     def test_serial_collective_rpc_all_ranks(self):
-        """``collective_rpc`` without *unique_reply_rank* collects
-        ``num_gpus`` responses.
+        """``collective_rpc`` without *unique_reply_rank* returns a single
+        response from rank 0 (only rank 0 has a result_mq).
         """
         engine, _, _, res_q = _make_engine(num_gpus=2)
 
-        # Pre-populate two results (simulating two workers replying)
+        # Pre-populate one result (only rank 0 replies via result_mq)
         res_q.put(_tagged_output("rank0"))
-        res_q.put(_tagged_output("rank1"))
 
         results = engine.collective_rpc("ping", args=("multi",))
 
-        assert len(results) == 2
+        # Only 1 response expected since only rank 0 has result_mq
+        assert len(results) == 1
         assert results[0].error == "rank0"
-        assert results[1].error == "rank1"
 
     def test_serial_add_req_then_collective_rpc(self):
         engine, _, req_q, res_q = _make_engine()

tests/entrypoints/test_omni_base_profiler.py

@@ -0,0 +1,263 @@
+"""Unit tests for OmniBase and AsyncOmni profiler methods."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestOmniBaseProfiler:
+    """Test suite for OmniBase profiler methods (start_profile, stop_profile)."""
+
+    @pytest.fixture
+    def mock_engine(self):
+        """Create a mock AsyncOmniEngine for testing."""
+        engine = MagicMock()
+        engine.num_stages = 3
+        engine.is_alive.return_value = True
+        engine.default_sampling_params_list = [MagicMock() for _ in range(3)]
+        engine.get_stage_metadata.side_effect = lambda i: {
+            "final_output_type": "text" if i == 0 else "audio",
+            "final_output": True,
+        }
+        engine.collective_rpc.return_value = [None, None, None]
+        return engine
+
+    @pytest.fixture
+    def omni_base_instance(self, mock_engine):
+        """Create an OmniBase instance with mocked dependencies."""
+        with (
+            patch("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", return_value=mock_engine),
+            patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", side_effect=lambda x: x),
+            patch("vllm_omni.entrypoints.omni_base.weakref.finalize"),
+        ):
+            from vllm_omni.entrypoints.omni_base import OmniBase
+
+            instance = OmniBase(model="test-model")
+            return instance
+
+    def test_start_profile_calls_collective_rpc(self, omni_base_instance, mock_engine):
+        """Test that start_profile calls collective_rpc with correct arguments."""
+        omni_base_instance.start_profile()
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(True, None),
+            stage_ids=None,
+        )
+
+    def test_start_profile_with_prefix(self, omni_base_instance, mock_engine):
+        """Test that start_profile passes profile_prefix to collective_rpc."""
+        omni_base_instance.start_profile(profile_prefix="test_trace")
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(True, "test_trace"),
+            stage_ids=None,
+        )
+
+    def test_start_profile_with_stages(self, omni_base_instance, mock_engine):
+        """Test that start_profile passes stages to collective_rpc."""
+        omni_base_instance.start_profile(stages=[0, 2])
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(True, None),
+            stage_ids=[0, 2],
+        )
+
+    def test_start_profile_with_prefix_and_stages(self, omni_base_instance, mock_engine):
+        """Test that start_profile passes both prefix and stages."""
+        omni_base_instance.start_profile(profile_prefix="my_prefix", stages=[1])
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(True, "my_prefix"),
+            stage_ids=[1],
+        )
+
+    def test_start_profile_returns_rpc_result(self, omni_base_instance, mock_engine):
+        """Test that start_profile returns the result from collective_rpc."""
+        expected_result = [{"stage_0": "started"}, {"stage_1": "started"}]
+        mock_engine.collective_rpc.return_value = expected_result
+
+        result = omni_base_instance.start_profile()
+
+        assert result == expected_result
+
+    def test_stop_profile_calls_collective_rpc(self, omni_base_instance, mock_engine):
+        """Test that stop_profile calls collective_rpc with correct arguments."""
+        omni_base_instance.stop_profile()
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(False, None),
+            stage_ids=None,
+        )
+
+    def test_stop_profile_with_stages(self, omni_base_instance, mock_engine):
+        """Test that stop_profile passes stages to collective_rpc."""
+        omni_base_instance.stop_profile(stages=[0])
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(False, None),
+            stage_ids=[0],
+        )
+
+    def test_stop_profile_returns_rpc_result(self, omni_base_instance, mock_engine):
+        """Test that stop_profile returns the result from collective_rpc."""
+        expected_result = [{"stage_0": "stopped"}, {"stage_1": "stopped"}]
+        mock_engine.collective_rpc.return_value = expected_result
+
+        result = omni_base_instance.stop_profile()
+
+        assert result == expected_result
+
+    def test_start_stop_profile_workflow(self, omni_base_instance, mock_engine):
+        """Test a typical start/stop profiling workflow."""
+        # Start profiling on specific stages
+        omni_base_instance.start_profile(profile_prefix="workflow_test", stages=[0, 1])
+
+        # Verify start was called correctly
+        mock_engine.collective_rpc.assert_called_with(
+            method="profile",
+            args=(True, "workflow_test"),
+            stage_ids=[0, 1],
+        )
+
+        # Reset mock to check stop call
+        mock_engine.collective_rpc.reset_mock()
+
+        # Stop profiling on the same stages
+        omni_base_instance.stop_profile(stages=[0, 1])
+
+        # Verify stop was called correctly
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(False, None),
+            stage_ids=[0, 1],
+        )
+
+    def test_start_profile_empty_stages_list(self, omni_base_instance, mock_engine):
+        """Test that start_profile handles empty stages list."""
+        omni_base_instance.start_profile(stages=[])
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(True, None),
+            stage_ids=[],
+        )
+
+    def test_stop_profile_empty_stages_list(self, omni_base_instance, mock_engine):
+        """Test that stop_profile handles empty stages list."""
+        omni_base_instance.stop_profile(stages=[])
+
+        mock_engine.collective_rpc.assert_called_once_with(
+            method="profile",
+            args=(False, None),
+            stage_ids=[],
+        )
+
+
+class TestOmniBaseProfilerSignatureConsistency:
+    """Test that profiler methods have consistent signatures with vLLM."""
+
+    def test_start_profile_signature(self):
+        """Verify start_profile has the expected signature parameters."""
+        import inspect
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+
+        sig = inspect.signature(OmniBase.start_profile)
+        params = list(sig.parameters.keys())
+
+        # Should have: self, profile_prefix, stages
+        assert "self" in params
+        assert "profile_prefix" in params
+        assert "stages" in params
+
+    def test_stop_profile_signature(self):
+        """Verify stop_profile has the expected signature parameters."""
+        import inspect
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+
+        sig = inspect.signature(OmniBase.stop_profile)
+        params = list(sig.parameters.keys())
+
+        # Should have: self, stages
+        assert "self" in params
+        assert "stages" in params
+
+    def test_start_profile_default_values(self):
+        """Verify start_profile has correct default parameter values."""
+        import inspect
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+
+        sig = inspect.signature(OmniBase.start_profile)
+
+        # profile_prefix should default to None
+        assert sig.parameters["profile_prefix"].default is None
+        # stages should default to None
+        assert sig.parameters["stages"].default is None
+
+    def test_stop_profile_default_values(self):
+        """Verify stop_profile has correct default parameter values."""
+        import inspect
+
+        from vllm_omni.entrypoints.omni_base import OmniBase
+
+        sig = inspect.signature(OmniBase.stop_profile)
+
+        # stages should default to None
+        assert sig.parameters["stages"].default is None
+
+
+class TestAsyncOmniProfilerSignatureConsistency:
+    """Test that AsyncOmni profiler methods have consistent signatures."""
+
+    def test_async_start_profile_signature(self):
+        """Verify AsyncOmni.start_profile has the expected signature parameters."""
+        import inspect
+
+        from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+        sig = inspect.signature(AsyncOmni.start_profile)
+        params = list(sig.parameters.keys())
+
+        # Should have: self, profile_prefix, stages
+        assert "self" in params
+        assert "profile_prefix" in params
+        assert "stages" in params
+
+    def test_async_stop_profile_signature(self):
+        """Verify AsyncOmni.stop_profile has the expected signature parameters."""
+        import inspect
+
+        from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+        sig = inspect.signature(AsyncOmni.stop_profile)
+        params = list(sig.parameters.keys())
+
+        # Should have: self, stages
+        assert "self" in params
+        assert "stages" in params
+
+    def test_async_start_profile_is_coroutine(self):
+        """Verify AsyncOmni.start_profile is an async method."""
+        import inspect
+
+        from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+        assert inspect.iscoroutinefunction(AsyncOmni.start_profile)
+
+    def test_async_stop_profile_is_coroutine(self):
+        """Verify AsyncOmni.stop_profile is an async method."""
+        import inspect
+
+        from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+        assert inspect.iscoroutinefunction(AsyncOmni.stop_profile)

vllm_omni/diffusion/data.py

@@ -21,6 +21,8 @@ from vllm_omni.diffusion.quantization import (
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 
 if TYPE_CHECKING:
+    from vllm.config import ProfilerConfig
+
     from vllm_omni.diffusion.quantization import DiffusionQuantizationConfig
 
 # Import after TYPE_CHECKING to avoid circular imports at runtime
@@ -476,6 +478,8 @@ class OmniDiffusionConfig:
     # Omni configuration (injected from stage config)
     omni_kv_config: dict[str, Any] = field(default_factory=dict)
 
+    profiler_config: "ProfilerConfig | dict[str, Any] | None" = None
+
     # Model-specific function for collecting CFG KV caches (set at runtime)
     cfg_kv_collect_func: Any | None = None
 
@@ -543,6 +547,11 @@ class OmniDiffusionConfig:
         initial_master_port = (self.master_port or 30005) + random.randint(0, 100)
         self.master_port = self.settle_port(initial_master_port, 37)
 
+        if isinstance(self.profiler_config, dict):
+            from vllm.config import ProfilerConfig
+
+            self.profiler_config = ProfilerConfig(**self.profiler_config)
+
         # Convert parallel_config dict/DictConfig to DiffusionParallelConfig
         # Use Mapping to handle both plain dicts and OmegaConf DictConfig
         if isinstance(self.parallel_config, Mapping):

vllm_omni/diffusion/diffusion_engine.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import os
 import threading
 import time
 from collections.abc import Iterable
@@ -309,124 +308,31 @@ class DiffusionEngine:
                     self.scheduler.pop_request_state(target_sched_req_id)
                     return output
 
-    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).
+    def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
+        """Start or stop torch profiling on all diffusion workers.
 
         Args:
-            trace_filename: Optional base filename (without extension or rank suffix).
-                            If None, generates one using current timestamp.
+            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.
         """
-        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=600)
-        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 ''}"
-            )
+        if is_start:
+            if profile_prefix is None:
+                profile_prefix = f"diffusion_{int(time.time())}"
+            logger.info(f"Starting diffusion profiling with prefix: {profile_prefix}")
         else:
-            logger.info("Profiling stopped — no trace files were collected from any rank.")
+            logger.info("Stopping diffusion profiling...")
 
-        if output_files["tables"]:
-            logger.debug(f"Collected {len(output_files['tables'])} profiling table(s)")
-
-        return output_files
+        try:
+            self.collective_rpc(method="profile", args=(is_start, profile_prefix))
+        except Exception as e:
+            action = "start" if is_start else "stop"
+            logger.error(f"Failed to {action} profiling on workers", exc_info=True)
+            if is_start:
+                raise RuntimeError(f"Could not {action} profiler: {e}") from e
 
     def _dummy_run(self):
         """A dummy run to warm up the model."""

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -204,20 +204,23 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         kwargs = kwargs or {}
 
         # Prepare RPC request message
+        # When unique_reply_rank is None, all workers must execute the RPC
+        # but only rank 0 can reply (it's the only one with a result_mq).
         rpc_request = {
             "type": "rpc",
             "method": method,
             "args": args,
             "kwargs": kwargs,
-            "output_rank": unique_reply_rank,
+            "output_rank": unique_reply_rank if unique_reply_rank is not None else 0,
+            "exec_all_ranks": unique_reply_rank is None,
         }
 
         try:
             # Broadcast RPC request to all workers via unified message queue
             self._broadcast_mq.enqueue(rpc_request)
 
-            # Determine which workers we expect responses from
-            num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus
+            # Only rank 0 has a result_mq, so we always expect exactly 1 response
+            num_responses = 1
 
             responses = []
             for _ in range(num_responses):

vllm_omni/diffusion/profiler/__init__.py

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

vllm_omni/diffusion/profiler/base.py

@@ -1,58 +0,0 @@
-# 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

@@ -1,126 +0,0 @@
-# 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/stage_diffusion_client.py

@@ -7,6 +7,7 @@ expects from any stage client.
 from __future__ import annotations
 
 import asyncio
+import time
 from typing import TYPE_CHECKING, Any
 
 from vllm.logger import init_logger
@@ -152,7 +153,27 @@ class StageDiffusionClient:
         """
         kwargs = kwargs or {}
 
-        if method in {"add_lora", "remove_lora", "list_loras", "pin_lora", "start_profile", "stop_profile"}:
+        # Handle profile method: inject stage_id into profile_prefix for diffusion stages
+        if method == "profile":
+            target = getattr(self._engine, method, None)
+            if target is None:
+                return {
+                    "supported": False,
+                    "todo": True,
+                    "reason": f"AsyncOmniDiffusion.{method} is not implemented",
+                }
+            # Extract is_start and profile_prefix from args
+            is_start = args[0] if args else True
+            profile_prefix = args[1] if len(args) > 1 else None
+            # Generate profile_prefix with stage_id if starting and no prefix provided
+            if is_start and profile_prefix is None:
+                profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
+            result = target(is_start, profile_prefix)
+            if timeout is not None:
+                return await asyncio.wait_for(result, timeout=timeout)
+            return await result
+
+        if method in {"add_lora", "remove_lora", "list_loras", "pin_lora"}:
             target = getattr(self._engine, method, None)
             if target is None:
                 return {

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -36,13 +36,13 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.diffusion.ipc import pack_diffusion_output_shm
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
-from vllm_omni.diffusion.profiler import CurrentProfiler
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.utils import RunnerOutput
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.profiler import OmniTorchProfilerWrapper, create_omni_profiler
 from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
 
 logger = init_logger(__name__)
@@ -83,6 +83,15 @@ 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,
+            )
         if not skip_load_model:
             self.load_model(load_format=self.od_config.diffusion_load_format)
             self.init_lora_manager()
@@ -174,15 +183,32 @@ class DiffusionWorker:
         """Generate output for the given requests."""
         return self.execute_model(request, 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)
+    def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
+        """Start or stop profiling for this GPU worker.
 
-    @classmethod
-    def stop_profile(cls) -> dict | None:
-        """Stop profiling and return the result dictionary."""
-        return CurrentProfiler.stop()
+        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.
+        """
+        if self.profiler is None:
+            logger.warning("Profiler not initialized, skipping profile(%s)", is_start)
+            return
+
+        if is_start:
+            from vllm_omni.profiler import OmniTorchProfilerWrapper
+
+            if isinstance(self.profiler, OmniTorchProfilerWrapper):
+                import time
+
+                filename = profile_prefix or f"diffusion_{int(time.time())}"
+                self.profiler.set_trace_filename(filename)
+            self.profiler.start()
+        else:
+            self.profiler.stop()
 
     def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput:
         """Execute a forward pass by delegating to the model runner."""

vllm_omni/engine/async_omni_engine.py

@@ -18,6 +18,7 @@ import time
 import uuid
 import weakref
 from collections.abc import Sequence
+from dataclasses import asdict
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
@@ -859,6 +860,15 @@ class AsyncOmniEngine:
                     "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
                     "quantization": kwargs.get("quantization", None),
                     "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
+                    **(
+                        {
+                            "profiler_config": asdict(kwargs["profiler_config"])
+                            if hasattr(kwargs["profiler_config"], "__dataclass_fields__")
+                            else kwargs["profiler_config"]
+                        }
+                        if kwargs.get("profiler_config") is not None
+                        else {}
+                    ),
                 },
                 "final_output": True,
                 "final_output_type": "image",

vllm_omni/entrypoints/async_omni.py

@@ -442,19 +442,30 @@ class AsyncOmni(EngineClient, OmniBase):
         async with self._pause_cond:
             return self._paused
 
-    async def start_profile(self, stages: list[int] | None = None) -> list[Any]:
+    async def start_profile(
+        self,
+        profile_prefix: str | None = None,
+        stages: list[int] | None = None,
+    ) -> list[Any]:
         """Start profiling specified stages.
 
-        TODO(AsyncOmni): normalize return payloads across LLM/diffusion stages.
+        Uses vLLM-compatible profile(is_start=True, profile_prefix) interface.
+
+        Args:
+            profile_prefix: Optional prefix for the trace file names.
+            stages: List of stage IDs to profile. If None, profiles all stages.
         """
-        return await self.collective_rpc(method="start_profile", stage_ids=stages)
+        return await self.collective_rpc(method="profile", args=(True, profile_prefix), stage_ids=stages)
 
     async def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
         """Stop profiling specified stages.
 
-        TODO(AsyncOmni): normalize return payloads across LLM/diffusion stages.
+        Uses vLLM-compatible profile(is_start=False) interface.
+
+        Args:
+            stages: List of stage IDs to profile. If None, stops all stages.
         """
-        return await self.collective_rpc(method="stop_profile", stage_ids=stages)
+        return await self.collective_rpc(method="profile", args=(False, None), stage_ids=stages)
 
     async def reset_mm_cache(self) -> None:
         """Reset the multi-modal cache for all stages.

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -440,28 +440,21 @@ class AsyncOmniDiffusion:
         )
         return all(results) if isinstance(results, list) else results
 
-    async def start_profile(self, trace_filename: str | None = None) -> None:
-        """Start profiling for the diffusion model.
+    async def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
+        """Start or stop profiling for the diffusion model.
 
         Args:
-            trace_filename: Optional base filename for trace files.
-                           If None, a timestamp-based name will be generated.
+            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.
         """
         loop = asyncio.get_event_loop()
         await loop.run_in_executor(
             self._executor,
-            self.engine.start_profile,
-            trace_filename,
-        )
-
-    async def stop_profile(self) -> dict:
-        """Stop profiling and return profiling results.
-
-        Returns:
-            Dictionary containing paths to trace and table files.
-        """
-        loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(
-            self._executor,
-            self.engine.stop_profile,
+            self.engine.profile,
+            is_start,
+            profile_prefix,
         )

vllm_omni/entrypoints/omni_base.py

@@ -292,6 +292,37 @@ class OmniBase:
     def close(self) -> None:
         self.shutdown()
 
+    def start_profile(
+        self,
+        profile_prefix: str | None = None,
+        stages: list[int] | None = None,
+    ) -> list[Any]:
+        """Start profiling specified stages.
+
+        Uses vLLM-compatible profile(is_start=True, profile_prefix) interface.
+
+        Args:
+            profile_prefix: Optional prefix for the trace file names.
+            stages: List of stage IDs to profile. If None, profiles all stages.
+
+        Returns:
+            List of results from each stage.
+        """
+        return self.engine.collective_rpc(method="profile", args=(True, profile_prefix), stage_ids=stages)
+
+    def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
+        """Stop profiling specified stages.
+
+        Uses vLLM-compatible profile(is_start=False) interface.
+
+        Args:
+            stages: List of stage IDs to profile. If None, stops all stages.
+
+        Returns:
+            List of results from each stage.
+        """
+        return self.engine.collective_rpc(method="profile", args=(False, None), stage_ids=stages)
+
     def _shutdown_base(self) -> None:
         if getattr(self, "_shutdown_called", False):
             return

vllm_omni/entrypoints/openai/api_server.py

@@ -118,18 +118,28 @@ router = APIRouter()
 profiler_router = APIRouter()
 
 
-def _should_enable_profiler_endpoints(args: Namespace) -> bool:
-    # Check upstream vLLM's profiler_config
-    profiler_config = getattr(args, "profiler_config", None)
-    if profiler_config is not None:
-        # profiler_config exists, check if profiler is set
-        profiler = getattr(profiler_config, "profiler", None)
-        if profiler is not None:
-            return True
-
-    # TODO: remove this env after refactoring torch profiler to CLI args
-    env_value = os.environ.get("VLLM_TORCH_PROFILER_DIR")
-    return env_value is not None
+def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool:
+    """Check if any stage has profiler_config set in its engine_args."""
+    if not stage_configs:
+        return False
+    for stage in stage_configs:
+        engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None)
+        if engine_args is None:
+            continue
+        profiler_config = (
+            engine_args.get("profiler_config")
+            if isinstance(engine_args, dict)
+            else getattr(engine_args, "profiler_config", None)
+        )
+        if profiler_config is not None:
+            profiler = (
+                profiler_config.get("profiler")
+                if isinstance(profiler_config, dict)
+                else getattr(profiler_config, "profiler", None)
+            )
+            if profiler is not None:
+                return True
+    return False
 
 
 class ProfileRequest(BaseModel):
@@ -292,8 +302,9 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         await omni_init_app_state(engine_client, app.state, args)
 
-        # Conditionally register profiler endpoints based on config or env var
-        if _should_enable_profiler_endpoints(args):
+        # Conditionally register profiler endpoints based on stage YAML configs
+        stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
+        if _should_enable_profiler_endpoints(stage_configs):
             logger.warning("Profiler endpoints are enabled. This should ONLY be used for local development!")
             app.include_router(profiler_router)
 

vllm_omni/platforms/interface.py

@@ -125,6 +125,16 @@ class OmniPlatform(Platform):
 
         os.environ.pop(cls.device_control_env_var, None)
 
+    @classmethod
+    def get_profiler_cls(cls) -> str:
+        """Get the profiler class for this platform.
+
+        Returns:
+            Fully qualified class path of the profiler.
+            Default returns the base OmniTorchProfilerWrapper.
+        """
+        return "vllm_omni.profiler.omni_torch_profiler.OmniTorchProfilerWrapper"
+
 
 class UnspecifiedOmniPlatform(OmniPlatform):
     _omni_enum = OmniPlatformEnum.UNSPECIFIED

vllm_omni/platforms/npu/platform.py

@@ -105,3 +105,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
     def get_device_total_memory(cls, device_id: int = 0) -> int:
         device_props = torch.npu.get_device_properties(device_id)
         return device_props.total_memory
+
+    @classmethod
+    def get_profiler_cls(cls) -> str:
+        return "vllm_omni.platforms.npu.profiler.NPUTorchProfilerWrapper"

vllm_omni/platforms/npu/profiler.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+
+from typing_extensions import override
+from vllm.config import ProfilerConfig
+from vllm.logger import init_logger
+
+from vllm_omni.profiler.omni_torch_profiler import (
+    OmniTorchProfilerWrapper,
+    TorchProfilerActivity,
+)
+
+logger = init_logger(__name__)
+
+
+class NPUTorchProfilerWrapper(OmniTorchProfilerWrapper):
+    """NPU-specific profiler using torch_npu.profiler.
+
+    Key differences from base profiler:
+    - Uses torch_npu.profiler instead of torch.profiler
+    - Different experimental_config options (AiCMetrics, profiler_level, etc.)
+    - Uses tensorboard_trace_handler for trace output
+    - No key_averages() support - requires offline parsing with
+      torch_npu.profiler.profiler.analyse()
+    """
+
+    @override
+    def _get_default_activities(self) -> list[TorchProfilerActivity]:
+        return ["CPU", "NPU"]
+
+    @override
+    def _create_profiler(
+        self,
+        profiler_config: ProfilerConfig,
+        activities: list[TorchProfilerActivity],
+    ):
+        import torch_npu
+
+        # Map activity names to torch_npu profiler activities
+        npu_activities = []
+        for activity in activities:
+            if activity == "CPU":
+                npu_activities.append(torch_npu.profiler.ProfilerActivity.CPU)
+            elif activity == "NPU":
+                npu_activities.append(torch_npu.profiler.ProfilerActivity.NPU)
+
+        # NPU-specific experimental config for detailed profiling
+        experimental_config = torch_npu.profiler._ExperimentalConfig(
+            export_type=torch_npu.profiler.ExportType.Text,
+            profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
+            msprof_tx=False,
+            aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
+            l2_cache=False,
+            op_attr=False,
+            data_simplification=True,
+            record_op_args=False,
+            gc_detect_threshold=None,
+        )
+
+        # Set up trace directory for NPU - tensorboard_trace_handler creates
+        # its own subdirectory structure. Use worker_name which includes stage_id.
+        npu_trace_dir = os.path.join(self._trace_dir, self._worker_name)
+        os.makedirs(npu_trace_dir, exist_ok=True)
+        self._trace_path = npu_trace_dir
+
+        return torch_npu.profiler.profile(
+            activities=npu_activities,
+            with_stack=False,
+            profile_memory=profiler_config.torch_profiler_with_memory,
+            # NOTE: torch_npu.profiler.with_modules is equivalent to
+            # torch.profiler.with_stack. The with_stack option in
+            # torch_npu.profiler introduces significant time overhead.
+            with_modules=profiler_config.torch_profiler_with_stack,
+            experimental_config=experimental_config,
+            # Use tensorboard_trace_handler directly - NPU profiler expects
+            # this specific handler format
+            on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(npu_trace_dir),
+        )
+
+    @override
+    def _on_stop_hook(self) -> None:
+        """NPU profiler doesn't support key_averages() - log offline parsing hint."""
+        if self.local_rank == 0:
+            logger.info(
+                "NPU profiler stopped. Use offline parsing to analyze: "
+                "from torch_npu.profiler.profiler import analyse; "
+                "analyse('%s')",
+                self._trace_path or self._trace_dir,
+            )

vllm_omni/platforms/npu/stage_configs/qwen3_omni_moe.yaml

@@ -22,6 +22,9 @@ stage_args:
       enable_prefix_caching: false
       hf_config_name: thinker_config
       tensor_parallel_size: 2
+      # profiler_config:
+      #   profiler: torch
+      #   torch_profiler_dir: ./perf
     final_output: true
     final_output_type: text
     is_comprehension: true

vllm_omni/platforms/npu/worker/base.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Base NPU worker class for vLLM-Omni with OmniProfiler support."""
+
+import time
+
+from vllm_ascend.worker.worker import NPUWorker
+
+
+class OmniNPUWorkerBase(NPUWorker):
+    """Base NPU worker for vLLM-Omni with OmniProfiler support.
+
+    This class replaces vllm-ascend's profiler with OmniProfiler for
+    unified profiling across all platforms.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Replace vllm-ascend's profiler with OmniProfiler
+        profiler_config = self.vllm_config.profiler_config
+        if profiler_config and profiler_config.profiler == "torch":
+            from vllm_omni.profiler import create_omni_profiler
+
+            stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
+            worker_name = f"stage{stage_id}_rank{self.rank}"
+            self.profiler = create_omni_profiler(
+                profiler_config=profiler_config,
+                worker_name=worker_name,
+                local_rank=self.local_rank,
+            )
+
+    def profile(self, is_start: bool = True, profile_prefix: str | None = None):
+        """Override to set trace filename before starting the profiler.
+
+        NPUWorker's profile() accepts profile_prefix, so we use it to generate
+        a descriptive trace filename for OmniProfiler.
+        """
+        if self.profiler is None:
+            raise RuntimeError(
+                "Profiling is not enabled. For diffusion models, set --profiler-config via CLI. "
+                "For omni models, add profiler_config to your stage config file."
+            )
+        if is_start:
+            from vllm_omni.profiler import OmniTorchProfilerWrapper
+
+            if isinstance(self.profiler, OmniTorchProfilerWrapper):
+                # Include stage_id and rank in default filename to distinguish
+                # traces from different stages profiling on the same local_rank
+                stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
+                filename = profile_prefix or f"stage{stage_id}_rank{self.rank}_{int(time.time())}"
+                self.profiler.set_trace_filename(filename)
+            self.profiler.start()
+        else:
+            self.profiler.stop()

vllm_omni/platforms/npu/worker/npu_ar_worker.py

@@ -2,13 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from vllm.v1.worker.workspace import init_workspace_manager
-from vllm_ascend.worker.worker import NPUWorker
 
+from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase
 from vllm_omni.platforms.npu.worker.npu_ar_model_runner import NPUARModelRunner
 from vllm_omni.worker.mixins import OmniWorkerMixin
 
 
-class NPUARWorker(OmniWorkerMixin, NPUWorker):
+class NPUARWorker(OmniWorkerMixin, OmniNPUWorkerBase):
     """NPU AR worker for thinker/talker stages in Omni model."""
 
     def init_device(self):

vllm_omni/platforms/npu/worker/npu_generation_worker.py

@@ -2,13 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from vllm.v1.worker.workspace import init_workspace_manager
-from vllm_ascend.worker.worker import NPUWorker
 
+from vllm_omni.platforms.npu.worker.base import OmniNPUWorkerBase
 from vllm_omni.platforms.npu.worker.npu_generation_model_runner import NPUGenerationModelRunner
 from vllm_omni.worker.mixins import OmniWorkerMixin
 
 
-class NPUGenerationWorker(OmniWorkerMixin, NPUWorker):
+class NPUGenerationWorker(OmniWorkerMixin, OmniNPUWorkerBase):
     """NPU generation worker for code2wav stage in Omni model."""
 
     def init_device(self):

vllm_omni/profiler/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from .omni_torch_profiler import OmniTorchProfilerWrapper, create_omni_profiler
+
+__all__ = ["OmniTorchProfilerWrapper", "create_omni_profiler"]

vllm_omni/profiler/omni_torch_profiler.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import os
+import subprocess
+import time
+from typing import Literal
+
+import torch
+from typing_extensions import override
+from vllm.config import ProfilerConfig
+from vllm.config.profiler import _is_uri_path
+from vllm.logger import init_logger
+from vllm.profiler.wrapper import WorkerProfiler
+
+logger = init_logger(__name__)
+
+# NPU has its custom profiler
+TorchProfilerActivity = Literal["CPU", "CUDA", "XPU", "NPU"]
+TorchProfilerActivityMap = {
+    "CPU": torch.profiler.ProfilerActivity.CPU,
+    "CUDA": torch.profiler.ProfilerActivity.CUDA,
+    "XPU": torch.profiler.ProfilerActivity.XPU,
+}
+
+
+class OmniTorchProfilerWrapper(WorkerProfiler):
+    """Base torch profiler wrapper with platform-agnostic functionality.
+
+    Provides common profiler features:
+    - Custom trace file naming with stage/rank info
+    - Background gzip compression via subprocess
+    - Returns trace file paths from get_results() for orchestrator collection
+
+    Subclasses can override hook methods for platform-specific behavior:
+    - _get_default_activities(): Return default activities for the platform
+    - _create_profiler(): Create platform-specific profiler instance
+    - _on_stop_hook(): Handle platform-specific post-stop logic
+    """
+
+    def __init__(
+        self,
+        profiler_config: ProfilerConfig,
+        worker_name: str,
+        local_rank: int,
+        activities: list[TorchProfilerActivity] | None = None,
+    ) -> None:
+        super().__init__(profiler_config)
+
+        if activities is None:
+            activities = self._get_default_activities()
+
+        self.local_rank = local_rank
+        self.profiler_config = profiler_config
+        self._worker_name = worker_name
+        self._trace_dir = profiler_config.torch_profiler_dir
+        self._use_gzip = profiler_config.torch_profiler_use_gzip
+        self._trace_filename: str | None = None
+        self._trace_path: str | None = None
+        self._table_path: str | None = None
+
+        if local_rank in (None, 0):
+            logger.info_once(
+                "Omni torch profiling enabled. Traces will be saved to: %s",
+                self._trace_dir,
+                scope="local",
+            )
+
+        self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
+        self.profiler = self._create_profiler(profiler_config, activities)
+
+    def _get_default_activities(self) -> list[TorchProfilerActivity]:
+        """Get default activities for this platform.
+
+        Override in subclasses to provide platform-specific defaults.
+        """
+        return ["CPU", "CUDA"]
+
+    def _create_profiler(
+        self,
+        profiler_config: ProfilerConfig,
+        activities: list[TorchProfilerActivity],
+    ):
+        """Create the profiler instance.
+
+        Override in subclasses for platform-specific profiler creation.
+        """
+        return torch.profiler.profile(
+            activities=[TorchProfilerActivityMap[a] for a in activities],
+            record_shapes=profiler_config.torch_profiler_record_shapes,
+            profile_memory=profiler_config.torch_profiler_with_memory,
+            with_stack=profiler_config.torch_profiler_with_stack,
+            with_flops=profiler_config.torch_profiler_with_flops,
+            on_trace_ready=self._on_trace_ready,
+        )
+
+    def set_trace_filename(self, filename: str) -> None:
+        """Set the trace filename before starting profiling.
+
+        Args:
+            filename: Base filename without extension or rank suffix.
+                      e.g. "stage_0_llm_1234567890"
+                      Can also be a full path (e.g. from diffusion engine).
+        """
+        self._trace_filename = filename
+
+    def _on_trace_ready(self, prof) -> None:
+        """Custom trace handler: export chrome trace with omni naming."""
+        rank = self.local_rank
+        filename = self._trace_filename or f"{self._worker_name}_{int(time.time())}"
+        # If filename already contains a directory, use as-is (e.g. from
+        # diffusion engine which builds full path). Otherwise join with trace_dir.
+        if os.path.dirname(filename):
+            json_file = f"{filename}_rank{rank}.json"
+        else:
+            json_file = os.path.join(self._trace_dir, f"{filename}_rank{rank}.json")
+
+        os.makedirs(os.path.dirname(json_file), exist_ok=True)
+
+        try:
+            prof.export_chrome_trace(json_file)
+            logger.info("[Rank %s] Trace exported to %s", rank, json_file)
+
+            if self._use_gzip:
+                try:
+                    subprocess.Popen(["gzip", "-f", json_file])
+                    logger.info(
+                        "[Rank %s] Triggered background compression for %s",
+                        rank,
+                        json_file,
+                    )
+                    self._trace_path = f"{json_file}.gz"
+                except Exception as compress_err:
+                    logger.warning(
+                        "[Rank %s] Background gzip failed to start: %s",
+                        rank,
+                        compress_err,
+                    )
+                    self._trace_path = json_file
+            else:
+                self._trace_path = json_file
+
+        except Exception as e:
+            logger.warning("[Rank %s] Failed to export trace: %s", rank, e)
+
+    @override
+    def _start(self) -> None:
+        self.profiler.start()
+
+    @override
+    def _stop(self) -> None:
+        """Stop profiler, export trace via on_trace_ready, and dump table."""
+        self.profiler.stop()
+        self._on_stop_hook()
+
+    def _on_stop_hook(self) -> None:
+        """Hook called after profiler.stop().
+
+        Override in subclasses for platform-specific post-stop handling.
+        Base implementation handles CUDA time total dump.
+        """
+        rank = self.local_rank
+
+        if self.profiler_config.torch_profiler_dump_cuda_time_total:
+            profiler_dir = self.profiler_config.torch_profiler_dir
+            sort_key = "self_cuda_time_total"
+            table = self.profiler.key_averages().table(sort_by=sort_key)
+
+            if not _is_uri_path(profiler_dir):
+                table_file = os.path.join(profiler_dir, f"profiler_out_{rank}.txt")
+                with open(table_file, "w") as f:
+                    print(table, file=f)
+                self._table_path = table_file
+
+            if rank == 0:
+                print(table)
+
+        if self.dump_cpu_time_total and rank == 0:
+            logger.info(self.profiler.key_averages().table(sort_by="self_cpu_time_total", row_limit=50))
+
+    @override
+    def annotate_context_manager(self, name: str):
+        return torch.profiler.record_function(name)
+
+    def get_results(self) -> dict:
+        """Return collected trace and table paths after stop."""
+        return {
+            "trace": self._trace_path,
+            "table": self._table_path,
+        }
+
+
+def create_omni_profiler(
+    profiler_config: ProfilerConfig,
+    worker_name: str,
+    local_rank: int,
+    activities: list[TorchProfilerActivity] | None = None,
+) -> OmniTorchProfilerWrapper:
+    """Factory function to create platform-specific profiler.
+
+    Uses the current platform's get_profiler_cls() to determine which
+    profiler class to instantiate.
+
+    Args:
+        profiler_config: Profiler configuration.
+        worker_name: Name of the worker.
+        local_rank: Local rank of the worker.
+        activities: Optional list of profiler activities.
+
+    Returns:
+        Platform-specific profiler instance.
+    """
+    from vllm.utils.import_utils import resolve_obj_by_qualname
+
+    from vllm_omni.platforms import current_omni_platform
+
+    profiler_cls_path = current_omni_platform.get_profiler_cls()
+    profiler_cls = resolve_obj_by_qualname(profiler_cls_path)
+    return profiler_cls(profiler_config, worker_name, local_rank, activities)

vllm_omni/worker/base.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import os
+import time
 
 import torch
 from vllm.logger import init_logger
@@ -24,8 +25,55 @@ class OmniGPUWorkerBase(GPUWorker):
     This class overrides determine_available_memory() to use per-process GPU
     memory tracking via pynvml, allowing multiple stages to initialize
     concurrently on the same GPU without memory accounting interference.
+
+    It also replaces vLLM's TorchProfilerWrapper with OmniTorchProfilerWrapper
+    for custom trace naming, background gzip, and trace path collection.
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Replace vLLM's profiler with platform-specific profiler
+        profiler_config = self.vllm_config.profiler_config
+        if profiler_config and profiler_config.profiler == "torch":
+            from vllm_omni.profiler import create_omni_profiler
+
+            stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
+            worker_name = f"stage{stage_id}_rank{self.rank}"
+            self.profiler = create_omni_profiler(
+                profiler_config=profiler_config,
+                worker_name=worker_name,
+                local_rank=self.local_rank,
+            )
+
+    def profile(self, is_start: bool = True, profile_prefix: str | None = None):
+        """Override to set trace filename before starting the profiler.
+
+        Args:
+            is_start: True to start profiling, False to stop.
+            profile_prefix: Optional prefix for trace filename (vLLM compat).
+
+        vLLM's profile() only passes is_start, so we generate a descriptive
+        trace filename here before delegating to the profiler.
+        """
+        if self.profiler is None:
+            raise RuntimeError(
+                "Profiling is not enabled. For diffusion models, set --profiler-config via CLI. "
+                "For omni models, add profiler_config to your stage config file."
+            )
+        if is_start:
+            from vllm_omni.profiler import OmniTorchProfilerWrapper
+
+            if isinstance(self.profiler, OmniTorchProfilerWrapper):
+                # Include stage_id and rank in default filename to distinguish
+                # traces from different stages profiling on the same local_rank
+                stage_id = getattr(self.vllm_config.model_config, "stage_id", 0)
+                filename = profile_prefix or f"stage{stage_id}_rank{self.rank}_{int(time.time())}"
+                self.profiler.set_trace_filename(filename)
+            self.profiler.start()
+        else:
+            self.profiler.stop()
+
     @torch.inference_mode()
     def determine_available_memory(self) -> int:
         """Process-scoped GPU memory profiling for concurrent stage initialization.