[Enhancement] add pytorch profiler ops and memory record

[bjf-frz]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Enhance the PyTorch Profiler to include memory timeline and operation table recordings.

Test Plan

When using online service, set --profiler-config as the startup parameter.

vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
  --omni \
  --port 8091 \
  --profiler-config '{
    "profiler": "torch",
    "torch_profiler_dir": "/tmp/vllm_profile_wan22_i2v",
    "torch_profiler_with_stack": true,
    "torch_profiler_with_flops": false,
    "torch_profiler_use_gzip": true,
    "torch_profiler_dump_cuda_time_total": true,
    "torch_profiler_record_shapes": true,
    "torch_profiler_with_memory": true,
    "ignore_frontend": false,
    "delay_iterations": 0,
    "max_iterations": 0
  }'

In the user end, curl the request like:

#!/usr/bin/env bash

curl -s -X POST http://localhost:8000/start_profile \
  -H "Content-Type: application/json" \
  -d '{}'

resp=$(curl -s -X POST http://localhost:8000/v1/videos \
  -F "prompt=一只兔子看向镜头" \
  -F "input_reference=@rabbit.png")

id=$(echo "$resp" | jq -r '.id')

while true; do
  status=$(curl -s http://localhost:8000/v1/videos/$id | jq -r '.status')
  [ "$status" = "completed" ] && break
  [ "$status" = "failed" ] && exit 1
  sleep 2
done

curl -s http://localhost:8000/v1/videos/$id/content -o output.mp4

curl -s -X POST http://localhost:8000/stop_profile \
  -H "Content-Type: application/json" \
  -d '{}'

echo "done"

Test Result

The profiling file gathered like:
20260403-072613_stage_0_diffusion_1775201173/
├── memory_snapshot_rank0.pickle
├── ops_rank0.xlsx
├── stacks_cpu_rank0.txt
├── stacks_cuda_rank0.txt
└── trace_rank0.json.gz

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[bjf-frz]

@david6666666 @gcanlin PTAL,thx.

[david6666666]

I found one blocking issue around the new Excel export dependency.

[david6666666]

Can other platform can adapt this feature? @bjf-frz @gcanlin

[bjf-frz]

Can other platform can adapt this feature? @bjf-frz @gcanlin

I modified this profiler function to include a graceful fallback: if the platform does not support it, it will be skipped, and a warning will be printed.

[hsliuustc0106]

add test result and highly recommended to update the docs

[bjf-frz]

add test result and highly recommended to update the docs

Profiling files include:
Python call stack archive (compressed in .gz format)
Memory usage timeline (pickle format)
Operator (op) table (.xlsx format)

[david6666666]

why use VLLM_TORCH_PROFILER_DIR, not args --profiler-dir

[bjf-frz]

In the offline service, this env arg is just used to initiate profiling. It does not point to the actual profiling file destination.

[david6666666]

new args need to explain

[bjf-frz]

done

[hsliuustc0106]

vllm now use

vllm serve meta-llama/Llama-3.1-8B-Instruct --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'

does it still use the envs?

[hsliuustc0106]

https://docs.vllm.ai/en/latest/contributing/profiling/#profile-with-pytorch-profiler

[bjf-frz]

vllm now use

vllm serve meta-llama/Llama-3.1-8B-Instruct --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'

does it still use the envs?

In the current implementation, enabling this environment argument is mandatory for profiling to be activated within the offline service.

[gcanlin]

That's the legacy code actually. I recommend to clean example instead of modifying the docs.

[hsliuustc0106]

?

[bjf-frz]

?

It may generate a file .codex when using codex, not only the .codex/ directory.