[Core] Add AR and diffusion backend config propagation

[shunwendongan]

Purpose

Partial implementation of #7924 Item 1. Refs #7924.

Complete the backend configuration/propagation slice discussed in the issue, without changing dispatch or registration APIs.

• Add first-class AR attention, MoE and linear fields to the deploy schema, and carry linear_backend through the AR and diffusion structured configs.
• Reuse upstream KernelConfig normalization/validation for MoE and linear choices. Preserve omitted/default/explicit auto behavior and existing Qwen3-Omni defaults.
• Pass diffusion kernel choices into the shared VllmConfig builder used by engine and workers, before layers are constructed. Do not let shared model defaults overwrite explicit diffusion choices.
• Give explicit first-class fields precedence over the same engine_extras keys, with a conflict warning. Preserve extras-only inputs and platform overlays, including ROCm attention overrides of legacy extras.

Previously, typed AR configuration rejected linear_backend as having no structured owner, and diffusion lacked the field and final kernel-config wiring. The regression tests follow both typed and legacy paths to the terminal config, not just the input dictionary.

Out of scope: facade/registration/discovery, CUSTOM/name-registry changes, declarative model defaults, FishSpeech, selector/per-role behavior, metadata/cache/CUDA-graph lifecycle, and broader validation or fallback redesign. Diffusion retains its existing attention selector and platform hooks; AR retains upstream attention handling.

Precedence contract

• Backend settings are stage-scoped; pipeline-level backend keys are not broadcast to stages.
• A non-null first-class stage backend wins over the same legacy engine_extras key, including nested platform extras. Different values produce a warning; extras-only inputs remain supported.
• Explicit first-class platform backends replace base legacy extras. The complete overlay is merged before backend normalization/duplicate removal, so nested extras cannot re-enter after the check.
• For typed diffusion stages, an explicit model_config.moe_backend / linear_backend wins over an explicit diffusion value, warning on different normalized values. Omitted model defaults do not overwrite explicit diffusion selections; explicit auto is retained.
• CLI/runtime precedence and upstream backend dispatch/fallback behavior remain unchanged.

Current-head update — 2026-10-14

Head: cd5db54695e69cdf04cd7206d05373bdfd4793e7. Includes upstream main 37f094a388ba1c79f83c129fda39797c1ebfc2a4.

• Resolved the diffusion config merge conflict by preserving both backend normalization and the upstream VAE encode fast-path validation.
• Compare MoE/linear extras through upstream KernelConfig normalization before emitting a conflict warning. Equivalent case/hyphen spellings no longer warn; conflicting or invalid shadowed extras still yield to the first-class field.
• Added the diffusion YAML attention assertion, global and stage-scoped CLI-over-YAML cases, independent backend selections per stage, and equivalent-spelling/platform-overlay regressions across typed and legacy config paths.
• Updated the existing FLUX smoke to use the repository's hardware-mark helper for one CUDA L4 card.

Local verification on Apple M1:

• Dependency-isolated warning regressions: 12 passed. The same checks against the pre-fix code produced 6 failures. This executes extracted production stage definitions and the repository regression function, using the official vLLM 0.31.0 KernelConfig source with a Pydantic dataclass wrapper; it does not import the full vLLM package.
• Existing isolated precedence, immutability, extras-only and explicit-field provenance checks: passed with that same upstream contract.
• Ruff check/format, compileall, git diff --check, SPDX, test marks, forbidden imports and the torch.cuda gate: passed across the PR's Python diff.
• Repository pytest could not collect in the selected local environment because the fixture imports require Torch; vLLM is also absent there. The added YAML/CLI/multi-stage cases need the project's CI environment. Full pytest and GPU/NPU runtime tests were not run for this head.

The AR MoE startup/prefill/decode and diffusion startup/request/real MoE-forward acceptance tests remain follow-ups, as suggested in the review. Previous-head pytest/GPU results below are historical evidence.

AI assistance: Codex implemented the sync and follow-ups, ran the local checks, and drafted this update.

Previous-head validation — 2026-10-10

Validated commit: d6d5ab233b13aa798a38b006211211c3ac4f5c5a, after merging main 874b8b3a.

Environment: Ubuntu/WSL2, Python 3.12.3, vLLM 0.31.0 (matching the updated CI Dockerfile), PyTorch 2.13.0+cu130, transformers 5.14.1, diffusers 0.40.0, pydantic 2.13.5. CUDA smoke ran on one RTX 3080 10 GB, compute capability 8.6; Windows NVIDIA driver 616.56.

