[Core] Optimize Qwen2.5-Omni Token2Wav buffer loading

[akshatvishu]

Resolves #7060.

Summary

find_all_registers() called state_dict() once for every named buffer. Qwen2.5-Omni made 219 calls during each startup, even though every call traversed the same module tree.

The loader now collects nonpersistent buffer names from _non_persistent_buffers_set and walks the named buffers once. The test checks that persistent child buffers remain included, nonpersistent buffers remain excluded, and state_dict() is not called.

The code follows the existing vLLM buffer loading pattern. vLLM-Omni also uses the same check in its diffusion weight loaders.

Benchmark

Item	Value
Model	Qwen/Qwen2.5-Omni-3B
GPU	1 AMD Instinct MI300X VF
GPU memory	192 GB
PyTorch	2.12.0+git6bbd260
ROCm	7.2.53211
vLLM	0.28.0+rocm723
vLLM-Omni	0.28.1.dev47

The server was started with:

MIOPEN_FIND_MODE=FAST vllm serve Qwen/Qwen2.5-Omni-3B --omni --host 127.0.0.1 --port 8091 --deploy-config /tmp/qwen2_5_omni_rocm_single_gpu.yaml

All three stages used device 0 with GPU memory utilization values of 0.35, 0.25, and 0.15. The benchmark used one warmup per version and five measured runs per version. The measured runs alternated between the baseline and optimized commits, and the model files were already cached.

Measurement	Baseline median	Optimized median	Change
Calls to state_dict()	219	0	Removed 219 calls
Buffer discovery	456.946 ms	1.564 ms	455.382 ms faster, 99.66%
Token2Wav weight loading	1099.526 ms	620.436 ms	479.090 ms faster, 43.57%
Startup to /health	105.738 s	104.322 s	1.416 s faster, 1.34%

The startup ranges overlapped, so the PR claims the direct reduction in Token2Wav loading work rather than a stable full startup improvement.

The logged GPU memory figures for stages 0 and 1 were unchanged. Stage 0 used 9.94 GiB for weights and non-PyTorch memory with 2.49 GiB of peak activation memory. Stage 1 used 1.95 GiB with 1.18 GiB of peak activation memory. No GPU memory improvement is claimed.

Test

2 passed, 14 warnings in 0.13s

Command:

.venv/bin/pytest tests/model_executor/models/qwen2_5_omni/test_weight_loading.py --run-level advanced_model -q

pre-commit run --files passed for both changed files.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to be related to model: qwen-omni.

Model owners: @amy-why-3459 @psv666

@akshatvishu, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[akshatvishu]

Self-review: The change is limited to persistent buffer discovery in Token2Wav and replaces 219 repeated state_dict() traversals with the _non_persistent_buffers_set pattern already used by vLLM. The regression test covers root and nested persistent buffers, excludes a nonpersistent buffer, and verifies that state_dict() is never called; both targeted tests and all applicable local pre-commit hooks passed. Five alternating ROCm runs, after one warmup per version, reduced median buffer discovery from 456.946 ms to 1.564 ms and Token2Wav weight loading from 1099.526 ms to 620.436 ms, while the server loaded the same 721 weights and found the same persistent buffer. The full startup ranges overlap, so the PR does not claim a stable end-to-end startup improvement; output quality was not compared because the change only alters how the same persistent buffer names are selected during loading.

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@akshatvishu this PR reads as a performance or value claim:

• title: [Core] Optimize Qwen2.5-Omni Token2Wav buffer loading

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

[akshatvishu]

@vllm-omni-review-bot All the claims, reason for bottleneck as well as A/B testing is already provided in the pr body itself

[linyueqian]

Correct and cheap. state_dict() excludes exactly the buffers whose names sit in their module's _non_persistent_buffers_set, so walking named_modules() once to collect those qualified names and filtering named_buffers() against the set reproduces the old membership test without rebuilding the state dict per buffer; it is the same pattern vLLM's own loaders and this repo's diffusion weight loaders already use. The test pins persistent root and nested buffers included, a non-persistent one excluded, and zero state_dict() calls. The A/B numbers in the description are from ROCm on an older vLLM, which is fine for a startup-time change of this shape.

Static read at 40ccac04; fork head, no PR code executed. pre-commit and DCO are green; the PR has no ready label so no general lane has run, and it needs one before merge.