Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 48bd050974 docs(review): enforce configuration provenance
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-09 22:39:57 +08:00
boatman fc8946fcf5 [Perf][MiniCPM-O-4.5]Add CUDA Graph For HiFTGenerator (#5869)
Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
2026-08-09 22:29:49 +08:00
amd-xiaoyu12 f01e4bdddc [Recipe] Document MiniMax H3 ROCm (gfx950) BF16 serving (#5723)
Signed-off-by: Xiao Yu <xiao.yu@amd.com>
Signed-off-by: amd-xiaoyu12 <xiao.yu@amd.com>
2026-08-09 22:18:24 +08:00
Weiming LiaoandHongsheng Liu 592637015f [BugFix][NPU] MiniMax-H3 RoPE crash: add the missing batch dim before the mindiesd kernel (#5896)
Signed-off-by: Weiming Liao <liaowm5@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-09 14:08:43 +00:00
Yueqian Lin ebe93f5d1e [Refactor][TTS] Derive TTS model detection from adapter metadata (#5682)
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
2026-08-09 21:30:05 +08:00
linzhenpl07andClaude Opus 4.8 8e2e9b6b53 [Feat] Cosmos3 session-memory port for UND text K/V (RFC #4480 Phase 0) (#4657)
Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-08-09 20:53:17 +08:00
Hongsheng Liu 9159cedbd7 docs: update vLLM-Omni architecture overview (#5914)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-09 09:16:18 +08:00
Caspian ChenandSYLAR 593b404539 [Quantization][Qwen2.5-Omni] Support thinker-only ModelOpt NVFP4 W4A4 checkpoints (#5073)
Signed-off-by: Caspian443 <scrisis843@gmail.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
2026-08-09 02:45:57 +08:00
Hongsheng Liu f663d120c7 [Misc] Add vLLM-Omni PR review skill (#5871)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-08 22:18:54 +08:00
Shenglei FuandClaude Opus 5 81b48e83ec [Model] Fail MiniMax H3 encoder load when a weight or fused shard is missing (#5824)
Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-08 13:14:00 +08:00
Hongsheng Liu fd0e92183a [Docs] [templates] Reorganize module design documentation (#5139)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-08 11:39:27 +08:00
SYLAR 44e38d7b38 [Diffusion] Fix DLO DP concurrent request execution (#5864)
Signed-off-by: lishunyang12 <lishunyang12@163.com>
2026-08-08 08:24:42 +08:00
Alex Brooks b15d4e6a64 [Frontend] Implement Batched Chat Completions (#5317)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-08-07 16:12:09 -06:00
NumberWan 18bf7e0ff9 [BugFix][CI] Project fa_deterministic into OmniDiffusionConfig fields (#5897)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-08-07 23:22:01 +08:00
NumberWan cfa6d66ed6 [BugFix][Nightly CI] Opt in FA deterministic for Qwen-Image accuracy (#5887)
Signed-off-by: NumberWan <wantszkin2003@gmail.com>
2026-08-07 17:15:22 +08:00
Mu GuanLin d8c35663b4 [Model][Feat] Support Minimax-H3 quality grading requests through dynamic loading/unloading with Cache-DiT (#5853)
Signed-off-by: mglyn <1203789601@qq.com>
2026-08-07 09:07:23 +00:00
NATURE e906f41dc8 [bugfix][CI] Fix Cache-DiT nested module discovery [issue 5879] (#5884)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-08-07 09:02:42 +00:00
dengyunyang f114113409 [Bugfix] Fix image num check error (#5838)
Signed-off-by: dengyunyang <584797741@qq.com>
2026-08-07 08:50:24 +00:00
RuQing Xu ab459b36c2 [Kernel] Refresh FlashInfer attention; Add quantized attention support (Blackwell QK16/V8) (#5344)
Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com>
2026-08-07 16:14:26 +08:00
linzhenpl07 24ede9e226 [Bugfix] Support GQA/MQA in the ring attention SDPA path (#5255)
Signed-off-by: linzhenpl07 <linzhenpl07@gmail.com>
2026-08-07 04:08:33 +00:00
145 changed files with 8093 additions and 760 deletions

No files matched your search

+4 -2
View File
@@ -47,15 +47,17 @@ steps:
- timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
mirror_hardwares: l4_1
- label: "Diffusion · Distributed Test"
- label: "Diffusion · Distributed & Attention Test"
source_file_dependencies:
- tests/diffusion/distributed/
- tests/diffusion/attention/test_attention_sp.py
- vllm_omni/diffusion/distributed/
- vllm_omni/diffusion/attention/parallel/
- vllm_omni/diffusion/attention/
- vllm_omni/diffusion/data.py
- vllm_omni/diffusion/models/
commands:
- timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model"
- timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda' --run-level "core_model"
mirror_hardwares: l4_4
- label: "Diffusion · Model Test"
+60 -55
View File
@@ -220,65 +220,70 @@ See `plan/voxcpm2_native_ar_design.md`.
### Steps
1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**;
partial integration causes hard-to-debug failures. This file is modified by every
model PR and is the most common source of rebase conflicts — see conflict note below.
1. **Write one adapter** under `vllm_omni/entrypoints/openai/tts_adapters/`.
`serving_speech.py` should not need an edit — detection, stage discovery and
dispatch are all derived from what the adapter declares.
**Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets):
Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`:
```python
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"}
@register_tts_adapter
class YourModelAdapter(ARTTSAdapter):
name = "your_model" # registry key + log label
stage_keys = frozenset({"your_stage_key"}) # the deploy yaml's model_stage
def validate(self, request) -> str | None:
if not request.input or not request.input.strip():
return "Input text cannot be empty"
return None
async def build(self, request, sampling_params_list, has_inline_ref_audio):
params = {"text": [request.input]}
if request.voice is not None:
params["voice"] = [request.voice]
return PreparedRequest(
prompt={"prompt": request.input},
tts_params=params,
model_type=self.name,
)
```
**Point 2** — union into `_TTS_MODEL_STAGES`:
```python
_TTS_MODEL_STAGES: set[str] = (
...
| _YOUR_MODEL_TTS_MODEL_STAGES
)
```
Then add the module to the import block at the bottom of
`tts_adapters/__init__.py` so it registers. That import line is the only shared
file a new model touches — which is also why the old rebase-conflict hotspot is
gone.
**Point 3** — model type detection in `_detect_tts_model_type()`:
```python
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
return "your_model"
```
> **Pure-diffusion TTS does not go through adapters yet.** Under
> `for_diffusion()`, `create_speech()` routes straight to
> `_create_diffusion_speech()` and never calls `validate()`/`build()`.
> `DiffusionTTSAdapter` is scaffolding with no production subclass, so logic
> placed in one would silently never run. Diffusion-engine models follow the
> existing diffusion path; wiring it through adapters is open work (#4855).
**Point 4** — validation dispatch in `_validate_tts_request()`:
```python
if self._tts_model_type == "your_model":
return self._validate_your_model_request(request)
```
**If a stage key alone cannot identify the model**, declare
`model_archs = frozenset({"YourModelForConditionalGeneration"})`; add
`arch_identifies_entry_stage = True` when the model owns no stage key at all
(Ming dense). For a rule that is not set membership, override `matches()` (see
`covo_audio.py`). For a genuine overlap with another adapter, give one an
explicit `detect_priority` — `test_tts_detection.py` fails on an unordered
overlap. Models that only serve speech in some topologies override
`stage_serves_speech()` (see `audex.py`).
**Point 5** — validation + parameter-builder methods:
```python
def _validate_your_model_request(self, request) -> str | None:
if not request.input or not request.input.strip():
return "Input text cannot be empty"
return None
**Reuse shared helpers via `self.ctx.server`** rather than reimplementing them:
`_resolve_ref_audio`, `_apply_uploaded_speaker`, `_validate_ref_audio_format`,
`_max_instructions_length`. Read a comparable adapter first — `fish_speech.py`
(voice cloning), `higgs_audio_v3.py` (parameter-heavy), `moss_tts.py` (family
sharing a base class).
def _build_your_model_params(self, request) -> dict:
params = {"text": [request.input]}
if request.voice is not None:
params["voice"] = [request.voice]
return params
```
Wire `_build_your_model_params` into `_create_tts_request()` alongside the other
model-specific param builders.
> **Do not add `self._tts_model_type == ...` branches to `serving_speech.py`.**
> Older models predate the adapter framework and still have them; they are being
> migrated out (RFC #4327, #4855). `tools/pre_commit/check_tts_adapter.py` is a
> ratchet on the remaining count and fails the commit if it grows. Behaviour that
> no adapter hook can express is a missing hook — propose it on the RFC.
> **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean
> instance attribute checked before `elif self._is_tts`, while all newer models
> (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by
> `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string
> pattern — do not add new `_is_*` flags.
> **Unused variable rule**: only extract fields in `_build_your_model_params` that
> are actually forwarded to the model. Unused extractions fail `ruff F841`.
> For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` →
> `prompt_text`), add them to the param builder and verify they reach the model call.
**Rebase conflict note**: when rebasing onto `main` after another model was merged,
`serving_speech.py` will conflict. Resolution: always keep *both* the upstream
model's additions and your own — never discard either side.
> **Unused variable rule**: only extract fields in `build()` that are actually
> forwarded to the model. Unused extractions fail `ruff F841`. For voice-cloning
> fields (`ref_audio` -> `prompt_audio_path`, `ref_text` -> `prompt_text`), add
> them to the params and verify they reach the model call.
2. **Handle model-specific parameters**:
- Voice cloning: `ref_audio` encoding and prompt injection
@@ -295,7 +300,7 @@ import base64
from pathlib import Path
def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
"""Build prompt with reference audio for voice cloning in serving_speech.py."""
"""Build prompt with reference audio for voice cloning, called from the adapter."""
audio_bytes = Path(ref_audio_path).read_bytes()
codes = codec.encode(audio_bytes) # Encode on CPU using model's codec (e.g., DAC)
token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()]
@@ -321,7 +326,7 @@ Classify the model's **CI priority** first (high / medium / low). High-priority
| Level | Location | Marker | CI pipeline | Notes |
|-------|----------|--------|-------------|-------|
| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, `serving_speech` validation — no GPU |
| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, adapter validation — no GPU |
| **L2** | `tests/e2e/online_serving/test_{slug}.py` | **`core_model` + `advanced_model`** (both on baseline smoke) + `tts` + `@hardware_test(...)` | `test-ready.yml` (`ready` label) | Default deploy smoke: single `/v1/audio/speech` or offline `OmniRunner` path |
| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `tts`) | `test-merge.yml` **or** merged into nightly TTS function job | Streaming, voice clone, batch/queue, async_chunk |
| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py`, optional offline expansion | `full_model` + `tts` | `test-nightly.yml` (`:full_moon: TTS · Function Test with L4`) | Feature matrix; perf → `tests/dfx/perf/tests/test_tts.json` |
@@ -366,7 +371,7 @@ See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers
### Deliverables
- Updated `serving_speech.py` with all 5 integration points (single commit)
- One adapter file under `tts_adapters/` plus its line in the package import block
- Client scripts and server launcher under `examples/online_serving/text_to_speech/<model>/`
- Gradio demo with streaming and voice cloning UI in the same dir
- E2E tests per **Test Case Writing (CI Levels)** above (priority tier determines L1–L4 scope)
@@ -514,8 +519,8 @@ Use this checklist when integrating a new TTS model:
- [ ] README.md written
### Phase 3: Online Serving
- [ ] All 5 `serving_speech.py` integration points added in one commit
- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841)
- [ ] Adapter written under `tts_adapters/` and registered in the import block
- [ ] Only extract params in `build()` that are forwarded to the model call (ruff F841)
- [ ] Prompt builder handles text input correctly
- [ ] Voice cloning works (if supported)
- [ ] All response formats work (wav, mp3, flac, pcm)
+238
View File
@@ -0,0 +1,238 @@
---
name: review-pr
description: Review pull requests and local branches for vllm-project/vllm-omni with a frozen snapshot, module-design ownership, feature-design overlays, targeted validation, and concise evidence-backed findings. Use for default, detailed, or repeat maintainer reviews; checking correctness, compatibility, tests, benchmarks, model additions, distributed changes, or breaking behavior; and identifying or explicitly requesting the most relevant code-owner reviewers. Use precheck-pr instead for an author's pre-submit self-check.
---
# Review vLLM-Omni Pull Requests
Review like a maintainer: direct, selective, and focused on issues that CI does
not prove. Prefer a few high-confidence findings over exhaustive commentary.
Zero findings is a valid result.
## Quality contract
Make every finding:
- **Correct:** prove a reachable failure, not a suspicion.
- **Prioritized:** lead with merge blockers and high-impact defects.
- **Actionable:** identify the smallest safe fix direction.
- **Evidence-based:** cite code, tests, docs, CI, or measurements.
- **Concise:** avoid review templates and repeated summaries.
- **Calibrated:** match severity to user and maintainer impact.
Do not report unrelated backlog, style already enforced by pre-commit, or a
missing test that would not protect changed behavior.
## Select the input and depth
Use `vllm-project/vllm-omni` as the base repository. Accept its forks and local
checkouts; use another skill for unrelated repositories.
| Input | Review surface |
| --- | --- |
| PR number or URL | Frozen PR metadata, full diff, and relevant threads. |
| Local branch/worktree | Frozen target-base SHA through committed, staged, unstaged, and in-scope untracked changes. |
| Pre-filled context | Reuse supplied metadata; fetch only missing facts and the full diff. |
Default to maintainer brevity. A detailed or audit request expands coverage and
lists `path:line` findings, but keeps the same confidence and severity bar.
## Reference guide
Load references in review order: process, one primary module contract, matching
feature designs, evidence checks, then delivery. Every file is linked directly
below; do not load unrelated module references.
Each concise reference links to the maintained
[vLLM-Omni documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/).
For branch-specific behavior, inspect the matching `docs/` file in the reviewed
checkout first; use the published latest docs for current guidance and discovery.
If docs and live code disagree, verify the code/tests and report the drift.
### Review process
| Reference | Read when |
| --- | --- |
| [review-execution.md](references/process/review-execution.md) | Every review; freeze inputs, inspect safely, and deliver against the same snapshot. |
| [general-checks.md](references/process/general-checks.md) | Every review; apply repository-wide correctness and evidence rules. |
| [design-contracts.md](references/process/design-contracts.md) | Every production review; resolve branch-local module and feature design status. |
| [review-routing.md](references/process/review-routing.md) | After the diff census; select one primary module and conditional overlays. |
### Primary module contract
| Reference | Read when |
| --- | --- |
| [entrypoints.md](references/modules/entrypoints.md) | Offline/CLI/API ingress, validation, rendering, streaming, or sessions change. |
| [configuration.md](references/modules/configuration.md) | Config construction, deploy/stage schema, defaults, registry, or topology changes. |
| [input-output-modality.md](references/modules/input-output-modality.md) | Requests, messages, serialization, output types, accumulation, or completion change. |
| [error-contracts.md](references/modules/error-contracts.md) | Error classification, fatality, propagation, sanitization, or rendering changes. |
| [engine-orchestration.md](references/modules/engine-orchestration.md) | Cross-stage routing, request state, output ordering, RPC correlation, or terminal convergence changes. |
| [stage-runtime.md](references/modules/stage-runtime.md) | Placement, startup, readiness, replica identity, affinity, membership, or shutdown changes. |
| [omni-connector.md](references/modules/omni-connector.md) | Cross-stage/process/device/node transport or synchronization changes. |
| [model-integration.md](references/modules/model-integration.md) | Registration, preprocessing, loading, runners, or model-specific execution changes. |
| [ar-runtime.md](references/modules/ar-runtime.md) | AR scheduling, request/cache state, adapters, workers, or upstream vLLM semantics change. |
| [diffusion.md](references/modules/diffusion.md) | Diffusion runtime, models, batching, parallelism, or offload changes. |
| [execution-platforms.md](references/modules/execution-platforms.md) | Hardware selection, capabilities, vendor workers, kernels, or patches change. |
| [cache-management.md](references/modules/cache-management.md) | Cache identity, reuse, validity, reset, eviction, or teardown changes. |
| [quantization.md](references/modules/quantization.md) | Quantization selection, checkpoint metadata, layer mapping, precision, or constraints change. |
| [observability.md](references/modules/observability.md) | Metrics, logs, units, labels, correlation, or lifecycle changes. |
| [profiling.md](references/modules/profiling.md) | Profiling instrumentation, traces, start/stop lifecycle, or overhead changes. |
| [benchmarking.md](references/modules/benchmarking.md) | Benchmark workload, metric calculation, CLI, or result metadata changes. |
### Feature-design overlays
| Reference | Read when |
| --- | --- |
| [runtime-stage-execution.md](references/features/runtime-stage-execution.md) | Disaggregated inference, async chunk/output/materialization, or prefix caching changes. |
| [communication.md](references/features/communication.md) | A concrete OmniConnector backend or its deployment contract changes. |
| [diffusion-acceleration.md](references/features/diffusion-acceleration.md) | Diffusion parallelism, attention, quantization, cache, batching, or offload changes. |
| [infrastructure-performance.md](references/features/infrastructure-performance.md) | Metrics infrastructure or documented speech optimization stacks change. |
### Evidence and quality checks
| Reference | Read when |
| --- | --- |
| [model-addition-checklist.md](references/checks/model-addition-checklist.md) | A model, architecture, loader, processor, registry, or stage config is added. |
| [perf-verification.md](references/checks/perf-verification.md) | The PR makes a latency, throughput, memory, or quality claim. |
| [config-provenance.md](references/checks/config-provenance.md) | Changed or newly load-bearing behavior reads config, CLI, environment, legacy, or default values. |
| [test-quality-evaluation.md](references/checks/test-quality-evaluation.md) | Tests change, are absent for risky code, or may not exercise production behavior. |
| [tests-docs-checklist.md](references/checks/tests-docs-checklist.md) | Coverage, CI markers, examples, user docs, or PR evidence need review. |
| [verification.md](references/checks/verification.md) | Hardware, a server, or a runnable affected path is available for active verification. |
### Delivery and reviewer coordination
| Reference | Read when |
| --- | --- |
| [maintainer-style-study.md](references/delivery/maintainer-style-study.md) | Findings are ready for concise maintainer-style delivery. |
| [review-requests.md](references/delivery/review-requests.md) | The user asks to identify, suggest, request, or ping code-owner reviewers. |
## Workflow
### 1. Freeze and report the snapshot
Pin the base and head before reading source or running validation. Within 60 seconds,
report the pinned head, CI, mergeability, and preliminary findings in the host
conversation. Do not wait for CI or post this update to GitHub.
If the target changes while fetching, discard the evidence and retry once. If
it changes again, report the churn and wait for a stable target.
For a trusted PR head, materialize the pinned head in an isolated detached
worktree. A worktree freezes identity but is not a security sandbox. Treat fork
heads as untrusted unless the user and environment policy explicitly establish
otherwise: execute them only in a disposable, secret-free sandbox with restricted
filesystem, network, and resources; without one, use static SHA-addressed reads
and CI evidence only. For a local review, freeze the committed, index, worktree,
and NUL-safe in-scope untracked contents. Follow
[review-execution.md](references/process/review-execution.md) for trust gates,
state fingerprints, and byte-for-byte staleness checks.
### 2. Build the diff census
Group files into production code, tests, docs, configuration, build/CI, and
generated artifacts. Map each changed production file and test group to the PR
goal. Compare the title/body claims with the actual diff; use linked issues only
when they define the contract or reproduction.
Mark unrelated scope and unexplained generated artifacts. Do not infer behavior
from the PR description without tracing the live code.
### 3. Route from the live behavior
Trace each claimed behavior through the changed producer to its live consumer,
then use [design-contracts.md](references/process/design-contracts.md) and
[review-routing.md](references/process/review-routing.md) to select one primary
module contract, a second only for a real documented cross-boundary call path,
and every matching feature-design and evidence overlay. Treat titles and paths
as hints; live behavior and the frozen head's current design metadata are
authoritative. For docs-, tests-, or CI-only changes, route to the production
contract they protect or use only the applicable evidence checks.
### 4. Run the blocker scan
Apply every category in [general-checks.md](references/process/general-checks.md) before
lower-priority comments.
For each changed or newly load-bearing value or behavior, trace:
```text
public ingress -> validation/defaulting -> producer -> transformations
-> stage/worker/connector boundary -> final consumer -> terminal cleanup
```
Treat unchanged context as in scope when the diff makes it authoritative for a
new consumer. When config, CLI, environment, legacy, or default sources can
control that value, apply
[config-provenance.md](references/checks/config-provenance.md) and prove the
documented source at the final sink with distinct sentinel values.
Cover every applicable offline/online, streaming/non-streaming, sync/async,
feature-on/off, topology, and compatibility path. Search bounded callers and
sibling implementations rather than assuming the changed hunk is the only path.
### 5. Apply module and feature contracts
Apply the reference set selected in step 3 and any matching repo-local skill.
Read the exact module and feature pages in the frozen head, including status,
ownership boundary, dependencies, candidate invariants, safe-change guide, and
promotion gate. Candidate or draft rules are questions, not blockers, unless
current code, tests, or policy enforce them. Inspect both sides of any config,
registry, serialization, connector, cache, or stage boundary.
When a diff adds or expands a helper, class, fallback, compatibility branch, or
public behavior, run a subtraction pass: remove out-of-scope behavior and check
whether each new abstraction can be deleted, merged, moved, or inlined.
### 6. Verify the changed path
Before each validation group, verify the frozen SHA plus the tracked, index,
untracked, and ignored-file fingerprint, or recreate a pristine snapshot. On a
trusted head or inside the required sandbox, run an import/version preflight,
then the narrowest relevant tests and low-cost static checks. Bind every result
to the head SHA, snapshot fingerprint, and environment fingerprint. Never run
imports, tests, builds, hooks, or repo-configurable tooling from an untrusted
head on the reviewer host.
- Treat CI as status evidence; inspect only the first overlapping failure.
- For docs-only changes, use diff hygiene, links/build checks, and bounded live
contract verification instead of dependency setup or pytest.
- For hardware-dependent paths, run available static/CPU checks and name the
exact GPU/NPU gap. Never simulate device evidence.
- For performance or accuracy claims, require comparable base/head runs with
the same environment, workload, warmup, repetitions, and quality criteria.
Stop when each changed semantic path has a supported finding or an explicit
no-issue conclusion. Do not search further only to increase confidence.
### 7. Consolidate and deliver
Verify each finding against the current diff, deduplicate by root cause, and
order by severity.
Re-read the remote head and reverify or recreate the pristine validation
snapshot immediately before delivery. If either changed, mark the review stale
and restart from the new snapshot.
Return findings first. Use
[maintainer-style-study.md](references/delivery/maintainer-style-study.md) to keep them
direct and brief. Each finding must include an exact `path:line`, trigger or
call path, current behavior, impact, and smallest fix direction. If there are
no findings, say so briefly and name material validation gaps.
Keep the review read-only unless the user explicitly authorizes posting. Do not
submit `APPROVE`, `COMMENT`, or `REQUEST_CHANGES`, add labels, edit code, or push
commits as an implied part of review.
### 8. Optionally request focused owner reviews
Only when the user asks to identify or request reviewers, read
[review-requests.md](references/delivery/review-requests.md). Rank path-matched
CODEOWNERS with the frozen module page's owners or required reviewers and
documented governance expertise; propose one to three focused reviewers with an
explicit contract rationale.
Identifying or suggesting reviewers is read-only. Requesting reviewers or
posting `@mention` comments changes external state and requires explicit user
authorization. When authorized, recheck the head, deduplicate existing
requests, and post at most one consolidated comment. Do not infer this
permission from a request to review the code.
@@ -0,0 +1,4 @@
interface:
display_name: "vLLM-Omni PR Review"
short_description: "Review vLLM-Omni PRs and route owner requests"
default_prompt: "Use $review-pr to review this vLLM-Omni pull request with owner-specific rules and evidence-backed findings."
@@ -0,0 +1,49 @@
# Configuration Provenance
Read this reference when changed or newly load-bearing behavior reads config,
CLI, environment, legacy, alias, or default values. Always use it for endpoint,
host, port, address, connector metadata, and `*_extra_config` consumers, even
when the schema or documentation is unchanged.
## Trace the source of truth
For each semantic value, record one compact ledger:
```text
documented input -> parsed field -> competing sources and precedence
-> runtime projection/transport -> final consumer
```
- Search the exact field, environment variable, default literal, and aliases in
code, tests, docs, examples, and deployment files.
- Treat unchanged producers and context lines as review scope when the diff
starts relying on them or makes their output authoritative.
- Resolve the intended precedence from active branch-local docs, enforced
code/tests, and compatibility policy. Report a documented input that is
silently ignored or reinterpreted at a live sink.
- Check every supported construction path and feature-off behavior without
inventing parity for unsupported paths.
## Require discriminating evidence
Give every competing source a distinct value so the test reveals which source
wins. Prefer non-default sentinels; for example:
```text
documented config port = 25301
environment/default port = 8998
expected final endpoint = http://10.0.0.8:25301
```
Assert the final URL, connection target, runtime argument, or other live sink,
not only an intermediate object. A test that changes the environment value to
the expected configured value proves formatting, not configuration propagation.
For a bug fix, show the discriminating test fails on the frozen base and passes
on the reviewed head when the environment permits execution.
## Finding bar
A supported non-default value reaching the wrong endpoint, device, topology, or
runtime behavior is a correctness or compatibility finding. A missing
discriminating test alone is a validation gap unless repository policy requires
that regression coverage for the change.
@@ -0,0 +1,121 @@
# Model Addition Checklist
Read this reference when a PR adds or expands a model architecture, pipeline,
loader, processor, registry entry, or stage configuration.
Official docs: [model contribution guides](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/),
[adding an omni model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_omni_model/),
[adding a diffusion model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model/),
[adding a TTS model](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_tts_model/),
[supported models](https://docs.vllm.ai/projects/vllm-omni/en/latest/models/supported_models/),
[diffusion features](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion_features/),
and [feature compatibility](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/feature_compatibility/).
## Match claims to the diff
- Reconcile every claimed architecture, modality, variant, feature, and example
with changed files and registry entries.
- Resolve config model types, architecture names, callable paths, exports, and
dependency extras against live symbols.
- Check sibling generations sharing an architecture name are disambiguated by
a stable config predicate rather than import order or an accidental default.
## Close the integration chain
Trace:
```text
public model id -> config/pipeline selection -> registry -> loader/processor
-> stage inputs -> model execution -> stage/public output
```
Verify optional dependencies fail with an actionable message, weight names and
dtypes map correctly, and each advertised serving mode reaches the production
dispatcher. Confirm outputs are non-empty and valid for their modality, shape,
sample rate, or response schema.
## Keep support documentation truthful
- Build an evidence-backed inventory of the new model's exact architecture and
checkpoint identifiers, modalities and tasks, hardware backends, serving
modes, and supported acceleration, parallelism, cache, offload, and
quantization features. Treat untested support as unknown, not supported.
- Update `docs/models/supported_models.md` with the architecture, model family,
example checkpoints, and only the backends validated by the PR.
- For a diffusion model, add it to the applicable ImageGen, VideoGen, or
AudioGen table in `docs/user_guide/diffusion_features.md`. Fill every feature
column according to its legend and document partial support or constraints.
- Update `docs/user_guide/feature_compatibility.md` when the model changes valid
feature combinations, configuration constraints, examples, or limitations,
and update any other branch-local support table affected by the claims.
- Cross-check every documented support mark against the registry, configuration,
production-path tests, and PR evidence; do not copy a sibling model's support
claims without validation.
## Ship the model recipe
- Require every new model PR to add or update one model-family recipe under
`recipes/<vendor>/` and its row in the `recipes/README.md` Available Recipes
table. Follow [`recipes/TEMPLATE.md`](../../../../../recipes/TEMPLATE.md) and
the layout rules in
[`recipes/README.md`](../../../../../recipes/README.md); the
external recipes repository is a structural reference, not a substitute.
- Include only tested task, serving-mode, and hardware configurations. Keep the
exact model identifiers, flags, commands, verification, feature support,
constraints, and limitations consistent with code, tests, examples, and the
support documentation.
- Cover each validated platform in its own hardware section. Use one recipe per
model family by default; justify a split when configurations cannot be kept
clear in one file, and do not add placeholder sections for untested platforms.
## Compare the reference implementation
Require every new model PR to include a head-versus-reference accuracy or
quality and performance comparison:
- Name the canonical implementation and pin its repository revision, package
version, configuration, and checkpoint. If it cannot run, explain why and use
the closest runnable reference; do not silently omit the comparison.
- Match inputs or dataset, seed, output settings, precision or quantization,
hardware, software, warmup, repetitions, synchronization, and timing and
memory scope. Include exact commands and environment details for both sides.
- Compare accuracy or output quality with a task-appropriate metric and explicit
tolerance. When no reliable metric exists, provide paired fixed-input outputs,
artifacts, and evaluation criteria, and explain expected numerical or
qualitative differences.
- Report end-to-end latency, task-appropriate throughput, peak device memory,
run-to-run variability, and head/reference ratios at equivalent output quality.
- Prefer a timing split for preprocessing or encoding, each model stage,
inter-stage transfer or communication, decoding or postprocessing, and output
assembly. Define the boundaries, account for overlap, reconcile the split with
end-to-end time, and explain any unavailable stage measurement.
Use [perf-verification.md](perf-verification.md) for the comparison protocol and
repository benchmark and profiling links.
## Remove accidental surface
Search bounded call sites for:
- inference-dead training `forward()` paths, unused factories/wrappers, never-set
branch keys, and immediately discarded parameters;
- duplicate payload strings, validation, or shape coercion across stages;
- compatibility aliases in brand-new code without an existing caller;
- private symbols re-exported as public API or module-level side effects.
Keep code only when it has a distinct live caller, invariant, or compatibility
contract. Prefer one typed producer-consumer schema to repeated string keys.
## Require proportionate evidence
- Run a representative production-path inference and assert output content, not
only process survival.
- Require profiling or A/B tables when the PR makes performance, memory,
precision, or quality claims, or when a suspected hot-path/device bug needs
that evidence. Do not impose fixed utilization or regression thresholds that
the repository does not define.
- Require focused registry/config/loading tests, a representative E2E path, and
user-facing model/usage documentation.
Use [diffusion.md](../modules/diffusion.md) for diffusion models and
[perf-verification.md](perf-verification.md) for quantitative claims.
@@ -0,0 +1,51 @@
# Performance and Accuracy Verification
Activate this reference when the PR adds a model or claims or intentionally
changes latency, throughput, memory, scaling, precision, or output quality.
Official docs: [profiling](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/profiling/),
[serving benchmarks](https://docs.vllm.ai/projects/vllm-omni/en/latest/cli/bench/serve/),
and [metrics](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/metrics/).
## Comparable A/B contract
Compare frozen base and head for an existing path. For a new model absent from
base, compare head with a pinned canonical reference implementation. Use the same:
- hardware, driver, software versions, model/checkpoint, and dependencies;
- input or dataset, seed, precision, batch/concurrency, topology, and feature
flags;
- warmup, measured repetitions, synchronization, timing scope, and memory
collection method.
Keep exact commands and report variability rather than a cherry-picked run.
Use the repository's benchmark when it covers the claim; otherwise use the
smallest reproducible workload that reaches the changed production path.
## Evidence table
| Dimension | Typical evidence |
| --- | --- |
| Latency | End-to-end latency and, when available, TTFT plus a defined per-stage and transfer-time split. |
| Throughput/scaling | Requests, tokens, frames, or audio duration per second across the claimed concurrency/topology. |
| Memory | Peak allocated/reserved device memory and OOM boundary if relevant. |
| Quality/accuracy | Repository metric or known-good output comparison with an explicit tolerance; paired samples when no metric exists. |
Every optimization needs correctness or quality evidence. A faster result on a
different workload, precision, seed, or topology is not a valid comparison.
## Reviewer verification ladder
1. Run base/head A/B on suitable hardware when affordable.
2. If only one side can run, verify the script and request a comparable pair.
3. If hardware or assets are unavailable, audit methodology, code path, and
contributor evidence; name the exact unverified claim.
Classify discrepancies before reporting: implementation regression, benchmark
bug, environmental drift, noise, or unsupported claim. Explain material
regressions against the PR's stated goal or repository contract; do not invent
universal percentage thresholds.
Report claimed and measured values together, bind them to the frozen SHAs and
environment, and keep unavailable hardware as a validation gap rather than a
fabricated pass.
@@ -0,0 +1,42 @@
# Test Quality Evaluation
Read this reference when tests change, a risky production path has no test, or
a test-only PR may not prove the intended behavior.
Official docs: [test system overview](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_system_overview/)
and [test writing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_writing_guide/).
## Static proof check
For each changed semantic path, ask whether the test would fail if that behavior
were reverted or broken. Check:
- assertions pin the contract, not only non-null output, logs, or “no crash”;
- the production dispatcher, registry, connector, or scheduler is reached
instead of the changed behavior being mocked away;
- fakes preserve relevant types, MRO, shapes, devices, async behavior, and
lifecycle transitions;
- configuration tests give competing config, CLI, environment, legacy, and
default sources distinct non-default values and assert the final consumer;
setting only the source read by the implementation does not prove precedence;
- seeds, ordering, timing, synchronization, external services, and numeric
tolerances are deterministic or explicitly controlled;
- normal, invalid, boundary, feature-off, failure/cancellation, and regression
paths are covered where the diff changes them;
- run-level and domain markers place hardware/model tests in the intended CI
lane without silently skipping the contract.
Map source symbols to tests with bounded `rg` searches; do not assume the test
tree mirrors production paths.
## Runtime check
Run the narrowest affected tests supported by the environment. Record skipped
hardware/model cases as gaps. Classify failures as code, test, infrastructure,
or flaky before turning them into findings; one passing rerun does not erase a
flakiness signal.
Keep grades or matrices internal. Report only concrete code bugs or the one or
two test defects that leave changed behavior materially unprotected. Use
[tests-docs-checklist.md](tests-docs-checklist.md) for CI and documentation
coverage.
@@ -0,0 +1,48 @@
# Tests and Documentation Checklist
Use this reference to connect changed behavior to coverage, CI selection,
examples, documentation, and reproducible PR evidence.
Official docs: [test writing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_writing_guide/),
[test execution guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_execution_guide/),
and [documentation guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/DOCS_GUIDE/).
## Coverage packet
For each high-risk change, record a compact internal packet:
```text
change and failure risk -> existing unit/E2E coverage -> uncovered boundary
-> smallest test and stable assertion that closes the gap
```
Prefer a regression test that fails on the frozen base for bug fixes. Check
numeric tolerances, deterministic inputs, realistic mocks, hardware/domain
markers, and Buildkite/CI wiring using the documentation present on the reviewed
branch. A skipped test is not a pass.
## Documentation sync
Require user-facing updates when the diff changes a model, feature, CLI/API,
config key, default, compatibility behavior, or supported platform. Verify:
- the actual docs navigation and support tables used by the branch;
- exact identifiers, defaults, constraints, known limitations, and migration;
- a minimal runnable example for newly exposed behavior;
- examples and docs agree on commands, parameters, output schema, and supported
modes;
- links and the bounded live contract build or resolve correctly.
Do not request docs for internal behavior with no user-visible contract.
## PR evidence
For contributor-run tests or benchmarks, require enough provenance to evaluate
the claim: hardware, relevant software versions, exact commands, inputs/model,
result, and approximate runtime. Performance or accuracy evidence must compare
base and head under equivalent conditions; use
[perf-verification.md](perf-verification.md).
Turn a missing test or document into a finding only when it protects or explains
changed behavior. Name the smallest useful addition rather than requesting a
broad coverage expansion.
@@ -0,0 +1,36 @@
# Active Verification
Read this reference when a runnable affected path and suitable local or remote
hardware are available. Never claim device evidence from static analysis or a
simulated backend.
Official docs: [installation](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/)
and [test execution guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_execution_guide/).
## Choose the narrowest level
| Available environment | Verification |
| --- | --- |
| Matching accelerator, model/assets, and server | Targeted unit/integration test plus one representative E2E request. |
| Matching accelerator without serving setup | Targeted model/runtime test and output or metric inspection. |
| CPU/static environment only | Import/version preflight, focused CPU tests, config/registry resolution, and static contract checks. |
| Missing dependencies or inaccessible hardware | Record the exact gap and evaluate contributor/CI evidence. |
## Execute safely
1. Bind the run to the frozen head SHA and record Python, platform, dependency,
model, device, and relevant environment fingerprints.
2. Run the smallest affected unit test before an E2E workload.
3. Exercise the changed public or production-dispatch path with bounded inputs.
4. Compare actual output, metrics, cleanup, and exit behavior with the PR claim.
5. For base/head claims, use the equivalent A/B protocol in
[perf-verification.md](perf-verification.md).
Avoid disturbing shared servers, ports, models, or accelerators without user
authorization. Do not install large dependencies, download gated assets, or
start remote workloads merely because hardware might exist.
Report the command, environment, result, and any skipped path. Classify failures
as change-induced, pre-existing, infrastructure, environment, or flaky before
making a finding. A hardware gap lowers confidence; it is not automatic proof
of a defect.
@@ -0,0 +1,105 @@
# Maintainer Review Style
Use this reference only when turning proved findings into comments. The baseline
maintainer samples favor short, direct comments and reserve long explanations
for architecture or non-obvious failure paths.
Official docs: [contributing and review process](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
## Evidence basis
This calibration used 928 PRs merged from 2026-05-10 through 2026-08-08.
Candidates were ranked by substantive human inline-review threads, then by
distinct reviewers and review rounds. Author self-comments, bots and AI
reviewers, and bodyless approvals were excluded. Counts selected the sample;
the practices below came from reading root comments and follow-up replies.
The main set contains the six leaders after filtering; three other high-ranked
PRs add core-runtime and benchmark coverage. All were merged.
| PR | Review boundary |
| --- | --- |
| [#2906](https://github.com/vllm-project/vllm-omni/pull/2906) | New TTS model, shared utilities, recipes, and request state. |
| [#3907](https://github.com/vllm-project/vllm-omni/pull/3907) | Full-duplex API, session lifecycle, concurrency, and realtime evidence. |
| [#2162](https://github.com/vllm-project/vllm-omni/pull/2162) | Diffusion world model, OpenPI serving, tests, examples, and docs. |
| [#3454](https://github.com/vllm-project/vllm-omni/pull/3454) | Diffusion model, generic CLI surface, loaders, and configuration scope. |
| [#4425](https://github.com/vllm-project/vllm-omni/pull/4425) | Structured configuration, typing, defaults, and compatibility. |
| [#5397](https://github.com/vllm-project/vllm-omni/pull/5397) | Distributed offload, collectives, failure handling, and load semantics. |
| [#4281](https://github.com/vllm-project/vllm-omni/pull/4281) | Composable parallelism, configuration contracts, and examples. |
| [#3628](https://github.com/vllm-project/vllm-omni/pull/3628) | Stage metrics, benchmark definitions, opt-in API behavior, and cleanup. |
| [#3855](https://github.com/vllm-project/vllm-omni/pull/3855) | Stage runtime, thread safety, startup, and compatibility. |
Historical comments calibrate delivery; they do not create merge policy. Live
code, tests, repository docs, and the current `CODEOWNERS` file remain
authoritative. When refreshing this study, promote a practice only when it
recurs across at least three PRs or is supported by current repository guidance.
## Recurring review behavior
- Start from a changed contract, not a style preference. State the exact
request, configuration, topology, or failure path that triggers the issue.
- Protect repository boundaries. Keep model-specific behavior in its domain,
move genuinely reusable logic to a shared module, and reuse existing loaders,
protocols, definitions, and config sources instead of bypassing them.
- Challenge unrelated scope and abstractions. Ask why a new flag, config, class,
helper, fallback, or file is needed when an existing path appears sufficient.
- Trace concurrency and terminal paths: overlapping requests, disconnect,
cancellation, partial failure, stale replies, collective schedule mismatch,
cleanup, and retry. Prefer an explicit startup error for unsupported
combinations over silently running with incomplete state or wrong weights.
- Preserve user-visible behavior across variants. An accepted CLI/API/config
field must be applied, explicitly rejected, or visibly warned about; compare
normal and headless, offline and online, and streaming and non-streaming paths.
- Ask for evidence at the affected boundary: a regression test using the real
tensor or protocol shape, correct pytest markers so CI runs it, comparable A/B
performance results, or an SLO-oriented latency/quality measurement.
- Keep interfaces explicit: avoid `Any` and opaque strings where a precise
union, enum, or typed envelope exists; remove dead code and duplicated
constants; validate requests before expensive execution.
- Document module scope and non-obvious invariants. Update examples, recipes,
supported-model/feature docs, and hardware assumptions when users rely on
them.
- When expertise crosses an ownership boundary, ask a named owner to inspect a
specific file or contract and explain why. Do not send a generic `PTAL` list.
- Keep non-blocking cleanup in a linked follow-up only after the current PR is
safe. If a revision fixes only part of a failure chain, continue the existing
thread with the remaining trigger instead of opening a duplicate finding.
## Write the finding
- Lead with the issue or decisive question; omit a review preamble.
- Name the concrete trigger and impact, then the smallest fix direction.
- Keep obvious fixes to one sentence. Use more detail only when the call path or
ownership boundary is otherwise unclear.
- Prefer one root-cause comment over repeated symptom comments.
- Hedge only when evidence is genuinely incomplete; mark that item as a
question or validation gap rather than overstating severity.
- Say briefly that no findings were found when the review is clean; zero
comments is a valid outcome when posting is authorized.
Useful forms:
```text
This leaves <state> allocated when <failure/cancellation path>. Release it in
the shared terminal cleanup path.
```
```text
Does this producer also update <consumer>? It still reads the old field here.
```
```text
With <configuration>, this bypasses <existing loader/validator>, so <impact>.
Reuse that path or reject this combination before startup.
```
```text
@<owner>, could you check <specific boundary>? This changes <contract> used by
<consumer>.
```
Use a priority label only for a proved high-impact defect. Avoid generic praise,
routine “Nit:” prefixes, decorative severity labels, audit templates, rule IDs,
comment-count narration, and “I left a few comments.” Keep approvals, review
events, reviewer requests, mentions, and GitHub replies subject to the user's
explicit authorization.
@@ -0,0 +1,99 @@
# Focused Code-Owner Review Requests
Read only when the user asks to identify, suggest, request, add, or ping
reviewers. Finding candidates is read-only; GitHub review requests and
`@mention` comments are external writes.
Owner sources, in order:
1. the reviewed head's
[CODEOWNERS](https://github.com/vllm-project/vllm-omni/blob/main/.github/CODEOWNERS)
for path ownership;
2. the active module design page's `owners` and `required_reviewers` metadata
for contract ownership and mandatory design review; and
3. documented
[governance expertise](https://docs.vllm.ai/projects/vllm-omni/en/latest/community/governance/)
for specialization and tie-breaking.
Resolve all three from the reviewed head when available because ownership
changes over time. Use published docs only to discover a page absent from an
older checkout, never to override the frozen head, and do not copy a static
owner list into review output. CODEOWNERS remains the path authority; module
metadata identifies contract experts and required design reviewers.
## Rank candidates
1. Freeze the PR head and collect changed files, author, requested reviewers,
submitted reviews, and existing comments.
2. Resolve the **last matching CODEOWNERS rule** for each changed file. Do not
union an earlier broad rule with a later specific rule.
3. Resolve the primary module and feature overlays through
[design-contracts.md](../process/design-contracts.md). Read the active module
page's status, `owners`, and `required_reviewers`; do not treat owners on an
archived, superseded, or unrelated page as current authority.
4. Group files by the selected primary owner and the exact module or feature
contract needing review. Weight production owners and live consumers above
tests/docs/CI owners and the repository-wide fallback. Include an applicable
`required_reviewers` entry even when it is not the path owner, and label why.
5. Use documented governance specialties only to choose among matching owners:
for example serving versus configuration, diffusion model versus cache or
parallelism, TTS versus generic model execution, or CI versus tests.
6. Exclude the PR author, bots, users who already completed the needed review,
and duplicate pending requests. Do not infer availability or invent
expertise from a username.
7. Select one reviewer for a single-owner change and at most three for a real
cross-module change. If more are plausible, rank them but do not ping a
broad group; ask the user before exceeding three.
If every specific owner is excluded, consider one repository-wide fallback
owner and label that choice as a fallback rather than a path match. Never fan
out to the full fallback list. Prefer a qualified owner who covers multiple
critical path groups over separate reviewers for each file. If CODEOWNERS is
missing or invalid, state that gap and use governance only as a labeled fallback.
Recent file history may break a tie only when CODEOWNERS and governance remain
ambiguous. Never use commit count alone to override an explicit current owner.
## Preview before writing
Show the user the exact candidate set and rationale:
```text
@owner — <owned paths or module role>; please verify <specific module or feature contract>.
```
Draft one consolidated comment tied to the frozen head:
```text
Requesting focused review for `<short-head-sha>`:
- @owner1: could you check the <module contract> in `<path/group>`?
- @owner2: could you check the <feature or cross-boundary contract> in `<path/group>`?
Relevant validation: <one concise result or exact remaining gap>.
```
Mention only owners being requested. Keep the reason specific enough that each
person knows what decision or invariant needs attention; do not post a generic
“please review” or a generated ownership audit.
## External-write gate
- “Who should review?”, “suggest reviewers”, or “draft a request” authorizes
only the preview.
- “Request/add these reviewers” authorizes the GitHub review-request mutation.
- “Ping/comment/tag the owners” authorizes the consolidated comment.
- When the user explicitly asks to request reviews **with an owner comment**,
both writes are authorized. Obey narrower wording when only one action is
requested.
Immediately before an authorized write, verify the PR is still open, the head
SHA is unchanged, and the candidate is not the author or already satisfied.
Report draft/WIP state and avoid notifications unless the user explicitly wants
reviewers requested before ready-for-review.
Use the GitHub connector for reviewer requests and the single top-level comment;
use `gh` only when connector coverage is unavailable. Do not submit an approval,
request-changes event, label, assignment, or multiple comments as part of this
action. After writing, reread the PR and report the reviewers requested and the
comment URL. If the head changed, discard the draft and rerank before posting.
@@ -0,0 +1,24 @@
# Communication Features
Use with the OmniConnector, stage-runtime, configuration, and I/O module
contracts. Load the concrete backend design from the reviewed head:
| Backend | Design |
| --- | --- |
| Mooncake store | [MooncakeStoreConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/mooncake_store_connector/) |
| Mooncake transfer engine | [MooncakeTransferEngineConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/mooncake_transfer_engine_connector/) |
| Mori transfer engine | [MoriTransferEngineConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/mori_transfer_engine_connector/) |
| Shared memory | [SharedMemoryConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/shared_memory_connector/) |
| Yuanrong store | [YuanrongConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/yuanrong_connector/) |
| Yuanrong transfer engine | [YuanrongTransferEngineConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_connectors/yuanrong_transfer_engine_connector/) |
## Feature checks
- Verify dependency, service, device/network, port, memory-pool, and environment
prerequisites with actionable startup errors.
- Match key/metadata construction, sender and receiver roles, allocation,
ownership, completion, and supported fast/serialized paths.
- Exercise timeout, retry, stale sender/buffer reclamation, partial transfer,
cancellation, and connector close semantics.
- Validate the smallest supported same-host or cross-node topology and document
backend limitations and operational diagnostics.
@@ -0,0 +1,26 @@
# Diffusion Acceleration Features
Use with the exact diffusion submodule plus platform, cache, quantization, or
benchmarking contracts selected by the live path.
| Feature family | Designs |
| --- | --- |
| Parallelism | [CFG](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel/), [Expert](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/expert_parallel/), [HSDP](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp/), [Pipeline](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/pipeline_parallel/), [Sequence](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel/), [Tensor](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel/), [VAE patch](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/vae_parallel/) |
| Attention | [Skip-Softmax](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/skip_softmax/) |
| Quantization | [Quantization](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/quantization/) |
| Cache | [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/teacache/) |
| Batching | [Diffusion Continuous Batching](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/diffusion_continuous_batching/) |
| Offload | [Distributed Layerwise Offload](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/distributed_layerwise_offload/) |
## Feature checks
- Verify feature selection, supported models/platforms/topologies, defaults,
dependencies, and every touched feature combination.
- Preserve feature-off and single-rank correctness; reject unsupported
combinations before collectives, loading, or execution begin.
- Trace shapes, shards, ranks, collectives, precision, cache validity, residency,
and per-request state through the selected runtime hooks.
- Compare output quality or accuracy against the disabled/reference path before
accepting latency or memory improvements.
- Require exact commands, topology, environment, warmup, repetitions, and
measurements for performance claims; update feature compatibility docs.
@@ -0,0 +1,22 @@
# Infrastructure and Performance Features
Load the matching design from the reviewed head:
- [Prometheus Metrics](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/metrics/)
for logger wrapping, gating, lifecycle, stage transfer, and metric definitions.
- [Speech Generation Performance Optimizations](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/qwen3_omni_tts_performance_optimization/)
for stage batching, CUDA graphs, async chunk, streaming, re-prefill, compile,
and deployment evidence.
## Feature checks
- Route instrumentation changes through observability or profiling and result
methodology through benchmarking and the performance evidence check.
- Preserve opt-in/default behavior, units, labels, aggregation, timing scope,
lifecycle resets, and bounded overhead.
- For stacked optimizations, measure each increment against the same baseline
and explain interactions, quality guards, regressions, and limitations.
- Require exact model, checkpoint, hardware, software, workload, topology,
concurrency, warmup, repetitions, and production commands.
- Update the design page when metric meaning, supported optimization stack, or
deployment guidance changes.
@@ -0,0 +1,24 @@
# Runtime and Stage Execution Features
Load the matching design from the reviewed head:
| Feature | Design | Primary module intersections |
| --- | --- | --- |
| Disaggregated inference | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/disaggregated_inference/) | Configuration, stage runtime, OmniConnector, I/O contracts |
| Async chunk | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/async_chunk/) | Model integration, orchestration, stage runtime, connector, I/O |
| Async diffusion output | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/async_diffusion_output/) | Diffusion runtime, I/O contracts, entrypoints |
| Async Omni output materialization | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/omni_async_output_materialization/) | AR runtime, model integration, I/O contracts, orchestration |
| Automatic prefix caching | [Design](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/prefix_caching/) | AR runtime, cache management, I/O contracts |
## Feature checks
- Confirm the feature's supported models, entrypoints, stage types, topology,
configuration, default, and compatibility limitations.
- Trace enablement from public/config ingress through every participating module
and verify the advertised production dispatch actually selects it.
- Preserve a correct disabled or fallback path and make unsupported combinations
fail or warn as documented.
- Test overlap, ordering, identity, backpressure, cancellation, failure, final
partial output, and cleanup at each asynchronous boundary.
- Update the feature page, support matrices, examples, and recipes when behavior
or coverage changes.
@@ -0,0 +1,22 @@
# Autoregressive Runtime
Primary design: [Autoregressive Runtime](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/ar_runtime/).
Use for AR request queues, scheduling state, token and memory budgets, KV/input
readiness, workers, runners, and upstream vLLM request/cache transitions.
## Contract checks
- Preserve upstream request-state, scheduling, cache, and preemption semantics
unless an Omni-specific difference is documented and tested.
- Convert modality-specific data through explicit input or output adapters; do
not inject it through unrelated scheduler state.
- Keep workers and model runners focused on assigned execution rather than
admission or cross-stage routing.
- Synchronize token, sequence, batch, and memory budgets with worker-visible
state and prevent starvation, double counting, or stale readiness.
- Preserve legal transitions and one terminal cleanup path across completion,
abort, timeout, preemption, worker failure, and shutdown.
Test request lifecycle, scheduling decisions, abort, cache state, adapters, and
every affected worker execution mode against the pinned upstream contract.
@@ -0,0 +1,24 @@
# Benchmarking
Primary design: [Benchmarking](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/benchmarking/).
Use when benchmark code, workloads, metric calculations, CLI benchmark
entrypoints, or result metadata change. Use the performance check separately
when a PR uses benchmark results as evidence.
## Contract checks
- Record model and revision, code SHA, hardware, software, configuration,
workload, precision, topology, concurrency, warmup, repetitions, and
measurement window.
- Validate successful and semantically acceptable outputs before accepting
timing, memory, or scaling results.
- Define every metric's unit, population, aggregation, synchronization, and
timing scope.
- Keep workload parsing and metric computation deterministic and reusable
across base and head.
- Reject comparisons that change material workload or runtime settings without
isolating and explaining that difference.
Test workload parsing and metric calculations with deterministic fixtures and
publish exact commands plus variability for quantitative claims.
@@ -0,0 +1,22 @@
# Cache Management
Primary design: [Cache Management](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/cache_management/).
Use for AR or diffusion cache identity, reusable state, validity, lifecycle,
reset, eviction, and teardown.
## Contract checks
- Include every model, request, input, timestep, precision, device, topology,
and feature property that can change a reused value in cache identity.
- Preserve a correct disabled and miss path; never reuse after validity inputs
change.
- Isolate request-scoped state across concurrent requests and define ownership
for model-scoped state.
- Make reset, invalidation, eviction, capacity, and teardown explicit on
completion, cancellation, failure, model reload, and shutdown.
- Keep cache optimizations behind the owning runtime's hooks rather than adding
a second request lifecycle.
Test hit, miss, invalidation, disabled behavior, concurrency, capacity, reset,
and teardown, with quality comparison when reuse is approximate.
@@ -0,0 +1,24 @@
# vLLM-Omni Configuration
Primary design: [vLLM-Omni Configuration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/vllm_omni_config/).
Use for config construction, deploy YAML, registries, schema, defaults,
aliases, CLI projection, stage config, and topology. Inspect the reviewed
page's status: while its contract remains deferred or draft, current code and
tests—not proposed precedence rules—are authoritative.
## Contract checks
- Resolve each supported structured, legacy, CLI, environment, default, alias,
and stage-override producer through its actual runtime consumer.
- Reject unknown or owner-mismatched fields explicitly; keep process-local
runtime objects out of transportable configuration.
- Verify default, explicit, and feature-off values reach every live factory and
consumer without silent reinterpretation.
- Treat stage count, placement, connector, parallelism, device, and memory
fields as one topology contract and reject impossible combinations early.
- Preserve supported construction-path parity until a documented migration
removes a path.
Require normalization and schema tests, a live-consumer assertion, and docs for
public keys, defaults, constraints, precedence, or migration behavior.
@@ -0,0 +1,33 @@
# Diffusion Module Family
Primary design: [Diffusion Modules](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/).
Use for diffusion-stage behavior, then load the exact submodule page from the
reviewed head:
| Submodule | Signals | Design |
| --- | --- | --- |
| Runtime | Admission, scheduling, execution, progress, outputs, cancellation, hooks, cleanup | [Diffusion Runtime](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/diffusion_runtime/) |
| Model Integration | Pipelines, registry, loader, adapters, shared layers, model-specific processing | [Diffusion Model Integration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/diffusion_model_integration/) |
| Continuous Batching | Compatibility, request/step batches, per-request progress, capacity | [Continuous Batching](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/continuous_batching/) |
| Parallelism | Rank topology, process groups, sharding, collectives, distributed execution | [Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/parallelism/) |
| Offloader | Residency, transfers, memory accounting, prefetch, teardown | [Offloader](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/diffusion/offloader/) |
## Family-wide checks
- Keep one scheduler-owned lifecycle from admission through completion,
cancellation, failure, and cleanup; optional features must use defined hooks.
- Select pipelines through registry and loader contracts instead of scattered
model-name conditionals; keep genuine model differences in model directories.
- Trace latent, conditioning, timestep, generator, shape, layout, dtype, device,
batch expansion, and output conversion through the actual consumer.
- Make batch compatibility and per-request isolation explicit; never associate
output by unstable batch position.
- Derive topology from validated configuration, define every sharding boundary,
keep collectives symmetric, and preserve a supported single-rank path.
- Track one residency owner per offloaded component, wait for transfer readiness,
preserve model state, bound retained copies, and clean up deterministically.
Load matching feature designs for cache, quantization, parallelism, continuous
batching, or offload. Require feature-off correctness, representative inference,
terminal-path tests, and quality evidence for optimizations.
@@ -0,0 +1,23 @@
# Engine Orchestration
Primary design: [Engine Orchestration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/engine_orchestration/).
Use for request-state creation, cross-stage routing, output ordering, companion
tracking, RPC correlation, cancellation propagation, and terminal convergence.
Orchestration does not own replica placement, stage-process startup, public
rendering, connector transport, or payload and error schemas.
## Contract checks
- Keep one owner for cross-stage routing; entrypoints, models, and stage clients
must not independently select the next logical stage.
- Preserve request identity and output order across concurrent stages and
correlate every control or RPC result to the originating request and wave.
- Make terminal state monotonic: no new forwarding after success, abort,
request failure, or fatal engine failure.
- Propagate cancellation and failure to every active stage and release
orchestrator-owned tasks, queues, companions, and request state.
- Preserve shutdown ordering across the orchestrator and stage runtime.
Test representative multi-stage routing, ordering, correlation, cancellation,
failure, and repeated shutdown.
@@ -0,0 +1,25 @@
# Entrypoints and Serving Boundaries
Primary design: [Entrypoints and Serving Boundaries](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/entrypoints/).
Use for offline APIs, CLI and server composition, OpenAI-compatible routes,
request validation, response conversion, streaming, sessions, and engine
handoff. Entrypoints adapt public contracts; they do not own cross-stage
routing, stage lifecycle, configuration precedence, payload schemas, or error
classification.
## Contract checks
- Normalize and validate public values once before engine submission.
- Preserve request identity, modality, ordering, and exactly one terminal event
across streaming, disconnect, cancellation, and failure.
- Translate internal outputs and errors explicitly for each advertised offline,
HTTP, SSE, WebSocket, or asynchronous-job path.
- Keep model-specific behavior behind a common adapter or processor instead of
adding model-name policy to generic entrypoints.
- Require API or CLI changes to explain necessity, alternatives, compatibility,
migration, and the supported route or transport matrix.
Test the affected public protocol, normalization, engine handoff, rendering,
and terminal outcomes. Route orchestration, I/O schema, and semantic error
changes to their module contracts.
@@ -0,0 +1,24 @@
# Error Classification, Propagation, and Rendering
Primary design: [Error Classification, Propagation, and Rendering](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/error_contracts/).
Use when a change classifies failures, crosses process boundaries, decides
fatal versus request-scoped impact, sanitizes context, or renders offline,
HTTP, SSE, WebSocket, or asynchronous-job errors. The I/O module owns the wire
schema; this module owns the meaning of its error fields.
## Contract checks
- Keep request-scoped failure isolated unless the underlying engine is no
longer usable.
- Preserve stable classification and request context across process and
protocol boundaries without exposing prompts, outputs, secrets, or internal
stack details.
- Keep cancellation distinguishable from internal failure.
- Terminate each public transport exactly once with consistent status, body,
event, finish reason, and metrics classification.
- Maintain compatibility across mixed or older producers only when the
reviewed contract defines that behavior.
Test table-driven classification, wire round trips, fatal/request-scoped
propagation, cancellation, sanitization, and every affected public renderer.
@@ -0,0 +1,23 @@
# Execution Platforms
Primary design: [Execution Platforms](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/execution_platforms/).
Use for hardware capability detection, worker selection, patches, kernels,
attention backends, and vendor-specific configuration or connectors.
## Contract checks
- Guard hardware behavior through platform selection or capability detection
and fail clearly when unsupported.
- Keep portable modules free of direct vendor imports when a common selector
exists; verify portable imports without the target accelerator.
- Override only behavior that differs from common or upstream implementations,
and compare sibling worker hooks when inheritance does not provide parity.
- Preserve dtype, layout, rank, stream, synchronization, workspace,
graph-capture, and resource-lifetime contracts.
- Do not let simplified CPU doubles hide platform MRO, initialization, or
dispatch failures.
Run portable checks first, then the smallest matching-device test in a fresh
environment. Report unavailable hardware as a validation gap, never simulated
device evidence.
@@ -0,0 +1,23 @@
# Input, Output, and Modality Contracts
Primary design: [Input, Output, and Modality Contracts](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/input_output_modality_contracts/).
Use for requests, prompt and modality types, message schemas, serialization,
output types, accumulation, completion, and compatibility shims crossing
entrypoint, orchestration, stage, connector, or model boundaries.
## Contract checks
- Give every boundary value an explicit modality and validated typed schema.
- Preserve request and stage identity across conversion, serialization,
streaming updates, batching, cancellation, and errors.
- Distinguish partial, final, empty, and failed output states for every modality.
- Verify producer and consumer agreement on shape, layout, dtype, device,
ownership, metadata, ordering, and optional or unknown fields.
- Keep internal objects out of public protocols; translate them at the owning
entrypoint.
- Preserve wire compatibility deliberately and give deprecated shims a bounded
migration path.
Require construction, round-trip, accumulation, completion, rejection, and
compatibility tests at each changed producer-consumer boundary.
@@ -0,0 +1,25 @@
# Model Integration
Primary design: [Model Integration](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/model_integration/).
Use for model registration, configuration adapters, preprocessing, checkpoint
loading, stage inputs, model-specific hooks, shared layers, runners, and model
execution. Use the diffusion family reference for diffusion models.
## Contract checks
- Declare explicitly how the model class, loader, input processor, stage
configuration, and runtime capabilities are selected.
- Preserve weight-name mapping, dtype/device placement, optional-dependency
errors, and checkpoint-specific loading semantics.
- Trace each typed input and output through preprocessing, runner state, model
call, and its actual consumer.
- Reuse upstream vLLM behavior when sufficient; document and test the precise
difference for every override or patch.
- Keep model code out of admission, cross-stage routing, public policy, and
terminal lifecycle ownership.
- Place reusable behavior in shared modules and test it with more than one
integration.
Require focused registry, loader, processor, runner, and representative model
execution tests. Apply the model-addition check for new or expanded support.
@@ -0,0 +1,21 @@
# Observability
Primary design: [Observability](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/observability/).
Use for metrics, logs, correlation fields, and cross-stage request visibility.
## Contract checks
- Preserve stable request and stage identity across entrypoints, orchestration,
stages, workers, errors, and terminal cleanup.
- Treat metric names, units, labels, aggregation, and lifecycle as public
contracts; document any semantic change.
- Keep labels bounded: never use request IDs, prompts, outputs, paths, or other
unbounded user values.
- Register and reset metrics exactly once across repeated startup and shutdown,
and keep disabled instrumentation inexpensive.
- Align emitted values with their actual timing and lifecycle owner rather than
reconstructing them from ambiguous client strings.
Test registration, label sets, units, resets, error classification, and
representative multi-stage emission.
@@ -0,0 +1,24 @@
# OmniConnector
Primary design: [OmniConnector](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/omni_connector/).
Use for connector contracts and model-agnostic data transfer across stages,
processes, devices, or nodes. Connector implementations transport and
synchronize data; they do not choose the next stage or implement model policy.
## Contract checks
- Prove producer and consumer agreement on request/stage identity, shape,
layout, dtype, placement, ownership, ordering, and completion.
- Keep backend selection and metadata/key construction explicit; do not silently
reinterpret or drop payloads across SHM, store, network, or vendor paths.
- Bound connect, put/get, retry, synchronization, and readiness waits with
backpressure, cancellation, and actionable errors.
- Prevent route, endpoint, port, key, and buffer collisions across repeated
startup, concurrent requests, and replica replacement.
- Release connections, handles, buffers, listener work, and stale sender state
on success, cancellation, failure, and shutdown.
Load the matching communication feature design for a concrete backend. Require
focused producer-consumer tests and the smallest runnable multi-process or
multi-node path.
@@ -0,0 +1,21 @@
# Profiling
Primary design: [Profiling](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/profiling/).
Use for opt-in instrumentation and trace collection across supported profiler
backends.
## Contract checks
- Keep expensive collection disabled by default and bounded by explicit start,
stop, timeout, and cleanup behavior.
- Preserve scheduling, generated outputs, resource ownership, and concurrency
semantics apart from documented synchronization overhead.
- Include stage, rank, worker, operation, and request correlation where useful
without embedding user payloads or secrets.
- Make repeated start/stop and partial failure safe, including multi-rank trace
naming and artifact ownership.
- Measure disabled and enabled overhead before making performance claims.
Test the disabled path, repeated lifecycle, failure cleanup, trace context, and
each affected backend.
@@ -0,0 +1,22 @@
# Quantization
Primary design: [Quantization](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/quantization/).
Use for method registration, checkpoint interpretation, layer replacement,
reduced precision, platform capability, and runtime constraints.
## Contract checks
- Validate checkpoint metadata, selected method, supported layers, device
capability, required dependencies, and parallel layout before loading.
- Preserve loader and weight-name mapping semantics, dtype, scale/granularity,
zero point, padding, and buffer handling through the real kernel consumer.
- Fail clearly for unsupported combinations; do not silently choose a different
method or precision.
- Keep online and pre-quantized checkpoint paths explicit and ensure patches or
fallbacks do not bypass model-owned loaders.
- Require numerical comparison against a defined reference with explicit
tolerance before treating performance results as valid.
Test registration, loading, representative layers, metadata, platform guards,
parallel combinations, and numerical correctness separately from performance.
@@ -0,0 +1,23 @@
# Stage Runtime and Replica Lifecycle
Primary design: [Stage Runtime and Replica Lifecycle](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/module/stage_runtime/).
Use for local or distributed placement, startup, readiness, stage pools,
replica identity and membership, selection, affinity, liveness, draining, and
stage-client or process lifecycle. Cross-stage request policy belongs to engine
orchestration; model scheduling belongs to the owning runtime.
## Contract checks
- Acquire, replace, and retire replicas through one stage-runtime lifecycle.
- Preserve stable replica identity and request affinity until termination or a
declared replica loss.
- Make readiness explicit before admission; surface partial-startup and
membership failures without leaving half-initialized capacity.
- Keep local and distributed semantics aligned for supported modes.
- Make drain, shutdown, and cleanup idempotent across repeated signals and
process failure.
Test startup, readiness, selection, affinity, membership changes, abort,
process loss, draining, and repeated shutdown on affected local/distributed
paths.
@@ -0,0 +1,34 @@
# Resolve Design Contracts
Use after freezing the review head and before selecting module or feature
references. The maintained [design index](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/)
defines two routing axes:
- **Module designs** own code boundaries and invariants. Select one primary
module and a second only when the live producer-consumer path crosses it.
- **Feature designs** describe behavior spanning modules. Load every feature
page whose implementation, default, support, or contract changes.
## Resolution order
1. Open `docs/design/index.md` in the reviewed head, not the reviewer's checkout
or published latest tree.
2. Match changed behavior against each module page's `primary_code_paths`,
`primary_path_exceptions`, ownership boundary, and `depends_on` metadata.
3. Read `status`, `architecture_state`, `last_verified_commit`, and any linked
in-flight work before treating prose as current behavior.
4. Use current code and tests when a page is draft, deferred, stale, absent, or
conflicts with the frozen implementation. Report material documentation
drift separately.
5. Treat candidate invariants and promotion gates as review questions until
repository policy, current code/tests, or owner-approved normative text
makes them enforceable. Do not turn a draft identifier into an automatic
blocker.
6. Use published latest docs only for discovery when the reviewed head lacks a
page; never import a newer contract into an older branch silently.
When behavior changes an active module or feature contract, require the same PR
to update that page. Require a new page only when repository policy, an accepted
RFC/issue, or explicit maintainer acceptance criteria do so; otherwise treat it
as a non-blocking design suggestion. Do not demand a second page when an existing
design already owns the exact contract.
@@ -0,0 +1,134 @@
# General Review Checks
Apply these checks to every PR after freezing the review surface. A pattern is
not a finding until the diff introduces or exposes a reachable trigger, impact,
and smallest safe fix. Do not report unrelated backlog.
Official docs: [contributing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/)
[feature compatibility](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/feature_compatibility/),
and [design documents](https://docs.vllm.ai/projects/vllm-omni/en/latest/design/).
## Contract and scope
- Match the title/body claims to every changed production path, test, config,
dependency, generated artifact, and user-visible document.
- Trace each changed value through public ingress, validation/defaulting,
producer, transformations, process or device boundaries, final consumer, and
terminal cleanup.
- Treat unchanged code as in scope when the diff makes its value newly
load-bearing for a changed consumer; validate that value's origin and contract
rather than reviewing only added lines.
- Check every applicable sync/async, offline/online, streaming/non-streaming,
feature-off, topology, and compatibility path without demanding unsupported
modes merely for symmetry.
## Module, API, and class design
Apply these rules to new or materially changed production code; do not require
unrelated cleanup from untouched files.
1. Start each source module with a concise module docstring that defines its
responsibility and boundaries and identifies its primary classes and key
functions. Place it after any required license header and before imports.
2. Give every added or materially changed class and function a clear docstring
covering its responsibility, parameters, return value, raised exceptions,
and non-obvious invariants. Use explicit typed parameters at owned API
boundaries; avoid opaque `**kwargs` and `Any`. Allow them only when an
upstream or compatibility interface requires them, and document and validate
the accepted keys and concrete types.
3. Move stable helpers shared by multiple features into an appropriately scoped
utility module. Do not let feature modules accumulate large collections of
unrelated private utility functions; keep a private helper local only when
it is single-purpose and tightly coupled to that feature.
4. Keep every method on a key class within that class's stated responsibility,
and document non-obvious design decisions and invariants. Before adding a
method, prefer reuse, consolidation, or extension of an existing abstraction;
split the class when new behavior creates a separate responsibility or makes
the class unnecessarily long.
5. Require changes to user-facing behavior to update `docs/`. When an existing
active module or feature design page owns the changed contract, update that
page. Require a new design page only when the reviewed branch's repository
policy, an accepted RFC/issue, or explicit maintainer acceptance criteria do
so; otherwise suggest it as non-blocking. For a major architectural change,
apply the branch's contributing policy for a design RFC or issue; the current
policy uses more than 500 changed production-code lines, excluding kernel,
data, config, and test changes. Resolve page status and authority through
[design-contracts.md](design-contracts.md), and do not demand a duplicate page
when an existing design owns the exact contract. When required, cover scope,
data flow, public contracts, failure behavior, alternatives, validation,
limitations, and rollout.
6. Require every performance PR to provide reproducible base-versus-head A/B
results under the same hardware, software, model, workload, precision,
topology, warmup, repetitions, and measurement scope. Include exact commands,
variability, and a correctness or quality guard; one-sided or non-comparable
numbers are incomplete evidence. For a new model absent from the frozen base,
compare head with its pinned canonical reference implementation under rule 12
instead. Follow
[perf-verification.md](../checks/perf-verification.md) for the full protocol.
7. Require every addition, removal, default change, or semantic change to a
user-facing inference API or CLI to explain why it is necessary, which
alternatives were considered, and why the chosen interface is preferable.
Record compatibility, deprecation or migration behavior, and the rationale
in the PR. Update the corresponding module and feature design documents when
they define the changed interface or workflow, or when policy requires one.
8. Require every bug-fix PR to add an automated regression test that reproduces
the defect on the frozen base and passes on the fixed head. Exercise the
production path or nearest stable boundary and assert the corrected behavior,
not only process survival. Record the exact base/head command and result, and
place hardware- or model-specific coverage in the correct CI lane. A manual
reproduction supplements but does not replace the regression test.
9. Require every new or materially changed recipe to follow
[`recipes/TEMPLATE.md`](../../../../../recipes/TEMPLATE.md) and the naming and
layout conventions in
[`recipes/README.md`](../../../../../recipes/README.md).
Complete every applicable summary, use-case, reference, hardware,
environment, exact-command, verification, and notes section with tested,
reproducible details; explain any intentional deviation in the PR.
10. Require every new model to verify its supported modalities, platforms,
execution modes, and features with production-path evidence. Update
`docs/models/supported_models.md` and every applicable feature support or
compatibility table; do not claim support that the PR does not validate.
Follow
[model-addition-checklist.md](../checks/model-addition-checklist.md) for the
model-specific documentation checks.
11. Require every new model PR to add or update its in-repository model-family
recipe under `recipes/` and its `recipes/README.md` index entry. Apply rule
9, and keep the recipe's model identifiers, tasks, serving modes, hardware,
commands, supported features, and limitations aligned with validated PR
evidence. An external recipe link does not replace the in-repository recipe.
12. Require every new model to compare accuracy or output quality and performance
against a pinned baseline or canonical reference implementation. Match the
checkpoint, inputs, precision, hardware, workload, and measurement method;
provide exact commands, results, and explanations for material differences.
Prefer a detailed timing split for each stage and explain when it is
unavailable. Follow
[model-addition-checklist.md](../checks/model-addition-checklist.md) and
[perf-verification.md](../checks/perf-verification.md) for the evidence
contract.
## Blocking risk scan
| Risk | Prove before reporting |
| --- | --- |
| Correctness | A live input reaches a wrong output, exception, silent drop, shape/dtype/device mismatch, or partially initialized state. |
| Compatibility | A supported API, config key, default, serialization schema, model identifier, or caller breaks without validation, migration, or deprecation. |
| Lifecycle | Allocation or startup lacks cleanup on partial failure, timeout, cancellation, shutdown, or repeated requests. |
| Concurrency/distribution | Ordering, identity, rank, world size, completion, timeout, backpressure, or collective assumptions fail on a supported path. |
| Cache/shared state | Identity omits a correctness-affecting input, or isolation, invalidation, disabled behavior, eviction, or cleanup is incomplete. |
| Async behavior | Blocking I/O, sleep, a lock across `await`, or unnecessary serialization harms a reachable async path. |
| Security/data handling | Secrets, user payloads, unsafe deserialization/eval/shell, unvalidated paths, or unbounded metric labels are reachable. |
| Validation evidence | Tests bypass the production dispatcher, mock away the change, use unrealistic types/MRO, or cannot fail when the defect returns. |
| User contract | User-visible model, feature, API, CLI, config, default, or compatibility behavior changes without accurate docs or examples. |
Search bounded callers and sibling implementations before reporting dead code,
duplication, or a missed consumer. Keep new abstractions only when they have a
distinct live caller, invariant, or compatibility need.
## Finding bar
Anchor each finding to the exact `path:line` that causes the behavior; it may be
unchanged when the diff makes it newly load-bearing, in which case name the
changed call path that brings it into scope. Include the trigger, impact, and
smallest fix direction. Treat pending CI, missing hardware, or unsupported
measurements as validation gaps unless the repository contract makes that
evidence a merge requirement.
@@ -0,0 +1,195 @@
# Review Execution and Delivery
Use this reference for every review. It defines snapshot collection, gates,
validation records, delivery permissions, maintainer tone, and re-review.
Official docs: [contributing guide](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/)
and [CI failure triage](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/failures/).
## Contents
- [Freeze the review surface](#freeze-the-review-surface)
- [Report status before analysis](#report-status-before-analysis)
- [Apply review gates](#apply-review-gates)
- [Run bounded validation](#run-bounded-validation)
- [Deliver maintainer-style findings](#deliver-maintainer-style-findings)
- [Re-review safely](#re-review-safely)
## Freeze the review surface
For a GitHub PR, read the base/head SHA before and after fetching metadata and
the diff:
```bash
gh api "repos/vllm-project/vllm-omni/pulls/<PR>" \
--jq '{base_sha: .base.sha, head_sha: .head.sha}'
REVIEW_FIELDS="number,url,title,body,isDraft,baseRefName,headRefName,mergeable,mergeStateStatus,statusCheckRollup,files"
gh pr view <PR> --repo vllm-project/vllm-omni \
--json "${REVIEW_FIELDS}"
gh pr diff <PR> --repo vllm-project/vllm-omni
gh api "repos/vllm-project/vllm-omni/pulls/<PR>" \
--jq '{base_sha: .base.sha, head_sha: .head.sha}'
```
Discard the snapshot if either SHA changed. Do not mix comments or validation
from different heads.
After the SHAs stabilize, fetch and materialize a trusted `head_sha` in an
isolated detached worktree. Run every source read, `rg` search, import, and test
from that snapshot, not the caller's checkout. A detached worktree provides
snapshot isolation, not a security boundary.
Treat a fork head as untrusted unless the user and environment policy explicitly
establish trust. Never run its imports, tests, builds, hooks, package setup, or
repo-configurable linters/plugins on the reviewer host. Execute them only in a
disposable sandbox or VM with no credentials or inherited secrets, no host agent
or service sockets, minimal read-only host mounts, disabled or explicitly
allowlisted network, resource/time limits, and destruction after the run. If
that boundary is unavailable, limit the review to the remote diff, SHA-addressed
reads such as `git show <head_sha>:<path>`, and existing CI evidence; report all
executable validation as a gap.
Before the first source read, record a pristine snapshot fingerprint outside
the reviewed worktree. Include `HEAD`, the NUL-delimited index entries and blob
IDs, and a NUL-safe manifest of every worktree entry except Git metadata,
including tracked, untracked, and ignored paths with file type, mode, symlink
target, and content hash. Before each validation group and delivery, recompute
and byte-compare the fingerprint as well as asserting `HEAD == head_sha`. If a
tool changes or creates any entry, discard affected evidence and recreate the
snapshot before another group; do not rely on `HEAD` alone or clean an unknown
worktree in place. If an exact snapshot cannot be created, use SHA-addressed
reads and report filesystem-dependent validation as a gap.
For a local branch/worktree, determine the target ref from the user, current PR,
or configured upstream; never infer it from the branch name. Resolve the target
once and include all task-owned worktree state:
```bash
git status --porcelain=v2 -z
git rev-parse HEAD
git rev-parse <target-ref>
git merge-base <target-base-sha> HEAD
git diff --stat <comparison-commit>
git diff --name-status <comparison-commit>
git diff --binary <comparison-commit>
git diff --cached --binary <comparison-commit>
git diff --binary
git ls-files --others --exclude-standard -z
```
Freeze the exact bytes of every in-scope untracked file from the NUL-delimited
list, plus a NUL-safe path, file-type, mode, and content-hash manifest. Also
fingerprint `HEAD`, the target and merge base, porcelain status, and both binary
index and worktree patches. Recording only untracked names is insufficient.
Keep these snapshots in the evidence packet and do not mutate the reviewed
checkout during a read-only review.
For local changes without PR metadata, mark the title/body as synthetic. Base
them on the user request and commit messages without inventing validation claims.
If no PR, branch, or usable worktree is identifiable, ask for a PR URL/number or
explicit base and head rather than guessing.
## Report status before analysis
Within 60 seconds of starting, and before source searches or tests, send this
host update:
```text
Pinned head: <SHA>
Base/comparison: <ref and SHA>
CI: <pass/fail/pending/not applicable>
Mergeability: <state/not applicable>
Preliminary findings: <brief finding or none yet>
```
Mark early findings as preliminary and continue. This is not a GitHub comment.
## Apply review gates
- Report draft/WIP state. Continue a local review when explicitly requested,
but do not publish review events without separate authorization.
- Record DCO, pre-commit, required CI, and mergeability. Pending or unknown gates
do not block source review.
- Treat a failed gate as its own evidence. Do not restate its formatting/lint
output as a new code-review finding.
- Open CI logs only when the first failing step overlaps the frozen diff or
blocks the verdict. Start with the first error, not the last cascade.
- Treat the PR body as navigation. Commands, benchmark tables, and claimed tests
become evidence only after their provenance and relevance are checked.
## Run bounded validation
Keep one evidence packet for files read, bounded searches, callers, tests, CI,
hardware, routes, and findings. Reuse it rather than refetching facts.
Before pytest, run a short import/version compatibility check. Record each
validation result with:
```text
repo, head SHA, command, result, Python/platform, dependency or lock fingerprint
```
After preflight, run targeted tests and low-cost static checks alongside source
inspection when possible. Imports, tests, builds, and linters may create caches
or rewrite files, so isolate each potentially mutating group in a disposable
snapshot or recreate the pinned snapshot before continuing. Map source symbols
to tests with bounded `rg` searches rather than assuming the test directory
mirrors production paths.
Classify failures as code, test, infrastructure, or flaky before reporting.
Skipped hardware tests are gaps, not passes. For available hardware, verify the
smallest representative unit/E2E path and compare actual output with PR claims.
## Deliver maintainer-style findings
Default to roughly 1-5 short comments, but treat that as calibration rather than
a quota: report every real blocker and report no finding when there is no issue.
Write each finding as:
```text
[P1] Short imperative title — path/to/file.py:<line>
<Trigger or call path>. <Current behavior and impact>. <Smallest fix direction>.
```
Use [maintainer-style-study.md](../delivery/maintainer-style-study.md) for
comment tone.
Keep rule IDs, grades, and audit matrices internal unless the user asks for the
complete audit.
Verify every inline `path:line` against the frozen diff before delivery. For a
PR, re-read the remote head and byte-compare the detached snapshot fingerprint,
or recreate it from `head_sha`, before asserting current lines. For a local
review, recompute and byte-compare the frozen `HEAD`, target, merge base, status,
index patch, worktree patch, and untracked-content manifest. Any mismatch makes
the review stale; discard the affected evidence and restart. Prefer one
root-cause comment to several symptom comments.
When the root cause is an unchanged line made newly load-bearing by the diff,
cite it in the consolidated review body instead of moving the comment to an
unrelated changed line merely to make it inline.
Local presentation is the default. Only explicit authorization permits GitHub
posting. When authorized, publish one consolidated final review after
validation; do not post preliminary or incremental comments. Submit a review
event (`APPROVE`, `COMMENT`, or `REQUEST_CHANGES`) only when the user explicitly
chooses that event.
Reviewer requests and owner `@mention` comments are separate external writes;
follow [review-requests.md](../delivery/review-requests.md) when the user asks
for them.
## Re-review safely
Freeze the new head and compare it with the previously reviewed SHA. Then:
1. inspect the delta and every file affected by conflict resolution or rebase;
2. revalidate earlier findings against current lines and behavior;
3. read unresolved, non-outdated threads and the author's evidence;
4. rerun only checks invalidated by the delta;
5. look for new regressions introduced by the fix;
6. do not repeat resolved or outdated comments.
If the target head changes during re-review, discard stale validation and start
again from the new snapshot.
@@ -0,0 +1,77 @@
# Module and Feature Review Routing
Read [design-contracts.md](design-contracts.md) first. Route from the frozen
head's live producer-consumer behavior, then confirm the result against its
`docs/design/` metadata. Titles and paths are evidence, not the deciding rule.
## Primary module contract
Choose one primary module and a second only when the live call path crosses a
documented dependency or primary-path exception.
| Module | Live behavior signals | Read |
| --- | --- | --- |
| Entrypoints | Offline/CLI/API ingress, validation, rendering, streaming, sessions | [entrypoints.md](../modules/entrypoints.md) |
| Configuration | Schema, defaults, deploy/stage construction, registry, topology | [configuration.md](../modules/configuration.md) |
| I/O and Modality | Requests, messages, serialization, output types, accumulation, completion | [input-output-modality.md](../modules/input-output-modality.md) |
| Error Contracts | Classification, fatality, propagation, sanitization, public rendering | [error-contracts.md](../modules/error-contracts.md) |
| Engine Orchestration | Cross-stage routing, request state, ordering, RPC correlation, terminal convergence | [engine-orchestration.md](../modules/engine-orchestration.md) |
| Stage Runtime | Placement, startup, readiness, replicas, affinity, membership, shutdown | [stage-runtime.md](../modules/stage-runtime.md) |
| OmniConnector | Cross-stage/process/device/node transport and synchronization | [omni-connector.md](../modules/omni-connector.md) |
| Model Integration | Registration, preprocessing, loading, runners, model-specific execution | [model-integration.md](../modules/model-integration.md) |
| AR Runtime | Scheduling, request/cache state, adapters, workers, upstream vLLM semantics | [ar-runtime.md](../modules/ar-runtime.md) |
| Diffusion Family | Diffusion runtime, models, batching, parallelism, or offload | [diffusion.md](../modules/diffusion.md) |
| Execution Platforms | Hardware selection, capabilities, vendor workers, kernels, patches | [execution-platforms.md](../modules/execution-platforms.md) |
| Cache Management | Cache identity, validity, reuse, eviction, reset, teardown | [cache-management.md](../modules/cache-management.md) |
| Quantization | Method/metadata selection, layer mapping, reduced precision, constraints | [quantization.md](../modules/quantization.md) |
| Observability | Metrics, logs, labels, units, correlation, lifecycle | [observability.md](../modules/observability.md) |
| Profiling | Opt-in instrumentation, traces, start/stop lifecycle, overhead | [profiling.md](../modules/profiling.md) |
| Benchmarking | Workloads, metric computation, benchmark CLI, result metadata | [benchmarking.md](../modules/benchmarking.md) |
For tests-, docs-, CI-, or recipe-only changes, route to the production module
whose contract they protect. When none exists, use the applicable evidence
check without inventing a production owner.
## Feature-design overlays
Load every category whose documented behavior, enablement, default, support, or
compatibility changes.
| Feature-design section | Read |
| --- | --- |
| Runtime and stage execution | [runtime-stage-execution.md](../features/runtime-stage-execution.md) |
| Communication and concrete connectors | [communication.md](../features/communication.md) |
| Diffusion acceleration | [diffusion-acceleration.md](../features/diffusion-acceleration.md) |
| Infrastructure and performance | [infrastructure-performance.md](../features/infrastructure-performance.md) |
## Evidence and change overlays
| Signal | Read | Optional repo-local skill |
| --- | --- | --- |
| New or expanded model, loader, processor, registry, or stage config | [model-addition-checklist.md](../checks/model-addition-checklist.md) | [`add-tts-model`](../../../add-tts-model/SKILL.md) or [`add-diffusion-model`](../../../add-diffusion-model/SKILL.md) |
| Latency, throughput, memory, scaling, precision, or quality claim | [perf-verification.md](../checks/perf-verification.md) | [`diffusion-perf-opt`](../../../diffusion-perf-opt/SKILL.md) for diffusion |
| Config/CLI/environment/default precedence; endpoint, host, port, address, connector metadata, or `*_extra_config` | [config-provenance.md](../checks/config-provenance.md) | None |
| Tests changed, absent for risky behavior, or test-only | [test-quality-evaluation.md](../checks/test-quality-evaluation.md) | [`vllm-omni-test`](../../../vllm-omni-test/SKILL.md) |
| CI, examples, docs, public behavior, or contributor evidence | [tests-docs-checklist.md](../checks/tests-docs-checklist.md) | None |
| Suitable hardware/server and runnable affected path | [verification.md](../checks/verification.md) | None |
| User asks who should review or requests owner notification | [review-requests.md](../delivery/review-requests.md) | None |
Use [`quantization`](../../../quantization/SKILL.md) for a quantization module
change and [`$vllm-omni-npu-model-runner-upgrade`](../../../vllm-omni-npu-upgrade/SKILL.md)
for an NPU execution-platform change when those skills are available.
For a bug fix, require a reachable reproduction and regression test. For a
refactor, prove parity and remove obsolete paths. For a feature, verify the
public/config contract, module integration, compatibility/default behavior,
production dispatch, feature design, and docs.
## Calibrate findings
- **P0:** security exposure, data corruption, or broad project unusability.
- **P1:** reachable runtime failure, wrong output, compatibility break, or
unsafe lifecycle in the changed behavior.
- **P2:** real non-blocking defect with a concrete future failure mode.
Treat draft candidate invariants, missing hardware, and unsupported claims as
questions or validation gaps unless current code, tests, or policy makes them a
merge requirement.
+7
View File
@@ -61,6 +61,13 @@ repos:
types: [python]
files: ^tests/
- id: check-tts-adapter-migration
name: Ratchet TTS per-model branches out of serving_speech.py
entry: python tools/pre_commit/check_tts_adapter.py
language: python
pass_filenames: false
files: ^vllm_omni/entrypoints/openai/(serving_speech\.py|tts_adapters/__init__\.py)$
# Keep `suggestion` last
- id: suggestion
name: Suggestion
+1 -1
View File
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/111-L8zF7A1j_YI_cR8JsblofdScdRr2f/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://deepwiki.com/vllm-project/vllm-omni"><b>DeepWiki</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing"><b>Slides</b></a> |
</p>
+1 -1
View File
@@ -63,5 +63,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
For more information, checkout the following:
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1aPj0OGl_-ZVoib-Qne5dGDAlrRFB-PdHl6E-EE99g8E/edit?usp=sharing)
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
+110 -63
View File
@@ -672,100 +672,147 @@ down and spawn a new one mid-module. A few rules that save real CI debugging tim
## Online Serving Integration
To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add
**all five** of the following integration points to
`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them
piecemeal causes partial-integration failures that are hard to debug.
To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, write
one adapter under `vllm_omni/entrypoints/openai/tts_adapters/`. The adapter declares
how your model is recognized and owns its request handling; `serving_speech.py` should
not need an edit.
### 1. Stage constant
### 1. Write the adapter
Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants:
Create `vllm_omni/entrypoints/openai/tts_adapters/your_model.py`:
```python
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"}
# SPDX-License-Identifier: Apache-2.0
"""YourModel serving adapter."""
from typing import TYPE_CHECKING
from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
if TYPE_CHECKING:
from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
@register_tts_adapter
class YourModelAdapter(ARTTSAdapter):
# The registry key, and the model-type string used in logs.
name = "your_model"
# The engine ``model_stage`` value(s) your deployment yaml sets. Detection,
# stage discovery and `/v1/audio/speech` routing all derive from this.
stage_keys = frozenset({"your_model_stage_key"})
def validate(self, request: "OpenAICreateSpeechRequest") -> str | None:
"""Return an error string, or None if the request is valid."""
if not request.input or not request.input.strip():
return "Input text cannot be empty"
return None
async def build(
self,
request: "OpenAICreateSpeechRequest",
sampling_params_list: list,
has_inline_ref_audio: bool,
) -> PreparedRequest:
"""Build the engine prompt and per-model parameters."""
params: dict[str, list] = {"text": [request.input]}
if request.voice is not None:
params["voice"] = [request.voice]
return PreparedRequest(
prompt={"prompt": request.input},
tts_params=params,
model_type=self.name,
)
```
### 2. Union into `_TTS_MODEL_STAGES`
Then add it to the import block at the bottom of
`vllm_omni/entrypoints/openai/tts_adapters/__init__.py` so it registers on import.
That import line is the only shared file a new model has to touch.
Add to the `_TTS_MODEL_STAGES` set union:
> **Pure-diffusion TTS does not go through adapters yet.** When the server is built
> by `OmniOpenAIServingSpeech.for_diffusion()`, `create_speech()` routes straight to
> `_create_diffusion_speech()` and never calls `validate()` or `build()`.
> `DiffusionTTSAdapter` is scaffolding for that migration and has no production
> subclass today, so validation or prompt construction placed in one would silently
> never run. If your model is served by the diffusion engine rather than
> `engine_client`, follow the existing diffusion path and say so on #4855 — wiring
> that path through the adapter layer is open work, not something to do inside a
> new-model PR.
### 2. If a stage key alone cannot identify your model
Two escape hatches, in order of preference:
```python
_TTS_MODEL_STAGES: set[str] = (
...
| _YOUR_MODEL_TTS_MODEL_STAGES
)
# The architecture is authoritative (e.g. your model has no dedicated
# model_stage value, or shares one with another model).
model_archs = frozenset({"YourModelForConditionalGeneration"})
# Set this too if your model owns NO stage key, so stage discovery has to find
# your AR entry stage by architecture.
arch_identifies_entry_stage = True
```
### 3. Model type detection
In `_detect_tts_model_type()`, add before the final `return None`:
If your rule is not a set-membership test, override `matches()` — see
`covo_audio.py`, which shares the generic `fused_thinker_talker` stage key with
non-CoVo deployments and confirms with the architecture:
```python
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
return "your_model"
@classmethod
def matches(cls, model_stage: str | None, model_arch: str | None) -> bool:
return super().matches(model_stage, model_arch) and bool(model_arch) and "CovoAudio" in model_arch
```
### 4. Request validation dispatch
If your model can match the same deployment as another adapter, give one of them an
explicit `detect_priority` (lower runs first). `tests/entrypoints/openai_api/
test_tts_detection.py` fails if two same-priority detectors can both match one input,
so an unordered overlap cannot be merged.
In `_validate_tts_request()`, add before the fallback `return`:
Models that are speech-capable only in some deployment topologies override
`stage_serves_speech()` — see `audex.py`, whose omni thinker is text-final unless the
speech decoder is deployed alongside it.
```python
if self._tts_model_type == "your_model":
return self._validate_your_model_request(request)
```
### 3. Reuse the shared helpers
### 5. Validation and parameter-builder methods
Adapters reach shared serving helpers through `self.ctx.server`: reference-audio
resolution (`_resolve_ref_audio`), uploaded-speaker handling
(`_apply_uploaded_speaker`), format validation (`_validate_ref_audio_format`), and
prompt-length limits (`_max_instructions_length`). Read a comparable adapter before
writing your own — `fish_speech.py` for voice cloning, `higgs_audio_v3.py` for
parameter-heavy models, `moss_tts.py` for a family sharing one base class.
Add two new methods:
### 4. Test it
```python
def _validate_your_model_request(
self, request: OpenAICreateSpeechRequest
) -> str | None:
"""Validate YourModel request. Returns an error string or None."""
if not request.input or not request.input.strip():
return "Input text cannot be empty"
return None
`tests/entrypoints/openai_api/test_tts_adapter.py` covers registry invariants and
`test_tts_detection.py` covers detection. Add your model to
`EXPECTED_MODEL_TYPES` in the former, and a case to the latter if you added a
`model_archs` or `matches()` rule.
def _build_your_model_params(
self, request: OpenAICreateSpeechRequest
) -> dict[str, Any]:
"""Build additional_information dict for YourModel."""
params: dict[str, Any] = {"text": [request.input]}
if request.voice is not None:
params["voice"] = [request.voice]
# Add any other model-specific fields here
return params
```
### Do not add branches to `serving_speech.py`
Then wire `_build_your_model_params` into the request-dispatch block in
`_create_tts_request()` (search for the equivalent `_build_*_params` call for an
existing model to find the right location). If the model supports voice cloning
(`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings
here too — follow any existing `_build_<model>_params` in `serving_speech.py` (e.g.
`_build_moss_tts_params` for the voice-cloning variant) for the pattern.
Older models predate the adapter framework and still have per-model branches in
`serving_speech.py`; they are being migrated out (RFC #4327, #4855). Do not copy
that pattern. `tools/pre_commit/check_tts_adapter.py` is a ratchet on the remaining
`self._tts_model_type == ...` comparisons and fails the commit if the count grows.
> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean
> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type`
> string pattern shown above. For new models, always use the string pattern — do not
> add new `_is_*` boolean flags.
> **Note on unused variables:** Only extract parameters in `_build_your_model_params`
> that you actually pass to the model's generate / `inference_stream` call. Extracting
> a variable without forwarding it will trigger a `ruff F841` pre-commit failure.
If your model genuinely needs behaviour that no adapter hook can express, that is a
missing hook — propose it on the RFC rather than adding a branch.
### Merge conflicts
`serving_speech.py` is modified by every new model PR and is the most common source of
rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution
is always to **keep both** the upstream model's additions and your own — never discard
either side. After resolving:
Your adapter is a new file, so it cannot conflict. The one shared line is your entry
in the import block at the bottom of `tts_adapters/__init__.py`, which conflicts only
when another model lands at the same time. The resolution is always to **keep both**
imports — never discard either side. After resolving:
```bash
git add vllm_omni/entrypoints/openai/serving_speech.py
git add vllm_omni/entrypoints/openai/tts_adapters/__init__.py
git rebase --continue
```
If you find yourself resolving a conflict inside `serving_speech.py`, something in
your model is still going through the legacy path; see the note above.
## Single-Stage Models
Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the
+406 -144
View File
@@ -1,6 +1,6 @@
# Architecture Overview
This document outlines the architectural design for vLLM-Omni.
This document outlines the architecture of vLLM-Omni.
<p align="center">
<picture>
@@ -9,194 +9,456 @@ This document outlines the architectural design for vLLM-Omni.
</picture>
</p>
# Goals
## Design goals
The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
The primary goal of vLLM-Omni is to provide a fast and easy-to-use inference
and serving engine for omni-modality models. vLLM-Omni extends vLLM's
text-oriented autoregressive (AR) runtime with stage-based execution for
non-textual outputs and non-autoregressive model components.
* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text.
* **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
* **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
* **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
The architecture is designed to:
* support text, image, audio, video, and action inputs and outputs;
* compose autoregressive, generation, and diffusion stages in one pipeline;
* reuse vLLM scheduling, cache, distributed-execution, and serving
primitives where they fit; and
* keep model topology, deployment placement, runtime lifecycle, transport,
and public API concerns in separate layers.
# Representative omni-modality models
## Model execution
According to analysis for current popular open-source models, most of them have the combination of AR+DiT. Specifically, they can be further categorized into 3 types below:
These four examples show how vLLM-Omni maps model components to stage-local
execution policies. Each summary covers the model-owned path, batching,
attention, parallelism, quantization, and primary serving targets. The profiles
are representative; a model may support multiple `PipelineConfig` and
`DeployConfig` combinations.
**DiT as a main structure, with AR as text encoder (e.g.: Qwen-Image)**
A powerful image generation foundation model capable of complex text rendering and precise image editing.
| Model | Model-owned path | vLLM-Omni stage view | Primary serving target |
| --- | --- | --- | --- |
| Qwen3-Omni | Thinker → Talker → Code2Wav | Three stages with asynchronous chunk hand-off | TTFT/TTFP, TPOT, E2EL, audio RTF, and concurrent throughput |
| HunyuanImage-3.0 | Multimodal AR understanding/reasoning plus image-generation DiT | AR-only, DiT-only, or split AR → DiT deployment | DiT E2EL/denoising latency and image throughput; TTFT/TPOT for AR tasks |
| MiniMax-H3 | Text encoder → task-specific FL2VA or Ref2VA DiT → video/audio VAE decoder | Three logical component/phase boundaries within one registered diffusion stage | Video/audio E2EL and media throughput; RTF when audio streaming is measured |
| Cosmos3 | Unified MoT reasoner tower plus diffusion generator tower | One diffusion pipeline materializes both towers inside the stage | Image/video/action E2EL and throughput; RTF for synchronized audio |
### Serving metric vocabulary
The target metric follows the output modality, not just the model name.
| Metric | Meaning in this document | Most useful for |
| --- | --- | --- |
| **TTFT** | Time to the first text token | AR reasoning or text output |
| **TPOT** | Time per output token after the first token | AR decode stages |
| **TTFP** | Time to the first streamed media packet, such as audio | Interactive speech and multimodal chat |
| **E2EL** | End-to-end latency from request admission to the final output | Image, video, audio, and action requests |
| **RTF** | Wall-clock processing time divided by generated audio duration; lower is faster than real time | Audio generation |
| **Throughput** | Requests, tokens, or media seconds produced per wall-clock second | Offline and concurrent serving |
### Qwen3-Omni: streaming AR pipeline
Qwen3-Omni illustrates a multi-stage streaming AR pipeline: Thinker → Talker →
Code2Wav. The [technical report](https://arxiv.org/abs/2509.17765) describes
the model components, and the figure shows the stage-level streaming schedule
used by vLLM-Omni.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png">
<img alt="Qwen-Image" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/dit-main-architecture.png" width=30%>
</picture>
<a href="https://github.com/vllm-project/vllm-omni/blob/main/docs/source/architecture/qwen3-omni-async-chunk.png">
<img alt="Qwen3-Omni Thinker, Talker, and Code2Wav streaming stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width="80%">
</a>
</p>
**AR as a main structure, with DiT as multi-modal generator (e.g. BAGEL)**
A unified multimodal comprehension and generation model, with cot text output and visual generation.
#### Model architecture by stage
| Stage | Model component | Stage output |
| --- | --- | --- |
| **Thinker** | AuT audio encoder, SigLIP2 vision encoder, and the 30B-total/3B-active Thinker MoE Transformer | Text tokens and conditioning for the Talker |
| **Talker** | 3B-total/0.3B-active Talker MoE Transformer with the MTP code-prediction path | Text tokens and audio codec tokens |
| **Code2Wav** | Approximately 200M-parameter ConvNet audio decoder | Waveform/audio chunks |
#### Representative stage execution
| Stage | Batching | Attention and execution | Parallelism | Quantization |
| --- | --- | --- | --- | --- |
| **Thinker** | Continuous batching with stage-local `max_num_seqs` and token budgets | vLLM KV-cached AR attention; CUDA Graph execution when the stage is not eager | Independent stage placement; use the AR runtime's configured tensor parallelism | Supported ModelOpt FP8/NVFP4 or AutoRound checkpoint paths target the Thinker; encoders remain BF16 in the documented paths |
| **Talker** | Continuous batching over autoregressive decode requests | KV-cached token decode with async-chunk output to Code2Wav | Independent stage placement; typically a separate GPU or replica from Thinker | BF16 baseline; no generic Talker quantization path is assumed |
| **Code2Wav** | Static/chunk batching for codec-to-waveform decode | Non-AR chunked audio decode; eager execution is commonly retained | Separate stage or colocated placement, depending on the deployment | BF16 baseline; keep codec/vocoder precision model-specific |
For this pipeline, the main serving target is interactive responsiveness:
TTFT for text, TTFP for the first audio packet, TPOT for autoregressive decode,
E2EL for the complete answer, and RTF/throughput for sustained audio serving.
Async chunking and streaming primarily reduce TTFP and improve stage overlap;
batching and CUDA Graphs primarily improve E2EL and throughput. See the
[Qwen3-Omni recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md)
and [async chunk design](feature/async_chunk.md) for the deployment details.
### HunyuanImage-3.0: shared multimodal model, split deployment choices
HunyuanImage-3.0 illustrates a shared multimodal model with flexible AR, DiT,
and split AR → DiT deployment. The official [technical
report](https://arxiv.org/abs/2509.23951) provides the model view; the tables
below map it to vLLM-Omni stages.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png">
<img alt="Bagel" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-main-architecture.png" width=30%>
</picture>
<a href="https://github.com/Tencent-Hunyuan/HunyuanImage-3.0/blob/main/assets/framework.png">
<img alt="HunyuanImage-3.0 framework for understanding, language modeling, and image generation" src="https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/framework.png" width="95%">
</a>
</p>
**AR+DiT (e.g. Qwen-Omni)**
A natively end-to-end omni-modal LLM for multimodal inputs (text/image/audio/video...) and outputs (text/audio...).
#### Model architecture by stage
| Stage or component | Model component | Stage output |
| --- | --- | --- |
| **AR understanding/reasoning** | Understanding and generation encoders feeding the Hunyuan-A13B decoder-only Transformer and text detokenizer | Text responses or image-generation conditioning |
| **Image generation** | Generation encoder, diffusion prediction path/Gen. Decoder, and the image VAE decode path | Image latents and final images |
The vLLM-Omni deployment may expose the AR path only, the DiT path only, or
both as an AR → DiT pipeline with a connector between them. That is a runtime
decomposition of the model's capabilities; it should not be read as a claim
that the published framework figure contains two unrelated backbones.
#### Representative stage execution
| Stage | Batching | Attention and execution | Parallelism | Quantization |
| --- | --- | --- | --- | --- |
| **AR** | Standard vLLM continuous batching when serving understanding/text tasks | KV-cached causal attention | Tensor/expert parallelism follows the selected AR deployment | Checkpoint-specific; do not infer DiT FP8 settings for the AR checkpoint |
| **DiT + VAE** | Request batching or step-wise batching; Hunyuan step batching with more than one request requires `TORCH_SDPA` | `TORCH_SDPA` is used because the model mixes causal and full attention | Validated profiles include TP4, TP2 + sequence parallelism, TP2 + CFG parallelism, and expert parallelism for MoE weights | Online FP8 and ModelOpt mixed FP8/NVFP4 are documented for the DiT; VAE precision stays separate |
Image generation is therefore measured mainly with per-image E2EL/denoising
latency, peak memory, and images-per-second throughput. TTFT and TPOT remain
meaningful for the optional AR understanding path, but they are not the right
headline metrics for the DiT path. See the
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
for the validated batching, attention, and quantization combinations.
### MiniMax-H3: conditioning, denoising, and decode
MiniMax-H3 illustrates logical component and phase boundaries between text
encoding, task-specific denoising, and media decoding:
`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
these components in one diffusion stage. These boundaries are useful for
profiling and offload; independent placement would require a future
multi-stage `PipelineConfig`. The stage mapping follows the implementation and
serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
notes that a detailed technical report is forthcoming.
#### Model architecture by component
| Logical component | Model component | Component output |
| --- | --- | --- |
| **Text Encoder** | Tokenizer/processor and Qwen3-VL text encoder; vision/audio references are prepared alongside the text condition | Packed multimodal conditioning for denoising |
| **DiT** | Task-specific `FL2VA` DiT for text/first-frame generation or `Ref2VA` DiT for mixed references | Video and audio latent tokens |
| **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
Both DiT partitions remain behind one task-selection point; `task` chooses
which partition runs. In the current implementation, the three logical
components execute inside one diffusion stage. This is different from
Qwen3-Omni's three independently scheduled stages, even though both models
produce synchronized audio and video-like outputs.
#### Representative component execution
| Logical component | Batching | Attention and execution | Parallelism | Quantization |
| --- | --- | --- | --- | --- |
| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
| **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
| **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
The primary target is video/audio E2EL and media throughput, with the logical
component boundaries making prompt encoding, denoising, and VAE decode costs
separately measurable. RTF is useful when the audio stream is evaluated
independently. TTFT and TPOT do not describe the main H3 generation path
because the user-visible output is diffusion media rather than a token stream.
The
[MiniMax-H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md)
contains the hardware-specific profiles and warmup requirements.
### Cosmos3: unified MoT reasoner and generator
Cosmos3 illustrates a unified Mixture-of-Transformers (MoT) reasoner and
diffusion generator. The official [technical
report](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf)
and NVIDIA's [technical explanation](https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/)
provide the two-tower model view used below.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png">
<img alt="Qwen-Omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/ar-dit-main-architecture.png" width=30%>
</picture>
<a href="https://developer.nvidia.com/blog/develop-physical-ai-reasoning-world-and-action-models-with-nvidia-cosmos-3/">
<img alt="Cosmos3 autoregressive reasoner and diffusion generator towers" src="https://developer-blogs.nvidia.com/wp-content/uploads/2026/05/image-11-e1780000686151.webp" width="95%">
</a>
</p>
# vLLM-Omni main architecture
#### Model architecture by stage
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png">
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-main-architecture.png" width=55%>
</picture>
</p>
| Stage or component | Model component | Stage output |
| --- | --- | --- |
| **Reasoner tower** | Autoregressive VLM path with causal attention over language and optional vision/action context; runs once per generation request | Discrete reasoning/context tokens |
| **Generator tower** | Diffusion path with full attention over the generator tokens and the reasoner K/V context; runs for each denoising step | Image, video, audio, or action tokens |
| **Shared pipeline** | `Cosmos3OmniDiffusersPipeline` plus modality encoders/decoders and VAE paths | Final image/video/audio or action response |
## Key Components
The reasoner and generator are model-owned components, not two independent
vLLM-Omni stages by default. The current Cosmos3 pipeline materializes both
inside one diffusion stage; model-level CPU offload can swap the reasoner and
generator components between their phases.
| Component | Description |
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
| **OmniRouter** | provide an intelligent router for Omni-modality requests dispatch |
| **EntryPoints** | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
| **AR** | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management |
| **Diffusion** | natively implemented and optimized using acceleration components |
| **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages |
#### Representative stage execution
Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
| Stage | Batching | Attention and execution | Parallelism | Quantization |
| --- | --- | --- | --- | --- |
| **Reasoner + generator pipeline** | Use request batching or step execution only when the selected pipeline advertises that capability; otherwise use `max_num_seqs=1` | Causal attention for reasoner tokens and full attention for diffusion tokens; the validated recipe uses a platform-selected aiter FlashAttention path with `TORCH_SDPA` fallback | Ulysses/context parallelism or tensor parallelism for the transformer; VAE tiling and layerwise/model-level offload address activation and weight memory | Online FP8 is supported for the diffusion pipeline; AutoRound/checkpoint-specific paths and quality validation remain separate from the BF16 baseline |
For Cosmos3, generation E2EL and output throughput are the primary targets.
Action-policy workloads add a control-loop latency target, while synchronized
audio workloads add RTF. TTFT/TPOT are appropriate for a reasoner-only service,
but not for the full diffusion generation path. See the
[Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
and [diffusion execution modes](../user_guide/diffusion/execution_modes.md)
for the current runtime choices.
## Why the stage-based architecture follows from these examples
The four representative pipelines show that omni-modality serving is not one
execution problem. Autoregressive decoding, diffusion denoising, multimodal
encoding, and media decoding have different scheduling, attention, parallelism,
memory, and quantization requirements. The architecture therefore needs a
unit that can be optimized and resourced independently while remaining part of
one request pipeline.
| Observation | Design implication | vLLM-Omni response |
| --- | --- | --- |
| AR, diffusion, encoder, and decoder components have different execution loops | One scheduler and one execution policy cannot fit every component | Specialized AR, generation, and diffusion stage runtimes |
| A model can be deployed as AR-only, DiT-only, or AR → DiT | Model structure must be separate from deployment topology | `PipelineConfig` describes logical stages and relationships; `DeployConfig` describes placement and resources |
| Text encoders, DiTs, and VAE decoders have different batch and memory behavior | Batching, attention, parallelism, and quantization must be stage-local | Each stage carries its own runtime and engine configuration |
| MiniMax-H3 can co-locate or logically split its text encoder and VAE decoder | A stage boundary should be configurable rather than forcing a process boundary | Logical stages can be co-located or assigned independent runtime placement |
| Qwen3-Omni transfers hidden states and codec chunks across stages | Large intermediate data needs typed transport and synchronization | `OmniConnector` transports payloads and KV-cache data without owning model routing |
| Requests span multiple stages and may be cancelled or produce ordered streaming output | Cross-stage request state needs a separate control plane | `Orchestrator` owns correlation, cancellation, routing state, and output ordering |
| Qwen3-Omni optimizes first-packet latency while image/video models optimize E2EL and throughput | Performance objectives must be model- and stage-aware | TTFT/TPOT/TTFP, E2EL, RTF, and throughput are measured according to the output modality |
This leads to a precise definition of a stage: it is a logical execution unit
with its own model component, execution policy, resource ownership, and
performance objectives. A stage is not necessarily a separate process. Cosmos3
keeps its reasoner and generator inside one diffusion pipeline, while
MiniMax-H3 exposes useful logical boundaries between text encoding, denoising,
and VAE decoding. Both are valid mappings of the same abstraction.
The resulting separation can be summarized as:
```text
Model structure What is computed and how components are related
↓
PipelineConfig Which logical stages exist and how they are connected
↓
Stage runtime policy How each stage batches, attends, parallelizes, and quantizes
↓
DeployConfig Where stages run, with which devices, replicas, and connectors
```
`AsyncOmniEngine` and `Orchestrator` coordinate the pipeline, `OmniConnector`
transports stage outputs, and `StageRuntime` materializes the selected
placement. This separation lets one logical model support multiple serving
topologies without coupling model code to process layout or deployment
resources.
## System architecture
The figure shows the main components of the stage-based runtime: entrypoints,
the shared `AsyncOmniEngine`, request orchestration, stage lifecycle
management, AR and diffusion execution modules, model/layer operations,
connector transport, and multimodal outputs. `AsyncOmniEngine` is the
composition root between the entrypoints and the orchestrator.
```mermaid
flowchart TB
subgraph omni["vLLM-Omni"]
subgraph entry["Entry points"]
api["API server"]
client["Omni / AsyncOmni"]
end
engine["AsyncOmniEngine<br/>composition root and background loop"]
orchestrator["Orchestrator<br/>request state and cross-stage routing"]
runtime["StageRuntime<br/>placement, replicas, readiness, and lifecycle"]
ar["AR module<br/>vLLM engine, scheduler, cache, executor, worker, and model runner"]
diffusion["Diffusion module<br/>scheduler, worker, model runner, and pipeline"]
ops["Model / layer / ops<br/>attention, parallelism, and quantization"]
connector["OmniConnector<br/>payload and KV transport plus synchronization"]
outputs["Multimodal outputs<br/>streaming chunks and final artifacts"]
api --> engine
client --> engine
engine --> orchestrator
engine --> runtime
orchestrator --> runtime
runtime --> ar
runtime --> diffusion
ar -. uses .-> ops
diffusion -. uses .-> ops
ar -. transports through .-> connector
diffusion -. transports through .-> connector
ar --> outputs
diffusion --> outputs
end
```
### Key components
| Component | Responsibility |
| --- | --- |
| **Entrypoints** | Translate offline, CLI, OpenAI-compatible, and duplex requests into engine operations and render outputs back to public protocols. |
| **Configuration resolution** | Combines pipeline topology, deployment settings, model metadata, and user overrides into a validated control-plane configuration. |
| **AsyncOmniEngine** | Owns engine composition, the background event loop, stage initialization, request submission, and output collection. |
| **Orchestrator** | Owns cross-stage request state, stage-to-stage routing, request correlation, cancellation, and output ordering. It does not own model selection or deployment placement. |
| **StageRuntime** | Expands logical stages into local or distributed replicas, starts stage clients and processes, and manages readiness, affinity, failure, and shutdown. |
| **AR runtime** | Extends vLLM's scheduler, KV-cache, worker, and model-runner path for omni-modality inputs and inter-stage outputs. |
| **Diffusion runtime** | Schedules and executes denoising workloads through diffusion executors, workers, pipelines, acceleration backends, and output materialization. |
| **OmniConnector** | Transports stage payloads and KV-cache data and provides synchronization. Connectors transport data; they do not choose the next logical stage. |
| **Multimodal outputs** | `MultimodalPayload` separates tensor content from metadata, while `OmniRequestOutput` carries pipeline and diffusion results through the common output path. |
## Configuration and runtime resolution
The control plane has five conceptual layers. Authoring inputs are resolved
once into a complete, transport-safe configuration before runtime processes are
started. This keeps stage topology and model capabilities distinct from
deployment placement and from process-local engine objects.
The figure shows only the primary hand-off object for each layer; the detailed
inputs, fields, and ownership rules are described below.
```mermaid
flowchart TB
layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
layer2["Layer 2 · Resolve once<br/>StageConfigFactory.create_from_model()<br/>VllmOmniConfig.from_pipeline_config()"]
layer3["Layer 3 · Transport-safe control plane<br/>VllmOmniConfig"]
layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
layer1 --> layer2 --> layer3 --> layer4 --> layer5
```
The single resolution boundary is implemented by
`StageConfigFactory.create_from_model()` and
`VllmOmniConfig.from_pipeline_config()` in
[`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
The legacy `stage_args` YAML path remains only for models that have not yet
migrated to `PipelineConfig` and `DeployConfig`.
In the typed path, each stage configuration derives from
`BaseVllmOmniStageConfig` and is specialized as
`VllmOmniARStageConfig`, `VllmOmniGenerationStageConfig`, or
`VllmOmniDiffusionStageConfig`. The request and engine-spec fields belong to the
control-plane boundary; the current implementation stores their equivalent
projections in the structured stage configuration and materializes
backend-specific engine objects during stage initialization.
The important ownership rules are:
1. `PipelineConfig` is the source of truth for stage topology, execution type,
model capabilities, and stage relationships.
2. `DeployConfig` describes placement, replicas, devices, connectors, and
deploy-time defaults; it does not redefine the model graph.
3. CLI and Python overrides are applied at the resolution boundary, with
per-stage overrides taking precedence over global values where supported.
4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
is runtime-private state, not a user configuration object.
5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
process that owns the corresponding engine.
## Main features
vLLM-Omni aims to be fast, flexible, and easy to use with the following features:
The feature surface is grouped to match the
[feature design documents](index.md#feature-design-documents). This page
summarizes each feature's architectural role; the linked design document is
the source for configuration, compatibility, and implementation details.
### Performance and Acceleration
### Runtime and stage execution
The framework achieves high performance through several optimization techniques:
* **Disaggregated inference:** Logical stages can run in separate processes,
devices, or nodes while the orchestrator preserves their declared
relationships. `OmniConnector` implementations transfer stage data and
control-plane metadata. See [Disaggregated Inference](feature/disaggregated_inference.md).
* **Asynchronous stage and output execution:** [Async Chunk](feature/async_chunk.md)
forwards partial stage outputs as they become available. [Async Diffusion
Output](feature/async_diffusion_output.md) overlaps device-to-host output
packing with the next diffusion request, while [Async Omni Output
Materialization](feature/omni_async_output_materialization.md) moves
CPU-side payload construction off the AR decode critical path.
* **Automatic prefix caching:** [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
reuses KV-cache-aligned stage outputs and multimodal tensors for requests
with common prefixes.
* **Efficient AR Support:** Leverages efficient KV cache management inherited from vLLM.
* **Pipelined Execution:** Uses pipelined stage execution overlapping to ensure high throughput.
* **Full Disaggregation:** Relies on the OmniConnector and dynamic resource allocation across stages.
* **Diffusion Acceleration:** Includes integrated support for diffusion acceleration. This is managed by the acceleration layer, which handles:
* **Cache:** Includes DBCache, TeaCache and third-party integration(e.g., [cache-dit](https://github.com/vipshop/cache-dit)).
* **Parallelism:** Supports TP, CP, USP, and CFG.
* **Attention:** Provides an interface for third-party integration (e.g., FA3, SAGE, MindIE-SD).
* **Quantization:** Supports various quantization implementations including FP8 and AWQ.
* **FusedOps:** Allows for custom and third-party integration.
### Communication
### Classifier-Free Guidance (CFG) Companion Flow
* **OmniConnector transport:** The connector contract carries tensors, KV-cache
data, and transport metadata across stage boundaries. The available
implementations cover shared memory and multi-node Mooncake, Mori, and
Yuanrong transports; see [Disaggregated Inference](feature/disaggregated_inference.md)
for the connector choices and configuration model.
vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
### Diffusion acceleration
### Flexibility and Usability
* **Request and step batching:** [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
defines request-batch and step-batch execution, scheduler admission, and the
common streaming output path.
* **Composable parallelism:** Diffusion stages can combine [CFG-Parallel](feature/cfg_parallel.md),
[Expert Parallel](feature/expert_parallel.md), [HSDP](feature/hsdp.md),
[Pipeline Parallel](feature/pipeline_parallel.md), [Sequence Parallel](feature/sequence_parallel.md),
[Tensor Parallel](feature/tensor_parallel.md), and [VAE Patch Parallelism](feature/vae_parallel.md)
according to the pipeline and hardware topology.
* **Attention and cache acceleration:** [Skip-Softmax](feature/skip_softmax.md),
[Cache-DiT](feature/cache_dit.md), and [TeaCache](feature/teacache.md)
provide backend and denoising-step optimizations without changing the
stage contract.
* **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
resolves per-pipeline or per-component quantization configurations, while
[Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
streams diffusion blocks from host memory within the existing parallel
topology.
vLLM-Omni is designed to be flexible and straightforward for users:
## Interfaces
* **Heterogeneous Pipeline Abstraction:** Manages complex model workflows effectively.
* **Hugging Face Integration:** Offers seamless integration with popular Hugging Face models.
* **Distributed Inference:** Supports tensor, pipeline, data, and expert parallelism.
* **Streaming Outputs:** Supports streaming outputs.
* **Unified API:** Provides a consistent and unified API interface compatible with vLLM.
* **OpenAI-compatible API Server:** Includes a FastAPI-based server for online serving that is compatible with the OpenAI API.
# Interface design
If you use vLLM, then you know how to use vLLM-Omni from Day 0:
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png">
<img alt="vLLM-Omni interface design" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/vllm-omni-user-interface.png" width=55%>
</picture>
</p>
Taking **Qwen3-Omni** as an example:
## Offline Inference
The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
The public interfaces map onto the same engine and stage boundaries:
```mermaid
flowchart LR
offline["Offline Python<br/>Omni.generate()"] --> engine["AsyncOmniEngine"]
online["OpenAI-compatible serving<br/>vllm serve ... --omni"] --> engine
duplex["Experimental duplex WebSocket<br/>/v1/duplex or realtime duplex"] --> engine
engine --> stages["Configured AR and diffusion stages"]
stages --> result["Streaming or final multimodal output"]
```
# Create an omni runtime with HF model name.
### Offline inference
The **Omni** class provides a Python interface for offline batched inference:
```python
from vllm_omni.entrypoints.omni import Omni
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
# Example prompts.
om_inputs = {"prompt": prompt,
"multi_modal_data": {
"video": video_frames,
"audio": audio_signal,
}}
om_inputs = {
"prompt": prompt,
"multi_modal_data": {
"video": video_frames,
"audio": audio_signal,
},
}
# Generate texts and audio from the multi-modality inputs.
outputs = omni.generate(om_inputs, sampling_params_list)
```
## Online Serving
Similar to vLLM, vLLM-Omni also provides a FastAPI-based server for online serving. Users can launch the server using the vllm serve command with the `--omni` flag:
### Online serving
```
The OpenAI-compatible server uses the same stage configuration and engine
boundaries:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
Users can send requests to the server using curl:
For example, a Qwen3-Omni chat request can contain text, image, audio, or
video content and a `sampling_params_list` for its configured stages. See the
[Qwen3-Omni serving example](../user_guide/examples/online_serving/qwen3_omni.md)
and the [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples)
for complete requests.
```
# prepare user content
user_content='[
{
"type": "video_url",
"video_url": {
"url": "'"$SAMPLE_VIDEO_URL"'"
}
},
{
"type": "text",
"text": "Why is this video funny?"
}
]'
sampling_params_list='[
'"$thinker_sampling_params"',
'"$talker_sampling_params"',
'"$code2wav_sampling_params"'
]'
mm_processor_kwargs="{}"
# send the request
curl -sS -X POST http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @- <<EOF
{
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"sampling_params_list": $sampling_params_list,
"mm_processor_kwargs": $mm_processor_kwargs,
"messages": [
{
"role": "system",
"content": [
{
"type": "text",
"text": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
}
]
},
{
"role": "user",
"content": $user_content
}
]
}
```
For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).
Some pipelines expose additional OpenAI-compatible endpoints, such as joint
video/audio generation. Endpoint support remains model-specific; consult the
relevant model guide before assuming that every OpenAI route applies to every
pipeline.
+65 -5
View File
@@ -20,11 +20,12 @@ extracts it into a shared, typed contract:
(matching the bespoke caches it replaces), and reports per-device byte usage
via `stats()` for observability.
Attention KV is *not* managed here: for DreamZero it is owned by the
AR-Diffusion engine's paged KV pool (PR
[#4534](https://github.com/vllm-project/vllm-omni/pull/4534)). The manager
covers the model's non-KV session state; integrating the engine's per-session
KV handle into the same byte accounting is RFC #4480 Phase 1.
Autoregressive generation KV is *not* managed here: for DreamZero it is owned
by the AR-Diffusion engine's paged KV pool (PR
[#4534](https://github.com/vllm-project/vllm-omni/pull/4534)). Cosmos3's fixed,
encode-once UND conditioning K/V is different and is covered by the dense
Cosmos3 adapter described below. Integrating the engine's per-session paged KV
handle into the same byte accounting is RFC #4480 Phase 1.
## Enabling
@@ -158,3 +159,62 @@ of `num_frame_per_block` frames instead: identical output, bounded memory, and
the per-step quantisation no longer scales with session length. Asking for more
frames than the ring retains raises rather than silently returning a shorter
window. The bespoke path is unchanged.
### Cosmos3
When `enable_session_state_manager` is enabled, the Cosmos3 adapter stores the
dense, encode-once UND text K/V for every transformer layer and active CFG
branch, together with that branch's GEN M-RoPE cosine/sine tensors. These are
GPU-resident tensors and can contribute to CUDA OOM. A useful description of
the retained storage on one rank is:
```text
sum_over_branches(
sum_over_layers(storage_bytes(K) + storage_bytes(V))
+ storage_bytes(freqs_gen_cos) + storage_bytes(freqs_gen_sin)
)
```
This is not a peak-memory upper bound: construction of the next branch can
temporarily overlap the already retained branch with UND activations, the new
full K/V result, RoPE tensors, and generator activations. K/V slices also keep
their underlying allocation alive, so physical storage can follow the padded
text length rather than only the exposed trimmed length. The manager's byte
accounting deliberately measures storage for this reason.
Branch placement changes the per-rank cost:
- no CFG retains one conditional branch;
- sequential CFG can retain both conditional and unconditional branches on the
same GPU;
- CFG-parallel retains one branch on each CFG rank;
- UND K/V is replicated across sequence-parallel ranks rather than divided by
the SP degree.
Consequently, risk increases with text batch/sequence length, media token count
(which sizes `freqs_gen`), sequential CFG, dtype width, and concurrent live
requests. Phase 0 records `nbytes:cuda:*` but does not enforce `byte_budget`, so
the manager cannot prevent an allocation from exhausting the device. The
count-based session limit is not an OOM guarantee.
The adapter transfers ownership by tensor reference; it does not clone K/V or
RoPE tensors. For a normal serial request this should not materially increase
tensor storage relative to Cosmos3's existing bespoke cond/uncond cache, which
retains the same branch values. It does add small CPU-side session/object
metadata. This claim is based on the reference flow in the implementation;
there is not yet a current-head GPU peak-memory A/B result for this PR.
Request-finally cleanup drops the manager-owned session, so branches do not
accumulate with completed request IDs. The transformer handoff field still
references the last loaded branch until `reset_cache()` runs at the start of
the next generation (or the pipeline is destroyed). Thus this is bounded
post-request residency, not growth per completed request, but manager
`stats()` after `drop_session()` no longer includes that last transformer-held
alias. A profiler that only reads manager statistics can therefore under-report
live process memory.
For OOM validation, measure complete cache-off and true flag-on generations in
otherwise identical fresh processes. Synchronize CUDA and report process-level
peak memory plus PyTorch peak allocated and peak reserved memory. Include the
period between requests when checking cleanup; do not infer safety from the
manager byte total, a single denoising step, or an operator-only profile.
+7
View File
@@ -97,3 +97,10 @@ th {
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
<sup>H3</sup> MiniMax H3 AMD GPU support is validated on gfx942 (MI300X) and
gfx950 (MI350) in BF16, using the AITER `FLASH_ATTN` diffusion attention backend.
Other AMD SKUs (e.g. MI325X) are not listed until their
own validation evidence is added. See the
[MiniMax H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for
per-architecture serving commands.
@@ -17,7 +17,7 @@ The full set of backends and their platform defaults is in the **Backend Options
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
| `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. |
@@ -84,7 +84,7 @@ A backend that needs configuration exposes it as a typed field on the spec:
When constructing `OmniDiffusionConfig` directly:
```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig
config = OmniDiffusionConfig(
diffusion_attention_config=AttentionConfig(
@@ -99,6 +99,27 @@ config = OmniDiffusionConfig(
A plain dict is also accepted and normalized to `AttentionConfig`.
#### FlashInfer QK16/V8 quantized attention
`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
```python
config = OmniDiffusionConfig(
diffusion_attention_config=AttentionConfig(
default=AttentionSpec(
backend="FLASHINFER_ATTN",
quant=AttnQuantSpec(
dtype_qk="bfloat16",
dtype_vo="fp8_e4m3",
),
),
),
...,
)
```
`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
## Platform Defaults
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
@@ -6,6 +6,7 @@
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Request-Scoped Quality](#request-scoped-quality-minimax-h3)
- [Acceleration Methods](#acceleration-methods)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
@@ -140,6 +141,43 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
--cache-config '{"Fn_compute_blocks": 1, "residual_diff_threshold": 0.12}'
```
## Request-Scoped Quality (MiniMax H3)
MiniMax H3 can switch Cache-DiT at the request boundary without enabling a
cache backend at server startup. Add a quality field to an online request:
```bash
# Native reference path for this request.
-F 'quality=lossless'
# H3's conservative Cache-DiT profile for this request.
-F 'quality=high'
```
For offline inference, set the same field on the sampling parameters:
```python
OmniDiffusionSamplingParams(
quality="high",
num_inference_steps=50,
extra_args={"task": "t2va", "duration": 5.0, "audio_flow_shift": 3.0},
)
```
| Server startup | Request quality | Behavior |
|---|---|---|
| `--cache-backend cache_dit` | omitted | Use or restore the startup Cache-DiT profile |
| `--cache-backend cache_dit` | `lossless` | Remove Cache-DiT hooks and run the reference path |
| `--cache-backend cache_dit` | `high` | Use or install H3's conservative Cache-DiT profile |
| no Cache-DiT | omitted or `lossless` | Run the reference path |
| no Cache-DiT | `high` | Install H3's conservative Cache-DiT profile |
The startup option is only needed when omitted-quality requests should use the
server-configured Cache-DiT profile by default.
See the [MiniMax H3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#request-scoped-quality)
for a complete request and measured trade-off.
---
## Acceleration Methods
@@ -0,0 +1,240 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.
Examples:
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
CUDA_VISIBLE_DEVICES=0,1 \
python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
--model /path/to/MiniMax-H3/FL2VA --mode request
"""
from __future__ import annotations
import argparse
import asyncio
import copy
import json
import multiprocessing
import time
from pathlib import Path
from typing import Any
import numpy as np
from vllm_omni.entrypoints.async_omni import AsyncOmni
DEFAULT_PROMPTS = (
"At night, three cats march into a bedroom playing tiny brass instruments, "
"then abruptly file out, with synchronized room ambience.",
"A paper boat crosses a rain-filled street while distant traffic and water sounds remain synchronized.",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
parser.add_argument(
"--mode",
choices=("dlo-dp2", "request"),
required=True,
help="DLO AllGather DP2 or ordinary non-DLO TP2 request mode",
)
parser.add_argument("--steps", type=int, default=2)
parser.add_argument("--duration", type=float, default=5.0)
parser.add_argument("--width", type=int, default=1344)
parser.add_argument("--height", type=int, default=768)
parser.add_argument("--batch-wait-ms", type=float, default=500.0)
parser.add_argument("--init-timeout", type=float, default=1800.0)
parser.add_argument("--output", type=Path)
return parser.parse_args()
def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
common: dict[str, Any] = {
"model": args.model,
"trust_remote_code": True,
"num_gpus": 2,
"usp": 1,
"ring": 1,
"vae_parallel_mode": "tile",
"vae_use_tiling": True,
"diffusion_attention_backend": "CUDNN_ATTN",
"request_batch_max_wait_ms": args.batch_wait_ms,
"enforce_eager": True,
"stage_init_timeout": args.init_timeout,
"init_timeout": args.init_timeout,
}
if args.mode == "dlo-dp2":
common.update(
tensor_parallel_size=1,
data_parallel_size=2,
text_encoder_tp_size=1,
vae_patch_parallel_size=1,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
dlo_resident_layers=0,
)
else:
common.update(
tensor_parallel_size=2,
data_parallel_size=1,
text_encoder_tp_size=2,
vae_patch_parallel_size=2,
enable_distributed_layerwise_offload=False,
)
return common
def sampling_params(
engine: AsyncOmni,
args: argparse.Namespace,
seed: int,
) -> list[Any]:
params = copy.deepcopy(engine.default_sampling_params_list)
diffusion = params[0]
diffusion.width = args.width
diffusion.height = args.height
diffusion.fps = 24
diffusion.num_inference_steps = args.steps
diffusion.seed = seed
diffusion.extra_args = {
"task": "t2va",
"duration": args.duration,
"aspect_ratio": "16:9",
"flow_shift": 12.0,
"audio_flow_shift": 3.0,
}
return params
async def generate_one(
engine: AsyncOmni,
args: argparse.Namespace,
*,
request_id: str,
prompt: str,
seed: int,
) -> Any:
final_output = None
async for output in engine.generate(
prompt=prompt,
request_id=request_id,
sampling_params_list=sampling_params(engine, args, seed),
):
if output.finished:
final_output = output
if final_output is None:
raise RuntimeError(f"{request_id} finished without an output")
return final_output
def output_summary(output: Any, args: argparse.Namespace) -> dict[str, Any]:
if not output.images:
raise RuntimeError(f"{output.request_id} returned no video")
frames = np.asarray(output.images[0])
audio = np.asarray(output.multimodal_output.get("audio"))
if frames.ndim != 4 or tuple(frames.shape[1:]) != (
args.height,
args.width,
3,
):
raise RuntimeError(f"{output.request_id} returned invalid video shape {tuple(frames.shape)}")
if audio.ndim != 3 or tuple(audio.shape[:2]) != (1, 2):
raise RuntimeError(f"{output.request_id} returned invalid audio shape {tuple(audio.shape)}")
if args.duration == 5.0 and args.height == 768 and args.width == 1344:
if tuple(frames.shape) != (124, 768, 1344, 3) or tuple(audio.shape) != (1, 2, 165600):
raise RuntimeError(
f"{output.request_id} default shape mismatch: video={tuple(frames.shape)}, audio={tuple(audio.shape)}"
)
return {
"request_id": output.request_id,
"frames_shape": list(frames.shape),
"audio_shape": list(audio.shape),
"peak_memory_mb": output.peak_memory_mb,
"stage_durations": output.stage_durations,
}
async def run(args: argparse.Namespace) -> dict[str, Any]:
engine = AsyncOmni(**engine_kwargs(args))
summary: dict[str, Any] = {
"mode": args.mode,
"engine_kwargs": engine_kwargs(args),
}
try:
if args.mode == "dlo-dp2":
started = time.perf_counter()
asymmetric = await asyncio.gather(
generate_one(
engine,
args,
request_id="invalid-empty",
prompt="",
seed=1001,
),
generate_one(
engine,
args,
request_id="invalid-peer",
prompt=DEFAULT_PROMPTS[0],
seed=1002,
),
return_exceptions=True,
)
summary["asymmetric_wave_s"] = time.perf_counter() - started
summary["asymmetric_errors"] = [
f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
]
if len(summary["asymmetric_errors"]) != 2:
raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")
started = time.perf_counter()
request_count = 2 if args.mode == "dlo-dp2" else 1
outputs = await asyncio.gather(
*(
generate_one(
engine,
args,
request_id=f"recovery-{index}",
prompt=DEFAULT_PROMPTS[index],
seed=2000 + index,
)
for index in range(request_count)
)
)
summary["valid_wave_s"] = time.perf_counter() - started
summary["outputs"] = [output_summary(output, args) for output in outputs]
finally:
started = time.perf_counter()
engine.close()
summary["shutdown_s"] = time.perf_counter() - started
children = multiprocessing.active_children()
summary["active_children"] = [{"name": child.name, "pid": child.pid} for child in children]
if children:
raise RuntimeError(f"Worker processes remain after shutdown: {children}")
return summary
def main() -> None:
args = parse_args()
summary = asyncio.run(run(args))
rendered = json.dumps(summary, indent=2, sort_keys=True)
print(f"E2E_RESULT {rendered}", flush=True)
if args.output is not None:
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(rendered + "\n", encoding="utf-8")
if __name__ == "__main__":
main()
+147 -5
View File
@@ -58,6 +58,10 @@ FlashAttention-4 extra:
uv pip install -e '.[fa4]'
```
On AMD ROCm, install without the `[fa4]` extra (FA4 is CUDA-only) and use
`--diffusion-attention-backend FLASH_ATTN`; see
[AMD ROCm (gfx942 / gfx950)](#amd-rocm-gfx942--gfx950).
`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
reference-video preparation and MP4 output.
@@ -387,6 +391,101 @@ incompatible with H3 layerwise offload because the offload path produces a
weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
parallelism and VAE tiling instead.
## AMD ROCm (gfx942 / gfx950)
The sections above describe the NVIDIA CUDA path. This section covers the AMD
ROCm path; use it instead of the CUDA commands when running on AMD Instinct
GPUs.
MiniMax H3 runs on AMD Instinct GPUs (gfx942 / gfx950) in BF16. Use
`--diffusion-attention-backend FLASH_ATTN`, which resolves to AITER packed varlen
attention on both architectures.
Select the device with `HIP_VISIBLE_DEVICES` (not `CUDA_VISIBLE_DEVICES`), drop the
CUDA-only `FLASHINFER_DISABLE_VERSION_CHECK`, and install without the `[fa4]` extra
(FA4 is CUDA-only). The VAE uses AITER GroupNorm on ROCm.
Install (ROCm wheel + source vLLM-Omni):
```bash
pip install "vllm==0.26.0+rocm723" \
--extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
VLLM_OMNI_TARGET_DEVICE=rocm pip install -e . --no-build-isolation
```
Prebuilt image: `vllm/vllm-omni-rocm:minimax-h3`. All tasks work out of the box:
the image bundles TorchCodec (for image+audio Ref2VA) and `ffmpeg` (for
video-reference Ref2VA).
### ROCm single GPU
Single GPU with model-level CPU offload keeps the Qwen3-VL encoder and DiT from
being co-resident:
```bash
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8091
HIP_VISIBLE_DEVICES=0 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni --host 0.0.0.0 --port "${PORT}" --trust-remote-code \
--num-gpus 1 --enable-cpu-offload \
--diffusion-attention-backend FLASH_ATTN
```
### ROCm four GPUs
The best-practice CUDA four-GPU configuration works on ROCm with the changes above.
It mirrors the CUDA command including Ulysses sequence parallelism, VAE patch
parallelism, and text-encoder tensor parallelism, with no CPU offload when the model
shards across the GPUs:
```bash
export MODEL="${MODEL_ROOT}/FL2VA"
export PORT=8091
HIP_VISIBLE_DEVICES=0,1,2,3 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
vllm serve "${MODEL}" \
--omni \
--host 0.0.0.0 \
--port "${PORT}" \
--trust-remote-code \
--num-gpus 4 \
--usp 4 \
--ring 1 \
--text-encoder-tp-size 4 \
--vae-patch-parallel-size 4 \
--vae-parallel-mode tile \
--vae-use-tiling \
--diffusion-attention-backend FLASH_ATTN
```
As on CUDA, H3 is CFG-distilled, so keep `--cfg-parallel-size` at 1, and the H3 VAE
supports only its native `tile` mode. `--text-encoder-tp-size` is validated on
gfx942; on gfx950 it was not exercised.
### Validated ROCm evidence
vLLM-Omni with MiniMax H3 support, BF16. gfx942 rows measured with the
`vllm/vllm-omni-rocm:minimax-h3` image; gfx950 rows measured with the
`0.26.0+rocm723` wheel (HIP 7.2).
| Workload | Configuration | Observed result |
|----------|---------------|-----------------|
| T2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 0.09 s, denoise 244.04 s, decode 4.15 s, 267.42 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
| FL2VA, 1344x768, 209 frames, 50 steps | 4x gfx942 (MI300X), FLASH_ATTN, USP4, text-enc TP4, VAE PP4 tile | encode 13.98 s, denoise 257.58 s, decode 4.11 s, 287.07 s client E2E; H.264 24 FPS + 32 kHz stereo AAC |
| T2VA, 832x480, ~4 s, 40 steps | 1x gfx950 (MI350), FLASH_ATTN, CPU offload | valid MP4 (H.264 + synced audio); ~0.73 s/denoise-step (~1.37 it/s), ~55 s client E2E incl. warmup |
gfx942 figures are the mean of three requests after one excluded warmup
(external evidence: vllm-project/recipes#732). gfx950 figures are
functional-correctness validations, not tuned throughput; the first request
includes lazy regional compilation. MI325X (gfx942) and other MI355X SKUs are not
listed until their own evidence is added.
## HTTP API examples
The following requests use the synchronous endpoint so the returned body can
@@ -574,10 +673,43 @@ default. `short_edge` controls the 768-pixel canvas and must be `768`.
outputs; the synchronous raw-MP4 endpoint returns the first output when more
than one is requested.
## Request-scoped quality
Add one of these fields to any HTTP request above. No Cache-DiT startup option
is required; H3 installs its conservative profile when a `high` request
arrives and removes it for a `lossless` request.
```bash
-F 'quality=lossless' # Native reference path for this request.
-F 'quality=high' # H3's conservative Cache-DiT profile.
```
Omitting `quality` preserves the startup default: it uses the reference path
normally, or the server-configured profile when the server was started with
`--cache-backend cache_dit`.
The following result was measured on 4× NVIDIA H200 with SP4, text-encoder
TP4, 1344×768, 124 frames, 24 FPS, and 50 inference steps. One full
`lossless` warmup was excluded, followed by three fixed prompt/seed pairs in
balanced switch order.
| `quality` | Median inference latency | Speedup | SSIM vs `lossless` | PSNR vs `lossless` | Expected trade-off |
|---|---:|---:|---:|---:|---|
| `lossless` | 85.49 s | 1.00× | 1.0000 | exact | Native reference path |
| `high` | 63.36 s | 1.35× | 0.9709 | 34.98 dB | Faster with measured same-seed deviation |
> [!NOTE]
> `high` selects a fixed Cache-DiT profile, but its cache hit rate and
> resulting latency/quality trade-off may vary by hardware, topology, and
> workload. The values above apply to this deployment and are not universal
> guarantees. `lossless` remains the exact reference path.
## Key parameters
| Parameter | Recommended value | Notes |
|-----------|-------------------|-------|
| `quality` | omitted or `lossless` | Request-level quality intent; `high` dynamically installs H3's conservative Cache-DiT profile |
| `extra_params.force_refresh_step_hint` | omitted | Optional positive 1-based denoising-step hint for an active Cache-DiT request; pair with `extra_params.force_refresh_step_policy`=`once` or `repeat` |
| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; selects the task-specific DiT |
| `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
| `fps` | `24` | H3 output FPS is fixed |
@@ -625,12 +757,22 @@ reduction.
- H3 currently executes one generation request per diffusion batch.
- The first regional-compile request is a warmup and should not be included in
steady-state performance measurements.
- The serving path accepts fewer references than the model supports. H3 documents up
to 9 images, 3 video clips, and 3 audio clips (12 files) per Omni Reference
request; the current vLLM-Omni path takes exactly one image plus one audio
reference, or one or more videos with no separate `audio_reference` (it uses the
source soundtracks).
- The 768 px short-edge mode is available for T2VA and FL2VA; 1344x768 is the
documented 16:9 request shape.
- `--cfg-parallel-size > 1` is rejected by design (CFG-distilled, no negative branch).
- VAE patch parallelism requires size 1 or the full DiT group size and supports the
H3 native `tile` mode only.
- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
pure Ulysses.
- Online FP8 is not compatible with layerwise offload.
- Image+audio Ref2VA accepts exactly one image and one audio reference.
- Video Ref2VA accepts one or more video files, but not an additional standalone
audio reference.
- VAE patch parallelism requires size 1 or the full DiT group size and supports
the H3 native `tile` mode only.
- Pure Ulysses still replicates the full DiT on every rank, so smaller-memory GPUs
cannot use `--usp N --tp 1` as a resident capacity path. Use DiT tensor parallelism
or model-level CPU offload; text-encoder TP alone is not sufficient.
## Additional resources
+3
View File
@@ -1805,6 +1805,7 @@ class TestQwen2_5OmniPipeline:
assert s.owns_tokenizer is True
assert s.engine_output_type == "latent"
assert s.requires_multimodal_data is True
assert s.hf_config_name == "thinker_config"
def test_talker(self):
p = resolve_pipeline_config("qwen2_5_omni")
@@ -1813,6 +1814,7 @@ class TestQwen2_5OmniPipeline:
s = p.get_stage(1)
assert isinstance(s, StagePipelineConfig)
assert s.input_sources == (0,)
assert s.hf_config_name == "talker_config"
assert s.sampling_constraints["stop_token_ids"] == [8294]
# thinker2talker was removed: qwen2_5_omni has no async_chunk support,
# so sync_process_input_func always wins and custom_process_input_func
@@ -1829,6 +1831,7 @@ class TestQwen2_5OmniPipeline:
assert s.execution_type == StageExecutionType.LLM_GENERATION
assert s.final_output_type == "audio"
assert s.engine_output_type == "audio"
assert s.hf_config_name == "thinker_config"
class TestQwen3TTSPipeline:
@@ -0,0 +1,192 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Equivalence and session-isolation tests for Cosmos3StateAdapter (RFC #4480).
These drive the adapter directly with tiny CPU tensors -- no model, no GPU --
and assert it stores/returns the per-branch UND K/V exactly and isolates
durable state between sessions. They do not cover the transformer's mutable
load-to-forward handoff.
"""
from __future__ import annotations
import pytest
import torch
from vllm_omni.experimental.world_models.adapters.state_cosmos3_adapter import (
Cosmos3StateAdapter,
)
from vllm_omni.experimental.world_models.session_state import SessionStateManager
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
LAYERS = 3
SEQ, DIM = 3, 4
def _fake_cached_kv(seed: int) -> list[tuple[torch.Tensor, torch.Tensor]]:
"""A list of per-layer (K, V); contents seeded so branches/sessions differ."""
g = torch.Generator().manual_seed(seed)
return [(torch.randn(1, SEQ, DIM, generator=g), torch.randn(1, SEQ, DIM, generator=g)) for _ in range(LAYERS)]
def _adapter(session_id: str, manager: SessionStateManager | None = None) -> Cosmos3StateAdapter:
# `manager is not None`, not `manager or ...`: an empty SessionStateManager is
# falsy (it defines __len__), so `or` would silently swap in a fresh manager.
return Cosmos3StateAdapter(session_id, manager if manager is not None else SessionStateManager())
def _assert_kv_equal(got: list | None, expected: list) -> None:
assert got is not None and len(got) == len(expected) == LAYERS
for (k, v), (ke, ve) in zip(got, expected):
torch.testing.assert_close(k, ke)
torch.testing.assert_close(v, ve)
class _FakeTransformer:
cached_kv = None
cached_freqs_gen = None
# ----------------------------- equivalence -----------------------------
def test_branch_starts_uninitialized() -> None:
a = _adapter("s0")
for is_neg in (False, True):
assert a.is_branch_initialized(is_neg) is False
assert a.get_branch_kv(is_neg) is None
def test_set_then_get_roundtrip() -> None:
a = _adapter("s0")
kv = _fake_cached_kv(1)
a.set_branch_kv(False, kv)
assert a.is_branch_initialized(False) is True
_assert_kv_equal(a.get_branch_kv(False), kv)
def test_cond_uncond_branches_isolated() -> None:
a = _adapter("s0")
kv_pos, kv_neg = _fake_cached_kv(1), _fake_cached_kv(2)
a.set_branch_kv(False, kv_pos)
a.set_branch_kv(True, kv_neg)
_assert_kv_equal(a.get_branch_kv(False), kv_pos)
_assert_kv_equal(a.get_branch_kv(True), kv_neg)
def test_encode_once_capture_is_noop_when_initialized() -> None:
a = _adapter("s0")
kv = _fake_cached_kv(1)
a.set_branch_kv(False, kv)
tr = _FakeTransformer()
tr.cached_kv = _fake_cached_kv(999) # different content
a.capture_from_transformer(tr, False)
_assert_kv_equal(a.get_branch_kv(False), kv) # still the first one
def test_capture_writes_when_uninitialized() -> None:
a = _adapter("s0")
tr = _FakeTransformer()
tr.cached_kv = _fake_cached_kv(5)
tr.cached_freqs_gen = (torch.ones(2), torch.zeros(2))
a.capture_from_transformer(tr, False)
_assert_kv_equal(a.get_branch_kv(False), _fake_cached_kv(5))
def test_capture_without_freqs_does_not_partially_initialize() -> None:
a = _adapter("s0")
tr = _FakeTransformer()
tr.cached_kv = _fake_cached_kv(5)
tr.cached_freqs_gen = None
with pytest.raises(RuntimeError, match="did not produce GEN RoPE frequencies"):
a.capture_from_transformer(tr, False)
assert a.is_branch_initialized(False) is False
def test_load_into_transformer_without_freqs_sets_kv() -> None:
a = _adapter("s0")
kv = _fake_cached_kv(1)
a.set_branch_kv(False, kv)
tr = _FakeTransformer()
a.load_into_transformer(tr, False)
_assert_kv_equal(tr.cached_kv, kv)
assert tr.cached_freqs_gen is None
def test_set_then_get_preserves_per_layer_values() -> None:
# The stored value must reflect the tensors passed in (per-layer).
a = _adapter("s0")
kv = _fake_cached_kv(3)
a.set_branch_kv(False, kv)
_assert_kv_equal(a.get_branch_kv(False), kv)
def test_load_into_transformer_restores_kv_and_freqs_gen() -> None:
a = _adapter("s0")
tr = _FakeTransformer()
tr.cached_kv = _fake_cached_kv(3)
tr.cached_freqs_gen = (torch.ones(2), torch.zeros(2))
a.capture_from_transformer(tr, False)
tr.cached_kv = None
tr.cached_freqs_gen = None
a.load_into_transformer(tr, False)
_assert_kv_equal(tr.cached_kv, _fake_cached_kv(3))
assert tr.cached_freqs_gen is not None
torch.testing.assert_close(tr.cached_freqs_gen[0], torch.ones(2))
def test_partial_branch_raises() -> None:
a = _adapter("s0")
a.set_branch_kv(False, _fake_cached_kv(1))
a._session.get("und_kv/1").reset() # type: ignore[union-attr]
with pytest.raises(RuntimeError, match="partially initialized"):
a.get_branch_kv(False)
# ----------------------------- reset -----------------------------
def test_reset_clears_both_branches() -> None:
a = _adapter("s0")
a.set_branch_kv(False, _fake_cached_kv(1))
a.set_branch_kv(True, _fake_cached_kv(2))
a.reset()
for is_neg in (False, True):
assert a.is_branch_initialized(is_neg) is False
assert a.get_branch_kv(is_neg) is None
# ----------------------- concurrency isolation (core) -----------------------
def test_two_sessions_do_not_clobber() -> None:
manager = SessionStateManager()
a, b = _adapter("A", manager), _adapter("B", manager)
kv_a, kv_b = _fake_cached_kv(10), _fake_cached_kv(20)
a.set_branch_kv(False, kv_a)
assert b.is_branch_initialized(False) is False # B not seeing A's state
b.set_branch_kv(False, kv_b)
_assert_kv_equal(a.get_branch_kv(False), kv_a)
_assert_kv_equal(b.get_branch_kv(False), kv_b)
def test_fresh_adapter_same_session_sees_prior_state() -> None:
manager = SessionStateManager()
kv = _fake_cached_kv(7)
_adapter("s", manager).set_branch_kv(False, kv)
again = Cosmos3StateAdapter("s", manager)
_assert_kv_equal(again.get_branch_kv(False), kv)
def test_adapter_survives_session_lru_eviction() -> None:
# An in-use adapter keeps its state even after the manager evicts its
# session id from the lookup table (LRU bounded by max_sessions).
manager = SessionStateManager(max_sessions=2)
a = _adapter("active", manager)
kv = _fake_cached_kv(11)
a.set_branch_kv(False, kv)
for i in range(3):
_adapter(f"other{i}", manager)
assert "active" not in manager
_assert_kv_equal(a.get_branch_kv(False), kv)
+16 -5
View File
@@ -31,6 +31,7 @@ import tempfile
import pytest
import torch
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.attention.parallel.allgather_kv import (
@@ -409,6 +410,7 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
torch.testing.assert_close(v_full, expected_value)
@hardware_test(res={"cuda": "L4"}, num_cards=4)
@pytest.mark.parametrize(
"test_model_cls",
[
@@ -416,10 +418,12 @@ def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa():
],
)
@pytest.mark.parametrize(
("ulysses_degree", "ring_degree", "allgather_degree"),
("ulysses_degree", "ring_degree", "allgather_degree", "num_kv_heads"),
[
pytest.param(2, 2, 1, id="ulysses-ring"),
pytest.param(1, 1, 2, id="allgather-kv"),
pytest.param(2, 2, 1, None, id="ulysses-ring"),
pytest.param(1, 1, 2, None, id="allgather-kv"),
pytest.param(1, 2, 1, 2, id="ring-gqa"),
pytest.param(1, 2, 1, 1, id="ring-mqa"),
],
)
@pytest.mark.parametrize("batch_size", [2])
@@ -445,6 +449,7 @@ def test_sequence_parallel(
seq_len: int,
num_heads: int,
head_size: int,
num_kv_heads: int | None,
):
"""Compare Ulysses/Ring and AllGather-KV SP against a single-rank run."""
sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree
@@ -476,6 +481,7 @@ def test_sequence_parallel(
seq_len,
num_heads,
head_size,
num_kv_heads,
dtype,
causal,
use_sync,
@@ -507,6 +513,7 @@ def test_sequence_parallel(
seq_len,
num_heads,
head_size,
num_kv_heads,
dtype,
causal,
use_sync,
@@ -603,6 +610,7 @@ def ulysses_attention_on_test_model(
seq_len: int,
num_heads: int,
head_size: int,
num_kv_heads: int | None,
dtype: torch.dtype,
causal: bool,
use_sync: bool,
@@ -658,10 +666,13 @@ def ulysses_attention_on_test_model(
cfg_parallel_size=1,
)
od_config = OmniDiffusionConfig(
od_config = OmniDiffusionConfig.from_kwargs(
model="test_model",
dtype=dtype,
parallel_config=parallel_config,
# This regression targets pytorch_attn_forward(). Do not let an
# installed FA/FA3 backend silently bypass the SDPA Ring path.
diffusion_attention_backend="TORCH_SDPA",
)
# Initialize model parallel
@@ -687,7 +698,7 @@ def ulysses_attention_on_test_model(
"head_size": head_size,
"hidden_size": hidden_size,
"causal": causal,
"num_kv_heads": None,
"num_kv_heads": num_kv_heads,
"scatter_idx": 2,
"gather_idx": 1,
"use_sync": use_sync,
+32 -1
View File
@@ -6,12 +6,14 @@ from __future__ import annotations
import os
import socket
import tempfile
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
from vllm_omni.diffusion.config import set_current_diffusion_config
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.parallel_state import (
@@ -39,6 +41,33 @@ def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
os.environ["MASTER_PORT"] = str(master_port)
@pytest.mark.core_model
@pytest.mark.cpu
def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
lambda default: "advanced_uaa",
)
sp_group = SimpleNamespace(
ulysses_group=None,
ulysses_world_size=2,
ulysses_rank=0,
ring_world_size=2,
)
strategy = UlyssesParallelAttention(
sp_group=sp_group,
scatter_idx=2,
gather_idx=1,
use_sync=False,
)
query = torch.zeros(1, 2, 8, 4)
key = torch.zeros(1, 2, 1, 4)
value = torch.zeros_like(key)
with pytest.raises(ValueError, match="does not support K/V head padding"):
strategy.pre_attention(query, key, value, None)
def _run_attention_case(
local_rank: int,
world_size: int,
@@ -241,7 +270,9 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
batch_size = 2
head_size = 8
seq_len = 10
num_heads = 3 # head_cnt not divisible by ulysses_degree=2 -> triggers head padding
# This is the positive Hybrid case: K/V heads must not require padding,
# which is unsupported by advanced_uaa when Ring is also enabled.
num_heads = 4
# Ensure ring ranks see equal post-Ulysses seq_len:
# rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
+80 -3
View File
@@ -11,6 +11,7 @@ This module tests the cache backend implementations:
- DiffusionCacheConfig: configuration dataclass
"""
from types import SimpleNamespace
from unittest.mock import Mock, patch
import cache_dit
@@ -21,6 +22,7 @@ from vllm_omni.diffusion.cache.cachedit import (
CacheDiTAdapterConfig,
CacheDiTBackend,
CacheDiTConfig,
cache_summary,
)
from vllm_omni.diffusion.cache.magcache import MagCacheBackend
from vllm_omni.diffusion.cache.selector import get_cache_backend
@@ -91,8 +93,8 @@ class TestCacheDiTBackend:
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
def test_enable_and_refresh_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
"""Combined pipelines enable and refresh Cache-DiT for every DiT."""
def test_lifecycle_applies_to_every_declared_dit(self, mock_cache_dit, mock_block_adapter):
"""Combined pipelines manage every exact Cache-DiT target."""
pipeline = Mock()
pipeline.__class__.__name__ = "MiniMaxH3Pipeline"
pipeline._dit_modules = ["transformer", "transformers_ref"]
@@ -102,18 +104,67 @@ class TestCacheDiTBackend:
transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
)
installed_adapters = [Mock(), Mock()]
mock_block_adapter.side_effect = installed_adapters
backend = CacheDiTBackend({"Fn_compute_blocks": 2})
backend.enable(pipeline)
backend.refresh(pipeline, num_inference_steps=20)
backend.disable(pipeline)
assert len(backend._refresh_funcs) == 2
assert mock_cache_dit.enable_cache.call_count == 2
assert mock_cache_dit.refresh_context.call_count == 2
assert {call.args[0] for call in mock_cache_dit.refresh_context.call_args_list} == {
pipeline.transformer,
pipeline.transformers_ref,
}
assert [call.args[0] for call in mock_cache_dit.disable_cache.call_args_list] == installed_adapters
assert backend._refresh_funcs == []
assert backend._cache_targets == []
assert not backend.is_enabled()
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
def test_summary_skips_uncached_nested_dit(self, mock_cache_dit):
"""Only DiT modules with an active Cache-DiT context are summarized."""
language_model = Mock()
language_model._is_cached = False
transformer = Mock()
transformer._is_cached = True
transformer.language_model = language_model
pipeline = SimpleNamespace(
_dit_modules=["transformer.language_model", "transformer"],
transformer=transformer,
)
cache_summary(pipeline)
mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
def test_enable_and_refresh_nested_declared_dit(self, mock_cache_dit, mock_block_adapter):
"""Dotted component paths resolve to their nested DiT module."""
transformer = Mock()
transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
block_forward_patterns={"blocks": ForwardPattern.Pattern_3}
)
pipeline = SimpleNamespace(
_dit_modules=["language_model.model"],
language_model=SimpleNamespace(model=transformer),
)
backend = CacheDiTBackend({"Fn_compute_blocks": 2})
backend.enable(pipeline)
backend.refresh(pipeline, num_inference_steps=20)
cache_summary(pipeline)
mock_cache_dit.enable_cache.assert_called_once()
mock_cache_dit.refresh_context.assert_called_once_with(
transformer,
num_inference_steps=20,
verbose=True,
)
mock_cache_dit.summary.assert_called_once_with(transformer, details=True)
@patch("vllm_omni.diffusion.cache.cachedit.backend.logger")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
@@ -246,6 +297,32 @@ class TestCacheDiTBackend:
assert mock_cache_dit.refresh_context.call_count == 2
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
def test_refresh_rearms_force_refresh_hint(self, mock_cache_dit, mock_block_adapter):
"""A once hint must be restored for every repeated request."""
pipeline = Mock()
pipeline.__class__.__name__ = "DiTPipeline"
pipeline.transformer = Mock()
pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
block_forward_patterns={"layers": ForwardPattern.Pattern_0}
)
config = DiffusionCacheConfig(
force_refresh_step_hint=7,
force_refresh_step_policy="once",
)
backend = CacheDiTBackend(config)
backend.enable(pipeline)
backend.refresh(pipeline, num_inference_steps=20)
backend.refresh(pipeline, num_inference_steps=20)
assert mock_cache_dit.refresh_context.call_count == 2
for call in mock_cache_dit.refresh_context.call_args_list:
refresh_config = call.kwargs["cache_config"]
assert refresh_config.force_refresh_step_hint == 7
assert refresh_config.force_refresh_step_policy == "once"
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
+2
View File
@@ -81,6 +81,8 @@ def test_cachedit_public_api_is_explicit():
"CacheDiTAdapterConfig",
"CacheDiTBackend",
"CacheDiTConfig",
"CacheDiTRequestSpec",
"RequestScopedCacheDiTRuntime",
"SensenovaCachedAdapter",
"cache_summary",
"enable_cache_for_dit",
+101
View File
@@ -0,0 +1,101 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
from vllm_omni.diffusion.cache.cachedit import runtime as runtime_module
from vllm_omni.diffusion.cache.cachedit.runtime import (
CacheDiTRequestSpec,
RequestScopedCacheDiTRuntime,
)
from vllm_omni.diffusion.data import DiffusionCacheConfig
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _spec(key: str, steps: int = 50) -> CacheDiTRequestSpec:
return CacheDiTRequestSpec(
installation_key=key,
cache_config=DiffusionCacheConfig(),
num_inference_steps=steps,
)
def _recording_backend(events):
class FakeBackend:
def __init__(self, config):
events.append(("create", config))
self.enabled = False
def enable(self, pipeline):
events.append(("enable", pipeline))
self.enabled = True
def refresh(self, pipeline, num_inference_steps):
events.append(("refresh", pipeline, num_inference_steps))
def disable(self, pipeline):
events.append(("disable", pipeline))
self.enabled = False
def is_enabled(self):
return self.enabled
return FakeBackend
def test_startup_cache_follows_omit_lossless_high_omit_state_machine(monkeypatch):
events = []
FakeBackend = _recording_backend(events)
monkeypatch.setattr(runtime_module, "CacheDiTBackend", FakeBackend)
pipeline = object()
config = DiffusionCacheConfig()
runtime = RequestScopedCacheDiTRuntime(pipeline)
startup_backend = FakeBackend(config)
startup_backend.enabled = True
runtime.adopt(startup_backend, installation_key="generic")
runtime.prepare(_spec("generic", 50)) # omitted: keep startup cache
runtime.prepare(None) # lossless: uninstall
runtime.prepare(None) # repeated lossless: no-op
runtime.prepare(_spec("high", 40)) # high: install H3 profile
runtime.prepare(_spec("high", 30)) # repeated high: refresh only
runtime.prepare(_spec("generic", 20)) # omitted: restore startup profile
assert [event[0] for event in events] == [
"create",
"refresh",
"disable",
"create",
"enable",
"refresh",
"refresh",
"disable",
"create",
"enable",
"refresh",
]
assert events[1][-1] == 50
assert events[5][-1] == 40
assert events[6][-1] == 30
assert events[10][-1] == 20
assert runtime.is_enabled
def test_high_installs_without_startup_cache_and_lossless_uninstalls(monkeypatch):
events = []
monkeypatch.setattr(runtime_module, "CacheDiTBackend", _recording_backend(events))
pipeline = object()
runtime = RequestScopedCacheDiTRuntime(pipeline)
runtime.prepare(_spec("high", 40))
runtime.prepare(None)
assert [event[0] for event in events] == [
"create",
"enable",
"refresh",
"disable",
]
assert events[2][-1] == 40
assert not runtime.is_enabled
+30 -4
View File
@@ -244,7 +244,7 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey
captured: dict[str, object] = {}
def rotary_position_embedding(x, cos, sin, **kwargs):
captured.update(x=x, cos=cos, sin=sin, kwargs=kwargs)
captured.update(cos=cos, sin=sin, kwargs=kwargs)
return x
monkeypatch.setitem(
@@ -256,10 +256,8 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey
cos = torch.randn(11, 96, dtype=torch.bfloat16)
sin = torch.randn(11, 96, dtype=torch.bfloat16)
actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
assert actual is x
assert captured["x"] is x
assert captured["cos"] is cos
assert captured["sin"] is sin
assert captured["kwargs"] == {
@@ -267,3 +265,31 @@ def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.Monkey
"head_first": False,
"fused": True,
}
def test_mindie_rope_restores_3d_layout(monkeypatch: pytest.MonkeyPatch) -> None:
"""mindiesd rotary_position_embedding only accepts 4D ``x``."""
captured: dict[str, torch.Tensor] = {}
def rotary_position_embedding(x, cos, sin, **kwargs):
# Mirror mindiesd's own input guard.
assert x.dim() == 4, f"mindiesd requires 4D x, got {x.dim()}"
captured.update(x=x)
return x
monkeypatch.setitem(
sys.modules,
"mindiesd",
types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding),
)
S, H, D = 11, 3, 96
x = torch.randn(S, H, D, dtype=torch.bfloat16)
cos = torch.randn(S, D, dtype=torch.bfloat16)
sin = torch.randn(S, D, dtype=torch.bfloat16)
actual = apply_rotary_emb_mindiesd(x, cos, sin, interleaved=False, half_head_dim=False)
assert captured["x"].shape == (1, S, H, D)
assert actual.shape == (S, H, D)
torch.testing.assert_close(actual, x, atol=0, rtol=0)
@@ -16,6 +16,7 @@ from PIL import Image
from torch import nn
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.experimental.world_models.session_state import SessionStateManager
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -425,6 +426,29 @@ def test_distilled_generation_rejects_robolab_policy(make_cosmos3_pipeline) -> N
pipeline._forward_robolab_policy(make_sampling_params(), None, 0.0)
def test_forward_threads_request_id_to_robolab(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
robolab_inputs = object()
captured: dict[str, Any] = {}
expected = object()
pipeline._build_robolab_policy_inputs = lambda sp, prompt, request_id: robolab_inputs
def fake_forward_robolab(sp, inputs, pipeline_start, session_id=None):
del sp, inputs, pipeline_start
captured["session_id"] = session_id
return expected
pipeline._forward_robolab_policy = fake_forward_robolab
request = SimpleNamespace(
prompts=["policy"],
sampling_params=make_sampling_params(),
request_id="robolab-request-7",
)
assert pipeline.forward(request) is expected
assert captured["session_id"] == "robolab-request-7"
@pytest.mark.parametrize(
("prompt", "sampling_params", "message"),
[
@@ -595,6 +619,7 @@ def _make_od_config(
return SimpleNamespace(
enable_cpu_offload=False,
enable_diffusion_pipeline_profiler=False,
enable_session_state_manager=False,
model="/nonexistent/model/path",
dtype=torch.float32,
flow_shift=None,
@@ -1673,6 +1698,32 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
pipeline._use_session_state = True
pipeline._memory_manager = SessionStateManager()
class FailingProgress:
def __iter__(self):
raise RuntimeError("progress failed before first step")
pipeline.progress_bar = lambda timesteps: FailingProgress()
with pytest.raises(RuntimeError, match="progress failed"):
pipeline.diffuse(
latents=torch.zeros(1, 2, 1, 1, 1),
timesteps=torch.tensor([7]),
cond_ids=_ids(2),
cond_mask=_mask(),
uncond_ids=_ids(1),
uncond_mask=_mask(),
guidance_scale=1.0,
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
session_id="request-that-fails",
)
assert "request-that-fails" not in pipeline._memory_manager
def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
pipeline = make_cosmos3_pipeline()
latents = torch.zeros(1, 2, 1, 1, 1)
@@ -1702,6 +1753,28 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
torch.testing.assert_close(result, torch.full_like(latents, 254.0))
def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None:
pipeline = make_cosmos3_pipeline()
pipeline._use_session_state = True
with pytest.raises(NotImplementedError, match="richer cache keys"):
pipeline.diffuse_transfer(
latents=torch.zeros(1, 2, 1, 1, 1),
timesteps=torch.tensor([7]),
cond_ids=_ids(2),
cond_mask=_mask(),
uncond_ids=_ids(1),
uncond_mask=_mask(),
guidance_scale=1.0,
control_guidance=1.0,
control_guidance_interval=None,
control_latents=[],
shared_kwargs={},
velocity_mask=torch.ones(1, 1, 1, 1, 1),
condition_latents=torch.zeros(1, 2, 1, 1, 1),
)
def test_diffuse_transfer_skips_idle_cfg_branches(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
latents = torch.zeros(1, 2, 1, 1, 1)
velocity_mask = torch.ones(1, 1, 1, 1, 1)
@@ -9,10 +9,75 @@ from unittest.mock import Mock
import numpy as np
import pytest
import torch
import torch.nn as nn
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pipeline = object.__new__(MiniMaxH3Pipeline)
torch.nn.Module.__init__(pipeline)
pipeline.partition = "fl2va"
pipeline.supported_tasks = frozenset({"t2va"})
pipeline.default_video_shift = 12.0
pipeline.default_audio_shift = 3.0
pipeline.device = torch.device("cpu")
pipeline.od_config = SimpleNamespace()
cache_spec = object()
quality_plan = SimpleNamespace(cache_dit=cache_spec)
pipeline._quality_policy = Mock()
pipeline._quality_policy.resolve.return_value = quality_plan
events = []
pipeline._cache_dit_runtime = SimpleNamespace(prepare=lambda spec: events.append(("prepare", spec)))
pipeline.encode_prompt = Mock(
return_value=(
torch.ones(1, 2),
torch.ones(1, dtype=torch.long),
)
)
def diffuse(**kwargs):
events.append(("diffuse", kwargs))
return torch.zeros(1), torch.zeros(1)
pipeline.diffuse = diffuse
pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
sampling = OmniDiffusionSamplingParams(
quality="high",
width=1344,
height=768,
fps=24,
num_frames=124,
num_inference_steps=50,
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
)
batch = DiffusionRequestBatch(
[
OmniDiffusionRequest(
prompt="quality boundary",
sampling_params=sampling,
request_id="quality-boundary",
)
]
)
output = pipeline.forward(batch)
assert events[0] == ("prepare", cache_spec)
assert events[1][0] == "diffuse"
pipeline._quality_policy.resolve.assert_called_once_with(
quality="high",
num_inference_steps=50,
extra_args={"task": "t2va", "aspect_ratio": "16:9"},
)
assert output.output == pipeline.decode.return_value
def test_pipeline_import_registry_and_component_discovery():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
from vllm_omni.diffusion.registry import (
@@ -1401,3 +1466,149 @@ def test_g4_reference_video_metadata_validation(field, value, message, tmp_path)
metadata[field] = value
with pytest.raises(ValueError, match=message):
_validate_reference_video_metadata(metadata, index=0, source=str(tmp_path / "reference.mp4"))
_ENCODER_HIDDEN = 4
_ENCODER_HEAD_DIM = 2
_ENCODER_NUM_HEADS = 2
_ENCODER_NUM_KV_HEADS = 1
_ENCODER_INTERMEDIATE = 4
# One distinct value per checkpoint tensor, none of them a parameter initializer
# (fused weights start as `torch.empty`, norms as ones). A uniform fill would
# prove only that every row was written, not that each source landed in its own
# row range, which is the invariant the fused loaders have to get right.
_SOURCE_FILL = {
"self_attn.q_proj.weight": 3.0,
"self_attn.k_proj.weight": 4.0,
"self_attn.v_proj.weight": 5.0,
"mlp.gate_proj.weight": 6.0,
"mlp.up_proj.weight": 7.0,
"input_layernorm.weight": 8.0,
}
_QKV_TARGET = "text_model.layers.0.self_attn.qkv_proj.weight"
_GATE_UP_TARGET = "text_model.layers.0.mlp.gate_up_proj.weight"
_NORM_TARGET = "text_model.layers.0.input_layernorm.weight"
def _one_layer_text_encoder():
"""Encoder stub holding both fused weights plus one plain parameter."""
from vllm_omni.diffusion.models.minimax_h3.encoder import (
MiniMaxH3Qwen3VLEncoder,
MiniMaxH3Qwen3VLMergedColumnParallelLinear,
MiniMaxH3Qwen3VLQKVParallelLinear,
MiniMaxH3Qwen3VLRMSNorm,
)
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _SingleRankEncoderGroup
group = _SingleRankEncoderGroup(0)
layer = nn.Module()
layer.self_attn = nn.Module()
layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
group,
hidden_size=_ENCODER_HIDDEN,
num_heads=_ENCODER_NUM_HEADS,
num_kv_heads=_ENCODER_NUM_KV_HEADS,
head_dim=_ENCODER_HEAD_DIM,
dtype=torch.float32,
)
layer.mlp = nn.Module()
layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
group,
input_size=_ENCODER_HIDDEN,
intermediate_size=_ENCODER_INTERMEDIATE,
dtype=torch.float32,
)
layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN)
encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
nn.Module.__init__(encoder)
encoder.text_model = nn.Module()
encoder.text_model.layers = nn.ModuleList([layer])
return encoder
def _write_text_encoder_checkpoint(path, *, omit=()):
"""Write a one-layer Qwen3-VL-named checkpoint, one fill value per source."""
import json
import safetensors.torch
prefix = "model.language_model.layers.0"
q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
kv_rows = _ENCODER_NUM_KV_HEADS * _ENCODER_HEAD_DIM
shapes = {
"self_attn.q_proj.weight": (q_rows, _ENCODER_HIDDEN),
"self_attn.k_proj.weight": (kv_rows, _ENCODER_HIDDEN),
"self_attn.v_proj.weight": (kv_rows, _ENCODER_HIDDEN),
"mlp.gate_proj.weight": (_ENCODER_INTERMEDIATE, _ENCODER_HIDDEN),
"mlp.up_proj.weight": (_ENCODER_INTERMEDIATE, _ENCODER_HIDDEN),
"input_layernorm.weight": (_ENCODER_HIDDEN,),
}
tensors = {f"{prefix}.{source}": torch.full(shape, _SOURCE_FILL[source]) for source, shape in shapes.items()}
for source in omit:
del tensors[f"{prefix}.{source}"]
safetensors.torch.save_file(tensors, str(path / "model.safetensors"))
index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")}
(path / "model.safetensors.index.json").write_text(json.dumps(index))
@pytest.mark.parametrize(
("omitted_sources", "expected_detail", "unreported_target"),
[
(("self_attn.q_proj.weight",), f"{_QKV_TARGET}: ['q']", _GATE_UP_TARGET),
(("self_attn.k_proj.weight",), f"{_QKV_TARGET}: ['k']", _GATE_UP_TARGET),
(("self_attn.v_proj.weight",), f"{_QKV_TARGET}: ['v']", _GATE_UP_TARGET),
(("mlp.gate_proj.weight",), f"{_GATE_UP_TARGET}: ['gate']", _QKV_TARGET),
(("mlp.up_proj.weight",), f"{_GATE_UP_TARGET}: ['up']", _QKV_TARGET),
(
("self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"),
f"{_QKV_TARGET}: ['q', 'k', 'v']",
_GATE_UP_TARGET,
),
],
)
def test_encoder_load_reports_missing_fused_source_shards(
tmp_path, omitted_sources, expected_detail, unreported_target
):
encoder = _one_layer_text_encoder()
_write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
with pytest.raises(RuntimeError) as excinfo:
encoder._load_weights(str(tmp_path))
message = str(excinfo.value)
assert expected_detail in message
assert unreported_target not in message
def test_encoder_load_rejects_unloaded_plain_param(tmp_path):
encoder = _one_layer_text_encoder()
_write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
encoder._load_weights(str(tmp_path))
def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
encoder = _one_layer_text_encoder()
_write_text_encoder_checkpoint(tmp_path)
encoder._load_weights(str(tmp_path))
params = dict(encoder.named_parameters())
q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
kv_rows = _ENCODER_NUM_KV_HEADS * _ENCODER_HEAD_DIM
qkv = params[_QKV_TARGET]
gate_up = params[_GATE_UP_TARGET]
slices = {
"self_attn.q_proj.weight": qkv[:q_rows],
"self_attn.k_proj.weight": qkv[q_rows : q_rows + kv_rows],
"self_attn.v_proj.weight": qkv[q_rows + kv_rows :],
"mlp.gate_proj.weight": gate_up[:_ENCODER_INTERMEDIATE],
"mlp.up_proj.weight": gate_up[_ENCODER_INTERMEDIATE:],
"input_layernorm.weight": params[_NORM_TARGET],
}
for source, rows in slices.items():
assert torch.equal(rows, torch.full_like(rows, _SOURCE_FILL[source])), source
@@ -0,0 +1,212 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
from vllm_omni.diffusion.data import DiffusionCacheConfig
from vllm_omni.errors import OmniClientError
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _quality_od_config(*, cache_backend: str = "cache_dit", **overrides):
values = {
"cache_backend": cache_backend,
"cache_config": DiffusionCacheConfig(),
"diffusion_attention_config": SimpleNamespace(default=None, per_role={}),
"enable_cpu_offload": False,
"enable_layerwise_offload": False,
"enable_distributed_layerwise_offload": False,
"enforce_eager": True,
"num_gpus": 4,
"quantization_config": None,
"parallel_config": SimpleNamespace(
allgather_degree=1,
cfg_parallel_size=1,
data_parallel_size=1,
pipeline_parallel_size=1,
ring_degree=1,
sequence_parallel_size=4,
tensor_parallel_size=1,
ulysses_degree=4,
use_hsdp=False,
),
}
values.update(overrides)
return SimpleNamespace(**values)
def _resolve_quality(policy, **overrides):
values = {
"quality": "high",
"num_inference_steps": 50,
"extra_args": None,
}
values.update(overrides)
return policy.resolve(**values)
def test_high_quality_policy_emits_cache_profile_without_deployment_gates():
from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
MiniMaxH3QualityPolicy,
)
policy = MiniMaxH3QualityPolicy(
_quality_od_config(
cache_backend="none",
enable_layerwise_offload=True,
enforce_eager=False,
num_gpus=1,
)
)
plan = _resolve_quality(policy, num_inference_steps=17)
spec = plan.cache_dit
assert spec is not None
assert spec.installation_key == "minimax_h3.high"
assert spec.num_inference_steps == 17
assert spec.cache_config.Fn_compute_blocks == 1
assert spec.cache_config.Bn_compute_blocks == 0
assert spec.cache_config.max_warmup_steps == 4
assert spec.cache_config.residual_diff_threshold == pytest.approx(0.04)
assert spec.cache_config.max_continuous_cached_steps == 1
assert spec.cache_config.enable_taylorseer is False
assert spec.cache_config.scm_steps_mask_policy is None
def test_h3_force_refresh_is_model_owned_and_request_scoped():
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
policy = MiniMaxH3QualityPolicy(_quality_od_config())
plan = _resolve_quality(
policy,
quality="high",
num_inference_steps=17,
extra_args={"force_refresh_step_hint": 7},
)
assert plan.cache_dit is not None
assert plan.cache_dit.installation_key == "minimax_h3.high:force_refresh=7:once"
assert plan.cache_dit.cache_config.force_refresh_step_hint == 7
assert plan.cache_dit.cache_config.force_refresh_step_policy == "once"
def test_h3_force_refresh_policy_supports_repeat_and_does_not_mutate_generic_config():
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
config = _quality_od_config()
plan = _resolve_quality(
MiniMaxH3QualityPolicy(config),
quality=None,
num_inference_steps=20,
extra_args={"force_refresh_step_hint": 5, "force_refresh_step_policy": "repeat"},
)
assert plan.cache_dit is not None
assert plan.cache_dit.installation_key == "minimax_h3.generic:force_refresh=5:repeat"
assert plan.cache_dit.cache_config is not config.cache_config
assert plan.cache_dit.cache_config.force_refresh_step_hint == 5
assert plan.cache_dit.cache_config.force_refresh_step_policy == "repeat"
assert config.cache_config.force_refresh_step_hint is None
assert config.cache_config.force_refresh_step_policy == "once"
restored = _resolve_quality(MiniMaxH3QualityPolicy(config), quality=None)
assert restored.cache_dit is not None
assert restored.cache_dit.installation_key == "minimax_h3.generic"
assert restored.cache_dit.cache_config is config.cache_config
@pytest.mark.parametrize(
("extra_args", "message"),
[
({"force_refresh_step_hint": True}, "positive integer"),
({"force_refresh_step_hint": 0}, "between 1"),
({"force_refresh_step_hint": 51}, "between 1"),
({"force_refresh_step_hint": 5, "force_refresh_step_policy": "always"}, "one of"),
({"force_refresh_step_policy": "repeat"}, "requires force_refresh_step_hint"),
],
)
def test_h3_force_refresh_rejects_invalid_values(extra_args, message):
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
with pytest.raises(OmniClientError, match=message):
_resolve_quality(
MiniMaxH3QualityPolicy(_quality_od_config()),
extra_args=extra_args,
)
@pytest.mark.parametrize(
("cache_backend", "quality"),
[
("cache_dit", "lossless"),
("none", None),
],
)
def test_h3_force_refresh_requires_active_cache_target(cache_backend, quality):
from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy
policy = MiniMaxH3QualityPolicy(_quality_od_config(cache_backend=cache_backend))
with pytest.raises(OmniClientError, match="require an active Cache-DiT"):
_resolve_quality(
policy,
quality=quality,
extra_args={"force_refresh_step_hint": 7},
)
@pytest.mark.parametrize(
("cache_backend", "quality", "expected_key"),
[
("cache_dit", None, "minimax_h3.generic"),
("cache_dit", "lossless", None),
("cache_dit", "high", "minimax_h3.high"),
("none", None, None),
("none", "lossless", None),
("none", "high", "minimax_h3.high"),
],
)
def test_model_policy_owns_request_cache_target(cache_backend, quality, expected_key):
from vllm_omni.diffusion.models.minimax_h3.quality_policy import (
MiniMaxH3QualityPolicy,
)
config = _quality_od_config(cache_backend=cache_backend)
policy = MiniMaxH3QualityPolicy(config)
plan = _resolve_quality(
policy,
quality=quality,
)
if expected_key is None:
assert plan.cache_dit is None
else:
assert plan.cache_dit is not None
assert plan.cache_dit.installation_key == expected_key
if quality is None:
assert plan.cache_dit.cache_config is config.cache_config
def test_h3_adopts_runner_installed_cache_dit_backend():
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
backend = object()
pipeline = object.__new__(MiniMaxH3Pipeline)
pipeline._cache_dit_runtime = SimpleNamespace(
adopt=lambda adopted, *, installation_key: setattr(
pipeline,
"adopted",
(adopted, installation_key),
),
is_enabled=True,
)
pipeline.adopt_cache_dit_backend(backend)
assert pipeline.adopted == (backend, "minimax_h3.generic")
assert pipeline.is_cache_dit_enabled()
+42 -3
View File
@@ -178,10 +178,9 @@ class TestAsyncOutputLoopLogic:
t = threading.Thread(target=proc._async_output_loop, daemon=True)
t.start()
# Give the loop time to process, then unblock and stop.
# Give the loop time to process, then send the shutdown sentinel.
time.sleep(0.5)
proc._running = False
proc._async_output_queue.put((DiffusionOutput(), "sentinel", None))
proc._async_output_queue.put(None)
t.join(timeout=2.0)
mock_pack.assert_called()
@@ -191,3 +190,43 @@ class TestAsyncOutputLoopLogic:
for c in proc.result_mq.enqueue.call_args_list
)
assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"
class TestWorkerProcShutdown:
def test_shutdown_stops_async_thread_and_releases_queues(self):
import threading
proc = _make_worker_proc(step_execution=False)
proc.worker = MagicMock()
proc.mq = MagicMock()
broadcast_mq = proc.mq
result_mq = proc.result_mq
thread = threading.Thread(target=proc._async_output_queue.get)
proc._async_output_thread = thread
thread.start()
proc.shutdown()
assert not thread.is_alive()
proc.worker.shutdown.assert_called_once_with()
broadcast_mq.shutdown.assert_called_once_with()
result_mq.shutdown.assert_called_once_with()
assert proc.mq is None
assert proc.result_mq is None
def test_shutdown_releases_queues_when_worker_shutdown_fails(self):
proc = _make_worker_proc(step_execution=True)
proc.worker = MagicMock()
proc.worker.shutdown.side_effect = RuntimeError("shutdown failed")
proc.mq = MagicMock()
broadcast_mq = proc.mq
result_mq = proc.result_mq
with pytest.raises(RuntimeError, match="shutdown failed"):
proc.shutdown()
broadcast_mq.shutdown.assert_called_once_with()
result_mq.shutdown.assert_called_once_with()
assert proc.mq is None
assert proc.result_mq is None
@@ -15,7 +15,10 @@ from vllm_omni.diffusion.data import (
DiffusionParallelConfig,
OmniDiffusionConfig,
)
from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
from vllm_omni.diffusion.model_metadata import (
HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -136,6 +139,30 @@ def test_task_type_roundtrip():
assert od.task_type == "model-defined-task"
def test_architecture_name_resolves_via_pipeline_class_fallback():
# Direct hit: the architecture name itself is the metadata table key.
qwen_od_config = OmniDiffusionConfig(
model="Qwen/Qwen-Image-Edit-2511",
model_class_name="QwenImageEditPlusPipeline",
)
qwen_od_config.update_multimodal_support()
assert qwen_od_config.supports_multimodal_inputs is True
assert qwen_od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
# Fallback: the HF architecture name differs from the internal pipeline class
# name, so ``get_diffusion_model_metadata`` must consult ``_DIFFUSION_MODELS``
# and resolve the metadata via the pipeline class name.
hunyuan_od_config = OmniDiffusionConfig(
model="tencent/HunyuanImage-3.0-Instruct",
model_class_name="HunyuanImage3ForCausalMM",
)
hunyuan_od_config.update_multimodal_support()
assert hunyuan_od_config.supports_multimodal_inputs is True
assert hunyuan_od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
def test_additional_config_roundtrip():
additional_config = {"torchair_graph_config": {"enabled": True}}
od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)
+49
View File
@@ -331,6 +331,55 @@ class TestRequestBatchCapability:
DiffusionEngine(od_config)
fake_executor_cls.assert_not_called()
def test_engine_allows_independent_dlo_dp_requests_for_single_request_pipeline(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(
model_class_name="SinglePipeline",
custom_pipeline_args=None,
streaming_output=False,
max_num_seqs=2,
request_batch_max_wait_ms=0,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
fake_executor = SimpleNamespace(
execute_request=mocker.Mock(return_value="per-request"),
execute_batch=mocker.Mock(return_value="batch"),
execute_step=mocker.Mock(return_value="step"),
)
fake_executor_cls = mocker.Mock(return_value=fake_executor)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
lambda *args, **kwargs: fake_executor_cls,
)
monkeypatch.setattr(
diffusion_engine_module.DiffusionModelRegistry,
"_try_load_model_cls",
lambda model_class_name: _SingleRequestPipeline,
)
engine = DiffusionEngine(od_config)
assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
assert engine.supports_request_batch is False
assert engine.dp_concurrent is True
assert engine.scheduler.max_num_running_reqs == 2
assert od_config.request_batch_max_wait_ms == 0
fake_executor_cls.assert_called_once_with(od_config)
@pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
self,
+173
View File
@@ -2,11 +2,19 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import contextlib
import multiprocessing as mp
import queue
import threading
import time
from multiprocessing import shared_memory
import numpy as np
import pytest
import torch
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
from vllm_omni.diffusion.ipc import (
_SHM_TENSOR_THRESHOLD,
DIFFUSION_RPC_RESULT_ENVELOPE,
@@ -30,6 +38,129 @@ def _cleanup_shm_handle(value: object) -> None:
_unpack_if_shm_handle(value)
def _result_queue_worker(connection, rank: int) -> None:
"""Send one nested NumPy result through a worker-owned MessageQueue."""
result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
try:
connection.send(result_mq.export_handle())
assert connection.recv() == "reader-ready"
frames = np.full(300_000, rank, dtype=np.float32)
audio = np.arange(300_000, dtype=np.float32) + rank
output = DiffusionOutput(
output={
"rank": rank,
"nested": {"frames": frames},
"audio": [audio],
}
)
pack_diffusion_output_shm(output)
connection.send(
[
output.output["nested"]["frames"]["name"],
output.output["audio"][0]["name"],
]
)
result_mq.enqueue(output)
assert connection.recv() == "shutdown"
finally:
result_mq.shutdown()
connection.close()
def test_per_worker_result_queues_release_nested_numpy_shm_and_processes() -> None:
"""Exercise the production per-worker queue/pump/SHM lifecycle."""
ctx = mp.get_context("spawn")
parent_connections = []
processes = []
result_mqs = []
shm_names = []
executor = None
try:
for rank in range(2):
parent_connection, child_connection = ctx.Pipe()
process = ctx.Process(
target=_result_queue_worker,
args=(child_connection, rank),
name=f"IPCResultWorker-{rank}",
)
process.start()
child_connection.close()
parent_connections.append(parent_connection)
processes.append(process)
for connection in parent_connections:
handle = connection.recv()
result_mqs.append(MessageQueue.create_from_handle(handle, 0))
connection.send("reader-ready")
shm_names.extend(connection.recv())
executor = object.__new__(MultiprocDiffusionExecutor)
executor._result_mqs = result_mqs
executor._result_mq = result_mqs[0]
executor._pump_running = False
executor._pump_stop = threading.Event()
executor._sync_result_buffer = queue.Queue()
executor._is_failed = False
executor._futures_lock = threading.RLock()
executor._rpc_futures = {}
executor._output_futures = {}
executor._completed_outputs = {}
executor._batch_split_map = {}
executor._start_result_pump()
deadline = time.monotonic() + 10.0
while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
time.sleep(0.01)
assert executor._sync_result_buffer.qsize() == 2
outputs = []
for _ in range(2):
output = executor._sync_result_buffer.get_nowait()
unpack_diffusion_output_shm(output)
outputs.append(output)
assert sorted(output.output["rank"] for output in outputs) == [0, 1]
for output in outputs:
rank = output.output["rank"]
assert output.output["nested"]["frames"][0] == rank
assert output.output["audio"][0][0] == rank
executor._pump_stop.set()
for thread in executor._result_pump_threads:
thread.join(timeout=3.0)
assert all(not thread.is_alive() for thread in executor._result_pump_threads)
finally:
if executor is not None:
executor._pump_stop.set()
for thread in getattr(executor, "_result_pump_threads", []):
thread.join(timeout=3.0)
for connection in parent_connections:
with contextlib.suppress(BrokenPipeError, EOFError, OSError):
connection.send("shutdown")
for process in processes:
process.join(timeout=10.0)
if process.is_alive():
process.terminate()
process.join(timeout=5.0)
for result_mq in result_mqs:
with contextlib.suppress(Exception):
result_mq.shutdown()
for connection in parent_connections:
connection.close()
assert all(process.exitcode == 0 for process in processes)
for name in shm_names:
try:
leaked = shared_memory.SharedMemory(name=name)
except FileNotFoundError:
continue
leaked.close()
pytest.fail(f"shared memory segment {name} still exists after unpack")
def test_diffusion_output_dict_tensors_round_trip_through_shm() -> None:
image = torch.arange(300_000, dtype=torch.float32)
video = torch.arange(300_000, dtype=torch.float32) * 2
@@ -83,6 +214,29 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
torch.testing.assert_close(output.output[1], frames[1])
def test_diffusion_output_numpy_array_round_trips_through_shm() -> None:
frames = np.arange(300_000, dtype=np.float32)
output = DiffusionOutput(output=frames)
pack_diffusion_output_shm(output)
assert output.output["__ndarray_shm__"] is True
unpack_diffusion_output_shm(output)
assert isinstance(output.output, np.ndarray)
np.testing.assert_array_equal(output.output, frames)
def test_numpy_object_array_stays_inline() -> None:
values = np.empty(200_000, dtype=object)
values[:] = "safe-to-pickle"
packed = _pack_value_if_large(values)
assert packed is values
def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
tensor = torch.arange(300_000, dtype=torch.float32)
envelope = {
@@ -108,6 +262,25 @@ def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
def test_rpc_result_envelope_dp_tagged_output_round_trips_through_shm() -> None:
frames = np.arange(300_000, dtype=np.float32)
envelope = {
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
"result": {"dp_rank": 1, "output": DiffusionOutput(output=frames)},
"rank_statuses": [{"rank": 1, "ok": True}],
}
packed = pack_diffusion_output_shm(envelope)
tagged = packed["result"]
assert tagged["dp_rank"] == 1
assert tagged["output"].output["__ndarray_shm__"] is True
unpacked = unpack_diffusion_output_shm(packed)
tagged = unpacked["result"]
assert tagged["dp_rank"] == 1
np.testing.assert_array_equal(tagged["output"].output, frames)
def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
first = torch.arange(_large_numel(torch.float32), dtype=torch.float32)
second = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + 1
@@ -162,6 +162,21 @@ def _fake_platform_for_peak_memory():
)
def test_request_scoped_cache_dit_lifecycle_is_pipeline_opt_in():
events = []
plain_pipeline = object()
request_scoped_pipeline = SimpleNamespace(
adopt_cache_dit_backend=lambda backend: events.append(backend),
is_cache_dit_enabled=lambda: True,
)
backend = object()
assert not model_runner_module.adopt_request_scoped_cache_dit(plain_pipeline, backend)
assert model_runner_module.adopt_request_scoped_cache_dit(request_scoped_pipeline, backend)
assert model_runner_module.is_request_scoped_cache_dit_enabled(request_scoped_pipeline)
assert events == [backend]
def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
@@ -419,6 +434,33 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
assert cache_backend.refresh_calls == [(runner.pipeline, 4, True)]
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("cache_dit_enabled", [False, True])
def test_execute_model_cache_summary_follows_pipeline_owned_cache_dit_state(monkeypatch, cache_dit_enabled):
runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
runner.pipeline.adopt_cache_dit_backend = lambda backend: None
runner.pipeline.is_cache_dit_enabled = lambda: cache_dit_enabled
req = _make_request()
cache_summary_calls = []
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
monkeypatch.setattr(
model_runner_module,
"cache_summary",
lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
)
output = DiffusionModelRunner.execute_model(runner, req)
assert output.output == "ok"
assert cache_summary_calls == ([(runner.pipeline, True)] if cache_dit_enabled else [])
@pytest.mark.core_model
@pytest.mark.cpu
def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch):
@@ -380,6 +380,63 @@ class TestRequestScheduler:
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
def test_batches_different_quality_levels_separately(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
high = scheduler.add_request(
_make_step_request(
"high",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=2,
quality="high",
),
)
)
scheduler.add_request(
_make_step_request(
"lossless",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=2,
quality="lossless",
),
)
)
first = scheduler.schedule()
assert _new_ids(first) == [high]
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
def test_batches_omitted_and_explicit_lossless_separately(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
omitted = scheduler.add_request(
_make_step_request(
"omitted",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=2,
),
)
)
scheduler.add_request(
_make_step_request(
"explicit",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=2,
quality="lossless",
),
)
)
first = scheduler.schedule()
assert _new_ids(first) == [omitted]
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
def test_batches_different_request_local_seed_together(self) -> None:
scheduler = RequestScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
@@ -930,6 +987,35 @@ class TestStepScheduler:
assert second.num_running_reqs == 1
assert second.num_waiting_reqs == 1
def test_step_batch_rejects_different_quality_levels(self) -> None:
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=2))
high = scheduler.add_request(
_make_step_request(
"high",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=4,
quality="high",
),
)
)
scheduler.add_request(
_make_step_request(
"lossless",
sampling_params=OmniDiffusionSamplingParams(
num_inference_steps=4,
quality="lossless",
),
)
)
first = scheduler.schedule()
assert _new_ids(first) == [high]
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
"""Multiple requests with the same LoRA (id + scale) co-batch."""
from vllm_omni.lora.request import LoRARequest
@@ -75,6 +75,32 @@ async def test_inline_dispatch_request_success(client, mock_engine):
assert output.request_id == "req-1"
@pytest.mark.asyncio
async def test_inline_requests_clone_shared_sampling_params(client, mock_engine):
requests = []
async def _step_streaming(request):
requests.append(request)
yield [OmniRequestOutput.from_diffusion(request_id=request.request_id, images=[MagicMock()])]
mock_engine.step_streaming = _step_streaming
shared = OmniDiffusionSamplingParams()
await client.add_request_async("req-1", "First prompt", shared)
await client.add_request_async("req-2", "Second prompt", shared)
for _ in range(10):
if len(requests) == 2:
break
await asyncio.sleep(0.01)
assert len(requests) == 2
assert requests[0].sampling_params is not requests[1].sampling_params
assert requests[0].sampling_params.guidance_scale_provided is False
assert requests[1].sampling_params.guidance_scale_provided is False
assert shared.guidance_scale is None
@pytest.mark.asyncio
async def test_inline_dispatch_request_streaming_success(client, mock_engine):
"""Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine"""
@@ -137,6 +163,20 @@ def test_inline_shutdown(client, mock_engine):
mock_engine.close.assert_called_once()
def test_inline_shutdown_runs_after_orchestrator_premark(client, mock_engine):
# The orchestrator pre-marks clients to suppress false worker-death errors
# before StagePool invokes the actual teardown.
client._shutting_down = True
client.shutdown()
mock_engine.close.assert_called_once()
assert client._shutdown_complete
client.shutdown()
mock_engine.close.assert_called_once()
def test_inline_registers_executor_failure_callback(client, mock_engine):
mock_engine.executor.register_failure_callback.assert_called_once()
+12
View File
@@ -137,6 +137,18 @@ class TestPackDiffusionOutputShmWithD2hStream:
d2h_stream=d2h,
)
@pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True])
def test_dp_tagged_output_path_with_stream(self, mocker, wrapped_in_rpc_envelope):
mock_pack_fields = mocker.patch(
"vllm_omni.diffusion.ipc._pack_diffusion_fields",
)
d2h = mocker.MagicMock()
result = DiffusionOutput()
tagged = {"dp_rank": 1, "output": result}
value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged
pack_diffusion_output_shm(value, d2h_stream=d2h)
mock_pack_fields.assert_called_once_with(result, d2h_stream=d2h)
def test_runner_output_path_with_stream(self, mocker):
mock_pack_fields = mocker.patch(
"vllm_omni.diffusion.ipc._pack_diffusion_fields",
@@ -276,6 +276,8 @@ class TestRequestModeDispatch:
def worker():
request = req_q.get(timeout=2)
assert "rpc_id" not in request
assert request["output_rank"] is None
assert request["collect_rank_status"] is False
wave_id = request["wave_id"]
for dp_rank in range(2):
res_q.put(
@@ -297,6 +299,181 @@ class TestRequestModeDispatch:
assert [result.error for result in results] == ["0", "1"]
@pytest.mark.parametrize(
("tp_size", "sp_size", "pp_size", "cfg_size", "expected_primary_ranks"),
[
(2, 1, 1, 1, [0, 2]),
(2, 2, 1, 1, [0, 4]),
(1, 2, 2, 1, [0, 4]),
(1, 1, 2, 2, [0, 4]),
],
)
def test_step_execution_reads_each_dp_primary_result_queue(
self,
tp_size,
sp_size,
pp_size,
cfg_size,
expected_primary_ranks,
):
dp_size = 2
num_gpus = dp_size * tp_size * sp_size * pp_size * cfg_size
executor, _, _ = _make_executor(num_gpus=num_gpus)
executor.od_config = SimpleNamespace(
step_execution=True,
parallel_config=SimpleNamespace(
data_parallel_size=dp_size,
tensor_parallel_size=tp_size,
sequence_parallel_size=sp_size,
pipeline_parallel_size=pp_size,
cfg_parallel_size=cfg_size,
),
)
request: dict = {}
executor._broadcast_mq = SimpleNamespace(enqueue=lambda value: request.update(value))
def result_queue(global_rank):
result_mq = MagicMock()
result_mq.dequeue.side_effect = lambda timeout=None: {
"dp_rank": global_rank,
"output": _tagged_output(str(global_rank)),
"wave_id": request["wave_id"],
}
return result_mq
# Every queue returns a valid, rank-tagged result. If the executor
# selects a non-primary queue, the returned rank exposes the mistake.
executor._result_mqs = [result_queue(rank) for rank in range(num_gpus)]
executor._result_mq = executor._result_mqs[0]
results = executor.collective_rpc(
"execute_model",
unique_reply_rank=None,
exec_all_ranks=True,
)
assert [result.error for result in results] == [str(rank) for rank in expected_primary_ranks]
for rank, result_mq in enumerate(executor._result_mqs):
assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
@pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor.collective_rpc = Mock()
scheduler_output = _make_sched_output("invalid", "valid")
scheduler_output.scheduled_new_reqs[0].req.prompt = empty_prompt
with pytest.raises(ValueError, match="non-empty prompt"):
executor.execute_request(scheduler_output)
executor.collective_rpc.assert_not_called()
def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
invalid_wave = _make_sched_output("invalid", "valid")
invalid_wave.scheduled_new_reqs[0].req.prompt = ""
with pytest.raises(ValueError, match="non-empty prompt"):
executor.execute_request(invalid_wave)
result = executor.execute_request(_make_sched_output("A", "B"))
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
executor.collective_rpc.assert_called_once()
def test_dlo_dp_allows_shared_default_denoise_steps(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
scheduler_output = _make_sched_output("A", "B")
for new_req in scheduler_output.scheduled_new_reqs:
new_req.req.sampling_params.num_inference_steps = None
result = executor.execute_request(scheduler_output)
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
executor.collective_rpc.assert_called_once()
@pytest.mark.parametrize(
("field", "value"),
[
("num_inference_steps", 2),
("guidance_scale", 7.5),
("width", 1024),
],
)
def test_dlo_dp_rejects_incompatible_collective_wave(self, field, value):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor.collective_rpc = Mock()
scheduler_output = _make_sched_output("A", "B")
setattr(scheduler_output.scheduled_new_reqs[1].req.sampling_params, field, value)
with pytest.raises(ValueError, match="compatible shape, CFG"):
executor.execute_request(scheduler_output)
executor.collective_rpc.assert_not_called()
def test_dlo_dp_partial_reply_times_out_and_fails_closed(self, monkeypatch):
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
executor, req_q, res_q = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor._sync_result_buffer = res_q
executor._fail_closed_on_dp_wave_timeout = Mock()
monkeypatch.setattr(executor_module, "_DLO_DP_WAVE_TIMEOUT_S", 0.05)
def reply_from_only_one_dp_rank():
request = req_q.get(timeout=2.0)
res_q.put(
{
"dp_rank": 0,
"output": _tagged_output("rank-0"),
"wave_id": request["wave_id"],
}
)
worker = threading.Thread(target=reply_from_only_one_dp_rank, daemon=True)
worker.start()
started = time.monotonic()
result = executor.execute_request(_make_sched_output("A", "B"))
elapsed = time.monotonic() - started
worker.join(timeout=2.0)
assert elapsed < 1.0
assert all("timed out" in output.result.error for output in result.runner_outputs)
executor._fail_closed_on_dp_wave_timeout.assert_called_once()
assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
# ───────────────── concurrent collective RPC ─────────────────
@@ -948,6 +1125,42 @@ class TestStageDiffusionClientErrorPropagation:
assert rpc_id not in client._pending_rpcs
class TestExecutorShutdownCleaner:
def test_worker_joins_share_one_global_deadline(self, monkeypatch):
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
class FakeProcess:
def __init__(self, name):
self.name = name
self.alive = True
self.terminated = False
self.join_timeouts = []
def is_alive(self):
return self.alive
def join(self, timeout):
self.join_timeouts.append(timeout)
if self.terminated:
self.alive = False
def terminate(self):
self.terminated = True
monotonic = Mock(side_effect=[100.0, 100.0, 110.0, 120.0, 120.0, 124.0])
monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
first = FakeProcess("worker-0")
second = FakeProcess("worker-1")
cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
cleaner()
assert first.join_timeouts == [15.0, 5.0]
assert second.join_timeouts == [5.0, 1.0]
assert first.terminated and second.terminated
assert not first.is_alive() and not second.is_alive()
# ───────── monitor thread & death sentinel integration tests ─────────
+50
View File
@@ -35,6 +35,7 @@ def _make_executor(step_execution=False):
executor._pump_stop = threading.Event()
executor._sync_result_buffer = queue.Queue()
executor._result_mq = MagicMock()
executor._result_mqs = []
executor._broadcast_mq = MagicMock()
executor._closed = False
executor._is_failed = False
@@ -157,6 +158,41 @@ class TestResultPumpDispatch:
retrieved = executor._sync_result_buffer.get_nowait()
assert isinstance(retrieved, DiffusionOutput)
def test_start_result_pump_reads_every_worker_queue(self):
executor = _make_executor()
messages = [DiffusionOutput(output="rank0"), DiffusionOutput(output="rank1")]
result_mqs = [MagicMock(), MagicMock()]
def make_dequeue(message):
emitted = False
def dequeue(timeout=None):
nonlocal emitted
if not emitted:
emitted = True
return message
executor._pump_stop.wait(0.01)
raise TimeoutError
return dequeue
for result_mq, message in zip(result_mqs, messages, strict=True):
result_mq.dequeue = make_dequeue(message)
executor._result_mqs = result_mqs
executor._result_mq = result_mqs[0]
executor._start_result_pump()
deadline = time.monotonic() + 2.0
while executor._sync_result_buffer.qsize() < 2 and time.monotonic() < deadline:
time.sleep(0.01)
executor._pump_stop.set()
for thread in executor._result_pump_threads:
thread.join(timeout=2.0)
assert len(executor._result_pump_threads) == 2
received = [executor._sync_result_buffer.get_nowait().output for _ in range(2)]
assert sorted(received) == ["rank0", "rank1"]
def test_compute_done_routes_to_rpc_future(self):
executor = _make_executor()
rpc_id = "42"
@@ -233,6 +269,20 @@ class TestResultPumpDispatch:
class TestShutdownCleansUpFutures:
"""Test that shutdown cancels pending futures."""
def test_shutdown_joins_result_pump_threads(self):
executor = _make_executor()
pump = threading.Thread(
target=executor._pump_stop.wait,
name="test-result-pump",
)
executor._result_pump_threads = [pump]
pump.start()
executor.shutdown()
assert not pump.is_alive()
assert executor._result_pump_threads == []
def test_shutdown_sets_exception_on_pending_futures(self):
executor = _make_executor()
+10 -2
View File
@@ -29,7 +29,7 @@ HEIGHT = 512
NUM_INFERENCE_STEPS = 20
TRUE_CFG_SCALE = 4.0
SEED = 42
SSIM_THRESHOLD = 0.94
SSIM_THRESHOLD = 0.97
PSNR_THRESHOLD = 30.0
MODEL_2512_ID = "Qwen/Qwen-Image-2512"
@@ -69,7 +69,15 @@ def _local_files_only(model: str) -> bool:
def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image:
server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
server_args = [
"--num-gpus",
"1",
"--stage-init-timeout",
"300",
"--init-timeout",
"900",
"--fa-deterministic",
]
with OmniServer(model, server_args, use_omni=True) as omni_server:
response = requests.post(
f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
@@ -136,3 +136,36 @@ def test_flux_2_dev(
}
openai_client.send_diffusion_request(request_config)
@pytest.mark.parametrize(
"omni_server",
_get_flux_2_dev_feature_cases(MODEL),
indirect=True,
)
def test_flux_2_dev_batched_chat_completions(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
):
"""Validate batched chat completions for diffusion models."""
messages = [
[{"role": "user", "content": PROMPT}],
[{"role": "user", "content": "A sunset over the ocean."}],
]
responses = openai_client.send_batched_chat_completions_http_request(
{
"json": {
"model": omni_server.model,
"messages": messages,
"num_inference_steps": 2,
"height": 512,
"width": 512,
"seed": 42,
},
},
)
assert responses and len(responses) == 1
resp = responses[0]
assert resp.success
choices = resp.json_body["choices"]
assert len(choices) == len(messages)
@@ -250,3 +250,101 @@ def test_completions_rejected_for_thinker_talker(omni_server, openai_client) ->
err_code=400,
)
assert not responses[0].success
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_batched_completions_text(omni_server, openai_client) -> None:
"""Ensure that we can make a batch chat completions request (text only)."""
responses = openai_client.send_batched_chat_completions_http_request(
{
"json": {
"model": omni_server.model,
"messages": [
[{"role": "user", "content": "What color is the sky? Answer in one word."}],
[{"role": "user", "content": "What is 2+2? Answer in one word."}],
],
"modalities": ["text"],
"max_tokens": 20,
},
},
)
assert responses
resp = responses[0]
assert resp.success
body = resp.json_body
choices = body["choices"]
assert len(choices) == 2
for choice in choices:
assert choice["message"]["content"]
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_batched_completions_audio_out(omni_server, openai_client) -> None:
"""Ensure that we can make a batch chat completions request (audio + text)."""
messages = [
[{"role": "user", "content": "Say hello."}],
[{"role": "user", "content": "Say goodbye."}],
]
num_messages = len(messages)
responses = openai_client.send_batched_chat_completions_http_request(
{
"json": {
"model": omni_server.model,
"messages": messages,
"modalities": ["text", "audio"],
"max_tokens": 200,
},
},
)
assert responses and len(responses) == 1
resp = responses[0]
assert resp.success
choices = resp.json_body["choices"]
assert len(choices) == num_messages
# Every message in the batch should contain both the audio and text content
for choice in choices:
assert choice["message"]["audio"]["data"]
assert choice["message"]["content"]
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize(
"sad_opts",
[
{"stream": "not-a-bool"}, # Fails validation
{"n": 2}, # vLLM Batched endpoints doesn't support n > 1
],
)
def test_batched_completions_with_bad_values(omni_server, openai_client, sad_opts) -> None:
"""Ensure that a bad values are correctly handled as 400s."""
messages = [
[{"role": "user", "content": "Say hello."}],
[{"role": "user", "content": "Say goodbye."}],
]
responses = openai_client.send_batched_chat_completions_http_request(
{
"json": {
"model": omni_server.model,
"messages": messages,
**sad_opts,
},
},
)
assert responses and len(responses) == 1
resp = responses[0]
assert not responses[0].success
assert resp.status_code == 400
+23 -1
View File
@@ -12,7 +12,7 @@ from unittest.mock import Mock
import pytest
from omegaconf import OmegaConf
from pydantic import ValidationError
from transformers import PretrainedConfig
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
from vllm.engine.arg_utils import EngineArgs
from vllm_omni.config.model import OmniModelConfig
@@ -244,6 +244,28 @@ def test_stage_specific_text_config_override():
assert omni_config.get_head_size() == talker_head_dim
def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size():
"""A stage without an embedding-size override keeps the default input width."""
vllm_config = EngineArgs().create_model_config()
vllm_config.hf_config = Qwen3OmniMoeConfig(enable_audio_output=True)
omni_config = OmniModelConfig.from_vllm_model_config(
vllm_config,
model_stage="talker",
hf_config_name="talker_config",
)
assert (
getattr(
omni_config.hf_config.talker_config,
"embedding_size",
None,
)
is None
)
assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size
def test_stage_configs_path_field():
"""OmniEngineArgs with stage_configs_path should construct without error."""
args = OmniEngineArgs(stage_configs_path="/some/path.yaml")
@@ -0,0 +1,150 @@
import asyncio
from unittest.mock import AsyncMock
import pytest
from openai.types.chat import ChatCompletionUserMessageParam
from openai.types.chat.chat_completion_audio import ChatCompletionAudio
from starlette.requests import Request
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
from vllm.entrypoints.openai.chat_completion.protocol import (
BatchChatCompletionRequest,
ChatCompletionResponse,
ChatCompletionResponseChoice,
ChatMessage,
)
from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
MSG_BATCH: list[list[ChatCompletionMessageParam]] = [
[ChatCompletionUserMessageParam(role="user", content="What color is the sky?")],
[ChatCompletionUserMessageParam(role="user", content="What is 2+2?")],
]
collapse = OmniOpenAIServingChatBatch._maybe_collapse_choices
# Helpers for testing creating a __new__ serving instance and ensuring IDs are correct
def _make_raw_request(headers: list[tuple[bytes, bytes]]):
scope = {
"type": "http",
"method": "POST",
"path": "/v1/chat/completions/batch",
"headers": headers,
}
return Request(scope)
def _make_completion(text):
return ChatCompletionResponse(
model="test-model",
choices=[
ChatCompletionResponseChoice(
index=0,
message=ChatMessage(role="assistant", content=text),
)
],
usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0),
)
def _make_handler(responses):
handler = OmniOpenAIServingChatBatch.__new__(OmniOpenAIServingChatBatch)
handler.create_chat_completion = AsyncMock(side_effect=responses)
return handler
# Helpers for creating choices of different modality types
def _text_choice(text="hello"):
return ChatCompletionResponseChoice(
index=0,
message=ChatMessage(role="assistant", content=text),
)
def _audio_choice():
audio = ChatCompletionAudio(id="a1", data="base64audio", expires_at=0, transcript="")
return ChatCompletionResponseChoice(
index=0,
message=ChatMessage(role="assistant", content=None, audio=audio),
)
def _image_choice():
content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}]
msg = ChatMessage.model_construct(role="assistant")
object.__setattr__(msg, "content", content)
return ChatCompletionResponseChoice.model_construct(
index=0,
message=msg,
)
def _get_subrequest_ids(handler):
return [call.args[0].request_id for call in handler.create_chat_completion.call_args_list]
@pytest.mark.parametrize("has_header", [True, False])
def test_subrequest_ids_are_unique(has_header: bool):
"""Ensure that a submitted batch request creates unique subrequest IDs"""
header = None if not has_header else b"user-123"
headers = [] if not header else [(b"x-request-id", header)]
handler = _make_handler([_make_completion("hello"), _make_completion("world")])
request = BatchChatCompletionRequest(
model="test-model",
messages=MSG_BATCH,
seed=42,
)
result = asyncio.run(handler.create_batch_chat_completion(request, _make_raw_request(headers)))
assert not isinstance(result, ErrorResponse)
assert result.id.startswith("chatcmpl-batch")
# Ensure request IDs are unique and that if we had a header, it's in all reqs
req_ids = _get_subrequest_ids(handler)
assert len(set(req_ids)) == len(req_ids) == 2
if header is not None:
assert all(header.decode("utf-8") in req_id for req_id in req_ids)
### Tests for ensuring that the (currently batch specific) choice collapsing is correct
def test_single_text_passthrough():
result = collapse([_text_choice()])
assert result.message.content == "hello"
def test_single_image_passthrough():
result = collapse([_image_choice()])
assert isinstance(result.message.content, list)
assert result.message.content[0]["type"] == "image_url"
def test_text_plus_audio():
result = collapse([_text_choice(), _audio_choice()])
assert result.message.content == "hello"
assert result.message.audio.data == "base64audio"
def test_audio_plus_text_order_independent():
result = collapse([_audio_choice(), _text_choice()])
assert result.message.content == "hello"
assert result.message.audio.data == "base64audio"
def test_image_plus_audio():
result = collapse([_image_choice(), _audio_choice()])
assert isinstance(result.message.content, list)
assert result.message.content[0]["type"] == "image_url"
assert result.message.audio.data == "base64audio"
def test_two_content_choices_raises():
with pytest.raises(ValueError, match="Multiple content choices cannot be set"):
collapse([_text_choice(), _text_choice()])
def test_three_choices_raises():
with pytest.raises(ValueError, match="got 3"):
collapse([_text_choice(), _audio_choice(), _text_choice()])
@@ -20,6 +20,10 @@ from vllm_omni.entrypoints.openai.serving_speech import (
OmniOpenAIServingSpeech,
)
# Serving-layer only: the engine is stubbed, so these run on the CPU lane
# alongside the other tests/entrypoints/openai_api suites.
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech:
serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
@@ -95,6 +95,7 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
messages=[],
modalities=["image"],
num_inference_steps=7,
quality="high",
size="768x512",
negative_prompt="avoid blur",
lora={"name": "adapter"},
@@ -110,6 +111,7 @@ def test_serving_boundary_normalizes_declared_root_and_nested_extras(serving_cha
}
assert diffusion_request_args["cfg_text_scale"] == 7.0
assert diffusion_request_args["num_inference_steps"] == 7
assert diffusion_request_args["quality"] == "high"
assert diffusion_request_args["size"] == "768x512"
assert diffusion_request_args["negative_prompt"] == "avoid blur"
assert diffusion_request_args["lora"] == {"name": "adapter"}
@@ -133,6 +135,19 @@ def test_unknown_root_extra_does_not_claim_canonical_extra(serving_chat):
assert "pipeline_option" not in diffusion_request_args
def test_serving_boundary_rejects_invalid_quality(serving_chat):
serving_chat._diffusion_extra_body_params = frozenset()
request = ChatCompletionRequest(
model="test",
messages=[],
modalities=["image"],
quality="medium",
)
with pytest.raises(ValueError, match="quality must be one of"):
serving_chat._normalize_diffusion_request_args(request)
def test_unregistered_cfg_scale_aliases_common_true_cfg_scale(serving_chat):
serving_chat._diffusion_extra_body_params = frozenset()
request = ChatCompletionRequest(
@@ -226,6 +241,7 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
negative_prompt="avoid blur",
cfg_scale=7.0,
true_cfg_scale=2.0,
quality="high",
extra_body={"extra_args": {"solver": "ddim"}},
)
@@ -239,6 +255,7 @@ def test_pure_consumer_preserves_defaults_and_separate_cfg_owners(serving_chat,
}
assert sampling_params.true_cfg_scale == 2.0
assert sampling_params.num_inference_steps == 7
assert sampling_params.quality == "high"
assert (sampling_params.height, sampling_params.width) == (512, 768)
assert captured["prompt"]["negative_prompt"] == "avoid blur"
assert response.error.message == "No output generated from AsyncOmni"
@@ -310,6 +327,7 @@ def test_mixed_consumer_keeps_root_common_args_with_nested_extras(serving_chat,
assert sampling_params_list[0].extra_args == {"target_h": 512, "target_w": 768}
diffusion_params = sampling_params_list[1]
assert diffusion_params.num_inference_steps == 7
assert diffusion_params.quality is None
assert (diffusion_params.height, diffusion_params.width) == (512, 768)
assert diffusion_params.extra_args == {"solver": "euler"}
assert captured["prompt"]["negative_prompt"] == "avoid blur"
@@ -3445,16 +3445,17 @@ class TestCosyVoice3Serving:
def test_cosyvoice3_model_type_detection(self, cosyvoice3_server):
assert cosyvoice3_server._tts_model_type == "cosyvoice3"
assert cosyvoice3_server._is_tts is True
assert cosyvoice3_server._is_cosyvoice3 is True
def test_cosyvoice3_stage_registered(self):
from vllm_omni.entrypoints.openai.serving_speech import (
_COSYVOICE3_TTS_MODEL_STAGES,
_TTS_MODEL_STAGES,
from vllm_omni.entrypoints.openai.tts_adapters import (
all_tts_stage_keys,
detect_tts_model_type,
resolve_adapter,
)
assert "cosyvoice3_talker" in _COSYVOICE3_TTS_MODEL_STAGES
assert "cosyvoice3_talker" in _TTS_MODEL_STAGES
assert "cosyvoice3_talker" in resolve_adapter("cosyvoice3").stage_keys
assert "cosyvoice3_talker" in all_tts_stage_keys()
assert detect_tts_model_type("cosyvoice3_talker", None) == "cosyvoice3"
def test_validate_cosyvoice3_empty_input(self, cosyvoice3_server):
request = OpenAICreateSpeechRequest(input="", ref_audio="data:audio/wav;base64,abc", ref_text="hello")
@@ -17,6 +17,7 @@ from pytest_mock import MockerFixture
from starlette.testclient import TestClient
from tests.helpers.media import generate_synthetic_image
from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
@@ -103,6 +104,55 @@ def _build_test_app(
class TestStreamingVideoOutputWebSocket:
"""WebSocket protocol tests for streaming generated video output."""
def test_sampling_params_quality_reaches_h3_pipeline_boundary(
self,
mocker: MockerFixture,
) -> None:
async def mock_generate(*_args, **_kwargs):
if False:
yield None
_app, handler, _engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[],
mock_generate=mock_generate,
)
request_data = handler._coerce_request_data(
{
"type": "session.start",
"prompt": "H3 streaming quality",
"sampling_params": {"quality": "high"},
}
)
request = VideoGenerationRequest(**request_data)
_prompt, sampling_params, _video_params = asyncio.run(handler._build_prompt_and_sampling_params(request))
assert sampling_params.quality == "high"
def test_omitted_quality_preserves_engine_default(self, mocker: MockerFixture) -> None:
async def mock_generate(*_args, **_kwargs):
if False:
yield None
_app, handler, engine_client = _build_test_app(
mocker=mocker,
streaming_chunks=[],
mock_generate=mock_generate,
)
engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams(quality="high")]
request_data = handler._coerce_request_data(
{
"type": "session.start",
"prompt": "H3 streaming default quality",
}
)
request = VideoGenerationRequest(**request_data)
_prompt, sampling_params, _video_params = asyncio.run(handler._build_prompt_and_sampling_params(request))
assert sampling_params.quality == "high"
def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture):
"""A full session delivers video.start, binary chunks, then session.done."""
@@ -0,0 +1,345 @@
# SPDX-License-Identifier: Apache-2.0
"""TTS model detection: registry-driven stage -> model-type resolution.
``serving_speech.py`` used to carry a hand-written 20-branch ladder mapping
``model_stage``/``model_arch`` to a model-type string, duplicating the
``stage_keys`` every adapter already declared. Detection is now derived from the
adapters. ``test_detection_matches_legacy_ladder`` pins the new implementation
against a verbatim copy of the ladder it replaced, over every input the ladder
could distinguish, so the migration is checkable rather than reviewable-by-eye.
"""
import itertools
import pytest
from vllm_omni.entrypoints.openai.tts_adapters import (
LEGACY_TTS_DETECTORS,
TTS_ADAPTER_REGISTRY,
all_tts_stage_keys,
detect_tts_model_type,
iter_tts_detectors,
resolve_adapter,
tts_entry_stage_archs,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# --------------------------------------------------------------------------
# The legacy ladder, copied verbatim from serving_speech.py at 67c54777b
# (the commit this refactor is based on), with the module constants inlined.
# Do not "clean up" this function: its value is being a frozen oracle.
# --------------------------------------------------------------------------
_LEGACY_MING_TTS_MODEL_ARCHS = {"MingTTSForConditionalGeneration"}
_LEGACY_VOXTRAL_TTS_MODEL_STAGES = {"audio_generation"}
_LEGACY_QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"}
_LEGACY_FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"}
_LEGACY_COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"}
_LEGACY_OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"}
_LEGACY_COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"}
_LEGACY_VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"}
_LEGACY_MING_TTS_MODEL_STAGES = {"ming_tts"}
_LEGACY_MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"}
_LEGACY_MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"}
_LEGACY_MOSS_TTS_LOCAL_MODEL_STAGES = {"moss_tts_local", "moss_tts_local_codec"}
_LEGACY_HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"}
_LEGACY_HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"}
_LEGACY_GLM_TTS_MODEL_STAGES = {"glm_tts"}
_LEGACY_STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"}
_LEGACY_INDEXTTS2_TTS_MODEL_STAGES = {"indextts2_talker"}
_LEGACY_AUDEX_TTS_MODEL_STAGES = {"audex_thinker", "audex_omni"}
_LEGACY_AUDEX_TTA_MODEL_STAGES = {"audex_tta_thinker"}
_LEGACY_TTS_MODEL_STAGES = (
_LEGACY_VOXTRAL_TTS_MODEL_STAGES
| _LEGACY_QWEN3_TTS_MODEL_STAGES
| _LEGACY_FISH_TTS_MODEL_STAGES
| _LEGACY_COSYVOICE3_TTS_MODEL_STAGES
| _LEGACY_OMNIVOICE_TTS_MODEL_STAGES
| _LEGACY_HIGGS_AUDIO_V2_TTS_MODEL_STAGES
| _LEGACY_HIGGS_V3_TTS_MODEL_STAGES
| _LEGACY_COVO_AUDIO_MODEL_STAGES
| _LEGACY_VOXCPM2_TTS_MODEL_STAGES
| _LEGACY_MING_TTS_MODEL_STAGES
| _LEGACY_MOSS_TTS_MODEL_STAGES
| _LEGACY_MOSS_TTS_FULL_MODEL_STAGES
| _LEGACY_MOSS_TTS_LOCAL_MODEL_STAGES
| _LEGACY_GLM_TTS_MODEL_STAGES
| _LEGACY_STEP_AUDIO2_TTS_MODEL_STAGES
| _LEGACY_INDEXTTS2_TTS_MODEL_STAGES
| _LEGACY_AUDEX_TTS_MODEL_STAGES
| _LEGACY_AUDEX_TTA_MODEL_STAGES
)
def _legacy_detect(model_stage, model_arch):
if model_arch == "VoxCPM2TalkerForConditionalGeneration":
return "voxcpm2"
if model_stage in _LEGACY_QWEN3_TTS_MODEL_STAGES:
return "qwen3_tts"
if model_stage in _LEGACY_VOXTRAL_TTS_MODEL_STAGES:
return "voxtral_tts"
if model_stage in _LEGACY_FISH_TTS_MODEL_STAGES:
return "fish_tts"
if model_stage in _LEGACY_COSYVOICE3_TTS_MODEL_STAGES:
return "cosyvoice3"
if model_stage in _LEGACY_OMNIVOICE_TTS_MODEL_STAGES:
return "omnivoice"
if model_stage in _LEGACY_COVO_AUDIO_MODEL_STAGES:
if model_arch and "CovoAudio" in model_arch:
return "covo_audio"
if model_stage in _LEGACY_VOXCPM2_TTS_MODEL_STAGES:
return "voxcpm2"
if model_stage in _LEGACY_MING_TTS_MODEL_STAGES:
return "ming_flash_omni_tts"
if model_arch in _LEGACY_MING_TTS_MODEL_ARCHS:
return "ming_tts"
if model_stage in _LEGACY_MOSS_TTS_MODEL_STAGES:
return "moss_tts_nano"
if model_stage in _LEGACY_MOSS_TTS_FULL_MODEL_STAGES:
return "moss_tts"
if model_stage in _LEGACY_MOSS_TTS_LOCAL_MODEL_STAGES:
return "moss_tts"
if model_stage in _LEGACY_HIGGS_AUDIO_V2_TTS_MODEL_STAGES:
return "higgs_audio_v2"
if model_stage in _LEGACY_HIGGS_V3_TTS_MODEL_STAGES:
return "higgs_audio_v3"
if model_stage in _LEGACY_GLM_TTS_MODEL_STAGES:
return "glm_tts"
if model_stage in _LEGACY_STEP_AUDIO2_TTS_MODEL_STAGES:
return "step_audio2"
if model_stage in _LEGACY_INDEXTTS2_TTS_MODEL_STAGES:
return "indextts2"
if model_stage in _LEGACY_AUDEX_TTS_MODEL_STAGES:
return "audex"
if model_stage in _LEGACY_AUDEX_TTA_MODEL_STAGES:
return "audex_tta"
return None
# Every ``model_stage`` value any in-tree pipeline emits, plus any stage key an
# adapter claims, so the parametrization below covers the real deployment space
# rather than a hand-picked subset. Refresh the pipeline half with:
# grep -rhoE 'model_stage="[^"]+"' --include='pipeline*.py' vllm_omni/model_executor/ \
# | sed 's/model_stage="//;s/"$//' | sort -u
# ``test_pipeline_stage_list_is_complete`` enforces the adapter half.
_PIPELINE_STAGES = [
"AR",
"ar",
"asr",
"audex_code2wav",
"audex_omni",
"audex_thinker",
"audex_tta_thinker",
"audex_xcodec",
"audio_generation",
"audio_tokenizer",
"audio_vae",
"aura",
"code2wav",
"cosyvoice3_code2wav",
"cosyvoice3_talker",
"dac_decoder",
"diffusion",
"dit",
"fish_speech_slow_ar",
"fused_thinker_talker",
"glm_tts",
"glm_tts_dit",
"higgs_audio_v2",
"higgs_audio_v3",
"indextts2_s2mel_decoder",
"indextts2_talker",
"latent_generator",
"llm",
"ming_tts",
"moss_tts",
"moss_tts_codec",
"moss_tts_local",
"moss_tts_local_codec",
"moss_tts_nano",
# Claimed by the OmniVoice adapter but emitted by no pipeline: OmniVoice is
# served through the diffusion path (``OmniOpenAIServingSpeech.for_diffusion``
# sets ``_tts_model_type`` directly and leaves ``_tts_stage`` None), so
# detection is never consulted for it. The key is vestigial in the adapter
# exactly as it was in the deleted ladder; kept here so detection is still
# exercised on it.
"omnivoice_generator",
"qwen3_tts",
"step_audio2_thinker",
"talker",
"thinker",
"token2audio",
"token2image",
"token2text",
"token2wav",
"tts",
]
_STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"]
_ARCHS = [
None,
"VoxCPM2TalkerForConditionalGeneration",
"MingTTSForConditionalGeneration",
"CovoAudioForConditionalGeneration",
"MyCovoAudioThing",
"Qwen3TTSForConditionalGeneration",
]
#: Architectures that no adapter claims. Paired with any stage key these are
#: inert, so the full cross product is meaningful for them.
_UNCLAIMED_ARCHS = [None, "Qwen3TTSForConditionalGeneration"]
#: Architectures a detector matches on. Pairing one of these with *another*
#: model's stage key describes a deployment no pipeline can produce (the
#: architecture comes from the checkpoint, the stage key from the pipeline
#: definition), and the two implementations deliberately disagree there —
#: see ``test_arch_matching_is_a_fallback_not_an_override``.
_CLAIMED_ARCH_CASES = [
# Ming dense: its AR stage is the generic ``llm``.
("llm", "MingTTSForConditionalGeneration"),
("audio_vae", "MingTTSForConditionalGeneration"),
(None, "MingTTSForConditionalGeneration"),
# Ming flash-omni owns the ``ming_tts`` stage key and must win it.
("ming_tts", "MingTTSForConditionalGeneration"),
# VoxCPM2: talker architecture plus its own stage key.
("latent_generator", "VoxCPM2TalkerForConditionalGeneration"),
(None, "VoxCPM2TalkerForConditionalGeneration"),
# CoVo shares ``fused_thinker_talker`` with non-CoVo fused deployments.
("fused_thinker_talker", "CovoAudioForConditionalGeneration"),
("fused_thinker_talker", "MyCovoAudioThing"),
("fused_thinker_talker", None),
("fused_thinker_talker", "Qwen3TTSForConditionalGeneration"),
]
_REACHABLE = list(itertools.product(_STAGES, _UNCLAIMED_ARCHS)) + _CLAIMED_ARCH_CASES
@pytest.mark.parametrize(("model_stage", "model_arch"), _REACHABLE)
def test_detection_matches_legacy_ladder(model_stage, model_arch):
"""Registry detection is identical to the ladder it replaced.
Covers every ``model_stage`` any in-tree pipeline emits (``_PIPELINE_STAGES``,
plus negatives) against every architecture no adapter claims, and each
claimed architecture against the stage keys it really ships with. Pairing a
claimed architecture with an unrelated model's stage key is excluded because
no pipeline can produce it, and is covered separately by
``test_arch_matching_is_a_fallback_not_an_override``.
"""
assert detect_tts_model_type(model_stage, model_arch) == _legacy_detect(model_stage, model_arch)
def test_arch_matching_is_a_fallback_not_an_override():
"""The one deliberate divergence from the ladder, pinned.
The ladder tested Ming's architecture at a fixed position partway down, so a
Ming-architecture model deployed under *another* model's stage key resolved
to ``ming_tts`` if that key was checked below Ming, and to the other model if
above. Detection now treats an architecture-only match as a fallback that
runs after every adapter owning a real stage key, so the stage key always
wins. That is both self-consistent and unreachable in practice: a
checkpoint's architecture and a pipeline's stage key cannot be mixed across
models.
"""
assert detect_tts_model_type("audex_omni", "MingTTSForConditionalGeneration") == "audex"
assert _legacy_detect("audex_omni", "MingTTSForConditionalGeneration") == "ming_tts"
# Where the two real Ming deployments live, they agree.
assert detect_tts_model_type("llm", "MingTTSForConditionalGeneration") == "ming_tts"
assert detect_tts_model_type("ming_tts", "MingTTSForConditionalGeneration") == "ming_flash_omni_tts"
def test_stage_keys_cover_legacy_stage_set():
"""No stage key was dropped when the module constants were deleted."""
assert all_tts_stage_keys() == frozenset(_LEGACY_TTS_MODEL_STAGES)
def test_pipeline_stage_list_is_complete():
"""``_PIPELINE_STAGES`` must still list every stage key adapters claim.
Guards the parametrization above from silently narrowing: a new adapter
whose stage key is missing here would be detected but never exercised.
Regenerate the list from the pipelines (see the comment beside it) when a
new model lands.
"""
missing = all_tts_stage_keys() - set(_PIPELINE_STAGES)
assert not missing, f"stage keys claimed by adapters but absent from _PIPELINE_STAGES: {sorted(missing)}"
def test_entry_stage_archs_is_ming_only():
"""Only Ming dense identifies its entry stage by architecture.
Widening this set would change which stage ``_find_tts_stage`` selects in
mixed deployments — notably VoxCPM2, which declares ``model_archs`` but is
found by its ``latent_generator`` stage key.
"""
assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"})
def test_every_detected_type_is_adapter_backed_or_declared_legacy():
"""Detection may only name a model that has an adapter or an explicit
:data:`LEGACY_TTS_DETECTORS` entry — never an undeclared string."""
legacy_names = {d.name for d in LEGACY_TTS_DETECTORS}
for stage, arch in itertools.product(_STAGES, _ARCHS): # full cross product on purpose
detected = detect_tts_model_type(stage, arch)
if detected is None:
continue
assert detected in TTS_ADAPTER_REGISTRY or detected in legacy_names, (
f"detection produced undeclared model type {detected!r} for stage={stage!r} arch={arch!r}"
)
def test_legacy_detectors_have_no_adapter():
"""A legacy entry that gained an adapter must be deleted from the list."""
for detector in LEGACY_TTS_DETECTORS:
assert resolve_adapter(detector.name) is None, (
f"{detector.name!r} now has an adapter; remove it from LEGACY_TTS_DETECTORS"
)
def test_no_ambiguous_detectors_at_equal_priority():
"""Two detectors that can match the same deployment must be ordered.
This is what makes ``detect_tts_model_type`` independent of adapter import
order: any real overlap has to be resolved by an explicit
``detect_priority``, not by whichever module happened to register first.
"""
detectors = iter_tts_detectors()
for left, right in itertools.combinations(detectors, 2):
if left.detect_priority != right.detect_priority:
continue
for stage, arch in itertools.product(_STAGES, _ARCHS):
assert not (left.matches(stage, arch) and right.matches(stage, arch)), (
f"{left.name!r} and {right.name!r} both match stage={stage!r} arch={arch!r} "
f"at priority {left.detect_priority}; give one an explicit detect_priority"
)
def test_detection_order_is_total_and_stable():
"""Resolution order must not depend on registry insertion order."""
keys = [(d.detect_priority, d.name) for d in iter_tts_detectors()]
assert keys == sorted(keys)
assert len(set(keys)) == len(keys)
def test_adapters_declare_at_least_one_match_rule():
"""An adapter with neither stage keys nor architectures is unreachable."""
for name, cls in TTS_ADAPTER_REGISTRY.items():
assert cls.stage_keys or cls.model_archs, f"adapter {name!r} declares no stage_keys and no model_archs"
def test_audex_omni_requires_speech_decoder():
"""The Audex omni thinker only serves speech alongside its decoder."""
audex = resolve_adapter("audex")
assert audex.stage_serves_speech("audex_omni", frozenset({"audex_omni", "audex_code2wav"})) is True
assert audex.stage_serves_speech("audex_omni", frozenset({"audex_omni"})) is False
# The dedicated TTS thinker is unconditional.
assert audex.stage_serves_speech("audex_thinker", frozenset({"audex_thinker"})) is True
def test_default_stage_serves_speech_is_unconditional():
for name, cls in TTS_ADAPTER_REGISTRY.items():
if name == "audex":
continue
for key in cls.stage_keys:
assert cls.stage_serves_speech(key, frozenset({key})) is True
@@ -868,6 +868,7 @@ def test_default_sampling_params_apply_to_video_requests(test_client, mocker: Mo
OmniDiffusionSamplingParams(
num_inference_steps=4,
guidance_scale=7.5,
quality="high",
generator_device="cpu",
enable_frame_interpolation=True,
frame_interpolation_exp=2,
@@ -890,6 +891,7 @@ def test_default_sampling_params_apply_to_video_requests(test_client, mocker: Mo
captured = engine.captured_sampling_params_list[0]
assert captured.num_inference_steps == 4
assert captured.guidance_scale == 7.5
assert captured.quality == "high"
assert captured.generator_device == "cpu"
assert captured.enable_frame_interpolation is True
assert captured.frame_interpolation_exp == 2
@@ -1445,6 +1447,7 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
def test_video_request_validation():
req = VideoGenerationRequest(prompt="test")
assert req.prompt == "test"
assert req.quality is None
assert req.generate_sound is False
assert req.sound_duration is None
assert VideoGenerationRequest(prompt="test", generate_sound=True, sound_duration=1.5).generate_sound is True
@@ -1464,6 +1467,8 @@ def test_video_request_validation():
VideoGenerationRequest(prompt="test", frame_interpolation_scale=0)
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", sound_duration=0)
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", quality="medium")
def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):
@@ -1975,6 +1980,7 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
"num_inference_steps": "30",
"guidance_scale": "6.5",
"seed": "42",
"quality": "high",
},
)
assert response.status_code == 200
@@ -1983,6 +1989,7 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture):
assert captured.num_inference_steps == 30
assert captured.guidance_scale == 6.5
assert captured.seed == 42
assert captured.quality == "high"
def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture):
@@ -2018,6 +2025,7 @@ def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocke
OmniDiffusionSamplingParams(
num_inference_steps=4,
guidance_scale=7.5,
quality="high",
enable_frame_interpolation=True,
frame_interpolation_exp=2,
frame_interpolation_scale=0.5,
@@ -2038,6 +2046,7 @@ def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocke
captured = engine.captured_sampling_params_list[0]
assert captured.num_inference_steps == 4
assert captured.guidance_scale == 7.5
assert captured.quality == "high"
assert captured.enable_frame_interpolation is True
assert captured.frame_interpolation_exp == 2
assert captured.frame_interpolation_scale == 0.5
+22
View File
@@ -1252,6 +1252,28 @@ class OpenAIClientHandler:
)
return [resp]
def send_batched_chat_completions_http_request(
self,
request_config: dict[str, Any],
*,
err_code: int | tuple[int, ...] | list[int] | None = None,
err_message: str | tuple[str, ...] | list[str] | None = None,
) -> list[HttpResponse]:
"""Post to batched chat completions."""
cfg = _merge_http_expectation_kwargs(
request_config,
err_code=err_code,
err_message=err_message,
)
r = self._post_json_endpoint("/v1/chat/completions/batch", cfg, default_timeout=120.0)
resp = self._http_response_from_requests(r)
assert_http_error(
resp,
err_code=cfg.get("err_code"),
err_message=cfg.get("err_message"),
)
return [resp]
def send_completions_http_request(
self,
request_config: dict[str, Any],
@@ -25,6 +25,23 @@ def test_from_params_returns_existing_omni_params() -> None:
assert converted is params
@pytest.mark.parametrize("quality", ["lossless", "high"])
def test_quality_accepts_supported_request_levels(quality: str) -> None:
params = OmniDiffusionSamplingParams(quality=quality)
assert params.quality == quality
def test_quality_defaults_to_model_owned_policy() -> None:
params = OmniDiffusionSamplingParams()
assert params.quality is None
def test_quality_rejects_unsupported_request_level() -> None:
with pytest.raises(ValueError, match="quality must be one of"):
OmniDiffusionSamplingParams(quality="medium")
def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> None:
params = _sampling_params(
seed=1234,
@@ -33,6 +50,7 @@ def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> Non
"height": 1024,
"width": 768,
"guidance_scale": 7.5,
"quality": "high",
},
)
@@ -43,6 +61,7 @@ def test_from_params_converts_sampling_params_seed_and_known_extra_args() -> Non
assert converted.height == 1024
assert converted.width == 768
assert converted.guidance_scale == 7.5
assert converted.quality == "high"
assert converted.extra_args == {}
@@ -105,6 +105,9 @@ class _FakeHiFT(nn.Module):
generated_source = speech[:, None]
return speech, generated_source
def inference(self, mel, source):
return self.forward(mel, source)
class _FakeToken2Wav:
def __init__(self):
@@ -0,0 +1,144 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
import torch.nn as nn
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
HiFTGenerator,
)
from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
HiFTGraphWrapper,
)
pytestmark = [pytest.mark.core_model]
class _F0Predictor(nn.Module):
def __init__(self) -> None:
super().__init__()
self.proj = nn.Conv1d(80, 1, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.proj(x).squeeze(1).abs()
class _DeterministicSineGen(nn.Module):
"""Remove source RNG so eager and replay compare only execution paths."""
def __init__(self, num_harmonics: int) -> None:
super().__init__()
self.num_harmonics = num_harmonics
def forward(self, f0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
shape = (*f0.shape[:-1], self.num_harmonics)
sine = f0.new_zeros(shape)
uv = f0.new_ones((*f0.shape[:-1], 1))
return sine, uv, sine
def _small_hift() -> HiFTGenerator:
hift = HiFTGenerator(
base_channels=32,
sampling_rate=24000,
upsample_rates=[8, 5, 3],
upsample_kernel_sizes=[16, 11, 7],
source_resblock_kernel_sizes=[7, 7, 11],
source_resblock_dilation_sizes=[[1, 3, 5]] * 3,
f0_predictor=_F0Predictor(),
)
hift.m_source.l_sin_gen = _DeterministicSineGen(hift.nb_harmonics + 1)
return hift.eval().cuda()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_hift_graph_replay_matches_eager_for_uncached_and_cached_shapes() -> None:
torch.manual_seed(0)
hift = _small_hift()
token2wav = SimpleNamespace(
hift=hift,
flow=SimpleNamespace(
encoder=SimpleNamespace(pre_lookahead_layer=SimpleNamespace(pre_lookahead_len=3)),
token_mel_ratio=2,
),
mel_cache_len=2,
source_cache_len=960,
)
wrapper = HiFTGraphWrapper(
token2wav,
connector_config={"codec_chunk_frames": 2, "codec_left_context_frames": 3},
capture_batch_sizes=[1],
)
wrapper.capture()
cases = (
(torch.randn(1, 80, 4, device="cuda"), torch.zeros(1, 1, 0, device="cuda")),
(torch.randn(1, 80, 6, device="cuda"), torch.randn(1, 1, 960, device="cuda")),
)
with torch.inference_mode():
for speech_feat, cache_source in cases:
expected_speech, expected_source = hift.inference(speech_feat, cache_source)
actual_speech, actual_source = wrapper.replay(speech_feat, cache_source)
torch.testing.assert_close(actual_speech, expected_speech, rtol=1e-4, atol=1e-5)
torch.testing.assert_close(actual_source, expected_source, rtol=1e-4, atol=1e-5)
class _FakeGraph:
def replay(self) -> None:
return None
def _fake_wrapper(monkeypatch: pytest.MonkeyPatch) -> HiFTGraphWrapper:
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
wrapper = object.__new__(HiFTGraphWrapper)
wrapper.capture_batch_sizes = [1]
wrapper.graph = {}
wrapper.static_speech_inputs = {}
wrapper.static_cache_source_inputs = {}
wrapper.static_magnitude_outputs = {}
wrapper.static_phase_outputs = {}
wrapper.static_cache_source_outputs = {}
wrapper.lazy_graph_count = 0
wrapper.max_lazy_graphs = 1
wrapper.decode_fn = Mock(return_value=(torch.tensor([[99.0]]), torch.tensor([[[98.0]]])))
wrapper.finalize_fn = lambda magnitude, phase: magnitude + phase
def capture(batch_size: int, num_frames: int, cache_len: int) -> None:
key = (batch_size, num_frames, cache_len)
wrapper.graph[key] = _FakeGraph()
wrapper.static_speech_inputs[key] = torch.zeros(batch_size, 80, num_frames)
wrapper.static_cache_source_inputs[key] = torch.zeros(batch_size, 1, cache_len)
wrapper.static_magnitude_outputs[key] = torch.ones(batch_size, 1, num_frames)
wrapper.static_phase_outputs[key] = torch.ones(batch_size, 1, num_frames)
wrapper.static_cache_source_outputs[key] = torch.ones(batch_size, 1, num_frames)
wrapper._capture = Mock(side_effect=capture)
return wrapper
def test_unseen_shape_is_lazily_captured(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _fake_wrapper(monkeypatch)
speech, source = wrapper.replay(torch.randn(1, 80, 7), torch.zeros(1, 1, 0))
wrapper._capture.assert_called_once_with(1, 7, 0)
assert wrapper.lazy_graph_count == 1
assert speech.shape == (1, 1, 7)
assert source.shape == (1, 1, 7)
wrapper.decode_fn.assert_not_called()
def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) -> None:
wrapper = _fake_wrapper(monkeypatch)
wrapper.lazy_graph_count = wrapper.max_lazy_graphs
speech_feat = torch.randn(1, 80, 9)
cache_source = torch.zeros(1, 1, 0)
result = wrapper.replay(speech_feat, cache_source)
wrapper._capture.assert_not_called()
wrapper.decode_fn.assert_called_once_with(speech_feat, cache_source)
assert result is wrapper.decode_fn.return_value
@@ -157,6 +157,7 @@ class TestDeployTopology:
assert connector["name"] == "SharedMemoryConnector"
assert connector["extra"]["codec_chunk_frames"] == 25
assert connector["extra"]["codec_left_context_frames"] == 3
assert connector["extra"]["enable_hift_graph"] is True
assert connector["extra"]["connector_get_max_wait_first_chunk"] == 3000
assert connector["extra"]["connector_get_max_wait"] == 300
expected_processor = "tts2code2wav_async_chunk" if deploy.async_chunk else "tts2code2wav_full_payload"
@@ -0,0 +1,54 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm_omni.model_executor.models.cosyvoice3.utils import mel_spectrogram
from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import _build_hift
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
flashcosyvoice_hifigan = pytest.importorskip("flashcosyvoice.modules.hifigan")
flashcosyvoice_audio = pytest.importorskip("flashcosyvoice.utils.audio")
def test_vendored_hift_matches_flashcosyvoice() -> None:
reference = flashcosyvoice_hifigan.HiFTGenerator().eval()
vendored = _build_hift().eval()
# Strict loading verifies that the vendored architecture remains compatible
# with the Step-Audio2 checkpoint produced for flashcosyvoice.
vendored.load_state_dict(reference.state_dict(), strict=True)
speech_feat = torch.randn(1, 80, 4)
cache_source = torch.zeros(1, 1, 0)
with torch.inference_mode():
torch.manual_seed(0)
expected_speech, expected_source = reference(speech_feat, cache_source)
torch.manual_seed(0)
actual_speech, actual_source = vendored.inference(speech_feat, cache_source)
torch.testing.assert_close(actual_source, expected_source)
torch.testing.assert_close(actual_speech, expected_speech)
def test_vendored_mel_spectrogram_matches_flashcosyvoice() -> None:
reference = flashcosyvoice_audio.mel_spectrogram
vendored = mel_spectrogram
audio = torch.randn(1, 48000)
kwargs = {
"n_fft": 1920,
"num_mels": 80,
"sampling_rate": 24000,
"hop_size": 480,
"win_size": 1920,
"fmin": 0,
"fmax": 8000,
}
with torch.inference_mode():
expected_spec = reference(audio, **kwargs)
actual_spec = vendored(audio, **kwargs)
torch.testing.assert_close(actual_spec, expected_spec)
+79
View File
@@ -0,0 +1,79 @@
# SPDX-License-Identifier: Apache-2.0
"""Tests for the TTS adapter migration ratchet.
The gate only earns its place if it actually catches the ways a per-model branch
gets reintroduced. These pin what it detects — and, just as importantly, the one
shape it is known to miss, so nobody mistakes it for an adversarial sandbox.
"""
from __future__ import annotations
import ast
import sys
from pathlib import Path
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tools" / "pre_commit"))
from check_tts_adapter import ( # noqa: E402
MAX_LEGACY_DETECTORS,
MAX_MODEL_TYPE_BRANCHES,
_legacy_detectors,
_model_type_branches,
main,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _count(src: str) -> int:
return len(_model_type_branches(ast.parse(src)))
@pytest.mark.parametrize(
("label", "src", "expected"),
[
("direct comparison", 'if self._tts_model_type == "foo": pass', 1),
("negated comparison", 'if self._tts_model_type != "foo": pass', 1),
("membership test", 'if self._tts_model_type in ("a", "b"): pass', 1),
("local alias", 'kind = self._tts_model_type\nif kind == "foo": pass', 1),
("getattr read", 'if getattr(self, "_tts_model_type") == "foo": pass', 1),
("getattr with default", 'if getattr(self, "_tts_model_type", None) == "foo": pass', 1),
("match statement counts each case", 'match self._tts_model_type:\n case "a": pass\n case _: pass', 2),
("alias then match", 'k = getattr(self, "_tts_model_type")\nmatch k:\n case "a": pass', 1),
("unrelated attribute is ignored", 'if self._other == "foo": pass', 0),
("unrelated local is ignored", 'kind = self._backend\nif kind == "foo": pass', 0),
],
)
def test_branch_detection(label, src, expected):
assert _count(src) == expected, label
def test_dict_dispatch_is_a_known_gap():
"""Documented limitation, asserted so it is a decision and not a surprise.
Table dispatch keyed on the model type reads as zero branches. The ratchet
guards against drift in the shape people actually write; catching every
indirection is a review question, not a lint one.
"""
assert _count('TABLE = {"foo": 1}\nvalue = TABLE[self._tts_model_type]') == 0
def test_legacy_detector_counting():
src = 'X = (LegacyDetector(name="a", stage_keys=frozenset({"s"})), LegacyDetector(name="b"))'
found = _legacy_detectors(ast.parse(src))
assert [name for _, name in found] == ["a", "b"]
def test_gate_passes_on_current_tree():
"""The recorded budgets must match reality, or the ratchet is decorative."""
assert main([]) == 0
def test_budgets_are_the_real_counts():
root = Path(__file__).resolve().parents[2]
serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py"
adapters = root / "vllm_omni" / "entrypoints" / "openai" / "tts_adapters" / "__init__.py"
assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES
assert len(_legacy_detectors(ast.parse(adapters.read_text()))) == MAX_LEGACY_DETECTORS
+181
View File
@@ -0,0 +1,181 @@
# SPDX-License-Identifier: Apache-2.0
"""Ratchet on the TTS adapter migration (RFC #4327, #4855).
Per-model behaviour is being moved out of ``serving_speech.py`` and into
``entrypoints/openai/tts_adapters/``. The migration has stalled before, because
nothing stopped a new model from adding one more ``self._tts_model_type == ...``
branch faster than the migration removed them.
This gate makes the two debt counters monotonically non-increasing:
1. ``self._tts_model_type`` comparisons in ``serving_speech.py``
2. entries in ``tts_adapters.LEGACY_TTS_DETECTORS`` (models with no adapter)
Lowering a budget after removing branches is expected and encouraged — the check
insists on it, so the recorded numbers cannot silently drift upward. Raising one
requires editing this file, which makes it a reviewable decision instead of an
invisible one.
Source-only (``ast``), so it runs in pre-commit without importing vllm.
"""
from __future__ import annotations
import argparse
import ast
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py")
ADAPTERS_INIT = Path("vllm_omni/entrypoints/openai/tts_adapters/__init__.py")
# Budgets. These may only go DOWN. See the module docstring.
MAX_MODEL_TYPE_BRANCHES = 29
MAX_LEGACY_DETECTORS = 1
_GUIDANCE = """
Move the per-model behaviour into its adapter under
``vllm_omni/entrypoints/openai/tts_adapters/`` instead of branching in the
shared serving module. ``docs/contributing/model/adding_tts_model.md`` walks
through it. If a branch genuinely cannot be expressed as an adapter hook, raise
the budget in this file and say why in the PR description.
""".strip()
_ATTR = "_tts_model_type"
def _reads_model_type(node: ast.AST) -> bool:
"""Whether ``node`` evaluates to the model-type discriminator."""
if isinstance(node, ast.Attribute) and node.attr == _ATTR:
return True
# getattr(self, "_tts_model_type") / getattr(self, "_tts_model_type", None)
return (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "getattr"
and len(node.args) >= 2
and isinstance(node.args[1], ast.Constant)
and node.args[1].value == _ATTR
)
def _local_aliases(tree: ast.AST) -> set[str]:
"""Local names bound directly to the model type, e.g. ``kind = self._tts_model_type``.
Shallow on purpose: one hop is what an author reaches for when a branch gets
long, which is the drift this ratchet exists to catch.
"""
aliases: set[str] = set()
for node in ast.walk(tree):
value = None
if isinstance(node, ast.Assign):
value = node.value
elif isinstance(node, ast.AnnAssign | ast.AugAssign):
value = node.value
if value is None or not _reads_model_type(value):
continue
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
for target in targets:
if isinstance(target, ast.Name):
aliases.add(target.id)
return aliases
def _model_type_branches(tree: ast.AST) -> list[tuple[int, str]]:
"""Every branch keyed on the model-type discriminator.
Counts direct comparisons, ``getattr`` reads, ``match`` subjects (one per
``case``), and comparisons against a local alias of the attribute. This is a
ratchet against drift, not an adversarial sandbox: dispatch tables and other
indirections still slip past, and are a review question.
"""
aliases = _local_aliases(tree)
def _is_subject(node: ast.AST) -> bool:
return _reads_model_type(node) or (isinstance(node, ast.Name) and node.id in aliases)
found: list[tuple[int, str]] = []
for node in ast.walk(tree):
if isinstance(node, ast.Compare):
if any(_is_subject(op) for op in [node.left, *node.comparators]):
found.append((node.lineno, ast.unparse(node)))
elif isinstance(node, ast.Match) and _is_subject(node.subject):
subject = ast.unparse(node.subject)
for case in node.cases:
found.append((case.pattern.lineno, f"match {subject}: case {ast.unparse(case.pattern)}"))
found.sort()
return found
def _legacy_detectors(tree: ast.AST) -> list[tuple[int, str]]:
"""Every ``LegacyDetector(...)`` construction."""
found: list[tuple[int, str]] = []
for node in ast.walk(tree):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "LegacyDetector":
name = "?"
for kw in node.keywords:
if kw.arg == "name" and isinstance(kw.value, ast.Constant):
name = str(kw.value.value)
found.append((node.lineno, name))
found.sort()
return found
def _check(label: str, path: Path, actual: int, budget: int, sample: list[str], budget_name: str) -> list[str]:
if actual > budget:
listing = "\n".join(f" {line}" for line in sample)
return [
f"{path}: {actual} {label} exceeds the budget of {budget}.\n{listing}\n\n{_GUIDANCE}\n",
]
if actual < budget:
return [
f"{path}: {actual} {label}, down from a budget of {budget}. "
f"Lower {budget_name} to {actual} in {Path(__file__).relative_to(REPO_ROOT)} "
f"so the ratchet holds the new ground.",
]
return []
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
# pre-commit passes the staged filenames; we always audit the same two files,
# so they are accepted and ignored.
parser.add_argument("filenames", nargs="*")
parser.parse_args(argv)
serving_path = REPO_ROOT / SERVING_SPEECH
adapters_path = REPO_ROOT / ADAPTERS_INIT
for path in (serving_path, adapters_path):
if not path.is_file():
print(f"check_tts_adapter: expected file not found: {path}", file=sys.stderr)
return 1
branches = _model_type_branches(ast.parse(serving_path.read_text()))
detectors = _legacy_detectors(ast.parse(adapters_path.read_text()))
errors = _check(
"`_tts_model_type` comparisons",
SERVING_SPEECH,
len(branches),
MAX_MODEL_TYPE_BRANCHES,
[f"line {lineno}: {src}" for lineno, src in branches],
"MAX_MODEL_TYPE_BRANCHES",
)
errors += _check(
"unmigrated model types in LEGACY_TTS_DETECTORS",
ADAPTERS_INIT,
len(detectors),
MAX_LEGACY_DETECTORS,
[f"line {lineno}: {name}" for lineno, name in detectors],
"MAX_LEGACY_DETECTORS",
)
for error in errors:
print(f"check_tts_adapter: {error}", file=sys.stderr)
return 1 if errors else 0
if __name__ == "__main__":
raise SystemExit(main())
+1 -16
View File
@@ -21,7 +21,6 @@ class RouteTarget(NamedTuple):
class OmniServingCapability(Enum):
"""Serving capabilities that pipelines can shut down."""
CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"}))
COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
@property
@@ -39,16 +38,6 @@ class EndpointRestriction:
reason: str
# Routes that are not supported for any model, but are supported in vLLM.
# This is only temporary to avoid 500s for batched chat completions.
UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = (
EndpointRestriction(
OmniServingCapability.CHAT_COMPLETIONS_BATCH,
"Batched chat completions are not yet supported by vLLM Omni.",
),
)
def build_rejection_handler(reason: str):
"""Build a rejection handler for a given endpoint for the provided reason."""
@@ -71,11 +60,7 @@ def shutdown_unsupported_routes(
"""
from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
# Generally these should not overlap since there is no point. If they do,
# we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages.
restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES)
for end_restrict in restricted_endpoints:
for end_restrict in endpoint_restrictions:
capability = end_restrict.capability
# Remove the route from the app
_remove_route_from_app(app, capability.path, capability.methods)
+8
View File
@@ -178,6 +178,14 @@ class OmniModelConfig(ModelConfig):
return override
return super().embedding_size
def get_inputs_embeds_size(self) -> int:
if self.hf_config_name is not None:
stage_config = getattr(self.hf_config, self.hf_config_name, None)
override = getattr(stage_config, "embedding_size", None)
if override is not None:
return override
return super().get_inputs_embeds_size()
def get_model_arch_config(self):
# For multi-stage omni models, use a stage-aware convertor so that
# only the correct stage's quantization config is surfaced.
+1
View File
@@ -568,6 +568,7 @@ class _DiffusionConfigProjection:
pin_cpu_memory: bool = True
diffusion_compile_granularity: Literal["regional", "full"] = "regional"
diffusion_compile_dynamic: bool = Field(default=True, strict=True)
fa_deterministic: bool = False
vae_use_slicing: bool = False
vae_use_tiling: bool = False
mask_strategy_file_path: str | None = None
+1
View File
@@ -394,6 +394,7 @@ class StageDeployConfig:
# Diffusion execution, cache, and VAE behavior.
diffusion_compile_granularity: str | None = None
diffusion_compile_dynamic: bool | None = None
fa_deterministic: bool | None = None
cache_backend: str | None = None
cache_config: dict[str, Any] | None = None
enable_cache_dit_summary: bool | None = None
+5
View File
@@ -14,6 +14,11 @@ connectors:
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 3
# HiFT capture buckets are derived from codec_chunk_frames and
# codec_left_context_frames. (#5869)
# A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true
stages:
- stage_id: 0
+5
View File
@@ -17,6 +17,11 @@ connectors:
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 3
# HiFT capture buckets are derived from codec_chunk_frames and
# codec_left_context_frames. (#5869)
# A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true
stages:
- stage_id: 0
+5
View File
@@ -17,6 +17,11 @@ connectors:
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 3
# HiFT capture buckets are derived from codec_chunk_frames and
# codec_left_context_frames. (#5869)
# A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true
stages:
- stage_id: 0
@@ -28,6 +28,11 @@ connectors:
connector_get_max_wait: 300
codec_chunk_frames: 25
codec_left_context_frames: 3
# HiFT capture buckets are derived from codec_chunk_frames and
# codec_left_context_frames. (#5869)
# A future initial_codec_chunk_frames shape will use lazy capture unless
# HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it.
enable_hift_graph: true
stages:
- stage_id: 0
@@ -15,6 +15,7 @@ from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
piecewise_attn,
)
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
logger = init_logger(__name__)
@@ -71,9 +72,20 @@ class FlashAttentionImpl(AttentionImpl):
self.causal = causal
self.softmax_scale = softmax_scale
self.qkv_layout = qkv_layout
cfg = get_current_diffusion_config_or_none()
self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False
if backend_kwargs:
logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
def _warn_fa_deterministic_non_dense(self, path: str) -> None:
if not self.fa_deterministic:
return
logger.warning_once(
"fa_deterministic=True is ignored on the %s FlashAttention path; "
"only the dense flash_attn_func path passes deterministic=True.",
path,
)
@staticmethod
def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
# FA3 may return (out, lse), FA2 returns out
@@ -250,6 +262,7 @@ class FlashAttentionImpl(AttentionImpl):
# Try piecewise attention
if full_attn_spans is not None:
self._warn_fa_deterministic_non_dense("piecewise")
logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
if flash_attn_func is not None:
attn_func = partial(
@@ -280,6 +293,7 @@ class FlashAttentionImpl(AttentionImpl):
if len(present_packed_keys) != len(packed_keys):
missing = sorted(set(packed_keys) - set(present_packed_keys))
raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
self._warn_fa_deterministic_non_dense("packed-varlen")
return self._forward_varlen_packed(
query,
key,
@@ -291,6 +305,7 @@ class FlashAttentionImpl(AttentionImpl):
)
if attention_mask is not None and torch.any(~attention_mask):
self._warn_fa_deterministic_non_dense("masked-varlen")
return self._forward_varlen_masked(
query,
key,
@@ -299,15 +314,16 @@ class FlashAttentionImpl(AttentionImpl):
)
if flash_attn_func is not None:
out = flash_attn_func(
query,
key,
value,
causal=self.causal,
softmax_scale=self.softmax_scale,
)
fa_kwargs = {
"causal": self.causal,
"softmax_scale": self.softmax_scale,
}
if self.fa_deterministic:
fa_kwargs["deterministic"] = True
out = flash_attn_func(query, key, value, **fa_kwargs)
return self._unwrap_flash_output(out)
self._warn_fa_deterministic_non_dense("dense-varlen-fallback")
return self._forward_varlen_dense(
query,
key,
@@ -1,7 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
import torch
from packaging.version import InvalidVersion, Version
from vllm.logger import init_logger
from vllm_omni.diffusion.attention.backends.abstract import (
@@ -13,7 +16,8 @@ from vllm_omni.diffusion.attention.backends.abstract import (
logger = init_logger(__name__)
try:
from flashinfer.prefill import single_prefill_with_kv_cache
import flashinfer
from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper
HAS_FLASHINFER = True
except Exception as e:
@@ -29,10 +33,8 @@ class FlashInferAttentionBackend(AttentionBackend):
@classmethod
def supports_attention_mask(cls) -> bool:
# FlashInfer single_prefill_with_kv_cache accepts a ``custom_mask`` kwarg
# (2D boolean, ``True`` = keep) for non-causal attention. See reviewer
# comment on #3079 and the flashinfer docs:
# https://docs.flashinfer.ai/generated/flashinfer.prefill.single_prefill_with_kv_cache.html
# Non-CuTe batch-prefill backends accept boolean custom masks. CuTe-DSL
# falls back to SDPA for nontrivial masks.
return True
@staticmethod
@@ -52,6 +54,27 @@ class FlashInferAttentionBackend(AttentionBackend):
class FlashInferAttentionImpl(AttentionImpl):
_QK_DTYPES = {torch.float16, torch.bfloat16}
_VO_DTYPES = {torch.float16, torch.bfloat16, torch.float8_e4m3fn}
@dataclass(frozen=True)
class _WrapperPlanKey:
batch_size: int
qo_len: int
kv_len: int
num_q_heads: int
num_kv_heads: int
head_dim_qk: int
head_dim_k: int
head_dim_vo: int
q_dtype: torch.dtype
k_dtype: torch.dtype
v_dtype: torch.dtype
o_dtype: torch.dtype
causal: bool
softmax_scale: float
has_custom_mask: bool
def __init__(
self,
num_heads: int,
@@ -60,10 +83,92 @@ class FlashInferAttentionImpl(AttentionImpl):
causal: bool = False,
num_kv_heads: int | None = None,
prefix: str = "",
backend_kwargs: dict | None = None,
**extra_impl_args,
) -> None:
self.causal = causal
self.softmax_scale = softmax_scale
self.device = torch.device("cuda", torch.accelerator.current_device_index())
backend_kwargs = backend_kwargs or {}
quant = backend_kwargs.get("quant") or {}
self.dtype_qk = self._check_dtype(quant.get("dtype_qk"), "dtype_qk", self._QK_DTYPES)
self.dtype_vo = self._check_dtype(quant.get("dtype_vo"), "dtype_vo", self._VO_DTYPES)
requested_backend = quant.get("flashinfer_backend", "auto")
if not HAS_FLASHINFER:
raise ImportError("FLASHINFER_ATTN backend requires flashinfer")
self._check_flashinfer_version()
self.flashinfer_backend = self._select_backend(requested_backend, self.device)
workspace_size = 0 if self.flashinfer_backend == "cute-dsl" else 128 * 1024 * 1024
self._workspace = torch.empty(
workspace_size,
device=self.device,
dtype=torch.uint8,
)
self._wrapper = BatchPrefillWithRaggedKVCacheWrapper(
self._workspace,
kv_layout="NHD",
backend=self.flashinfer_backend,
)
self._qo_indptr: torch.Tensor | None = None
self._kv_indptr: torch.Tensor | None = None
self._plan_key: FlashInferAttentionImpl._WrapperPlanKey | None = None
if self.dtype_qk is not None or self.dtype_vo is not None:
logger.info_once(
"FLASHINFER_ATTN dtype override: Q/K=%s, V=%s.",
self.dtype_qk,
self.dtype_vo,
)
logger.info_once(
"FLASHINFER_ATTN initialized backend=%s on %s.",
self.flashinfer_backend,
self.device,
)
def _check_flashinfer_version(self) -> None:
if self.dtype_qk == self.dtype_vo:
return
try:
flashinfer_version = Version(flashinfer.__version__)
except (AttributeError, InvalidVersion):
return
if flashinfer_version <= Version("0.6.15"):
raise RuntimeError(
f"FlashInfer {flashinfer_version} is too old for reliable mixed "
f"QK/V dtype attention (Q/K={self.dtype_qk}, V={self.dtype_vo}); "
"install flashinfer >= 0.6.16rc1."
)
@staticmethod
def _select_backend(requested_backend: str, device: torch.device) -> str:
if requested_backend != "auto":
return requested_backend
major, _minor = torch.cuda.get_device_capability(device)
if major >= 10:
return "cute-dsl"
if major >= 9:
return "fa3"
return "fa2"
@classmethod
def _check_dtype(
cls,
dtype: torch.dtype | str | None,
option_name: str,
allowed: set[torch.dtype],
) -> torch.dtype | None:
if dtype is None:
return None
if isinstance(dtype, str):
dtype = torch.float8_e4m3fn if dtype == "fp8_e4m3" else getattr(torch, dtype)
if dtype not in allowed:
choices = ", ".join(sorted(str(item) for item in allowed))
raise ValueError(f"Unsupported {option_name}={dtype}; expected one of: {choices}")
return dtype
@staticmethod
def _pack_mask_for_flashinfer(
@@ -128,6 +233,126 @@ class FlashInferAttentionImpl(AttentionImpl):
)
return impl.forward_cuda(query, key, value, attn_metadata)
@torch.compiler.disable
def _plan_wrapper(
self,
key: _WrapperPlanKey,
flat_mask: torch.Tensor | None,
) -> None:
self._wrapper.plan(
self._qo_indptr,
self._kv_indptr,
key.num_q_heads,
key.num_kv_heads,
key.head_dim_qk,
head_dim_vo=key.head_dim_vo,
custom_mask=flat_mask,
causal=key.causal,
sm_scale=key.softmax_scale,
q_data_type=key.q_dtype,
# FlashInfer versions newer than 0.6.15 can dispatch K and V with
# independent runtime dtypes even though plan() names this K/V.
kv_data_type=key.k_dtype,
o_data_type=key.o_dtype,
)
@torch.compiler.disable
def _make_indptr(self, batch_size: int, sequence_length: int) -> torch.Tensor:
return (
torch.arange(
batch_size + 1,
device=self.device,
dtype=torch.int32,
)
* sequence_length
)
def _ensure_plan(
self,
key: _WrapperPlanKey,
flat_mask: torch.Tensor | None,
) -> None:
key_changed = key != self._plan_key
if not key_changed and flat_mask is None:
return
if key_changed:
self._qo_indptr = self._make_indptr(key.batch_size, key.qo_len)
self._kv_indptr = self._make_indptr(key.batch_size, key.kv_len)
# A custom mask's contents may change without its shape changing, so
# copy/re-plan it for every masked invocation.
self._plan_wrapper(key, flat_mask)
self._plan_key = key
def _run_batch_prefill(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
custom_mask: torch.Tensor | None,
) -> torch.Tensor:
if query.device != self.device or key.device != self.device or value.device != self.device:
raise ValueError(
"FLASHINFER_ATTN inputs must remain on the layer initialization "
f"device {self.device}; got Q={query.device}, K={key.device}, "
f"V={value.device}"
)
batch_size, qo_len, num_q_heads, head_dim_qk = query.shape
kv_len = key.shape[1]
num_kv_heads = key.shape[2]
head_dim_k = key.shape[3]
head_dim_vo = value.shape[3]
q = query.reshape(batch_size * qo_len, num_q_heads, head_dim_qk)
k = key.reshape(batch_size * kv_len, num_kv_heads, head_dim_k)
v = value.reshape(batch_size * kv_len, num_kv_heads, head_dim_vo)
if self.dtype_qk is not None:
q = q.to(self.dtype_qk)
k = k.to(self.dtype_qk)
if self.dtype_vo is not None:
v = v.to(self.dtype_vo)
flat_mask = None
if custom_mask is not None:
if custom_mask.dim() == 2:
custom_mask = custom_mask.unsqueeze(0).expand(batch_size, -1, -1)
flat_mask = custom_mask.contiguous().view(-1)
self._ensure_plan(
FlashInferAttentionImpl._WrapperPlanKey(
batch_size=batch_size,
qo_len=qo_len,
kv_len=kv_len,
num_q_heads=num_q_heads,
num_kv_heads=num_kv_heads,
head_dim_qk=head_dim_qk,
head_dim_k=head_dim_k,
head_dim_vo=head_dim_vo,
q_dtype=q.dtype,
k_dtype=k.dtype,
v_dtype=v.dtype,
o_dtype=query.dtype,
causal=self.causal,
softmax_scale=self.softmax_scale,
has_custom_mask=flat_mask is not None,
),
flat_mask,
)
out = self._run_wrapper(q, k, v)
out = out.reshape(batch_size, qo_len, num_q_heads, head_dim_vo)
return out.to(query.dtype) if out.dtype != query.dtype else out
@torch.compiler.disable
def _run_wrapper(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
) -> torch.Tensor:
return self._wrapper.run(query, key, value)
def forward_cuda(
self,
query: torch.Tensor,
@@ -135,18 +360,7 @@ class FlashInferAttentionImpl(AttentionImpl):
value: torch.Tensor,
attn_metadata: AttentionMetadata | None = None,
) -> torch.Tensor:
if not HAS_FLASHINFER:
raise ImportError(
"FLASHINFER_ATTN backend requires flashinfer. "
"Install it or set DIFFUSION_ATTENTION_BACKEND to another backend."
)
# Try the custom_mask path; if the mask can't be packed into the
# (qo_len, kv_len) layout FlashInfer expects, fall back to SDPA
# instead of risking an illegal-memory-access crash in the kernel.
# Input layout is (B, S, H, D); FlashInfer dense prefill takes (S, H, D).
batch_size = query.shape[0]
custom_mask = None
if attn_metadata is not None and attn_metadata.attn_mask is not None:
try:
@@ -156,28 +370,14 @@ class FlashInferAttentionImpl(AttentionImpl):
qo_len=query.shape[1],
kv_len=key.shape[1],
)
except ValueError as e:
logger.debug("Falling back to SDPA for mask path: %s", e)
except ValueError as error:
logger.debug("Falling back to SDPA for mask path: %s", error)
return self._sdpa_fallback(query, key, value, attn_metadata)
# FlashInfer cannot combine causal masking with a custom_mask; rather
# than silently dropping the explicit mask (diverging from SDPA), let
# SDPA handle the causal+mask case correctly.
if custom_mask is not None and self.causal:
logger.debug("causal=True with explicit attn_mask; deferring to SDPA")
return self._sdpa_fallback(query, key, value, attn_metadata)
outputs = []
for b in range(batch_size):
kwargs: dict = {
"sm_scale": self.softmax_scale,
"causal": self.causal,
"return_lse": False,
}
if custom_mask is not None:
kwargs["custom_mask"] = custom_mask if custom_mask.dim() == 2 else custom_mask[b]
out = single_prefill_with_kv_cache(query[b], key[b], value[b], **kwargs)
outputs.append(out)
if custom_mask is not None and self.flashinfer_backend == "cute-dsl":
logger.debug("CuTe DSL does not support custom masks; deferring to SDPA")
return self._sdpa_fallback(query, key, value, attn_metadata)
if batch_size == 1:
return outputs[0].unsqueeze(0)
return torch.stack(outputs, dim=0)
return self._run_batch_prefill(query, key, value, custom_mask)
@@ -83,6 +83,19 @@ def pytorch_attn_forward(
k = k.transpose(1, 2)
v = v.transpose(1, 2)
# GQA/MQA support. Ring attention only ever *communicates* K/V with H_kv
# heads -- that is precisely the benefit of GQA -- so expanding here does
# not change communication volume. The expansion is needed because the aten
# SDPA ops below are required for the LSE output that ring accumulation
# consumes, and they do not accept `enable_gqa`.
if k.shape[1] != q.shape[1]:
num_heads, num_kv_heads = q.shape[1], k.shape[1]
if num_heads % num_kv_heads != 0:
raise ValueError(f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA/MQA.")
n_rep = num_heads // num_kv_heads
k = k.repeat_interleave(n_rep, dim=1)
v = v.repeat_interleave(n_rep, dim=1)
if op_type == "flash":
out, lse = _scaled_dot_product_flash_attention(
q,
@@ -204,6 +204,25 @@ class UlyssesParallelAttention:
attn_metadata: AttentionMetadata | None,
):
mode = get_ulysses_mode(default="strict")
ulysses_world_size = self._sp_group.ulysses_world_size
# advanced_uaa pads non-divisible head counts before the Ulysses
# all-to-all. Padding K/V is not valid in hybrid Ulysses+Ring: the Ring
# GQA/MQA path would repeat the padded zero heads as real K/V heads.
# Reject this layout until K/V replication is performed before Ulysses.
if mode == "advanced_uaa" and self._sp_group.ring_world_size > 1:
for name, tensor in (("key", key), ("value", value)):
kv_head_cnt = int(tensor.shape[2])
if kv_head_cnt % ulysses_world_size != 0:
raise ValueError(
"ulysses_mode='advanced_uaa' with hybrid Ulysses+Ring "
"does not support K/V head padding. "
f"{name}_head_cnt={kv_head_cnt}, "
f"ulysses_degree={ulysses_world_size}. "
"Use ring_degree=1, choose a compatible ulysses_degree, "
"or replicate K/V heads before Ulysses."
)
joint_tensor_query = joint_tensor_key = joint_tensor_value = None
joint_strategy = "front"
joint_len = 0
@@ -226,7 +245,6 @@ class UlyssesParallelAttention:
# Slice joint_query for this Ulysses rank
# joint_query is (B, S, H, D). We split H (dim 2).
ulysses_world_size = self._sp_group.ulysses_world_size
ulysses_rank = self._sp_group.ulysses_rank
joint_head_cnt = int(joint_tensor_query.shape[-2])
joint_orig_head_cnt = joint_head_cnt
+6
View File
@@ -20,6 +20,10 @@ from vllm_omni.diffusion.cache.cachedit.model_specific import (
from vllm_omni.diffusion.cache.cachedit.model_specific import (
register_custom_dit_enablers as _register_custom_dit_enablers,
)
from vllm_omni.diffusion.cache.cachedit.runtime import (
CacheDiTRequestSpec,
RequestScopedCacheDiTRuntime,
)
_register_custom_dit_enablers()
@@ -29,6 +33,8 @@ __all__ = [
"CacheDiTAdapterConfig",
"CacheDiTBackend",
"CacheDiTConfig",
"CacheDiTRequestSpec",
"RequestScopedCacheDiTRuntime",
"SensenovaCachedAdapter",
"cache_summary",
"enable_cache_for_dit",
Loaded 100 of 145 files, more files were not shown because too many files have changed in this diff. Show more