[Enhancement]modify profiling.md (#3051)

Signed-off-by: bjf-frz <frz123db@gmail.com>

docs/contributing/model/adding_diffusion_model.md

@@ -813,16 +813,16 @@ class WanTransformer3DModel(nn.Module):
 
 ---
 
-### Diffusion Timing (Performance Profiling)
+### Diffusion Pipeline Profiler (Performance Profiling)
 When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
 vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
 
 !!! info
-      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
+      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
 
 This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
 
-**Enabling Diffusion Timing**
+**Enabling Diffusion Pipeline Profiler**
 
 
 Enable timing by setting:
@@ -843,7 +843,7 @@ If not specified, the default targets are used:
 **Adding DiffusionPipelineProfilerMixin to a Pipeline**
 To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
 ```python
-from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.profiler import DiffusionPipelineProfilerMixin
 
 class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     # Optional: Specify custom timing targets
@@ -862,7 +862,9 @@ class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         ...
 
         # initialize timing profiler
-        self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
 ```
 The mixin dynamically wraps selected methods and records their execution time during inference.
 
@@ -906,9 +908,9 @@ tokenizer.forward
 
 When enabled, timing logs appear like this:
 ```
-[DiffusionTiming] text_encoder.forward took 0.018s
-[DiffusionTiming] diffuse took 2.412s
-[DiffusionTiming] vae.decode took 0.063s
+[DiffusionPipelineProfiler] text_encoder.forward took 0.018s
+[DiffusionPipelineProfiler] diffuse took 2.412s
+[DiffusionPipelineProfiler] vae.decode took 0.063s
 ```
 These measurements help identify bottlenecks during model adaptation and optimization
 

docs/contributing/profiling.md

@@ -1,15 +1,17 @@
-# Profiling vLLM-Omni
+# Profiling Diffusion Models
 
 > **Warning:** Profiling is for development and debugging only. It adds significant overhead and should not be enabled in production.
 
-vLLM-Omni supports two profiler backends through `profiler_config`:
+Diffusion profiling supports two backends through `profiler_config`:
 
-- `torch`: detailed CPU/CUDA traces written to `torch_profiler_dir`
+- `torch`: detailed CPU/CUDA traces, operator tables, and optional memory snapshots
 - `cuda`: low-overhead CUDA range control for NVIDIA Nsight Systems (`nsys`)
 
-## 1. Configure Profiling
+## 1. Configure `profiler_config`
 
-Use the same `profiler_config` shape everywhere:
+Use `profiler_config` to enable profiling for a diffusion model. For diffusion usage, pass it directly to `Omni(...)` or `vllm serve`.
+
+Minimal torch-profiler config:
 
 ```yaml
 profiler_config:
@@ -21,54 +23,114 @@ Supported fields:
 
 | Field | Description |
 |---|---|
-| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. |
-| `torch_profiler_dir` | Output directory for torch traces. Required when `profiler: torch`. |
+| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. |
+| `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. |
+| `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. |
+| `torch_profiler_record_shapes` | Record input shapes and add a `by_shape` sheet to `ops_rank*.xlsx`. |
+| `torch_profiler_with_stack` | Record call stacks, add a `by_stack` sheet to `ops_rank*.xlsx`, and export `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt`. |
+| `torch_profiler_with_memory` | Enable memory profiling and attempt to dump `memory_snapshot_rank*.pickle`. The pickle is only generated when the current backend supports memory history and snapshot APIs. |
+| `torch_profiler_with_flops` | Enable FLOPs collection in `torch.profiler`. This does not add a separate output file. |
+| `torch_profiler_dump_cuda_time_total` | Export an additional text summary `profiler_out_<rank>.txt` sorted by `self_cuda_time_total`. |
 | `delay_iterations` | Number of worker iterations to skip before profiling starts. |
 | `max_iterations` | Maximum number of worker iterations to capture before auto-stop. |
+| `wait_iterations` | Torch-profiler wait iterations before warmup. |
 | `warmup_iterations` | Torch-profiler warmup iterations. |
 | `active_iterations` | Torch-profiler active iterations. |
-| `wait_iterations` | Torch-profiler wait iterations before warmup. |
 
-For multi-stage omni pipelines, put `profiler_config` under the target stage's `engine_args`.
+### Minimal configurations by output
 
-```yaml
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    engine_args:
-      profiler_config:
-        profiler: torch
-        torch_profiler_dir: ./perf
-```
-
-For single-stage diffusion usage, pass `profiler_config` directly to `Omni(...)` or `vllm serve`.
-
-## 2. Profiling Omni Pipelines
-
-It is usually best to profile only the stages you need.
+Only collect trace output:
 
 ```python
-# Profile all stages.
-omni.start_profile()
-
-# Profile selected stages only.
-omni.start_profile(stages=[0, 2])
-...
-omni.stop_profile(stages=[0, 2])
+profiler_config = {
+    "profiler": "torch",
+    "torch_profiler_dir": "./perf",
+}
 ```
 
-Always stop the same stage set that you started. If only some stages have `profiler_config`, pass an explicit `stages=[...]` list instead of relying on the default "all stages" behavior.
+Outputs:
 
-Examples:
+- `trace_rank*.json`
+- `ops_rank*.xlsx` with a `summary` sheet
 
-1. [Qwen2.5-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
-2. [Qwen3-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
+Collect compressed trace output:
 
-## 3. Profiling Single-Stage Diffusion
+```python
+profiler_config = {
+    "profiler": "torch",
+    "torch_profiler_dir": "./perf",
+    "torch_profiler_use_gzip": True,
+}
+```
 
-Single-stage diffusion models use the same `start_profile()` / `stop_profile()` controls, but you must provide `profiler_config` explicitly.
+Outputs:
 
-### PyTorch profiler
+- `trace_rank*.json.gz`
+- `ops_rank*.xlsx` with a `summary` sheet
+
+Collect trace and full operator tables:
+
+```python
+profiler_config = {
+    "profiler": "torch",
+    "torch_profiler_dir": "./perf",
+    "torch_profiler_record_shapes": True,
+    "torch_profiler_with_stack": True,
+}
+```
+
+Outputs:
+
+- `trace_rank*.json`
+- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
+- `stacks_cpu_rank*.txt`
+- `stacks_cuda_rank*.txt`
+
+Collect trace, operator tables, and memory snapshots:
+
+```python
+profiler_config = {
+    "profiler": "torch",
+    "torch_profiler_dir": "./perf",
+    "torch_profiler_record_shapes": True,
+    "torch_profiler_with_stack": True,
+    "torch_profiler_with_memory": True,
+}
+```
+
+Outputs:
+
+- `trace_rank*.json`
+- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
+- `stacks_cpu_rank*.txt`
+- `stacks_cuda_rank*.txt`
+- `memory_snapshot_rank*.pickle` when supported by the current backend
+
+### Full torch-profiler configuration
+
+If you want to enable the commonly used torch-profiler options together:
+
+```python
+profiler_config = {
+    "profiler": "torch",
+    "torch_profiler_dir": "./perf",
+    "torch_profiler_use_gzip": False,
+    "torch_profiler_record_shapes": True,
+    "torch_profiler_with_stack": True,
+    "torch_profiler_with_memory": True,
+    "torch_profiler_with_flops": False,
+    "torch_profiler_dump_cuda_time_total": False,
+    "delay_iterations": 0,
+    "max_iterations": 0,
+    "wait_iterations": 0,
+    "warmup_iterations": 0,
+    "active_iterations": 0,
+}
+```
+
+## 2. Profiling Diffusion with PyTorch Profiler
+
+Single-stage diffusion models use `start_profile()` / `stop_profile()` controls. The profiler only writes artifacts after profiling has been started and then stopped.
 
 ```python
 from vllm_omni import Omni
@@ -86,7 +148,29 @@ omni.start_profile()
 omni.stop_profile()
 ```
 
-### Nsight Systems (`nsys`)
+For diffusion offline example scripts under `examples/offline_inference/`, pass `--profiler-config` as a JSON object. The script enables profiling when this argument is set and wraps generation with `start_profile()` / `stop_profile()`.
+
+Example:
+
+```bash
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+  --image input.jpg \
+  --prompt "A cat playing with yarn" \
+  --profiler-config '{
+    "profiler": "torch",
+    "torch_profiler_dir": "./perf",
+    "torch_profiler_record_shapes": true,
+    "torch_profiler_with_stack": true
+  }'
+```
+
+Examples:
+
+1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
+2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
+
+## 3. Profiling Diffusion with Nsight Systems (`nsys`)
 
 For Nsight Systems, use `profiler: cuda` and wrap the process with `nsys profile`.
 
@@ -108,30 +192,15 @@ profiler_config={"profiler": "cuda"}
 
 Then call `start_profile()` before the requests you want to capture and `stop_profile()` after them. The diffusion worker processes open and close the CUDA capture range themselves, so `nsys` sees the actual GPU work instead of only the parent process.
 
-Examples:
-
-1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
-2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
-
 ## 4. Profiling Online Serving
 
-When any stage has `profiler_config.profiler` set, the server exposes:
+When `profiler_config.profiler` is set for a diffusion model, the server exposes:
 
 - `POST /start_profile`
 - `POST /stop_profile`
 
 ### Start the server
 
-Multi-stage omni serving:
-
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B \
-  --omni \
-  --port 8091
-```
-
-(The default deploy config at `vllm_omni/deploy/qwen2_5_omni.yaml` is loaded automatically. Pass `--deploy-config /path/to/custom.yaml` to override.)
-
 Single-stage diffusion serving with torch profiler:
 
 ```bash
@@ -144,21 +213,17 @@ vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
     "torch_profiler_with_stack": true,
     "torch_profiler_with_flops": false,
     "torch_profiler_use_gzip": true,
-    "torch_profiler_dump_cuda_time_total": true,
+    "torch_profiler_dump_cuda_time_total": false,
     "torch_profiler_record_shapes": true,
     "torch_profiler_with_memory": true,
     "delay_iterations": 0,
-    "max_iterations": 0
+    "max_iterations": 0,
+    "wait_iterations": 0,
+    "warmup_iterations": 0,
+    "active_iterations": 0
   }'
 ```
 