The official vLLM 0.31.0 package was installed in a separate target directory and used via PYTHONPATH; these are real pytest runs, not a stub harness. The commands below used the existing Python environment with PYTHONPATH=/home/shunwendongan/.venvs/vllm-omni-pr8092-031.

Commands and results

python -m pytest -q tests/config/test_backend_selection.py -m 'core_model and cpu' --tb=short
# 30 passed, 16 warnings in 8.82s

python -m pytest -q -rs --tb=short \
  tests/config/test_backend_selection.py \
  tests/config/test_omni_config.py \
  tests/config/test_config_factory.py \
  tests/engine/test_stage_engine_args.py \
  tests/diffusion/diffusion_kv/test_initialization.py \
  tests/diffusion/diffusion_kv/test_kv_connector.py \
  tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py::TestRegistry::test_deploy_yaml_uses_native_flashinfer
# 946 passed, 2 skipped, 22 warnings in 159.59s

python -m pytest -s -q -rs --tb=short \
  tests/diffusion/models/flux/test_backend_selection.py -m 'local_model and cuda'
# 1 passed, 15 warnings in 5.16s

The two regression skips are existing Qwen3-Omni cases requiring an HF config. The runs emit upstream deprecation warnings and the checkout's missing generated _version.py warning.

The CUDA component smoke executes 15 real vLLM linear layers, observes torch reaching the real GEMM selector, checks default_unquantized_gemm and two SDPAImpl attention layers, and checks finite output of shape (1, 4, 64). It uses a tiny native FLUX transformer with random weights and synthetic inputs; no checkpoint download or numerical/dispatch implementation replacement.

At d6d5ab23 on 2026-10-10, both precedence review threads were resolved and the PR retained approval. GitHub reported MERGEABLE; Python 3.11/3.12 wheel builds, pre-commit, DCO and Read the Docs passed on that head. Buildkite results were absent. These results apply to that previous head.

These checks validate configuration propagation and a CUDA component execution path, not full AR generation, image generation/quality, performance, or NPU/ROCm runtime support. AI assistance: Codex assisted implementation, review, test execution, and this PR update.

Historical validation on vLLM 0.29.0 (not current-head evidence)

These results predate the latest precedence fixes and main sync. They are retained as historical evidence; current-head results are listed above.

Test Plan

vLLM Version: 0.29.0

vLLM-Omni base: 67aa30c9 (rebased before validation)

Environment: Ubuntu/WSL2, Python 3.12.3, PyTorch 2.13.0+cu130 (CUDA 13.0), transformers 5.14.1, diffusers 0.40.0, pydantic 2.13.5. GPU: one RTX 3080 10 GB, SM86.

Configuration/engine regression, existing defaults, platform overlays, and affected Higgs deploy coverage:

pytest -q -rs --tb=short \
  tests/config/test_backend_selection.py \
  tests/config/test_omni_config.py \
  tests/config/test_config_factory.py \
  tests/engine/test_stage_engine_args.py \
  tests/diffusion/diffusion_kv/test_initialization.py \
  tests/diffusion/diffusion_kv/test_kv_connector.py \
  tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py::TestRegistry::test_deploy_yaml_uses_native_flashinfer

The new config tests use core_model + cpu and are collected by the existing L1 sweep. Quick focused command:

pytest -q tests/config/test_backend_selection.py -m 'core_model and cpu'

Opt-in, offline CUDA component smoke (no checkpoint download):

pytest -s -q tests/diffusion/models/flux/test_backend_selection.py -m 'local_model and cuda'

Test Result

• Configuration/engine regression: 810 passed, 2 skipped. The skips are existing Qwen3-Omni pipeline cases requiring an HF config.
• The focused config tests on the unmodified base: 21 failed, 2 passed; on this branch, the full regression above includes all 23 passing.
• CUDA smoke: 1 passed. Tiny native FLUX uses one joint block and one single block, BF16, seed 42, random weights and synthetic embeddings. All 15 vLLM linear layers execute, with default_unquantized_gemm, and both attention layers select SDPAImpl; output (1, 4, 64) is finite. A spy observes torch reaching the real GEMM selector without replacing its behavior, since torch and auto may resolve identically on Ampere.
• Pre-commit: all applicable hooks pass except mypy-3.10. On Windows, both the unmodified base and this branch produce the same 88 diagnostics in the affected existing files (compared after removing line numbers); no suppressions or allowlist changes were added. git diff --check passes. Full diff-scoped precheck-pr review completed with this baseline warning.

