[shunwendongan]

[Bugfix] Preserve generic diffusion deploy config with FLUX runtime coverage

Follow-up to #8092. Head: a7821c4cf6d36958dbd2525e3c8f49dc0da78c73; base: 8da036a81e5a5680e6ff0b77cb759e45da9b4194.

Contributor intent (required human answers)

The following public-safe summaries reflect the human contributor's explicitly confirmed intent for this follow-up. The public technical reference is the maintainer's acceptance-test recommendation on #8092.

Scope against the #8092 follow-up recommendation

This PR implements the diffusion startup/request part of the recommendation and the associated configuration regressions. The AR MoE deploy/prefill/decode run and the diffusion-MoE real-forward run remain separate follow-up work: tiny FLUX is dense, so this PR does not claim MoE coverage.

The configuration portions are covered by regression tests for resolved diffusion attention selection, backend-specific CLI-over-YAML precedence, different backend values per stage, and equivalent canonical spellings such as TORCH/torch without a false engine_extras conflict warning.

1. What outcome do you want to achieve?

Complete the diffusion startup/request follow-up suggested when #8092 was merged, as an independent contribution.

2. What is the highest-leverage variable for that outcome?

Runtime acceptance coverage: extend the existing backend configuration/component tests to a deployed worker that actually serves requests.

3. How will this work move that variable?

Add small native FLUX startup/request tests for both mp and uni, and fix the deploy YAML configuration loss exposed while preparing that validation.

4. Whose customer or user needs does this address, and what are those needs?

vLLM-Omni maintainers and users need evidence that requested diffusion backends reach the worker and execute, rather than only surviving configuration parsing. This contribution supplies that follow-up validation in a small test workload.

Purpose

Follow up on #8092 with deployment-to-request coverage for dense tiny FLUX.

Before this change, generic diffusion fallback did not pass the supplied deploy YAML into its typed builder. A YAML-only regression requesting linear_backend=torch, SDPA, eager mode, and uni instead observed auto, no explicit attention selection, False, and no explicit executor. Supplying those values again through kwargs passed, which isolated the missing deploy path.

The resolver now supplies that path to the generic typed builder. The builder merges the complete deploy config through VllmOmniConfig.from_pipeline_config before deriving generic defaults; terminal dtype/cache/parallel defaults are not promoted to explicit CLI overrides. YAML/platform/caller precedence, model output metadata, and generic parallel/device/sampling/extras behavior are regression-tested. The existing no-YAML path is retained. Engine aliases enter the merge only in canonical stage form, so quantization and auxiliary encoder overrides do not reappear as conflicting raw global kwargs.

The new test_backend_runtime.py uses the shared async runner lifecycle and the existing tiny FLUX builder. A worker extension observes the real selector, Torch GEMM/SDPA calls, and finite transformer outputs while delegating computation to the original implementations. Backend, dtype, eager, and executor settings come only from deploy YAML. Each executor handles two sequential requests with counters reset before each request. PID assertions distinguish the multiprocessing worker from this single-stage, single-replica inline uni worker.

The existing synthetic FLUX backend component test now passes its local temporary directory as the model path instead of the fake model ID unused. This keeps the component smoke offline and allows it to reach the real backend computation; no production model-path behavior is changed.

The configuration module design page is a deferred draft and identifies current code/tests as authoritative. This fixes an existing deployment input path without adding fields or changing the public support matrix. The attention selection contract and user guides are unchanged because selection policy and kernels are unchanged.

Test Plan

vLLM Version: 0.31.0

vLLM-Omni Commit: a7821c4cf6d36958dbd2525e3c8f49dc0da78c73, based on 8da036a81e5a5680e6ff0b77cb759e45da9b4194. Validation ran on the corresponding pre-commit working tree; all 13 outgoing file hashes were verified unchanged before committing.

Environment: WSL2 Ubuntu 24.04.4, Python 3.12.3, NVIDIA RTX 3080 10 GiB (sm_86), driver 616.56, torch 2.13.0+cu130 / CUDA runtime 13.0, Triton 3.7.1, transformers 5.14.1, pytest 9.1.1.

Runtime workload: generated random tiny FLUX weights, default/native weight loading, bfloat16, eager, batch 1, 128×128, 2 steps, guidance 0, maximum sequence length 64, seed 42. Configuration/tokenizer inputs were staged from the official FLUX.1-schnell revision 741f7c3ce8b383c54771c7003378a50191e9efe9; model processes ran with HF_HUB_OFFLINE=1 and TRANSFORMERS_OFFLINE=1. Full FLUX weights are not required.

python -m pytest tests/model_tests/diffusion/test_backend_runtime.py \
  -vv -s --run-level core_model -n 0

python -m pytest tests/diffusion/models/flux/test_backend_selection.py \
  -vv -s --tb=short --run-level core_model -n 0

