[Config] Remove invalid LLM-only engine_args from diffusion stage configs

[ianliuy]

Purpose

Fix for: #2563

Remove dead engine_args fields from diffusion stage configs (stage_type: diffusion). These fields were copy-pasted from LLM stage configs and are silently dropped by OmniDiffusionConfig.from_kwargs().

Note: Generation/AR stages (worker_type: generation, worker_type: ar) use OmniEngineArgs where these fields are actively consumed they are intentionally left unchanged.

Changes

Diffusion config cleanup (11 YAMLs, 51 lines)

Main configs (9 files in vllm_omni/model_executor/stage_configs/):

File	Fields removed
bagel.yaml	gpu_memory_utilization, engine_output_type, enable_prefix_caching, max_num_batched_tokens, tensor_parallel_size
bagel_multiconnector.yaml	same 5
bagel_think.yaml	same 5
bagel_single_stage.yaml	same 5
bagel_usp2.yaml	same 5
hunyuan_image_3_moe.yaml	4 (no top-level tensor_parallel_size)
hunyuan_image3_moe_dit.yaml	4
hunyuan_image3_moe_dit_2gpu_fp8.yaml	4
omnivoice.yaml	2 (gpu_memory_utilization, engine_output_type)

Test configs (2 files in tests/e2e/offline_inference/stage_configs/):

File	Fields removed
bagel_mooncake_ci.yaml	same 5 + load_format (OmniDiffusionConfig uses diffusion_load_format)
bagel_sharedmemory_ci.yaml	same 5 + load_format

Regression test (new file)

tests/test_diffusion_config_fields.py scans all YAML configs (main + test dirs) and asserts diffusion stage engine_args only contain valid OmniDiffusionConfig fields.

Allowlisted fields consumed outside the dataclass:

• model_stage stage init layer
• model_arch diffusion model class resolution
• quantization mapped to quantization_config via backwards-compat in from_kwargs()

Why these fields are safe to remove

OmniDiffusionConfig.from_kwargs() explicitly filters unknown fields:

\\python
valid_fields = {f.name for f in fields(cls)}
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
return cls(**filtered_kwargs)
\\

No behavioral change all removed fields were already silently dropped at runtime.

Test Plan

• Regression test passes
• Existing CI tests pass (no behavioral change)

[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.

[hsliuustc0106]

let me help run the ci tests

[hsliuustc0106]

I think you also need to make changes to the yamls under tests

[ianliuy]

Thanks for running the CI and the feedback @hsliuustc0106!

Updated here's what changed:

1. Fixed the CI failure (regression test)

The test was flagging model_arch and quantization as invalid. Both are actually consumed:

• quantization mapped to quantization_config via backwards-compat in from_kwargs() (data.py L679-682)
• model_arch consumed by the stage init layer for model class resolution

Added both to the test's allowlist.

2. Cleaned test YAMLs (14 files, 65 lines)

Removed the same dead fields from diffusion/generation stages in:

Directory	Files
tests/e2e/stage_configs/	dynin_omni_ci, mimo_audio_ci, qwen2_5_omni_ci, qwen3_omni_ci
tests/e2e/stage_configs/rocm/	qwen2_5_omni_ci, qwen3_omni_ci
tests/e2e/stage_configs/xpu/	qwen2_5_omni_ci, qwen3_omni_ci
tests/e2e/offline_inference/stage_configs/	bagel_mooncake_ci, bagel_sharedmemory_ci
tests/e2e/offline_inference/stage_configs/npu/	qwen2_5_omni_ci
tests/dfx/perf/stage_configs/	qwen3_omni, qwen3_tts
tests/dfx/stability/stage_configs/	qwen3_omni

3. Rebased on latest main

Total: 24 files changed, +56 104.

[hsliuustc0106]

please check whether the ci failure is related to this PR

[ianliuy]

Pushed another update narrowed the scope after investigating more carefully:

Only diffusion-stage configs are cleaned. The test YAMLs with worker_type: generation (qwen3, qwen2.5, dynin, mimo, etc.) use OmniEngineArgs where these fields are actively consumed, so those are left unchanged.

Changes in this push:

• Removed dead fields from 2 diffusion test YAMLs (bagel_mooncake_ci, bagel_sharedmemory_ci)
• Also removed load_format: dummy from their diffusion stages (also dead OmniDiffusionConfig uses diffusion_load_format instead)
• Extended the regression test to scan tests/**/*.yaml recursively, still only flagging stage_type: diffusion stages

Total: 12 files, +68 51.

[princepride]

Why we remove tensor_parallel_size? Seems someone just edit tensor_parallel_size place, right?

[ianliuy]

Good observation! Yes tensor_parallel_size was effectively "moved" to parallel_config.tensor_parallel_size when DiffusionParallelConfig was introduced in PR #189. Since then, OmniDiffusionConfig no longer has a top-level tensor_parallel_size field, so from_kwargs() silently drops it (data.py L692-694). The value here is also 1, which matches the DiffusionParallelConfig default. Issue #2635 also tracks this inconsistency.

[princepride]

Same here

[ianliuy]

Same reasoning as above.

[princepride]

If future users add new special fields, how should we maintain this UT, considering these new fields might only apply to a specific model?

[ianliuy]

Good question. Two cases:

• If the new field is added to the OmniDiffusionConfig dataclass, the test picks it up automatically via fields(OmniDiffusionConfig) no changes needed.
• If it's consumed outside the dataclass (like model_stage by the stage init layer, or quantization via from_kwargs() backwards-compat), add it to the allowlist in this test.

I'll add a comment in the test to document this but before I do, I'd love to hear your thoughts on whether this approach works for you.

[princepride]

I think we should maintain it, correct me if I am wrong

[ianliuy]

These are safe to remove this is a stage_type: diffusion stage, and OmniDiffusionConfig has neither field. from_kwargs() silently drops them (data.py L692-694). For engine_output_type specifically, extract_stage_metadata() also hardcodes it to None for all diffusion stages (stage_init_utils.py L171). Audio routing is handled by the stage-level final_output_type: audio (which is preserved) and the SupportAudioOutput interface on OmniVoicePipeline.

[lishunyang12]

LGTM

[lishunyang12]

LGTM, verified the dead-field claims against data.py / stage_init_utils.py. cc @princepride for re-review.

[lishunyang12]

nit: worth a short block comment above the allowlist explaining the maintenance policy — entries here are fields consumed outside OmniDiffusionConfig (e.g. by extract_stage_metadata or the backwards-compat path in from_kwargs). Saves the next maintainer a git blame.