[Bugfix] Give LLM stages typed owners for vLLM LoRA engine args

[MinhaoLi0318]

Purpose

vLLM LoRA cannot be enabled on any AR/generation stage on main. This is a regression since #5140 (first released in v0.29.0rc1), which made the typed VllmOmniConfig resolver the runtime path.

What broke

Every input path for vLLM's LoRA engine args (enable_lora, max_lora_rank, max_loras, max_cpu_loras, lora_target_modules, fully_sharded_loras, default_mm_loras, ...) is rejected during config resolution: the top-level Omni/AsyncOmni kwarg, stage_overrides, stage_N_* keys, deploy YAML stage keys / engine_extras, and the serve CLI.

ValueError: Stage 0 (llm_ar) has explicit engine argument(s) with no structured config owner: enable_lora, max_lora_rank

Because a stage cannot enable LoRA, a lora_request passed to AsyncOmni.generate (forwarded since #5374) fails in vLLM with LoRA is not enabled!. On v0.28.0 the legacy stage-config path passed these keys through, so existing deploy YAMLs that set enable_lora / max_lora_rank on stage 0 (for example in downstream projects such as UniRL, which pins v0.28.0) are rejected on main.

Repro (CPU, no model weights)

python - <<'EOF'
from vllm_omni.config import VllmOmniConfig
from vllm_omni.config.pipeline_registry import resolve_pipeline_config

pipeline = resolve_pipeline_config("qwen3_tts", None)
cfg = VllmOmniConfig.from_pipeline_config(
    pipeline, cli_overrides={"stage_0_enable_lora": True, "stage_0_max_lora_rank": 64}
)
print(cfg.stage_by_id(0).lora_config)
EOF

• main (c2cd13c): raises the ValueError above.
• this PR: prints OmniStageLoRAConfig(max_lora_rank=64, max_loras=1, ...).

Root cause

These fields live in vllm.config.LoRAConfig, not ModelConfig, and no structured stage config owns them, so _validate_stage_engine_override_ownership and _validate_global_stage_cli_ownership reject them. #7390 gave typed owners to other CLI flags and left enable_lora / max_loras for a follow-up because they need their own owner. The question of how to handle them is still open in #7352. This PR is that follow-up. It is also needed for #8196 (lora_request on Omni.generate).

Changes

• OmniStageLoRAConfig inherits vllm.config.LoRAConfig, like the existing load/cache/scheduler/parallel stage configs. BaseVllmOmniStageConfig gets an optional lora_config.
• The owned fields are enable_lora plus a field map generated from LoRAConfig with _upstream_engine_field_map, aliasing target_modules to lora_target_modules. They are added to _LLM_STAGE_ENGINE_FIELDS only.
• As in vLLM, a stage carries lora_config only when enable_lora is set; other LoRA options are ignored otherwise. The construction validator only checks max_cpu_loras >= max_loras (non-positive treated as unset); platform-dependent checks stay in the engine process. Like the other inherited stage configs, it stays msgpack transport-safe (covered by the existing transport test).
• _project_omni_stage_engine_args (vllm_omni/engine/stage_init_utils.py) projects enable_lora=True and the explicitly set LoRA fields back to EngineArgs.
• Global vs stage-scoped input:
  • Single-stage pipeline: top-level LoRA args apply to that stage.
  • Multi-stage pipeline: top-level vLLM LoRA args are rejected with a message that names the LLM stage ids and suggests --stage-overrides. Broadcasting would also enable LoRA on talker/codec stages that do not support it, while request LoRA only reaches stage 0.
  • Pipeline with a diffusion stage (e.g. bagel): a global max_cpu_loras keeps its existing diffusion meaning and is not applied to LLM stages. An LLM stage sets its value with a stage-scoped key.
• Diffusion stages are unchanged: they still reject enable_lora and keep lora_path / lora_scale / lora_backend / max_cpu_loras on diffusion_config.
• Docs: a "vLLM LoRA on LLM stages" section in docs/configuration/stage_configs.md, the owner/alias notes in docs/design/module/vllm_omni_config.md, and a LoRA design page, docs/design/feature/lora.md (supported scope and boundaries, per-stage LoRAConfig, global vs. stage-scoped options, which stage a request adapter reaches, what is not supported), linked from the docs nav and both pages.
• vllm_omni/config/__init__.py exports OmniStageLoRAConfig; the SPDX header was added by the check-spdx-header pre-commit hook because the file had none.

Usage

omni = Omni(
    model,
    stage_overrides={"0": {"enable_lora": True, "max_lora_rank": 64, "max_loras": 2}},
)

vllm serve <model> --omni --stage-overrides '{"0": {"enable_lora": true, "max_lora_rank": 64}}'

stages:
  - stage_id: 0
    enable_lora: true
    max_lora_rank: 64

Out of scope

• Routing a request's LoRA to stages other than stage 0.
• The legacy stage-config runtime paths (being removed in [Core] Remove legacy stage config runtime paths #8077).
• Diffusion LoRA behavior.
• AsyncOmni.add_lora / remove_lora still send the RPC to every stage. With LoRA enabled on one stage only, targeting those calls at specific stages is a separate change.

Notes for merging

• [Core] Add AR and diffusion backend config propagation #8092 also edits the _ModelEngineOverrides / stage config region in omni_config.py; expect a small textual conflict.
• test_stage_cli_field_selection_defers_ownership_validation_until_sources_are_merged in tests/config/test_omni_config.py conflicts textually with [Core] Remove legacy stage config runtime paths #8077. When resolving, keep the cpu_offload_gb example; enable_lora is now owned and no longer an example of an unowned field. The new engine test only uses _project_omni_stage_engine_args, which [Core] Remove legacy stage config runtime paths #8077 keeps.

AI assistance: I used Claude Code to help investigate the regression, draft the code, tests and docs, review the diff, and draft this description. I reviewed every change and ran the checks below myself.

Test Plan

vLLM Version: 0.31.0 (source on PYTHONPATH, CPU, VLLM_TARGET_DEVICE=cpu), torch 2.13 CPU, macOS. No GPU was available.

vLLM-Omni Commit: based on c2cd13c

New tests:

• tests/config/test_omni_config.py: stage-scoped LoRA from stage CLI keys, deploy engine_extras and deploy YAML on qwen3_tts and bagel; global args on a single-stage pipeline; global args rejected on multi-stage pipelines; global max_cpu_loras on a mixed AR + DiT pipeline; options ignored without enable_lora; vLLM LoRA rejected on a diffusion stage; the validator does not resolve deferred fields. Existing export / public-field / keyword-only / msgpack tests now include OmniStageLoRAConfig.
• tests/engine/test_stage_engine_args.py: test_stage_lora_reaches_projected_engine_args_only_on_owning_stage, plus projection test updates.

Commands:

# targeted
pytest -p no:cacheprovider -n 4 tests/config tests/engine/test_stage_engine_args.py \
  tests/entrypoints/test_async_omni_diffusion_config.py tests/diffusion/test_diffusion_config_propagation.py

# broader CPU suite, run on this branch and on unmodified c2cd13c00
pytest -p no:cacheprovider -q -n 4 -m 'core_model and cpu' \
  --ignore=tests/engine/test_async_omni_engine_abort.py tests/config tests/engine tests/entrypoints

# lint / pre-commit hooks on changed files
ruff check <changed .py files>          # ruff 0.14.10
ruff format --check <changed .py files>
typos <changed files>
markdownlint-cli2 docs/configuration/stage_configs.md docs/design/module/vllm_omni_config.md docs/design/feature/lora.md
# plus the local check_spdx_header, check_forbidden_imports, check_torch_cuda, check_test_marks hooks

I also ran a local CPU script that goes through prepare_stage_config_inputs -> load_and_resolve_stage_configs -> _project_omni_stage_engine_args -> vLLM 0.31.0 create_engine_config, with a tiny Qwen2 config standing in for the model.

Test Result

• Targeted: 1116 passed, 4 failed, 2 skipped. The 4 failures also fail on unmodified c2cd13c in the same environment: test_injected_aligner_survives_typed_config_and_runtime_projections, test_frontend_log_stats_flag_is_not_an_unowned_stage_argument[True] / [False], test_mps_stays_in_runtime_instead_of_engine_arguments.
• Broader CPU suite: this branch 4745 passed, 13 failed, 6 skipped, 1 xfailed; c2cd13c 4726 passed, 13 failed, 6 skipped, 1 xfailed. The 13 failed test ids are identical on both trees.
• Mutation check: removing the multi-stage rejection makes the 5 rejection cases fail; removing the LoRA projection makes the projection test fail.
• ruff check / format, typos, markdownlint, the mypy-3.10 hook and the local hooks listed above pass (mypy: 0 errors on both this branch and c2cd13c).
• CPU end-to-end script, resolved through vLLM create_engine_config:
  • qwen2_5_omni with stage_overrides={"0": {"enable_lora": True, "max_lora_rank": 64, "max_loras": 2}}: only stage 0 gets LoRAConfig(max_lora_rank=64, max_loras=2, max_cpu_loras=2); stages 1 and 2 have no LoRA config.
  • bagel with stage 0 enable_lora: AR stage 0 gets a LoRAConfig; DiT stage 1 has no enable_lora and keeps max_cpu_loras=1.
  • Single-stage qwen2_5_omni_thinker_only with a global enable_lora works.
  • Global enable_lora / max_lora_rank on qwen2_5_omni and bagel raise the new error; on single-stage diffusion wan2_2_ti2v it is still rejected.
  • Deploy YAML stage keys (max_lora_rank: 32, lora_target_modules: [o_proj]) work; a deploy enable_lora: true overridden to false by stage_overrides gives no LoRA config.

Nothing was run on GPU. I will do the end-to-end GPU check (serving a LoRA adapter on stage 0 and sending a lora_request) together with #8196.

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.

Note

To avoid agent spam, vLLM-Omni limits contributors without write access to 5 open PRs. If this blocks well-intentioned critical work, please contact a committer to request bypass list access.
(anything written below this line will be removed by GitHub Actions)

[vllm-omni-review-bot]

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

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

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

@MinhaoLi0318, 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 routing record

Assigned Strict on zcode (GLM-5.3-Flash) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 24af0c7ede97 produced:

• Priority: high. Prompt maintainer attention is suggested.

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

[MinhaoLi0318]

Verified on a GPU together with #8196 (vLLM 0.31.0, 1x A100 80GB): Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_overrides={"0": {"enable_lora": True, "max_lora_rank": 32, "max_loras": 1}}) builds LoRAConfig(max_lora_rank=32, ...) on stage 0 only, and a rank-32 adapter loads and changes the stage-0 output. Script and output: #8196 (comment).

[MinhaoLi0318]

Self-review (author). What I checked before asking for review:

• Scope: 83 added production lines in 3 files (config/omni_config.py, engine/stage_init_utils.py, and the OmniStageLoRAConfig export plus SPDX header in config/__init__.py); the rest is tests and docs.
• Off by default: a stage gets lora_config only when enable_lora is set, and _project_omni_stage_engine_args writes LoRA fields back to EngineArgs only in that case, so a stage that does not set enable_lora gets the same engine args as before. Diffusion stages select no vLLM LoRA fields, still reject enable_lora, and keep their own max_cpu_loras / lora_* handling.
• Global args: on a multi-stage pipeline, global LoRA args raise an error that names the LLM stage ids. I checked the 68 registered pipelines: none is multi-stage with only diffusion stages, so that list is never empty. A global max_cpu_loras next to a diffusion stage keeps its diffusion meaning.
• Against current main: re-checked today on c55c7f7e6, 34 commits past the base. The branch merges cleanly. On the merge result, the targeted suite from the description gives 1160 passed and 6 failed. The same 6 tests fail on unmodified c55c7f7e6: the 4 listed in the description, plus 2 in tests/entrypoints/test_async_omni_diffusion_config.py.
• GPU: verified together with [Frontend] Forward lora_request through Omni.generate #8196 on 1x A100 (the comment above): only stage 0 gets LoRAConfig(max_lora_rank=32), and a rank-32 adapter loads and changes the stage-0 output.

@lishunyang12 @tzhouam, as config owners, could you take a look when you have time? The main open question is the one from #7352: does a typed LoRAConfig subclass on LLM stages, with global LoRA args rejected on multi-stage pipelines, fit the direction you want for VllmOmniConfig?

AI assistance: drafted with Claude Code; I reviewed the diff and re-ran the checks above before posting.

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: c806bc69-904d-435e-a5d6-5e6e69619dfb) — check zcode login and the CLI version; retrying strict/zcode/GLM-5.3-Flash in 120s (try ).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 6fa3de45-a7dd-4c72-8e08-c53acffe447f) — check zcode login and the CLI version; retrying strict/zcode/GLM-5.3-Flash in 600s (try ).