python -m pytest -q --tb=short \
  tests/config/test_backend_selection.py \
  tests/config/test_config_factory.py \
  tests/entrypoints/test_utils.py \
  tests/helpers/tests/test_runtime_async_omni.py \
  tests/helpers/tests/test_async_diffusion_client.py \
  tests/entrypoints/test_async_omni_diffusion_config.py \
  -m 'core_model and cpu' -n 0 \
  --deselect tests/entrypoints/test_async_omni_diffusion_config.py::test_legacy_default_stage_build_accepts_engine_adapter_metadata \
  --deselect 'tests/entrypoints/test_async_omni_diffusion_config.py::test_default_diffusion_factory_preserves_engine_quantization[typed]'

Test Result

• Original YAML regression before the fix: 1 failed, 55 passed; the explicit-kwargs control passed. Expected/observed values are listed above.
• CPU suite on the updated base and final candidate code: 711 passed, 2 deselected.
• Native runtime on the updated base: 2 passed, four real requests total (mp and uni, two each). Each request executed all 29 tracked unquantized transformer linear layers, all 4 tracked SDPA attention layers, and 2 transformer forwards with finite [1, 64, 64] outputs; completed image outputs had the requested count and dimensions.
• Negative propagation control: expected failure, 1 failed / 1 deselected. In an isolated worktree, forcing the worker KernelConfig.linear_backend to auto after successful startup produced assert 'auto' == 'torch' at the intended worker-backend assertion. The outgoing diff has no mutation. After restoring and verifying the editable candidate source, the final fresh-process rerun was 2 passed, with four successful real requests again.
• Existing native tiny FLUX baseline on the final updated candidate: 1 passed, 3 subtests passed (text-to-image, determinism, multi-output).
• Existing synthetic FLUX backend component smoke after the local model-path correction: 1 passed, with 15 real default_unquantized_gemm linear layers, 2 SDPAImpl attention layers, and finite (1, 4, 64) output. The original fake-model failure was reproduced against the original base before changing the test input.
• mypy-3.10 via the repository wrapper in an isolated Linux environment with exactly the declared hook dependencies: passed, 2 production files + 11 test files.
• Ruff 0.14.10 check/format, typos, whitespace/debug/SPDX, test marks, forbidden imports and new torch.cuda call checks: passed on all 13 candidate files. These hooks ran on Windows; mypy ran separately in Linux because the pre-existing os.getuid reference is not a Windows API.

The two explicitly deselected tests use uncached fake model IDs (unused and generic-diffusion) and fail with LocalEntryNotFoundError offline. Both failures were separately reproduced with the original HEAD factory/resolver loaded in memory. They were not changed or counted as passes.

The existing CUDA ready-pipeline sweep includes the new module's core_model/cuda marks and path; no CI YAML change is needed. These results are local validation, not an official Buildkite/GitHub CI result. The full project CUDA ready sweep was not run locally. AR MoE, diffusion MoE, alternative hardware, image quality, and performance are not validated by this dense tiny test.

Review and AI assistance

OpenAI Codex assisted with source inspection, implementation, local test execution, and the PR description. The contribution follows the human contributor's confirmed #8092 follow-up intent. The human contributor reviewed and approved the final diff and public description and authorized submission; agent reviews are not presented as human review.

