From 7217557950320d825b5c158ea6a2469166e02a4b Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Tue, 24 Mar 2026 15:25:52 +0800 Subject: [PATCH] [Refactor] Unify torch profiler for omni and diffusion models (#2099) Signed-off-by: gcanlin 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 | 144 +++++++--- .../offline_inference/qwen3_omni/end2end.py | 19 +- tests/conftest.py | 39 +++ .../test_multiproc_engine_concurrency.py | 11 +- tests/entrypoints/test_omni_base_profiler.py | 263 ++++++++++++++++++ vllm_omni/diffusion/data.py | 9 + vllm_omni/diffusion/diffusion_engine.py | 134 ++------- .../diffusion/executor/multiproc_executor.py | 9 +- vllm_omni/diffusion/profiler/__init__.py | 9 +- vllm_omni/diffusion/profiler/base.py | 58 ---- .../diffusion/profiler/torch_profiler.py | 126 --------- vllm_omni/diffusion/stage_diffusion_client.py | 23 +- .../diffusion/worker/diffusion_worker.py | 44 ++- vllm_omni/engine/async_omni_engine.py | 10 + vllm_omni/entrypoints/async_omni.py | 21 +- vllm_omni/entrypoints/async_omni_diffusion.py | 29 +- vllm_omni/entrypoints/omni_base.py | 31 +++ vllm_omni/entrypoints/openai/api_server.py | 39 ++- vllm_omni/platforms/interface.py | 10 + vllm_omni/platforms/npu/platform.py | 4 + vllm_omni/platforms/npu/profiler.py | 93 +++++++ .../npu/stage_configs/qwen3_omni_moe.yaml | 3 + vllm_omni/platforms/npu/worker/base.py | 56 ++++ .../platforms/npu/worker/npu_ar_worker.py | 4 +- .../npu/worker/npu_generation_worker.py | 4 +- vllm_omni/profiler/__init__.py | 6 + vllm_omni/profiler/omni_torch_profiler.py | 221 +++++++++++++++ vllm_omni/worker/base.py | 48 ++++ 28 files changed, 1063 insertions(+), 404 deletions(-) create mode 100644 tests/entrypoints/test_omni_base_profiler.py delete mode 100644 vllm_omni/diffusion/profiler/base.py delete mode 100644 vllm_omni/diffusion/profiler/torch_profiler.py create mode 100644 vllm_omni/platforms/npu/profiler.py create mode 100644 vllm_omni/platforms/npu/worker/base.py create mode 100644 vllm_omni/profiler/__init__.py create mode 100644 vllm_omni/profiler/omni_torch_profiler.py diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 8dbe4c0f6..7a2e64f13 100644 --- a/docs/contributing/profiling.md +++ b/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/) diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 639a6187c..793908519 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py index f2d866a58..0d90f639d 100644 --- a/tests/conftest.py +++ b/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): diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 0465a1932..adb8dc338 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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() diff --git a/tests/entrypoints/test_omni_base_profiler.py b/tests/entrypoints/test_omni_base_profiler.py new file mode 100644 index 000000000..0c1ddc6a5 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index f50f2c0ea..98d8e8c83 100644 --- a/vllm_omni/diffusion/data.py +++ b/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): diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 1b36bb0ef..ff0f753b4 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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