This is not an AR generation, full image-generation, image-quality, or performance benchmark. No new backend or faster-kernel claim is made. NPU/ROCm runtime behavior is not hardware-validated here; upstream linear filtering/fallback remains in effect (including unquantized ROCm dispatch not consulting this flag in vLLM 0.29).

AI assistance: Codex assisted implementation, test execution, code review, and documentation/PR drafting. The results above are actual command outputs, with limitations called out explicitly.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/stage_runtime.md, docs/design/module/diffusion/index.md.

Module owners: @david6666666 @Isotr0py @princepride

Routing: @david6666666 via module of the changed files, module named in the PR description, CODEOWNERS; @Isotr0py via module of the changed files, module named in the PR description, CODEOWNERS; @princepride via module of the changed files, 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 cd5db54695e6 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[shunwendongan]

Self-review: I rechecked the rebased diff against the latest main and verified that the change is limited to Item 1 configuration/projection. I reviewed the typed and legacy AR/diffusion paths, upstream KernelConfig normalization, first-class versus engine_extras precedence, platform overlay compatibility, and the existing Higgs deploy assertions. The focused regression suite passes (819 passed, 2 existing HF-config skips), and the opt-in RTX 3080 tiny FLUX smoke passes with 15 real linear layers and 2 SDPA attention layers. No registration, selector, FishSpeech, model-default, metadata/cache/CUDA-graph, or fallback redesign is included. The local mypy warning is pre-existing and matches the rebased baseline; no suppressions were added.

[hsliuustc0106]

One more [P3], not anchored to a diff line: load_deploy_config only reads a fixed list of top-level keys (stage_config.py:835-861), so linear_backend: / moe_backend: / attention_backend: hoisted to the deploy top level are silently ignored. That's a pre-existing pattern for stage fields, but these are exactly the keys the new doc section introduces, and the surrounding guidance ("pipeline-wide settings ... are declared at the top level and applied to every stage") makes a hoist a natural first attempt for a multi-stage pipeline. Consider either a near-miss warning for these keys at top level or a one-line note in the doc that the backend fields are stage-only.

Reviewed at 84158a8. No correctness blockers found: the KernelConfig reuse (defaults, normalization, validation) matches upstream v0.29.0, both typed and legacy paths were traced to the terminal config, the removed worker-side MoE override has no remaining dependents, create_base_diffusion_vllm_config is the only diffusion VllmConfig construction site, and the promoted attention_backend YAML keys keep flowing through both builders (legacy via the asdict loop, typed via the model-config projection). The three P3s above are precedence-consistency gaps rather than functional breakage.

[vllm-omni-review-bot]

Omni ReviewBot review

PR description

This change promotes AR attention_backend, moe_backend, and linear_backend to first-class deploy fields and carries the shared MoE/linear choices through typed AR and diffusion configs. Those values are normalized with upstream KernelConfig and, on the diffusion path, written into the shared VllmConfig builder before layers are constructed. Explicit first-class fields beat the same engine_extras keys after a warning, while omitted or extras-only inputs and existing Qwen3-Omni / Higgs deploy defaults keep working.

Change flow

flowchart TD
  A["[EXISTING] Deploy YAML, extras, platform overlay"]:::existing
  B["[NEW] StageDeployConfig attention/moe/linear fields"]:::new
  C["[CHANGED] Typed Omni configs plus KernelConfig normalize"]:::changed
  D["[CHANGED] Legacy and typed engine-args projection"]:::changed
  E["[CHANGED] create_base_diffusion_vllm_config kernel wiring"]:::changed
  F["[EXISTING] VllmConfig.kernel_config before layer build"]:::existing
  A --> B --> C --> D --> E --> F
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

No actionable findings.

[hsliuustc0106]

@shunwendongan gentle ping — this high-priority PR has seen no updates for 4 days and review feedback is waiting. Could you share a status or ETA when you get a chance? Thanks!

[shunwendongan]

Thanks for the detailed review. I addressed all three P3 items.

1. Platform overlays using nested engine_extras can no longer reintroduce
   attention_backend, moe_backend, or linear_backend after the first-class
   precedence check. The legacy backend key is removed from engine_extras,
   and typed and legacy regression tests cover the behavior.

2. For diffusion typed stages, when model_config and diffusion_config both
   explicitly provide different moe_backend or linear_backend values,
   model_config takes precedence and a UserWarning is emitted. The explicit
   field provenance is retained in the diffusion projection.

3. The stage configuration documentation now states that backend fields are
   stage-scoped and are not accepted as pipeline-level top-level settings.
   Users must repeat them under each stage.

The focused source checks pass locally. The full pytest command still needs to
be rerun in the project environment with the matching vLLM dependencies.