[vllm-omni-review-bot]

Omni ReviewBot attempt record

Review attempt ended as failed (step 'review' (agent.review_diff): unhandled error: RuntimeError: zcode exited 1 without a result event: statusCode: undefined } Error: Turn execution failed (traceId: 729536bc-fd89-42ac-9898-126f6b74f458) — check zcode login and the CLI version; falling back to direct/cursor/auto).

[vllm-omni-review-bot]

Omni ReviewBot review

0 actionable finding(s).

CI at 1e1a77be86f8 (2026-10-10T06:41:21.431901+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Note: The assigned review arm strict/zcode/GLM-5.3-Flash 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

LLM stages can take vLLM LoRA engine arguments again. OmniStageLoRAConfig owns enable_lora and the LoRAConfig fields that also exist on EngineArgs (target_modules aliased to lora_target_modules). A stage stores that config only when enable_lora is set, and projection writes enable_lora=True plus the explicitly set fields back into that stage's engine args. On a multi-stage pipeline, top-level vLLM LoRA arguments are rejected so they are not broadcast onto talker or codec stages; a global max_cpu_loras next to a diffusion stage still sizes only the diffusion cache. Diffusion stages still reject enable_lora and keep their own LoRA fields.

Change flow

flowchart LR
  A["[EXISTING] CLI, stage overrides, deploy YAML"]:::existing
  B["[CHANGED] Stage ownership filter"]:::changed
  C["[NEW] OmniStageLoRAConfig"]:::new
  D["[CHANGED] EngineArgs projection"]:::changed
  E["[EXISTING] Diffusion LoRA config"]:::existing
  A --> B
  B --> C
  C --> D
  B --> E
  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.

---

_{🤖 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!}

[MinhaoLi0318]

@lishunyang12 Could you review this one when you have time? It fixes LoRA config on LLM stages: on main, enable_lora on an LLM stage fails with no structured config owner, and #8196 (lora_request forwarding) depends on it.

• The design point to check is in vllm_omni/config/omni_config.py: OmniStageLoRAConfig is set only on LLM stages that set enable_lora, and on multi-stage pipelines global LoRA args are rejected instead of being broadcast to talker/codec stages.
• End-to-end check (BAGEL-7B-MoT, 1x A100, run with this PR applied): [Frontend] Forward lora_request through Omni.generate #8196 (comment). GitHub checks pass on 1e1a77b, and it still merges cleanly onto current main; Buildkite needs ready.
• Downstream, verl-omni carries a local workaround (_OmniLoRALoadConfig in utils/vllm_omni/patch.py) that its patch-debt ledger lists as gated on this PR: [RFC][tracking] Upstream patch debt: audited ledger of temporary compatibility patches verl-project/verl-omni#445

Thanks!

AI assistance: drafted with Claude Code; I reviewed it before posting.

[hsliuustc0106]

@Acerak01-fy PTAL

[hsliuustc0106]

can you add or update the lora feature design docs?

[Acerak01-fy]

LGTM.

[Acerak01-fy]

Regarding the revision of the doc, it is recommended to update the supported range and boundaries of LORA.

[MinhaoLi0318]

@hsliuustc0106 @Acerak01-fy Added docs/design/feature/lora.md in 24af0c7. There was no LoRA design page for the LLM stages before, only the diffusion user guide. The page covers:

• supported scope and boundaries for both paths: which stage models can take vLLM LoRA (classes implementing SupportsLoRA), how it is enabled, adapters per request, limits and runtime adapter management;
• how OmniStageLoRAConfig is built per stage and projected to the vLLM engine args;
• where LoRA options may be set (global vs. stage-scoped, and why multi-stage pipelines reject the global form);
• which stage a request-level LoRARequest reaches (the entry stage; offline Omni.generate forwarding is [Frontend] Forward lora_request through Omni.generate #8196);
• what is not supported, including diffusion adapter stacking, which [Diffusion][LoRA] Support stacking multiple LoRA adapters #9058 proposes.

It is linked from vllm_omni_config.md, stage_configs.md and the docs nav.

AI assistance: drafted with Claude Code; I reviewed the doc change and this reply.