• Local affected-code, examples-policy, and simplification review performed.
• Human contributor has reviewed the final diff and public description.
• Human contributor authorized the signed-off commit, push, and PR submission.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-18T01:41:56.586487Z	a7821c4	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[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 belong to: docs/design/module/vllm_omni_config.md, docs/design/module/stage_runtime.md, docs/design/module/entrypoints.md, docs/design/module/diffusion/index.md.

Module owners: @alex-jw-brooks @lishunyang12 @david6666666

Routing: @alex-jw-brooks via module of the changed files, module named in the PR description, CODEOWNERS; @lishunyang12 via module of the changed files, module named in the PR description, CODEOWNERS; @david6666666 via module named in the PR description

@shunwendongan, 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.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 4e067548b980 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[BLANKETusers]

Supersedes my earlier review on this PR, which was posted in Chinese. GitHub does not allow editing or deleting a submitted review body, so the findings are re-stated here in English with P0/P1/P2 severities. The findings are unchanged.

Independent review (frozen head 4e067548b980)

Scope: vllm_omni/config/config_factory.py, vllm_omni/config/resolver.py, plus the added/modified tests and CI wiring. Snapshot: head 4e067548b9805b2f1849c63b308a8cdd0da3d8d7, mergeable=CONFLICTING, all checks green (build 3.11/3.12, pre-commit, CodeQL, DCO, Analyze, readthedocs).

---

P0 — Validation evidence is not bound to the current head, and the config ownership audit was not redone after the rebase

• Location: PR description "Test Plan" / vllm_omni/config/config_factory.py:793-912
• Trigger path: The description states Head: a7821c4cf6d36958dbd2525e3c8f49dc0da78c73 and base: 8da036a81e5a5680e6ff0b77cb759e45da9b4194, and claims "all 13 outgoing file hashes were verified unchanged before committing". The actual PR head is now 4e067548b9805b2f1849c63b308a8cdd0da3d8d7 and contains two Merge branch 'main' commits; GitHub reports mergeable=CONFLICTING.
• Current behavior: The runtime acceptance runs in the Test Plan (tests/model_tests/diffusion/test_backend_runtime.py FLUX mp/uni, tests/config/test_backend_selection.py, tests/diffusion/models/flux/test_backend_selection.py) were all produced from the pre-rebase working tree, not from the reviewed revision.
• Impact: This PR changes exactly the deploy → typed projection field ownership and precedence contract. If either upstream main merge added dataclass fields, default factory parameters, or pipeline-wide fields, the auto-merge result must be treated as a new semantic diff and the field sets plus final per-stage values re-compared — otherwise the conclusion "the deploy YAML is no longer dropped" may hold only on the old base.
• Minimal fix: Update the description's head/base to the frozen revision, re-run the Test Plan's three commands on the current base and report the selector used with the output, and re-walk the config ownership audit for _create_typed_diffusion_from_deploy after the rebase.

---

P1 — The deploy dtype materialization duplicates an existing fallback and pushes dtype into the highest-precedence opaque engine_extras layer

• Location: vllm_omni/config/config_factory.py:811-812

  if deploy.dtype is not None:
      stage_deploy.engine_extras.setdefault("dtype", deploy.dtype)

• Trigger path: resolve_omni_config → _create_typed_diffusion_from_deploy writes the YAML top-level dtype into stage 0's engine_extras before VllmOmniConfig.from_pipeline_config runs _apply_platform_overrides.
• Current behavior: In the legacy builder, engine_extras is merged last and therefore wins over everything else (stage_config.py:1258-1269: pipeline-wide fields → per-stage fields → engine_args.update(ds.engine_extras)). The typed projection already has an equivalent fallback (omni_config.py:2273-2274: if "dtype" not in diffusion_kwargs and deploy.dtype is not None), so these lines are redundant for the final value. The comment's claim that "platform and explicit caller inputs still overlay it normally" holds only because _apply_platform_overrides happens to write the same engine_extras["dtype"] key (stage_config.py:1138-1139) — not because the layer precedence actually works that way.
• Impact: A pipeline-wide field is promoted to a stage-owned opaque extra that outranks every other deploy layer. If a platform overlay ever expresses the same setting through a first-class field, or a platform stage carries no explicit dtype, the materialized value silently defeats the platform value — while the comment implies it cannot.
• Minimal fix: Drop these two lines and rely on the existing omni_config.py:2273-2274 fallback; if the materialization must stay, materialize after the platform overlay is merged, or have the platform branch overwrite the same key explicitly with an assertion.

---

P2 — Probe teardown lacks the hasattr guard, so the test's finally assertion masks the real failure

• Location: tests/helpers/diffusion_backend_probe.py:194-205
• Trigger path: DiffusionBackendProbe.stop_backend_runtime_probe() iterates self._probe_linear_originals without a guard, while shutdown() in the same file (lines 207-211) guards the identical call with hasattr.
• Current behavior: If load_model fails early (probe attributes never installed), test_backend_runtime.py:146-148 runs finally: assert await engine.collective_rpc("stop_backend_runtime_probe", ...), which raises AttributeError and replaces the real model-loading exception.
• Impact: A failure in this test surfaces a misleading stack trace in CI, increasing triage cost — and it is the exact anti-pattern this repo's review guidance forbids (assertions must bind to real behavior).
• Minimal fix: Add the same hasattr guard used by shutdown() inside stop_backend_runtime_probe() (or make the finally block clean up only, without asserting).

---

Verified as sound

• Regression coverage: tests/config/test_backend_selection.py:124-163 parametrizes over the deploy-yaml case (no CLI kwargs at all) and asserts each of the four previously-wrong values (linear_backend, attention, enforce_eager, distributed_executor_backend), keeping an explicit-kwargs-control comparison. This satisfies the requirement that a bugfix pins the original failure path.
• CI wiring is complete: the new runtime test carries core_model / diffusion / cuda/L4/cards_1 marks and is selected by -m 'core_model and cuda' at .buildkite/cuda/test-ready.yml:97; the new config regressions carry core_model + cpu and are selected by the tests/ step at .buildkite/cuda/test-ready.yml:34; diffusion_tiny_model's source-file dependencies cover tests/model_tests/diffusion/. No guard step is declared but never run.
• The probe observes rather than substitutes: it wraps and still delegates to the original _gemm_impl / attention forward, is revertible, and its assertions bind to real call counts plus transformer output shape and finiteness.
• No duplicate findings: this PR has no inline review comments, and the Codex review ended failed on usage limits without producing findings.