[shunwendongan]

Addressed all three review items in commit fbacd4a2a493ac14ef2dd5247046886c1ca2afde:

1. Platform overlays now re-apply backend-field precedence after merging nested engine_extras. attention_backend, moe_backend, and linear_backend remain first-class, conflicting legacy values are removed with a warning, and regression coverage includes all three fields, both payload builders, key-order independence, and input immutability.
2. _DiffusionConfigProjection now retains explicit-field provenance. Typed diffusion projection warns when model_config and diffusion_config explicitly provide different normalized MoE/linear backends; model_config remains the winning value. Equivalent normalized values do not warn.
3. stage_configs.md now states that backend fields are stage-scoped and documents typed diffusion precedence.

Validation on Apple M1 (PyTorch 2.12.1, MPS available, CUDA unavailable):

• CPU dependency-isolated review harness: PASS. It uses pydantic 2.11.10 and a pinned vLLM 0.29.0 KernelConfig contract stub; it covers syntax, three-field precedence, key-order independence, input immutability, extras-only behavior, provenance, and the production conflict branch.
• Ruff checks, format checks, Python compileall, and git diff --check: PASS.
• Focused pytest collection could not start because this checkout has no vllm package; no CUDA/NPU/model runtime test was run.

AI assistance: Codex helped draft the implementation, tests, documentation, and this summary. I reviewed the diff and validation results.

[hsliuustc0106]

Both precedence P3s are fixed at this head — _normalize_stage_backend_fields in StageDeployConfig.__post_init__ now normalizes promoted backends and strips the duplicate legacy keys from engine_extras after nested-extras merging, which closes the platform-overlay re-entry path generally (and covers the top-level hoist case by the same mechanism). Explicit-diffusion provenance via _TrackExplicitConfigFields resolves the model-vs-diffusion silent win. Branch conflicts with main — rebase owed before merge. Thanks for the fast turnaround on the ping.

[shunwendongan]

@hsliuustc0106 Synced the branch with main at 874b8b3a via merge commit d6d5ab23. The branch includes that main revision, and the merge completed without conflicts. Both backend precedence fixes, explicit diffusion provenance, and the stage-only configuration documentation are preserved.

Local checks on the updated PR diff passed: Ruff check/format, compileall, git diff --check, and the SPDX, test-mark, forbidden-import, and torch.cuda hook scripts. The isolated CPU configuration check also passed for platform precedence, input immutability, extras-only behavior, and explicit provenance. That check uses the real configuration branches with Pydantic 2 and a vLLM 0.29 KernelConfig stub.

