Compare commits
20
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
48bd050974 | ||
|
|
fc8946fcf5 | ||
|
|
f01e4bdddc | ||
|
|
592637015f | ||
|
|
ebe93f5d1e | ||
|
|
8e2e9b6b53 | ||
|
|
9159cedbd7 | ||
|
|
593b404539 | ||
|
|
f663d120c7 | ||
|
|
81b48e83ec | ||
|
|
fd0e92183a | ||
|
|
44e38d7b38 | ||
|
|
b15d4e6a64 | ||
|
|
18bf7e0ff9 | ||
|
|
cfa6d66ed6 | ||
|
|
d8c35663b4 | ||
|
|
e906f41dc8 | ||
|
|
f114113409 | ||
|
|
ab459b36c2 | ||
|
|
24ede9e226 |
No files matched your search
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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()
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
@@ -81,6 +81,8 @@ def test_cachedit_public_api_is_explicit():
|
||||
"CacheDiTAdapterConfig",
|
||||
"CacheDiTBackend",
|
||||
"CacheDiTConfig",
|
||||
"CacheDiTRequestSpec",
|
||||
"RequestScopedCacheDiTRuntime",
|
||||
"SensenovaCachedAdapter",
|
||||
"cache_summary",
|
||||
"enable_cache_for_dit",
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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 ─────────
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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())
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
Reference in new issue
Block a user