-Useful optional fields:
-
-- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files.
-- `torch_profiler_record_shapes`: export `by_shape` operator views.
-- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history.
-- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`.
-
 Single-stage diffusion serving with Nsight Systems:
 
 ```bash
@@ -176,30 +241,37 @@ nsys profile \
 
 ### Control capture
 
+Example profiling flow for an online Qwen-Image request:
+
 ```bash
-# Start profiling on all profiled stages.
+# Start profiling.
 curl -X POST http://localhost:8091/start_profile
 
-# Start profiling on selected stages.
-curl -X POST http://localhost:8091/start_profile \
+# Send a Qwen-Image generation request while profiling is active.
+curl http://localhost:8091/v1/images/generations \
   -H "Content-Type: application/json" \
-  -d '{"stages": [0]}'
+  -d '{
+    "model": "Qwen/Qwen-Image",
+    "prompt": "A red vintage bicycle parked beside a quiet canal at sunset"
+  }'
 
-# Stop profiling.
+# Stop profiling and flush profiler artifacts.
 curl -X POST http://localhost:8091/stop_profile
 ```
 
-For mixed-stage pipelines, use explicit `stages` and pass the same stage list to both endpoints.
+## 5. Diffusion Pipeline Profiler
 
-## 5. Analyze Results
+For lightweight per-stage pipeline timing such as `vae.decode` or `diffuse`, see [Diffusion Pipeline Profiler](model/adding_diffusion_model.md#diffusion-pipeline-profiler-performance-profiling). That utility logs stage durations only and does not generate torch-profiler artifacts such as `trace.json`, Excel tables, or memory snapshots.
 
-Torch profiler output:
+## 6. Analyze Results
+
+Torch-profiler output:
 
 - Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz`
 - Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets
 - Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled
 - Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend
-- Optional aggregated CUDA-time tables under the same session directory
+- Optional CUDA-time text summary: `profiler_out_<rank>.txt` when `torch_profiler_dump_cuda_time_total` is enabled
 
 CUDA profiler / Nsight Systems output:
 
@@ -211,4 +283,4 @@ Recommended viewers:
 - `nsys stats <report>.nsys-rep` for CLI summaries
 - Nsight Systems GUI for CUDA kernel timelines
 
-vLLM-Omni reuses the vLLM profiling infrastructure where possible. For the upstream reference, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
+For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -44,9 +44,11 @@ Layered RGBA output:
 
 import argparse
 import asyncio
+import json
 import os
 import time
 from pathlib import Path
+from typing import Any
 
 import torch
 from PIL import Image
@@ -58,6 +60,16 @@ from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
+    return config
+
+
 # ===========================
 # Argument Parser
 # ===========================
@@ -99,6 +111,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--vae-use-slicing", action="store_true")
     parser.add_argument("--vae-use-tiling", action="store_true")
     parser.add_argument("--enable-cpu-offload", action="store_true")
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
 
     return parser.parse_args()
 
@@ -158,12 +176,13 @@ async def main():
         enable_cpu_offload=args.enable_cpu_offload,
         diffusion_load_format="dummy",
         custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},
+        profiler_config=args.profiler_config,
     )
 
     print(">>> Pipeline loaded successfully")
 
     # ---- Profiling + Info ----
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.profiler_config is not None
     print(f"\n{'=' * 60}")
     print("Generation Configuration")
     print(f"Model: {args.model}")

examples/offline_inference/image_to_image/image_edit.py

@@ -87,9 +87,11 @@ For more options, run:
 """
 
 import argparse
+import json
 import os
 import time
 from pathlib import Path
+from typing import Any
 
 import torch
 from PIL import Image
@@ -101,6 +103,16 @@ from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
+    return config
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
     parser.add_argument(
@@ -360,6 +372,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
     return parser.parse_args()
 
 
@@ -425,11 +443,11 @@ def main():
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        profiler_config=args.profiler_config,
     )
     print("Pipeline loaded")
 
-    # Check if profiling is requested via environment variable
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.profiler_config is not None
 
     # Time profiling for generation
     print(f"\n{'=' * 60}")

examples/offline_inference/image_to_video/image_to_video.py

@@ -33,9 +33,10 @@ Usage:
 """
 
 import argparse
-import os
+import json
 import time
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import PIL.Image
@@ -48,6 +49,16 @@ from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
+    return config
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
     parser.add_argument(
@@ -194,6 +205,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
     return parser.parse_args()
 
 
@@ -274,8 +291,7 @@ def main():
             "rel_l1_thresh": 0.2,
         }
 
-    # Check if profiling is requested via environment variable
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.profiler_config is not None
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
@@ -300,6 +316,7 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        profiler_config=args.profiler_config,
     )
 
     if profiler_enabled:

examples/offline_inference/text_to_image/text_to_image.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
-import os
+import json
 import time
 from pathlib import Path
 from typing import Any
@@ -30,6 +30,16 @@ def is_nextstep_model(model_name: str) -> bool:
     return False
 
 
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
+    return config
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
     parser.add_argument(
@@ -254,6 +264,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
     parser.add_argument(
         "--log-stats",
         action="store_true",
@@ -333,8 +349,7 @@ def main():
         enable_expert_parallel=args.enable_expert_parallel,
     )
 
-    # Check if profiling is requested via environment variable
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.profiler_config is not None
 
     # Prepare LoRA kwargs for Omni initialization
     lora_args: dict[str, Any] = {}
@@ -375,6 +390,7 @@ def main():
         "mode": "text-to-image",
         "log_stats": args.log_stats,
         "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
+        "profiler_config": args.profiler_config,
         "init_timeout": args.init_timeout,
         "stage_init_timeout": args.stage_init_timeout,
         **lora_args,

examples/offline_inference/text_to_video/text_to_video.py

@@ -2,9 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
-import os
+import json
 import time
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import torch
@@ -44,6 +45,16 @@ def _detect_preset(model: str) -> dict:
     return _MODEL_PRESETS["wan"]
 
 
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
+    return config
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Generate a video from a text prompt. "
@@ -185,6 +196,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable diffusion pipeline profiler to display stage durations.",
     )
+    parser.add_argument(
+        "--profiler-config",
+        type=parse_profiler_config,
+        default=None,
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+    )
     parser.add_argument(
         "--quantization",
         type=str,
@@ -248,8 +265,7 @@ def main():
         enable_expert_parallel=args.enable_expert_parallel,
     )
 
-    # Check if profiling is requested via environment variable
-    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    profiler_enabled = args.profiler_config is not None
 
     omni_kwargs = dict(
         model=args.model,
@@ -263,6 +279,7 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        profiler_config=args.profiler_config,
     )
     if args.boundary_ratio is not None:
         omni_kwargs["boundary_ratio"] = args.boundary_ratio

vllm_omni/profiler/omni_torch_profiler.py

@@ -474,10 +474,13 @@ class OmniTorchProfilerWrapper(WorkerProfiler):
             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")
+                table_file = os.path.join(
+                    self._ensure_session_dir(),
+                    f"profiler_out_{rank}.txt",
+                )
                 with open(table_file, "w") as f:
                     print(table, file=f)
-                self._table_path = table_file
+                self._artifact_paths["profiler_out"] = table_file
 
             if rank == 0:
                 print(table)