The full pytest suite was not rerun locally: my M1 environment lacks vLLM and has no CUDA/NPU device. The sync includes upstream NPU/XPU vLLM 0.31 alignment (#8588/#8528) and the AMD AITER cache fix (#8527); fresh CI results are needed to confirm their effect on the earlier failures.

AI assistance: Codex assisted with the upstream sync, local checks, and drafting this reply.

[shunwendongan]

@hsliuustc0106 The new head d6d5ab23 has passed GitHub Actions, DCO, and Read the Docs, but it has no Buildkite status contexts yet, so GitHub still reports mergeStateStatus: BLOCKED. The ready label was already present before this head was pushed, and the Buildkite pipeline is label-triggered. Could you please re-apply the ready label (or manually retrigger the Buildkite PR checks) for the current head? I cannot remove/re-add the label with my contributor permissions.

[shunwendongan]

Current-head validation and review follow-up for d6d5ab23 (2026-10-10):

• Ran the real configuration/engine regression with official vLLM 0.31.0, matching the updated CI Dockerfile: 946 passed, 2 skipped in 159.59s. The skips are the existing Qwen3-Omni cases requiring an HF config.
• Focused backend/precedence suite: 30 passed in 8.82s.
• Offline CUDA component smoke on RTX 3080 10 GB (sm_86): 1 passed in 5.16s. All 15 real linear layers ran the requested Torch GEMM path, both attention layers selected SDPA, and output (1, 4, 64) was finite. This is component execution evidence, not serving, image-quality, or performance validation.
• Verified the fixes for both review findings in the current source and resolved the two review threads. The contract is now explicit in the PR description: first-class stage fields beat legacy extras after the platform merge; explicit typed model backends beat conflicting explicit diffusion values with a warning, while model defaults do not overwrite explicit diffusion choices.
• Updated the PR description to separate current-head vLLM 0.31.0 results from historical vLLM 0.29.0 results. Commands and environment are included there.

The current head is mergeable and retains approval. The listed GitHub checks are successful; Buildkite results are still absent on this head. The Oct 9 maintainer-trigger request remains pending.

AI assistance: Codex performed this verification and drafted the update; the test counts above are actual command results.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 120s (try 2 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; retrying strict/cursor/cursor-grok-4.6-high in 600s (try 3 of 3)).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (failed; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot review

Changes since the previous review

• 0 new inline finding(s); 0 finding(s) below.

CI at d6d5ab233b13 (2026-10-10T16:40:01.268490+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Note: The assigned review arm strict/cursor/cursor-grok-4.6-high could not complete this review, so it was produced by the fallback arm direct/cursor/auto. It is excluded from the routing experiment.

Full review analysis

PR description

Stage deploy config now accepts attention_backend, moe_backend, and linear_backend as real stage fields, and both the legacy engine-args builder and the typed stage projection carry MoE and linear choices through to the final config. Upstream KernelConfig normalizes and validates those two values, while a shared precedence step drops the same keys from engine_extras when a first-class value is set, including after a platform overlay merges nested extras. Diffusion no longer patches MoE onto kernel_config after worker setup; the shared VllmConfig builder receives both kernel choices before layers are built, and an explicit typed model value still beats a different explicit diffusion value with a warning.

Change flow

flowchart LR
  yaml["[EXISTING] Stage deploy settings"]:::existing
  deploy["[CHANGED] StageDeployConfig backend fields"]:::changed
  norm["[NEW] First-class extras precedence"]:::new
  builders["[CHANGED] Legacy and typed projection"]:::changed
  kernel["[CHANGED] Diffusion VllmConfig kernel_config"]:::changed
  worker["[CHANGED] Worker drops late MoE write"]:::changed
  yaml --> deploy --> norm --> builders --> kernel --> worker
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[Bounty-hunter]

Suggestions / Follow-ups

• tests/config/test_backend_selection.py:L71 / tests/diffusion/models/flux/test_backend_selection.py:L38 — All three backend dimensions have configuration coverage, but runtime coverage is limited to diffusion linear/SDPA in a directly constructed FLUX component without MoE. For functional acceptance, add a small AR MoE run through deploy/worker startup and prefill/decode, observing actual linear, attention and MoE implementations or documented fallback. Also add a diffusion startup/request smoke; diffusion MoE claims need a real MoE forward. These can be follow-up acceptance tests for this config-only PR.
• tests/config/test_backend_selection.py:L93-L99 — The diffusion YAML branch asserts only linear/MoE values. Add an assertion for the resolved attention selection, plus backend-specific CLI-over-YAML and different-backends-per-stage cases.
• vllm_omni/config/stage_config.py:L49-L50 — Comparing normalized first-class values with raw extras produces a false conflict warning for equivalent inputs such as "TORCH" in both locations. Compare canonical MoE/linear spellings and add a regression case. This is a non-blocking P3.

[Bounty-hunter]

please fix conflicts

[shunwendongan]

@Bounty-hunter @hsliuustc0106 Synced with main 37f094a3 and pushed cd5db546. GitHub now reports MERGEABLE and retains APPROVED.

Addressed the configuration follow-ups:

• MoE/linear extras are compared using canonical upstream KernelConfig values, so equivalent case/hyphen spellings do not produce a conflict warning. Invalid shadowed extras still warn and are discarded.
• Added the diffusion YAML attention assertion, global/stage-scoped CLI-over-YAML tests, independent backend choices per stage, and equivalent-spelling regressions including platform overlays.
• Migrated the existing FLUX smoke to the current hardware-mark helper.

Local M1 verification: 12 isolated warning regressions passed using official vLLM 0.31.0 KernelConfig source and extracted production definitions; the pre-fix code fails 6 of those checks. Existing isolated precedence/provenance checks and all diff-scoped Ruff, compileall, SPDX, test-mark, forbidden-import, torch.cuda, and whitespace checks passed. Repository pytest cannot collect in the selected local environment without the Torch/vLLM dependency stack, so the YAML/CLI/multi-stage suite and full runtime acceptance need project CI. No GPU/NPU runtime test was rerun for this head; the PR description distinguishes the earlier head's results.

The AR MoE startup/prefill/decode and diffusion startup/request/real MoE-forward tests remain follow-up acceptance work, as suggested.

New GitHub Actions have been queued and DCO passed. buildkite/vllm-omni is still absent for this head despite the existing ready label. Please re-apply ready or manually trigger Buildkite for cd5db546; my contributor account has no label-management permission.

AI assistance: Codex implemented the updates, ran the local verification and drafted this reply.

[Bounty-hunter]

LGTM