Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 2a47d901a0 docs: align quantization overview
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-12 08:56:19 +08:00
182 changed files with 1234 additions and 15828 deletions

No files matched your search

-12
View File
@@ -227,18 +227,6 @@ steps:
"
mirror_hardwares: h100_1
- label: "Omni · MiniCPM-o 4.5 Perf Test"
timeout_in_minutes: 60
source_file_dependencies:
- tests/dfx/perf/tests/test_minicpmo_4_5_ready.json
- vllm_omni/deploy/minicpmo_4_5.yaml
- vllm_omni/model_executor/models/minicpmo_4_5/
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
mirror_hardwares: h100_1
- label: "TTS · VoxCPM2 Test"
source_file_dependencies:
- tests/e2e/online_serving/test_voxcpm2_tts.py
-9
View File
@@ -17,15 +17,6 @@ steps:
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
commands:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and npu and A3 and omni' --run-level 'core_model'
- label: "Omni · MiniCPM-o 4.5 Perf Test"
timeout_in_minutes: 60
mirror_hardwares: a3_npu_2
env:
VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
commands:
- export BENCHMARK_DIR=tests/dfx/perf/results
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and npu and A3 and omni'
- group: ":card_index_dividers: TTS Model Test"
key: ready-tts-group
+1 -23
View File
@@ -40,11 +40,9 @@ steps:
commands:
- "bash .buildkite/release/scripts/publish-release-wheels-pypi.sh"
# Auto-skip when NIGHTLY=1 (scheduled); otherwise require manual unblock.
- block: "Unblock to build release Docker images"
depends_on: ~
key: block-build-release-images
if: build.env("NIGHTLY") != "1"
- group: "Build release Docker images"
key: "build-release-images"
@@ -65,7 +63,7 @@ steps:
--progress plain \
-f docker/Dockerfile.cuda .
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
# Fallback default tag (no arch suffix) if arm64 build fails.
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-omni-release-repo:$BUILDKITE_COMMIT"
env:
@@ -88,26 +86,6 @@ steps:
env:
DOCKER_BUILDKIT: "1"
# NIGHTLY=1 only: promote ECR images to DockerHub nightly tags, then prune.
- label: "Publish nightly multi-arch image to DockerHub"
depends_on:
- build-release-image-x86
- build-release-image-arm64
if: build.env("NIGHTLY") == "1"
key: publish-nightly-images-dockerhub
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/release/scripts/publish-release-images.sh nightly"
- "bash .buildkite/release/scripts/cleanup-nightly-builds.sh"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- block: "Publish release images to DockerHub"
key: block-publish-release-images
depends_on:
@@ -1,147 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
#
# Clean up old nightly builds from DockerHub (vllm/vllm-omni), keeping the
# newest 14 tags with the nightly- prefix.
set -ex
TAG_PREFIX="nightly-"
REPO="vllm/vllm-omni"
REPO_API_URL="https://hub.docker.com/v2/repositories/${REPO}/tags"
echo "Cleaning up tags with prefix: $TAG_PREFIX in repository: $REPO"
if [ -z "$DOCKERHUB_TOKEN" ]; then
echo "Error: DOCKERHUB_TOKEN environment variable is not set"
exit 1
fi
if [ -z "$DOCKERHUB_USERNAME" ]; then
echo "Error: DOCKERHUB_USERNAME environment variable is not set"
exit 1
fi
echo "Getting DockerHub bearer token..."
set +x
BEARER_TOKEN=$(curl -s -X POST \
-H "Content-Type: application/json" \
-d "{\"username\": \"$DOCKERHUB_USERNAME\", \"password\": \"$DOCKERHUB_TOKEN\"}" \
"https://hub.docker.com/v2/users/login" | jq -r '.token')
set -x
if [ -z "$BEARER_TOKEN" ] || [ "$BEARER_TOKEN" = "null" ]; then
echo "Error: Failed to get DockerHub bearer token"
exit 1
fi
get_all_tags() {
local page=1
local all_tags=""
# Hub `name` is a partial match (not a strict prefix). Use it to narrow the
# list, then enforce startswith() client-side. Exhaust API pages; do not
# stop merely because a page has no matching tags after filtering.
local name_query
name_query=$(printf '%s' "$TAG_PREFIX" | jq -sRr @uri)
while true; do
set +x
local response
response=$(curl -s -H "Authorization: Bearer $BEARER_TOKEN" \
"$REPO_API_URL?page=$page&page_size=100&name=${name_query}")
set -x
local page_count
page_count=$(echo "$response" | jq -r '.results | length')
if [ -z "$page_count" ] || [ "$page_count" = "0" ] || [ "$page_count" = "null" ]; then
break
fi
local tags
tags=$(echo "$response" | jq -r --arg prefix "$TAG_PREFIX" \
'.results[] | select(.name | startswith($prefix)) | "\(.last_updated)|\(.name)"')
if [ -n "$tags" ]; then
all_tags="$all_tags$tags"$'\n'
fi
page=$((page + 1))
done
echo "$all_tags" | sort -r | cut -d'|' -f2
}
delete_tag() {
local tag_name="$1"
echo "Deleting tag: $tag_name"
local delete_url="https://hub.docker.com/v2/repositories/${REPO}/tags/$tag_name"
set +x
local response http_code body
response=$(curl -s -w "\n%{http_code}" -X DELETE \
-H "Authorization: Bearer $BEARER_TOKEN" "$delete_url")
set -x
http_code=$(printf '%s\n' "$response" | tail -n1)
body=$(printf '%s\n' "$response" | sed '$d')
if [ "$http_code" -ge 200 ] && [ "$http_code" -lt 300 ]; then
echo "Successfully deleted tag: $tag_name (HTTP ${http_code})"
return 0
fi
local detail=""
if echo "$body" | jq -e '.detail' > /dev/null 2>&1; then
detail=$(echo "$body" | jq -r '.detail')
elif [ -n "$body" ]; then
detail="$body"
else
detail="(empty response body)"
fi
echo "Error: Failed to delete tag $tag_name (HTTP ${http_code}): ${detail}"
return 1
}
echo "Fetching all tags from DockerHub..."
all_tags=$(get_all_tags)
if [ -z "$all_tags" ]; then
echo "No tags found to clean up"
exit 0
fi
total_tags=$(echo "$all_tags" | wc -l)
echo "Found $total_tags tags"
tags_to_keep=14
tags_to_delete=$((total_tags - tags_to_keep))
if [ $tags_to_delete -le 0 ]; then
echo "No tags need to be deleted (only $total_tags tags found, keeping $tags_to_keep)"
exit 0
fi
echo "Will delete $tags_to_delete old tags, keeping the newest $tags_to_keep"
tags_to_delete_list=$(echo "$all_tags" | tail -n +$((tags_to_keep + 1)))
if [ -z "$tags_to_delete_list" ]; then
echo "No tags to delete"
exit 0
fi
echo "Deleting old tags..."
delete_failures=0
while IFS= read -r tag; do
if [ -n "$tag" ]; then
if ! delete_tag "$tag"; then
delete_failures=$((delete_failures + 1))
fi
sleep 1
fi
done <<< "$tags_to_delete_list"
if [ "$delete_failures" -gt 0 ]; then
echo "Cleanup finished with ${delete_failures} deletion failure(s)"
exit 1
fi
echo "Cleanup completed successfully"
@@ -1,87 +1,50 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
#
# Publish Docker images from ECR staging to DockerHub.
#
# Usage:
# publish-release-images.sh # release mode (default)
# publish-release-images.sh release
# publish-release-images.sh nightly
#
# Release tags: latest, v<version>
# Nightly tags: nightly, nightly-<commit>
# Publish release Docker images from ECR to DockerHub.
# Pulls per-arch images, tags with `latest` and versioned tags, pushes them,
# then creates and pushes a multi-arch manifest.
set -euo pipefail
MODE="${1:-release}"
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
echo "ERROR: release-version metadata not set"
exit 1
fi
COMMIT="$BUILDKITE_COMMIT"
DOCKERHUB_REPO="vllm/vllm-omni"
ECR_REPO="public.ecr.aws/q9t5s3a7/vllm-omni-release-repo"
COMMIT="$BUILDKITE_COMMIT"
PRIMARY_TAGS=()
EXTRA_MANIFEST_TAGS=()
case "${MODE}" in
release)
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
echo "ERROR: release-version metadata not set"
exit 1
fi
PRIMARY_TAGS=("latest" "v${RELEASE_VERSION}")
;;
nightly)
# Commit-pinned manifest reuses nightly-{x86_64,aarch64} arch images.
PRIMARY_TAGS=("nightly")
EXTRA_MANIFEST_TAGS=("nightly-${COMMIT}")
;;
*)
echo "Usage: $0 {release|nightly}"
exit 2
;;
esac
echo "========================================"
echo "Publishing ${MODE} images to ${DOCKERHUB_REPO}"
echo "Publishing release images v${RELEASE_VERSION}"
echo " Commit: ${COMMIT}"
echo " Primary tags: ${PRIMARY_TAGS[*]}"
if [ ${#EXTRA_MANIFEST_TAGS[@]} -gt 0 ]; then
echo " Extra manifests: ${EXTRA_MANIFEST_TAGS[*]}"
fi
echo "========================================"
# Login to ECR to pull staging images
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
SRC_X86="${ECR_REPO}:${COMMIT}-x86_64"
SRC_ARM="${ECR_REPO}:${COMMIT}-aarch64"
docker pull "${ECR_REPO}:${COMMIT}-x86_64"
docker pull "${ECR_REPO}:${COMMIT}-aarch64"
docker pull "${SRC_X86}"
docker pull "${SRC_ARM}"
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:latest-x86_64"
docker tag "${ECR_REPO}:${COMMIT}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
docker push "${DOCKERHUB_REPO}:latest-x86_64"
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64"
publish_manifest() {
local tag="$1"
local arch_tag_base="$2"
docker manifest rm "${DOCKERHUB_REPO}:${tag}" || true
docker manifest create \
"${DOCKERHUB_REPO}:${tag}" \
"${DOCKERHUB_REPO}:${arch_tag_base}-x86_64" \
"${DOCKERHUB_REPO}:${arch_tag_base}-aarch64"
docker manifest push "${DOCKERHUB_REPO}:${tag}"
}
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:latest-aarch64"
docker tag "${ECR_REPO}:${COMMIT}-aarch64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
docker push "${DOCKERHUB_REPO}:latest-aarch64"
docker push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
for tag in "${PRIMARY_TAGS[@]}"; do
docker tag "${SRC_X86}" "${DOCKERHUB_REPO}:${tag}-x86_64"
docker tag "${SRC_ARM}" "${DOCKERHUB_REPO}:${tag}-aarch64"
docker push "${DOCKERHUB_REPO}:${tag}-x86_64"
docker push "${DOCKERHUB_REPO}:${tag}-aarch64"
publish_manifest "${tag}" "${tag}"
done
ARCH_BASE="${PRIMARY_TAGS[0]}"
for tag in "${EXTRA_MANIFEST_TAGS[@]}"; do
publish_manifest "${tag}" "${ARCH_BASE}"
done
docker manifest rm "${DOCKERHUB_REPO}:latest" || true
docker manifest rm "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" || true
docker manifest create "${DOCKERHUB_REPO}:latest" "${DOCKERHUB_REPO}:latest-x86_64" "${DOCKERHUB_REPO}:latest-aarch64"
docker manifest create "${DOCKERHUB_REPO}:v${RELEASE_VERSION}" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-x86_64" "${DOCKERHUB_REPO}:v${RELEASE_VERSION}-aarch64"
docker manifest push "${DOCKERHUB_REPO}:latest"
docker manifest push "${DOCKERHUB_REPO}:v${RELEASE_VERSION}"
echo ""
echo "Successfully published ${MODE} images"
echo "Successfully published release images for v${RELEASE_VERSION}"
+1 -4
View File
@@ -1,6 +1,6 @@
---
name: precheck-pr
description: Self-check your branch before creating a PR — catch dead code, prevent new model-specific Python examples, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
---
# PR Pre-Check
@@ -65,8 +65,6 @@ Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type
**Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope.
**Also run the Examples-Policy check on every PR** using [references/examples-policy.md](references/examples-policy.md). Inspect only Python paths introduced by the diff. Treat a new model-, checkpoint-, vendor-, or family-specific Python example as blocking; do not report pre-existing example debt.
### Step 5: Print Report
```
@@ -79,7 +77,6 @@ Pre-check report for <branch>
───────────────── ──────
PR title format ✓
Code quality ⚠ 1 broad except, 2 Any hints
Examples policy ✗ new model-specific Python example
PR desc integrity ✓
Registry/config ✓
Dead code ⚠ 2 warnings
@@ -47,19 +47,6 @@ Diff-scoped sweep for five fragility patterns that are pervasive in the existing
Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row.
---
## All PRs: Examples Policy (run on every PR, quick & full)
Inspect Python files introduced under `examples/` using
[examples-policy.md](examples-policy.md). This is a diff-scoped ratchet: existing
model-specific examples may be edited or removed without triggering the check.
- [ ] **No new model-specific Python examples:** the diff adds, copies, or renames no Python file whose path or behavior is specific to one model, checkpoint, vendor, or model family. ✗ for any violation. Route inference through a shared task/protocol entrypoint, model request behavior through `vllm_omni/model_extras`, and commands through task docs or `recipes/`.
- [ ] **New generic example code is justified:** any new task-neutral Python file under `examples/` is an actual user-facing entrypoint rather than reusable production logic, a benchmark, an app, a setup tool, or a reproducer. ⚠ when the placement is questionable.
Roll the result into the report as a separate **Examples policy** dimension.
### Conventions (eyeball — no grep)
Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check:
@@ -1,68 +0,0 @@
# Examples Python Policy
Prevent new model-specific Python examples without forcing cleanup of the
existing backlog. Run this check on every PR in quick and full modes.
## Scope
Set the merge base as in the main workflow, then list Python paths introduced
by the PR, including additions, copies, and renames:
```bash
BASE="$(git merge-base HEAD origin/main)"
git diff --name-status --diff-filter=ACR --find-renames --find-copies \
"${BASE}"...HEAD -- 'examples/**/*.py'
```
Inspect destination paths for copies and renames. Existing model-specific files
that are only modified or deleted are grandfathered and must not be reported.
## Blocking classification
Mark ✗ when a newly introduced Python file does any of the following:
- Names or scopes a path to one model, checkpoint, vendor, or model family.
- Hard-codes a model identifier, model-specific prompt, request shape, output
conversion, or launch configuration behind an otherwise generic filename.
- Duplicates an existing task/protocol runner to demonstrate one model.
- Adds a per-model helper below a generic task hub such as `text_to_speech/`.
Do not infer safety from a generic filename alone. For example, a file named
`text_to_audio.py` is still model-specific if it only implements one model's
contract.
## Acceptable additions
A new Python example may pass when it is a genuinely model-neutral,
user-facing task or protocol entrypoint. It must accept model selection through
configuration and keep model-specific behavior out of the script.
Use ⚠ when the file is model-neutral but reusable logic appears misplaced under
`examples/`. Route code by purpose:
| Code purpose | Preferred location |
|---|---|
| Model prompt, defaults, request/output adaptation | `vllm_omni/model_extras/` or production model modules |
| Runnable model command and validation evidence | Task documentation or `recipes/` |
| UI or interactive application | `apps/` |
| Benchmark or evaluator | `benchmarks/` |
| Regression reproducer | `tests/` |
| Download, conversion, or setup utility | `tools/` |
If a model cannot use an existing shared runner, report the missing generic
capability. Do not solve the gap by adding a model-specific Python example.
## Report format
Add one row to the pre-check report:
```text
Examples policy ✓ no new Python example paths
Examples policy ⚠ generic helper belongs in tools/
Examples policy ✗ examples/offline_inference/new_model/end2end.py is model-specific
Examples policy ✗ examples/offline_inference/x_to_y_model_name.py is model-specific
```
For a blocker, name the path, the model-specific behavior, and the appropriate
shared runner or destination. Do not propose deleting unrelated existing
examples as part of the PR.
-8
View File
@@ -97,7 +97,6 @@ If docs and live code disagree, verify the code/tests and report the drift.
| [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. |
| [examples-policy.md](../precheck-pr/references/examples-policy.md) | The PR adds, copies, or renames Python under `examples/`; apply the canonical policy shared with `precheck-pr`. |
### Delivery and reviewer coordination
@@ -153,13 +152,6 @@ contract they protect or use only the applicable evidence checks.
Apply every category in [general-checks.md](references/process/general-checks.md) before
lower-priority comments.
If the diff census contains an added, copied, or renamed Python path under
`examples/`, read and apply the canonical
[examples policy](../precheck-pr/references/examples-policy.md). Treat a new
model-specific Python example as blocking. Do not flag model-specific example
debt that the PR only modifies or removes, and do not run the rest of the
author-oriented `precheck-pr` workflow.
For each changed value or behavior, trace:
```text
+11 -10
View File
@@ -158,6 +158,7 @@ def _build_omni_kwargs(args, quantization=None):
def _generate_image(omni, args, prompt, seed):
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
@@ -175,7 +176,7 @@ def _generate_image(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
req_out = outputs[0] if isinstance(outputs, list) else outputs
req_out = OmniRequestOutput.unwrap_result(outputs)
if not req_out.images:
raise ValueError("Could not extract image output from result.")
img = req_out.images[0]
@@ -205,15 +206,15 @@ def _generate_video(omni, args, prompt, seed):
elapsed = time.perf_counter() - start
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput):
first = outputs[0]
elif isinstance(outputs, OmniRequestOutput):
first = outputs
else:
raise ValueError("Could not extract video frames from output.")
if hasattr(first, "images"):
frames = first.images[0] if first.images else None
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images[0]
else:
raise ValueError("Could not extract video frames from output.")
+1 -1
View File
@@ -300,7 +300,7 @@ def benchmark(args: argparse.Namespace) -> None:
output_idx = 0
for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
if stage_outputs.final_output_type == "image":
request_output = stage_outputs
request_output = stage_outputs.request_output
request_id = getattr(request_output, "request_id", "")
images = getattr(request_output, "images", [])
@@ -131,7 +131,7 @@ def build_sampling_params() -> list[SamplingParams]:
def save_stage_output(stage_outputs, output_dir: str) -> None:
output = stage_outputs
output = stage_outputs.request_output
request_id = output.request_id
if stage_outputs.final_output_type == "text":
text_output = output.outputs[0].text
@@ -104,7 +104,7 @@ def main(args):
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -228,7 +228,7 @@ async def main():
if not outputs:
raise ValueError("No output produced from omni.generate()")
first_out = outputs[0]
first_out = outputs[0].request_output
req_out: OmniRequestOutput = first_out
# Verify trajectory data (from custom pipeline)
@@ -741,7 +741,7 @@ def make_mmu_prompt(
def iter_mm_outputs(outputs: list[Any]):
for omni_out in outputs:
req_out = omni_out
req_out = getattr(omni_out, "request_output", None)
req_list = req_out if isinstance(req_out, list) else [req_out]
for item in req_list:
if item is None:
@@ -303,7 +303,7 @@ def main():
)
img_idx = 0
for req_output in omni_outputs:
ro = req_output
ro = getattr(req_output, "request_output", None)
stage_id = getattr(req_output, "stage_id", None)
# AR stage text — each CompletionOutput.text is already a delta when
@@ -621,7 +621,7 @@ def main():
images = getattr(output, "images", None)
if images:
break
req_out = output
req_out = getattr(output, "request_output", None)
images = getattr(req_out, "images", None) if req_out is not None else None
if images:
break
@@ -646,8 +646,8 @@ def main():
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
if frames.is_pipeline_output and frames is not None:
inner_output = frames
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]
@@ -1,117 +0,0 @@
"""
Offline inference example: MammothModa2 image summarization (single AR stage).
Example:
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
--model path/to/MammothModa2-Preview \
--deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
--image /path/to/input.jpg \
--question "Please summarize the content of this image."
"""
from __future__ import annotations
import argparse
import os
from PIL import Image
from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
parser.add_argument(
"--max-tokens",
type=int,
default=512,
help="Max new tokens to generate.",
)
parser.add_argument("--temperature", type=float, default=0.2)
parser.add_argument("--top-p", type=float, default=0.9)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--trust-remote-code", action="store_true")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
return parser.parse_args()
def build_prompt(system: str, question: str) -> str:
return (
f"<|im_start|>system\n{system}<|im_end|>\n"
"<|im_start|>user\n"
"<|vision_start|><|image_pad|><|vision_end|>"
f"{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
def main() -> None:
args = parse_args()
if not os.path.exists(args.image):
raise FileNotFoundError(f"Image file not found: {args.image}")
pil_image = Image.open(args.image)
image_data = convert_image_mode(pil_image, "RGB")
prompt = build_prompt(args.system, args.question)
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
trust_remote_code=args.trust_remote_code,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
try:
sp = SamplingParams(
temperature=float(args.temperature),
top_p=float(args.top_p),
top_k=-1,
max_tokens=int(args.max_tokens),
seed=int(args.seed),
detokenize=True,
)
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
# Consume it inside the try block so the worker isn't closed early.
outputs = list(
omni.generate(
[
{
"prompt": prompt,
"multi_modal_data": {"image": image_data},
"additional_information": {"omni_task": ["chat"]},
}
],
[sp],
)
)
finally:
omni.close()
lines: list[str] = []
for stage_outputs in outputs:
ro = stage_outputs
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
lines.append("answer:\n")
lines.append(text.strip() + "\n")
lines.append("\n")
print("\n".join(lines))
if __name__ == "__main__":
main()
@@ -297,7 +297,7 @@ def main(args):
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
output = stage_outputs
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -347,7 +347,7 @@ def main(args):
omni.start_profile(stages=args.profiler_stages)
for stage_outputs in omni.generate(prompts, sampling_params_list):
output = stage_outputs
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -1,229 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Offline speech-to-speech with NVIDIA-NemotronLabs-VoiceChat-11B.
Feeds a 16 kHz user utterance (plus a spoken-style system prompt) through the
3-stage nemotron_voicechat pipeline and writes the agent's reply as text and a
22.05 kHz WAV.
Usage:
python end2end.py \
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--wav /path/to/user_question_16k.wav \
--output-dir results/nemotron_voicechat
The system prompt defaults to the NeMo reference default. The text tokenizer
resolves from the ``nvidia/NVIDIA-Nemotron-Nano-9B-v2`` HF id (or the
``NEMOTRON_VOICECHAT_LLM_PATH`` env var for air-gapped runs).
Sizing the input WAV: the timeline is frame-locked, so the reply budget IS the
input duration — the acoustic channel trails the text channel, and a reply
that has not finished when the frames run out is truncated SILENTLY (no error;
ASR of the WAV will just disagree with the emitted text). Leave generous
trailing silence after the question (a question ending at ~4.5 s truncated in
an 8 s WAV but completed cleanly in a 16 s one). The script warns when the
text channel is still emitting non-PAD tokens near the last frame.
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import numpy as np
import soundfile as sf
# Warn if the text channel is still speaking within this many final frames
# (12 frames = ~1 s) — the frame-locked reply likely got cut off.
TRUNCATION_GUARD_FRAMES = 12
# The NeMo reference offline script's default system prompt.
DEFAULT_SYSTEM_PROMPT = (
"You are an AI voice assistant developed by NVIDIA. "
"Your name is NVIDIA Voice Chat. "
"Answer in a spoken, conversational style rather than a written one. "
"Do not repeat the same sentence over and over again. "
"Start the conversation by greeting the user."
)
def load_wav_16k_mono(path: str) -> np.ndarray:
wav, sr = sf.read(path, dtype="float32")
if wav.ndim > 1:
wav = wav.mean(axis=1)
if sr != 16000:
from vllm.multimodal.audio import resample_audio_scipy
wav = resample_audio_scipy(wav, orig_sr=sr, target_sr=16000)
return wav.astype(np.float32)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", required=True, help="NVIDIA-NemotronLabs-VoiceChat-11B directory")
parser.add_argument("--wav", required=True, help="user speech (any sr; resampled to 16 kHz mono)")
parser.add_argument("--output-dir", default="results/nemotron_voicechat")
parser.add_argument("--system-prompt", default=DEFAULT_SYSTEM_PROMPT)
parser.add_argument(
"--stage-configs-path",
default=None,
help=(
"Deploy yaml override. Pass vllm_omni/deploy/nemotron_labs_voicechat_streaming.yaml "
"for the async-chunk streaming pipeline (stages overlap; audio starts decoding "
"before the thinker finishes). Default: the bundled non-streaming yaml (bit-exact "
"parity path)."
),
)
args = parser.parse_args()
if not Path(args.wav).is_file():
raise FileNotFoundError(f"Input WAV not found: {args.wav}")
ckpt = Path(args.checkpoint)
raw_cfg = json.loads((ckpt / "config.json").read_text())
stt_cfg = raw_cfg["model"]["stt"]["model"]
from transformers import AutoTokenizer
from vllm.sampling_params import SamplingParams
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
compute_acoustic_frame_count,
)
tok_ref = os.environ.get("NEMOTRON_VOICECHAT_LLM_PATH") or stt_cfg.get(
"pretrained_llm", "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
)
tokenizer = AutoTokenizer.from_pretrained(tok_ref, trust_remote_code=False)
bos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("bos_token", "<s>"))
eos_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("eos_token", "</s>"))
pad_id = tokenizer.convert_tokens_to_ids(stt_cfg.get("pad_token", "<SPECIAL_12>"))
# NeMo encode_system_prompt: [bos] + tokens + [eos].
prompt_ids = [bos_id] + tokenizer.encode(args.system_prompt, add_special_tokens=False) + [eos_id]
wav = load_wav_16k_mono(args.wav)
n_frames = compute_acoustic_frame_count(stt_cfg, int(wav.shape[0]))
print(f"input: {wav.shape[0]} samples @16kHz -> {n_frames} acoustic frames; prompt {len(prompt_ids)} tokens")
inputs = {
# vLLM prompt = system prompt + one placeholder position (acoustic
# frame 0); the thinker replaces every position with fused embeddings.
"prompt_token_ids": prompt_ids + [pad_id],
"additional_information": {
# Plain float list: raw ndarrays do not survive the engine-core
# message serialization.
"nvc_audio": wav.astype(np.float32).tolist(),
"nvc_sr": 16000,
"nvc_prompt_token_ids": prompt_ids,
"nvc_expected_frames": n_frames,
},
}
sampling_params_list = [
# thinker: greedy, frame-locked generation length. NO min_tokens: the
# tokenizer's eos_token is <SPECIAL_12> — the frame-locked PAD/silence
# token — and vLLM's min_tokens masks "EOS" to -inf, which would forbid
# silence and force the agent to babble through the whole utterance.
SamplingParams(temperature=0.0, max_tokens=n_frames, ignore_eos=True, seed=0),
# talker: placeholder loop, stops on the stage's stop token. 16383:
# the 1-token placeholder prompt + max_tokens must fit the stage's
# max_model_len (16384).
SamplingParams(temperature=0.0, max_tokens=16383, stop_token_ids=[1], detokenize=False, seed=0),
# code2wav: single decode step.
SamplingParams(temperature=0.0, max_tokens=1, detokenize=False, seed=0),
]
omni_kwargs: dict = {}
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
omni = Omni(model=str(ckpt), trust_remote_code=True, **omni_kwargs)
outputs = omni.generate(inputs, sampling_params_list)
out_dir = Path(args.output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
stem = Path(args.wav).stem
agent_text = None
audio = None
sr = 22050
for stage_outputs in outputs:
# One entry per final stage (text from the thinker, audio from code2wav);
# mirrors the qwen2_5_omni offline example's collection loop.
output = getattr(stage_outputs, "request_output", stage_outputs)
final_type = getattr(stage_outputs, "final_output_type", None)
completions = getattr(output, "outputs", None) or []
if not completions:
continue
completion = completions[0]
if final_type in (None, "text"):
text = getattr(completion, "text", None)
if text:
agent_text = text
token_ids = getattr(completion, "token_ids", None)
if token_ids:
(out_dir / f"{stem}_text_tokens.json").write_text(json.dumps(list(map(int, token_ids))))
# Frame lock: the reply budget IS the input length. If the text
# channel is still mid-sentence when the frames run out, the
# spoken answer is silently truncated — warn instead.
tail = list(map(int, token_ids))[-TRUNCATION_GUARD_FRAMES:]
if any(t != pad_id for t in tail):
print(
f"WARNING: the text channel emits non-PAD tokens within the last "
f"{TRUNCATION_GUARD_FRAMES} frames ({TRUNCATION_GUARD_FRAMES * 80} ms) — the reply "
"likely ran out of acoustic frames and is cut off mid-sentence. "
"Pad the input WAV with more trailing silence and rerun."
)
# Best-effort: dump the function-channel debug timeline if the
# engine surfaced the request info on this output.
info = getattr(output, "additional_information", None)
func_timeline = info.get("nvc_function_tokens") if isinstance(info, dict) else None
if func_timeline is not None:
values = func_timeline.tolist() if hasattr(func_timeline, "tolist") else list(func_timeline)
(out_dir / f"{stem}_function_tokens.json").write_text(json.dumps([int(v) for v in values]))
print(f"function-channel timeline: {len(values)} entries")
if final_type in (None, "audio"):
# multimodal_output is a MultimodalPayload (Mapping) — duck-type it.
mm = getattr(completion, "multimodal_output", None)
get = getattr(mm, "get", None)
if not callable(get):
continue
candidate = None
for key in ("audio", "model_outputs", "wav", "waveform"):
value = get(key)
if value is not None:
candidate = value
break
if isinstance(candidate, list):
candidate = candidate[0] if candidate else None
if candidate is not None:
arr = np.asarray(
candidate.float().cpu() if hasattr(candidate, "cpu") else candidate, dtype=np.float32
).reshape(-1)
if arr.size > 0:
audio = arr
sr_val = get("sr")
if sr_val is not None:
sr_item = sr_val[0] if isinstance(sr_val, list) else sr_val
sr = int(sr_item.item() if hasattr(sr_item, "item") else sr_item)
if agent_text is not None:
# The frame-locked text channel emits PAD tokens on silent frames;
# strip them (NeMo's post-inference detokenizer does the same).
pad_token = stt_cfg.get("pad_token", "<SPECIAL_12>")
agent_text = " ".join(agent_text.replace(pad_token, " ").split())
text_path = out_dir / f"{stem}_output.txt"
text_path.write_text(agent_text + "\n")
print(f"agent text -> {text_path}\n{agent_text}")
else:
print("WARNING: no agent text in outputs")
if audio is not None and audio.size > 0:
wav_path = out_dir / f"{stem}_output.wav"
sf.write(wav_path, audio, sr)
print(f"agent audio -> {wav_path} ({audio.size / sr:.2f}s @ {sr} Hz)")
else:
raise RuntimeError("No audio produced by the code2wav stage.")
if __name__ == "__main__":
main()
@@ -392,7 +392,7 @@ def main(args):
total_requests = len(prompts)
processed_count = 0
for stage_outputs in omni_generator:
output = stage_outputs
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -367,7 +367,7 @@ def main(args):
print(f"query type: {args.query_type}")
for stage_outputs in omni_generator:
output = stage_outputs
output = stage_outputs.request_output
if stage_outputs.final_output_type == "text":
request_id = output.request_id
text_output = output.outputs[0].text
@@ -237,7 +237,7 @@ async def run_single_request(
sampling_params_list=sampling_params_list,
output_modalities=output_modalities,
):
output = omni_output
output = omni_output.request_output
if omni_output.final_output_type == "text":
if stage_0_first_output_ts is None:
stage_0_first_output_ts = time.perf_counter()
@@ -26,6 +26,7 @@ import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
@@ -351,8 +352,8 @@ def main():
else:
print("[Profiler] No valid profiling data returned.")
# omni.generate() returns a list for sync calls; unwrap single-result list.
output = result[0] if isinstance(result, list) else result
# Extract output from result
output = OmniRequestOutput.unwrap_result(result)
if not output.images:
raise ValueError("No video frames found in OmniRequestOutput.")
@@ -345,7 +345,7 @@ def main(args):
# Process outputs from each stage
for stage_idx, stage_outputs in enumerate(omni_outputs):
# request_output may be a single object or a list
raw_output = stage_outputs
raw_output = stage_outputs.request_output
outputs_list = raw_output if isinstance(raw_output, list) else [raw_output]
if stage_outputs.final_output_type == "text":
@@ -43,7 +43,6 @@ from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
@@ -389,9 +388,9 @@ def main():
raise ValueError("No output generated from omni.generate()")
output = outputs[0]
if not isinstance(output, OmniRequestOutput) or not output:
if not hasattr(output, "request_output") or not output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
request_output = output
request_output = output.request_output
if not hasattr(request_output, "multimodal_output"):
raise ValueError("No multimodal_output found in request_output")
@@ -115,7 +115,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
)
images_outputs = []
for output in outputs:
req_out = output
req_out = output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
images = req_out.images
@@ -656,7 +656,7 @@ def main():
images = getattr(output, "images", None)
if images:
break
req_out = output
req_out = getattr(output, "request_output", None)
images = getattr(req_out, "images", None) if req_out is not None else None
if images:
break
@@ -372,10 +372,18 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
Streaming requires `async_chunk: true` in the stage config.
### Word Timestamps
Word-level timestamps are currently a serving-path feature: launch
`vllm-omni serve` with `--forced-aligner` and request `word_timestamps`
(see `examples/online_serving/text_to_speech/README.md`). An offline
example is not provided in this release.
Generate a WAV offline and a JSON sidecar with word-level timestamps from
`Qwen/Qwen3-ForcedAligner-0.6B`:
```bash
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
--text "Hello world." \
--output-dir /tmp/qwen3_tts_timestamps
```
The script writes `qwen3_tts_word_timestamps.wav` and
`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
### Batched decoding
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
@@ -168,7 +168,7 @@ def run_e2e():
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro = output
ro = output.request_output
if ro is None:
print("No request_output found.")
continue
@@ -179,7 +179,7 @@ def main(args):
t_start = time.perf_counter()
for stage_outputs in omni.generate(inputs):
request_output = stage_outputs
request_output = stage_outputs.request_output
if request_output is None or not request_output.outputs:
continue
_save_wav(
@@ -215,7 +215,7 @@ async def main_streaming(args):
sample_rate = None
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = stage_output.outputs[0].multimodal_output
mm = stage_output.request_output.outputs[0].multimodal_output
if not stage_output.finished:
t_now = time.perf_counter()
sample_rate, n = _collect_audio_chunk(mm, all_audio_chunks, sample_rate)
@@ -137,7 +137,7 @@ def main(args) -> None:
params_list = sampling_params
for stage_outputs in omni.generate(inputs, params_list):
for i, req_output in enumerate(stage_outputs):
for i, req_output in enumerate(stage_outputs.request_output):
for j, out in enumerate(req_output.outputs):
mm = out.multimodal_output
if mm is None:
@@ -135,7 +135,7 @@ def run_e2e():
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro = output
ro = output.request_output
if ro is None:
print("No request_output found.")
continue
@@ -391,7 +391,7 @@ def main(args):
for batch_start in range(0, len(inputs), batch_size):
batch = inputs[batch_start : batch_start + batch_size]
for stage_outputs in omni.generate(batch):
output = stage_outputs
output = stage_outputs.request_output
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
@@ -412,7 +412,7 @@ async def main_streaming(args):
t_prev = t_start
chunk_idx = 0
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = stage_output.outputs[0].multimodal_output
mm = stage_output.request_output.outputs[0].multimodal_output
if not stage_output.finished:
t_now = time.perf_counter()
audio = mm.get("audio")
@@ -116,7 +116,7 @@ def main(args: Any) -> None:
prompt = _build_custom_voice_input(args)
final_output = None
for stage_outputs in omni.generate([prompt]):
final_output = stage_outputs
final_output = stage_outputs.request_output
if final_output is None:
raise RuntimeError("Qwen3-TTS did not produce an output.")
@@ -38,7 +38,7 @@ def _stage_durations_from_output(outputs) -> dict[str, float]:
durations = getattr(omni_out, "stage_durations", None) or {}
if durations:
return dict(durations)
ro = omni_out
ro = getattr(omni_out, "request_output", None)
if ro is not None:
return dict(getattr(ro, "stage_durations", None) or {})
return {}
@@ -99,7 +99,7 @@ def extract_audio(outputs) -> tuple[np.ndarray, int]:
out = outputs[0]
mm = getattr(out, "multimodal_output", None)
if mm is None:
ro = out
ro = getattr(out, "request_output", None)
if ro is not None:
mm = getattr(ro, "multimodal_output", None)
if mm is None and ro.outputs:
@@ -390,7 +390,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = result
inner = getattr(result, "request_output", None)
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
@@ -575,8 +575,8 @@ def main():
if frames.multimodal_output and "audio" in frames.multimodal_output:
audio = frames.multimodal_output["audio"]
audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate)
if frames.is_pipeline_output and frames is not None:
inner_output = frames
if frames.is_pipeline_output and frames.request_output is not None:
inner_output = frames.request_output
if isinstance(inner_output, OmniRequestOutput):
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
audio = inner_output.multimodal_output["audio"]
@@ -138,7 +138,7 @@ def _extract_peak_memory_mb(result: Any) -> float:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = result
inner = getattr(result, "request_output", None)
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
@@ -605,13 +605,6 @@ python qwen3_tts/streaming_speech_client.py \
The client writes one PCM file per sentence and a matching
`sentence_XXX_timestamps.json` sidecar.
Non-streaming requests can also ask for timestamps: pass
`"word_timestamps": true` to `POST /v1/audio/speech` and read the
`X-Word-Timestamps` response header (JSON list of `{word, start_ms, end_ms}`,
ASCII-escaped). Past 4 KB the header is replaced by
`X-Word-Timestamps-Omitted: oversize; bytes=<n>; limit=4096` and the audio
still returns — use the WebSocket path for long transcripts.
To *see* the alignment instead of reading a JSON sidecar, run the
word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
```bash
-1
View File
@@ -111,7 +111,6 @@ plugins:
- "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building
- "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building
- "vllm_omni.experimental" # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building
- "vllm_omni.model_executor.models.nemotron_voicechat.nemo_vendored" # vendored NeMo code; upstream docstrings trip griffe strict warnings
- mkdocstrings:
handlers:
python:
-13
View File
@@ -206,12 +206,6 @@ ignore = [
# These files keep upstream conventions (star imports, long lines, math-symbol
# argument names) so they stay easy to diff against upstream releases.
"vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
# NemotronVoiceChat inference modules vendored from NVIDIA-NeMo/Speech (branch
# nemotron-labs-voicechat). Numerics-bearing code is kept byte-close to
# upstream to stay diffable and parity-safe, so upstream style (long lines,
# legacy typing, whitespace, unused-import shims) is exempted. TID251 and the
# F-correctness rules still apply.
"vllm_omni/model_executor/models/nemotron_voicechat/nemo_vendored/**" = ["E501", "E721", "E741", "F401", "F841", "I001", "N803", "N814", "UP006", "UP007", "UP008", "UP011", "UP032", "UP035", "W291", "W293"]
"vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines
[tool.ruff.lint.flake8-tidy-imports.banned-api]
@@ -320,10 +314,6 @@ extend-ignore-identifiers-re = [
".*[Nn][Dd]3.*",
".*updator.*",
"byt5.*",
# NemotronVoiceChat vendored triton kernel: dout/DOUT abbreviate d_out
# (output feature dim), not "doubt".
".*dout.*",
".*DOUT.*",
]
extend-ignore-re = [
"\\.\\.\\.n,d->\\.\\.\\.nd",
@@ -344,9 +334,6 @@ byt5 = "byt5"
# "seeked" is the standard HTML media element DOM event name (fires when a
# seek completes), used by the word-timestamps Gradio demo. Not a typo.
seeked = "seeked"
# Ryan Leary: author credited in the NemotronVoiceChat vendored NeMo
# features.py copyright header. A surname, not a typo of "leery".
Leary = "Leary"
# IndexTTS 2.5 language codes and speech-emotion-recognition event token.
fo = "fo"
ba = "ba"
-113
View File
@@ -1,113 +0,0 @@
# NemotronLabs VoiceChat 11B
> Offline speech-to-speech: a 16 kHz user utterance in, the agent's spoken reply
> (text + 22.05 kHz WAV) out, on a 3-stage vLLM-Omni pipeline.
## Summary
- Vendor: NVIDIA
- Model: [`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Task: Speech-to-speech voice chat. The model runs a frame-locked 12.5 Hz
timeline: a Conformer + NemotronH hybrid-Mamba thinker emits one text token
per acoustic frame, a Gemma3-1B EAR-TTS talker turns the text timeline into
31-quantizer RVQ code stacks, and an RVQ-VAE codec decodes them to audio.
- Mode: Offline single-turn inference (batch=1). Online/duplex serving,
function calling, and batch>1 are documented follow-ups.
- Maintainer: [`@yuekaizhang`](https://github.com/yuekaizhang)
## When to use this recipe
Use this recipe to run a single-turn voice-chat exchange with
`NVIDIA-NemotronLabs-VoiceChat-11B` on one GPU: you provide a user utterance as
a WAV file (any sample rate; it is resampled to 16 kHz mono) plus an optional
spoken-style system prompt, and get back the agent's reply as text and a
22.05 kHz WAV. The integration is NeMo-free at runtime — the perception
Conformer, EAR-TTS talker, and RVQ-VAE codec are vendored, dependency-stripped
NeMo modules (`nemo_vendored/`), so no `nemo_toolkit` install is needed.
## References
- Offline example:
[`examples/offline_inference/nemotron_voicechat/end2end.py`](../../examples/offline_inference/nemotron_voicechat/end2end.py)
- Model modules (thinker / talker / code2wav / vendored NeMo):
[`vllm_omni/model_executor/models/nemotron_voicechat/`](../../vllm_omni/model_executor/models/nemotron_voicechat/)
- Staged pipeline config:
[`vllm_omni/deploy/nemotron_labs_voicechat.yaml`](../../vllm_omni/deploy/nemotron_labs_voicechat.yaml)
- Upstream model card:
[`nvidia/NVIDIA-NemotronLabs-VoiceChat-11B`](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B)
- Reference implementation: NVIDIA-NeMo/Speech, branch `nemotron-labs-voicechat`
## Pipeline
| stage | arch | dtype | role |
|---|---|---|---|
| 0 thinker | `NemotronVoiceChatThinkerForConditionalGeneration` (LLM_AR) | fp32 | WAV + system prompt -> frame-locked text-token timeline (+ function channel) |
| 1 talker | `NemotronVoiceChatTalker` (LLM_AR) | fp32 | text timeline -> 31-quantizer RVQ code stacks (one per 80 ms frame) |
| 2 code2wav | `NemotronVoiceChatCode2Wav` (LLM_GENERATION) | fp32 | RVQ-VAE decode -> 22.05 kHz PCM |
Stages 0/1 default to fp32 for exact parity with the NeMo reference
implementation (greedy decoding matches it token for token on the acceptance
fixture). The deploy yaml documents a ~2x-faster bf16 thinker option whose
output stayed within one word of the reference in testing.
## Hardware Support
## GPU
### 1x H100 80GB
#### Environment
- OS: Linux
- Python: 3.12
- vLLM version: 0.26.0
- vLLM-Omni version or commit: this PR / current `main`
#### Command
```bash
# Tokenizer: the checkpoint ships no HF tokenizer; it resolves from the
# nvidia/NVIDIA-Nemotron-Nano-9B-v2 HF id automatically. For air-gapped runs,
# point NEMOTRON_VOICECHAT_LLM_PATH at a local snapshot of that repo instead.
python examples/offline_inference/nemotron_voicechat/end2end.py \
--checkpoint /path/to/NVIDIA-NemotronLabs-VoiceChat-11B \
--wav /path/to/user_question.wav \
--output-dir results/nemotron_voicechat
```
#### Verification
```bash
ls results/nemotron_voicechat
# <stem>_output.txt the agent reply as text
# <stem>_output.wav the agent reply as 22.05 kHz audio
# <stem>_text_tokens.json the frame-locked text-token timeline
```
The reply text should read as a coherent spoken-style answer to the question in
the input WAV, and the WAV should transcribe to (approximately) the same text
with any ASR model.
#### Notes
- Memory usage: the shipped yaml runs all three stages on one GPU
(`gpu_memory_utilization` 0.62 / 0.12 / 0.06); peak usage is dominated by the
fp32 thinker. The fp32 default has a hard floor of roughly 43 GB of thinker
weights alone (9B backbone + 587M `embed_tokens` + 587M `function_head` +
0.6B Conformer), so 48 GB cards cannot run it — use the bf16 thinker option
documented in the deploy yaml on anything smaller than an 80 GB part.
- Input sizing: the timeline is frame-locked, so the reply budget IS the input
duration. The acoustic channel trails the text channel; if the WAV does not
carry enough trailing silence for the reply to finish, the spoken answer is
truncated silently. Leave generous trailing silence (a question ending at
~4.5 s truncated in an 8 s WAV but completed cleanly in 16 s); the offline
example warns when the text channel is still speaking near the last frame.
- Key flags: sampling is greedy end to end. The thinker is frame-locked —
`max_tokens` equals the acoustic frame count with `ignore_eos=True`. Do NOT
set `min_tokens` on the thinker: the tokenizer's EOS token is also the
frame-locked PAD/silence token, so masking it forces the model to babble
instead of pausing.
- The talker's `max_tokens` is 16383 (its stage prompt is one placeholder
token, and the stage context is 16384).
- Known limitations: batch=1 only; offline single turn (no online/duplex
serving yet); the function-call channel is decoded but not yet acted on.
-1
View File
@@ -66,7 +66,6 @@ recipes/
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
| [`NVIDIA/NemotronLabs-VoiceChat.md`](./NVIDIA/NemotronLabs-VoiceChat.md) | Offline speech-to-speech voice chat (11B, frame-locked 12.5 Hz, 3-stage thinker/talker/code2wav) | 1x H100 80GB |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
@@ -1,92 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
"""Offline Daily-Omni sources: local mirrors must not fall back to the Hub id."""
import json
import os
import tarfile
from pathlib import Path
import pytest
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
daily_omni_local_qa_json,
daily_omni_local_videos_dir,
resolve_daily_omni_local_root,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
def _write_qa_json(root: Path) -> Path:
qa = root / "qa.json"
qa.write_text(json.dumps([{"Question": "q", "Answer": "A", "video_id": "vid0"}]), encoding="utf-8")
return qa
def test_resolve_local_root_returns_none_for_hub_id() -> None:
assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
assert resolve_daily_omni_local_root("") is None
assert resolve_daily_omni_local_root(None) is None
def test_plain_local_mirror(tmp_path: Path) -> None:
_write_qa_json(tmp_path)
(tmp_path / "Videos" / "vid0").mkdir(parents=True)
root = resolve_daily_omni_local_root(str(tmp_path))
assert root == tmp_path.resolve()
assert daily_omni_local_qa_json(root) == (tmp_path / "qa.json").resolve()
assert daily_omni_local_videos_dir(root) == root / "Videos"
def test_hf_cache_dir_resolves_to_snapshot(tmp_path: Path) -> None:
cache_dir = tmp_path / "datasets--liarliar--Daily-Omni"
snapshot = cache_dir / "snapshots" / "deadbeef"
snapshot.mkdir(parents=True)
_write_qa_json(snapshot)
(cache_dir / "refs").mkdir()
(cache_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
root = resolve_daily_omni_local_root(str(cache_dir))
assert root == snapshot.resolve()
assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
def test_hf_cache_dir_without_refs_picks_newest_snapshot(tmp_path: Path) -> None:
cache_dir = tmp_path / "datasets--liarliar--Daily-Omni"
older = cache_dir / "snapshots" / "aaa"
newer = cache_dir / "snapshots" / "bbb"
older.mkdir(parents=True)
newer.mkdir(parents=True)
os.utime(older, (1, 1))
os.utime(newer, (10, 10))
assert resolve_daily_omni_local_root(str(cache_dir)) == newer.resolve()
def test_videos_tar_is_extracted(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("HF_HOME", str(tmp_path / "hf_home"))
root = tmp_path / "mirror"
(root / "Videos" / "vid0").mkdir(parents=True)
(root / "Videos" / "vid0" / "vid0_video.mp4").write_bytes(b"fake")
with tarfile.open(root / "Videos.tar", "w") as tf:
tf.add(root / "Videos", arcname="Videos")
# Only the tarball is left in the mirror.
for child in (root / "Videos").rglob("*"):
if child.is_file():
child.unlink()
(root / "Videos" / "vid0").rmdir()
(root / "Videos").rmdir()
videos_dir = daily_omni_local_videos_dir(root)
assert videos_dir is not None
assert (videos_dir / "vid0" / "vid0_video.mp4").is_file()
# Second call reuses the extraction instead of redoing it.
assert daily_omni_local_videos_dir(root) == videos_dir
def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
root = resolve_daily_omni_local_root(str(tmp_path))
assert daily_omni_local_qa_json(root) is None
assert daily_omni_local_videos_dir(root) is None
+19 -7
View File
@@ -2,8 +2,9 @@
The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
OpenBMB's reported Daily-Omni score when frames and 1s audio segments arrive strictly
interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and no Qwen system prompt.
These tests pin that contract so a refactor cannot silently regress the score.
interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and an empty MiniCPM
system message (not the Qwen system prompt). These tests pin that contract so a refactor
cannot silently regress the score.
vllm stubs are installed by tests/benchmarks/conftest.py before collection.
"""
@@ -98,13 +99,22 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (7, 7), (3, 5), (1, 1)])
def test_uniform_sample_indices_matches_numpy_linspace(n: int, k: int) -> None:
@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (1000, 64), (7, 7), (3, 5), (1, 1)])
def test_uniform_sample_indices_matches_openbmb_midpoint(n: int, k: int) -> None:
"""Pin OpenBMB midpoint-bin sampling (not np.linspace endpoints)."""
got = _uniform_sample_indices(n, k)
if n <= k:
assert got == list(range(n))
else:
assert got == [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
gap = n / k
expected = [int(i * gap + gap / 2) for i in range(k)]
assert got == expected
# Endpoints differ from linspace for typical Daily-Omni downsample sizes.
linspace = [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
if n == 1000 and k == 64:
assert got[0] == 7 and got[-1] == 992
assert linspace[0] == 0 and linspace[-1] == 999
assert got != linspace
def test_uniform_sample_indices_non_positive_k() -> None:
@@ -213,14 +223,16 @@ def test_minicpm_interleave_missing_local_video_returns_none(qa_json, tmp_path)
# ---------------------------------------------------------------------------
def test_minicpm_interleave_omits_qwen_system_message(monkeypatch, qa_json, video_dir) -> None:
def test_minicpm_interleave_keeps_empty_system_message(monkeypatch, qa_json, video_dir) -> None:
"""Official MiniCPM omni mode still sends role=system with empty content."""
ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
_stub_extract(monkeypatch, ds, [])
parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
messages = ds._build_daily_omni_openai_messages(parts, "q?", {"A": "a"})
assert [m["role"] for m in messages] == ["user"]
assert [m["role"] for m in messages] == ["system", "user"]
assert messages[0]["content"][0]["text"] == ""
def test_qwen_pack_mode_keeps_system_message(qa_json, video_dir) -> None:
@@ -1,71 +0,0 @@
import pytest
import vllm_omni.config.pipeline_registry # noqa: F401 (populate registry)
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
from vllm_omni.config.stage_config import (
DeployConfig,
StageExecutionType,
)
from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
POOLING_OUTPUT_DECODER_PATH,
TIMESTAMPS_MODALITY,
)
from vllm_omni.utils.forced_aligner import _DEFAULT_CONFIG_PATH, inject_forced_aligner_stage
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_inject_appends_pooling_aligner_stage():
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
deploy = DeployConfig()
n0 = len(pipeline.stages)
ext_pipeline, ext_deploy = inject_forced_aligner_stage(
pipeline, deploy, {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}
)
# A new pooling stage is appended; the original pipeline is unchanged.
assert len(pipeline.stages) == n0
assert len(ext_pipeline.stages) == n0 + 1
# The original deploy config is not mutated either (double-injection safe).
assert len(deploy.stages) == 0
assert len(ext_deploy.stages) == 1
aligner = ext_pipeline.stages[-1]
# Pooling stage reuses LLM_AR (+ runner=pooling), no dedicated LLM_POOLING.
assert aligner.execution_type == StageExecutionType.LLM_AR
assert aligner.final_output is True
assert aligner.final_output_type == TIMESTAMPS_MODALITY
assert aligner.input_sources == (n0 - 1,)
assert aligner.requires_multimodal_data is True
assert aligner.owns_tokenizer is True
# The deploy stage carries the separate aligner checkpoint + decoder hook.
ds = ext_deploy.stages[-1]
assert ds.engine_extras["model"] == "/models/Qwen3-ForcedAligner-0.6B"
assert ds.engine_extras["runner"] == "pooling"
assert ds.engine_extras["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH
def test_inject_from_config_only():
# --forced-aligner-config alone (model in the YAML) must also inject.
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
deploy = DeployConfig()
n0 = len(pipeline.stages)
ext_pipeline, ext_deploy = inject_forced_aligner_stage(
pipeline, deploy, {"forced_aligner": None, "forced_aligner_config": str(_DEFAULT_CONFIG_PATH)}
)
assert len(ext_pipeline.stages) == n0 + 1
assert ext_pipeline.stages[-1].final_output_type == TIMESTAMPS_MODALITY
assert ext_deploy.stages[-1].engine_extras["model"] == "Qwen/Qwen3-ForcedAligner-0.6B"
def test_inject_noop_without_forced_aligner():
pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
deploy = DeployConfig()
ext_pipeline, ext_deploy = inject_forced_aligner_stage(pipeline, deploy, {})
assert ext_pipeline is pipeline
assert len(ext_deploy.stages) == 0
-21
View File
@@ -513,7 +513,6 @@ def test_sub_config_fields_match_structured_scopes():
"enable_prefix_caching",
"disable_hybrid_kv_cache_manager",
"mm_processor_cache_gb",
"mamba_ssm_cache_dtype",
}
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
"max_num_seqs",
@@ -779,26 +778,6 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
def test_from_pipeline_config_routes_ltx2_model_extra(tmp_path):
deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
deploy_path.write_text(
"\n".join(
[
"pipeline: dreamzero",
"async_chunk: false",
"stages:",
" - stage_id: 0",
" extras:",
" ltx2_use_channels_last_3d: true",
]
)
)
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
assert stage.diffusion_config.extras["ltx2_use_channels_last_3d"] is True
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
@@ -27,7 +27,6 @@ _MIXIN_UPDATE_HELPERS = (
"_attach_finished_request_sets",
"_attach_scheduler_stats",
"_capture_omni_connector_output",
"_maybe_decode_pooling_output",
)
@@ -139,7 +138,6 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
pending_stop_after_extraction=set(),
waiting_for_transfer_free=set(),
_new_prompt_len_snapshot={},
_pooling_output_decoder=None,
finished_req_ids=set(),
finished_req_ids_dict=defaultdict(set),
kv_cache_manager=SimpleNamespace(take_events=lambda: None),
@@ -269,72 +267,3 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
assert output_by_id["bad"].new_token_ids == []
assert output_by_id["good"].new_token_ids == [8]
np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
return SimpleNamespace(
sampled_token_ids=[[]],
logprobs=None,
prompt_logprobs_dict={},
pooler_output=[pooler_tensor],
num_nans_in_logits=None,
kv_connector_output=None,
cudagraph_stats=None,
req_id_to_index={"req": 0},
routed_experts=None,
)
def _make_pooling_stub(decoder) -> tuple[_Request, SimpleNamespace]:
request = _Request("req")
request.sampling_params = SimpleNamespace(num_logprobs=None)
request.pooling_params = SimpleNamespace()
scheduler = _make_scheduler_stub([request])
scheduler._pooling_output_decoder = decoder
scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None))
_bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
return request, scheduler
def test_pooling_decode_success_emits_decoded_payload() -> None:
import torch
decoded = torch.tensor([[1, 2]], dtype=torch.int32)
request, scheduler = _make_pooling_stub(lambda logits, req, hf: decoded)
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"req": 1},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
outputs = OmniARScheduler.update_from_output(
scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
)
(output,) = outputs[0].outputs
assert output.finish_reason is FinishReason.STOP
assert output.pooling_output is decoded
def test_pooling_decode_failure_finishes_request_with_error() -> None:
import torch
def raising_decoder(logits, req, hf):
raise ValueError("boom")
request, scheduler = _make_pooling_stub(raising_decoder)
scheduler_output = SimpleNamespace(
num_scheduled_tokens={"req": 1},
scheduled_spec_decode_tokens={},
num_invalid_spec_tokens=0,
)
outputs = OmniARScheduler.update_from_output(
scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
)
(output,) = outputs[0].outputs
# Decoder failure is a request error (500), not an empty success.
assert output.finish_reason is FinishReason.ERROR
assert output.pooling_output is None
assert "pooling output decode failed" in str(request.stop_reason)
@@ -120,9 +120,6 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
("MiMoAudioModel", "code2wav"),
("Qwen3TTSCode2Wav", "code2wav"),
("MiniCPMO45Code2Wav", "code2wav"),
# nemotron_voicechat: sync-path RVQ-VAE decode waits for the
# talker's full [frames, 31] code-stack payload.
("NemotronVoiceChatCode2Wav", "code2wav"),
("CosyVoice3Model", "cosyvoice3_code2wav"),
# audex: sync-path TTS decoder (async_chunk: false deploys) and
# the TTA XCodec1 stage (always sync full-payload).
@@ -1,80 +0,0 @@
[
{
"test_name": "test_minicpmo_4_5_ready",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100",
"npu": "A3"
},
"num_cards": 1
}
},
"core_model",
"omni"
],
"server_params": {
"model": "openbmb/MiniCPM-o-4_5",
"extra_cli_args": [
"--deploy-config",
"vllm_omni/deploy/minicpmo_4_5.yaml",
"--trust-remote-code"
]
},
"benchmark_params": [
{
"dataset_name": "seed-tts",
"dataset_path": "zhaochenyang20/seed-tts-eval",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [
10,
40
],
"max_concurrency": [
1,
4
],
"no_oversample": true,
"trust_remote_code": true,
"seed_tts_locale": "en",
"extra_body": {
"modalities": [
"text",
"audio"
],
"chat_template_kwargs": {
"enable_thinking": false,
"use_tts_template": true
}
},
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
"baseline": {
"H100": {
"request_throughput": [
0.7803,
1.1008
],
"mean_ttft_ms": [
166.9764,
307.7133
],
"mean_e2el_ms": [
1361.1682,
3800.6332
],
"mean_audio_ttfp_ms": [
635.8099,
1615.6591
],
"mean_audio_rtf": [
0.3151,
0.8568
]
}
}
}
]
}
]
@@ -125,7 +125,7 @@ def _extract_images(output: OmniRequestOutput) -> list:
return output.images
# When the output comes from the orchestrator pipeline, images may be
# nested inside request_output.
inner = output
inner = getattr(output, "request_output", None)
if inner is not None and hasattr(inner, "images") and inner.images:
return inner.images
return []
+2 -2
View File
@@ -56,10 +56,10 @@ def test_teacache(model_name: str):
# Extract images from request_output['images']
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not isinstance(first_output, OmniRequestOutput) or not first_output:
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output
req_out = first_output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
@@ -131,7 +131,7 @@ def _run_inference(
elapsed_ms = (time.time() - start) * 1000
return InferenceResult(
images=outputs[0].images,
images=outputs[0].request_output.images,
elapsed_ms=elapsed_ms,
)
finally:
@@ -49,10 +49,10 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
def _extract_single_image(outputs) -> Image.Image:
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not isinstance(first_output, OmniRequestOutput) or not first_output:
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output
req_out = first_output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
@@ -231,128 +231,6 @@ def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.Mon
assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize("split_dim", ["height", "width"])
@pytest.mark.parametrize(
("full_extent", "world_size"),
[
pytest.param(5, 4, id="four-ranks-empty-tail"),
pytest.param(20, 8, id="production-eight-ranks-empty-tail"),
pytest.param(39, 2, id="two-ranks-partial-tail"),
],
)
def test_reshard_from_trimmed_extent_clamps_and_reconstructs(
monkeypatch: pytest.MonkeyPatch,
split_dim: str,
full_extent: int,
world_size: int,
):
local_extent = (full_extent + world_size - 1) // world_size
shape = (1, 1, 1, full_extent, 1) if split_dim == "height" else (1, 1, 1, 1, full_extent)
x = torch.arange(full_extent, dtype=torch.float32).reshape(shape)
outputs = []
for rank in range(world_size):
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group, rank=rank: (rank, world_size))
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=full_extent,
local_input_extent=local_extent,
split_dim=split_dim,
rank=rank,
world_size=world_size,
)
)
try:
out = wan_spatial_shard.reshard_from_trimmed_extent(
x,
local_extent=local_extent,
split_dim=split_dim,
group=object(),
)
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
start = rank * local_extent
valid_extent = min(local_extent, max(0, full_extent - start))
expected = torch.zeros(local_extent, dtype=torch.float32)
if valid_extent:
expected[:valid_extent] = torch.arange(start, start + valid_extent, dtype=torch.float32)
assert out.shape[_spatial_dim_for_test(split_dim)] == local_extent
assert torch.equal(out.flatten(), expected)
outputs.append(out)
reconstructed = torch.cat(outputs, dim=_spatial_dim_for_test(split_dim))
reconstructed = reconstructed.narrow(_spatial_dim_for_test(split_dim), 0, full_extent)
assert torch.equal(reconstructed, x)
def _spatial_dim_for_test(split_dim: str) -> int:
return -2 if split_dim == "height" else -1
@pytest.mark.core_model
@pytest.mark.cpu
def test_reshard_from_trimmed_extent_rejects_rank_mismatch(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 2))
x = torch.zeros((1, 1, 1, 1, 4), dtype=torch.float32)
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=4,
local_input_extent=2,
split_dim="width",
rank=0,
world_size=2,
)
)
try:
with pytest.raises(RuntimeError, match="group_rank=1, context_rank=0"):
wan_spatial_shard.reshard_from_trimmed_extent(
x,
local_extent=2,
split_dim="width",
group=object(),
)
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
@pytest.mark.core_model
@pytest.mark.cpu
def test_attention_wrapper_rejects_spatial_extent_change(monkeypatch: pytest.MonkeyPatch):
class ShrinkingWanAttentionBlock(torch.nn.Module):
def forward(self, x):
return x[..., :-1]
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 2))
monkeypatch.setattr(
wan_spatial_shard,
"all_gather_along_dim",
lambda x, *, group, dim, dst=None: torch.cat([x, x], dim=dim),
)
module = ShrinkingWanAttentionBlock()
wan_spatial_shard._patch_attention_block(module, group=object(), split_dim="width")
x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32)
token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set(
wan_spatial_shard.SpatialShardContext(
input_extent=4,
local_input_extent=2,
split_dim="width",
rank=0,
world_size=2,
)
)
try:
with pytest.raises(
RuntimeError,
match=r"local_extent=2, gathered_extent=4, trimmed_extent=4, output_extent=3",
):
module(x)
finally:
wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token)
@pytest.mark.core_model
@pytest.mark.cpu
def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
+2 -2
View File
@@ -29,10 +29,10 @@ def test_diffusion_model(tmp_path: Path, omni_runner):
raise ValueError("Empty outputs from Omni.generate()")
first_output = outputs[0]
assert first_output.final_output_type == "image"
if not isinstance(first_output, OmniRequestOutput) or not first_output:
if not hasattr(first_output, "request_output") or not first_output.request_output:
raise ValueError("No request_output found in OmniRequestOutput")
req_out = first_output
req_out = first_output.request_output
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
raise ValueError("Invalid request_output structure or missing 'images' key")
return req_out.images
@@ -11,100 +11,6 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestLTXChannelsLast3D:
@pytest.mark.parametrize(
("extras", "expected"),
[({}, False), ({"ltx2_use_channels_last_3d": True}, True)],
)
def test_channels_last_opt_in_is_ltx2_model_extra(self, extras, expected):
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
assert _ltx2_use_channels_last_3d(SimpleNamespace(extras=extras)) is expected
def test_channels_last_opt_in_rejects_non_boolean_model_extra(self):
from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_channels_last_3d
with pytest.raises(TypeError, match="ltx2_use_channels_last_3d"):
_ltx2_use_channels_last_3d(SimpleNamespace(extras={"ltx2_use_channels_last_3d": "true"}))
@pytest.mark.parametrize(
("causal", "kernel_size"),
[
(True, 3),
(False, 3),
(True, (3, 1, 1)),
(True, 1),
],
)
def test_layout_preserving_causal_conv_matches_original(self, causal, kernel_size):
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
_enable_channels_last_3d,
_weight_is_channels_last_3d,
)
torch.manual_seed(0)
conv = LTX2VideoCausalConv3d(4, 5, kernel_size).to(dtype=torch.float32).eval()
hidden_states = torch.randn(1, 4, 5, 6, 7, dtype=torch.float32)
# Pointwise Conv3d storage is valid in both contiguous formats because
# all spatial/temporal kernel dimensions are singleton.
if kernel_size != 1:
assert not _weight_is_channels_last_3d(conv)
with torch.no_grad():
expected = conv(hidden_states.clone(), causal=causal)
converted, patched = _enable_channels_last_3d(conv)
assert (converted, patched) == (1, 1)
assert _weight_is_channels_last_3d(conv)
with torch.no_grad():
actual = conv(hidden_states.clone(), causal=causal)
torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4)
def test_temporal_pad_replicates_edges_in_channels_last_3d(self):
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
_temporal_pad_channels_last_3d,
)
torch.manual_seed(0)
hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
actual = _temporal_pad_channels_last_3d(hidden_states, left=2, right=1)
expected = torch.cat(
[
hidden_states[:, :, :1].repeat(1, 1, 2, 1, 1),
hidden_states,
hidden_states[:, :, -1:],
],
dim=2,
)
assert actual.is_contiguous(memory_format=torch.channels_last_3d)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_contiguous_weight_keeps_original_padding_path(self, monkeypatch):
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
import vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 as ltx2_vae
conv = LTX2VideoCausalConv3d(4, 5, 3).to(dtype=torch.float32).eval()
hidden_states = torch.randn(1, 4, 3, 2, 2, dtype=torch.float32)
def fail_channels_last_pad(*_args, **_kwargs):
raise AssertionError("contiguous weights must retain the original padding path")
monkeypatch.setattr(ltx2_vae, "_temporal_pad_channels_last_3d", fail_channels_last_pad)
with torch.no_grad():
expected = conv(hidden_states, causal=True)
actual = ltx2_vae._layout_preserving_causal_conv_forward(conv, hidden_states, causal=True)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
class TestLTXOutputRank:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
@@ -1,26 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
from vllm_omni.diffusion.layers.norm import RMSNorm
from vllm_omni.diffusion.models.minimax_h3.encoder import MiniMaxH3Qwen3VLRMSNorm
def test_qwen3_vl_rmsnorm_uses_common_fused_rmsnorm_contract() -> None:
"""Qwen keeps BF16 gamma while native fallback accumulates in FP32."""
eps = 1e-6
norm = MiniMaxH3Qwen3VLRMSNorm(hidden_size=4, eps=eps, dtype=torch.bfloat16)
norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
x = torch.tensor(
[[1.0, -2.0, 3.0, -4.0], [-5.0, 6.0, -7.0, 8.0]],
dtype=torch.bfloat16,
)
x_fp32 = x.float()
expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype)
assert isinstance(norm, RMSNorm)
assert norm.weight.dtype == torch.bfloat16
assert set(norm.state_dict()) == {"weight"}
torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)
@@ -122,8 +122,8 @@ def test_cpu_offload_diffusion_model(model_name: str):
print(f"No offload peak memory: {no_offload_peak_memory} MB")
if model_name == "stabilityai/stable-audio-open-1.0":
audio_offload = output_offload[0].multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2)
# Set platform-specific VRAM saving thresholds to account
@@ -20,10 +20,18 @@ IMAGE_VIDEO_MODELS = {
MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
MODEL_MARKS = {
"riverclouds/qwen_image_random": pytest.mark.core_model,
"stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
}
LAYERWISE_MODEL_CASES = [
pytest.param(
"stabilityai/stable-audio-open-1.0",
marks=[pytest.mark.full_model, pytest.mark.diffusion],
id="stable_audio",
),
pytest.param(
"riverclouds/qwen_image_random",
marks=[pytest.mark.core_model, pytest.mark.diffusion],
id="qwen_image_random",
),
]
AUDIO_MODEL_PARAMS = {
"runner_params": {},
@@ -93,9 +101,8 @@ def run_inference(
return peak, output
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
@pytest.mark.parametrize("model_name", list(MODELS.keys()))
@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES)
def test_layerwise_offload_diffusion_model(model_name: str):
"""Test that layerwise offloading reduces GPU memory usage.
@@ -127,8 +134,8 @@ def test_layerwise_offload_diffusion_model(model_name: str):
print(f"No offload peak memory: {no_offload_peak_memory} MB")
if model_name == "stabilityai/stable-audio-open-1.0":
audio_offload = output_offload[0].multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].multimodal_output.get("audio")
audio_offload = output_offload[0].request_output.multimodal_output.get("audio")
audio_no_offload = output_no_offload[0].request_output.multimodal_output.get("audio")
# Match the sibling cpu-offload test's tolerance: layerwise offload moves
# blocks across the PCIe bus on a side stream, which can perturb cuBLAS
# algorithm selection and produce ~ULP-level drift larger than 1e-3.
@@ -94,8 +94,8 @@ def _generate_single_stage_image(
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
else:
assert isinstance(first_output, OmniRequestOutput) and first_output
request_output = first_output
assert hasattr(first_output, "request_output") and first_output.request_output
request_output = first_output.request_output
if isinstance(request_output, list):
req_out = request_output[0]
else:
@@ -88,8 +88,8 @@ def _generate_single_stage_image(
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
else:
assert isinstance(first_output, OmniRequestOutput) and first_output
request_output = first_output
assert hasattr(first_output, "request_output") and first_output.request_output
request_output = first_output.request_output
if isinstance(request_output, list):
req_out = request_output[0]
else:
@@ -147,10 +147,10 @@ def _generate_single_stage_video(
first = outputs[0]
# Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput).
# Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output).
frames: Any = None
if isinstance(first, OmniRequestOutput) and isinstance(first, list):
inner = first[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and inner.images:
frames = inner.images[0]
if frames is None and hasattr(first, "images") and first.images:
@@ -234,8 +234,8 @@ def _generate_bagel_image(
if images := getattr(req_output, "images", None):
generated_image = images[0]
break
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_outputs = req_output
if hasattr(req_output, "request_output") and req_output.request_output:
stage_outputs = req_output.request_output
if not isinstance(stage_outputs, list):
stage_outputs = [stage_outputs]
for stage_out in stage_outputs:
@@ -250,8 +250,8 @@ def _generate_bagel_image(
# Check LLM stage output — should have finish_reason=stop (not length)
for req_output in omni_outputs:
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_outputs = req_output
if hasattr(req_output, "request_output") and req_output.request_output:
stage_outputs = req_output.request_output
if not isinstance(stage_outputs, list):
stage_outputs = [stage_outputs]
for stage_out in stage_outputs:
@@ -310,7 +310,7 @@ def _generate_image(omni, config: QualityTestConfig):
first = outputs[0]
if hasattr(first, "images") and first.images:
return first.images[0], peak_mem
inner = first
inner = first.request_output
if inner is not None and hasattr(inner, "images") and inner.images:
return inner.images[0], peak_mem
raise ValueError("Could not extract image from output.")
@@ -341,15 +341,15 @@ def _generate_video(omni, config: QualityTestConfig):
)
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
if isinstance(outputs, list) and isinstance(outputs[0], OmniRequestOutput):
first = outputs[0]
elif isinstance(outputs, OmniRequestOutput):
first = outputs
else:
raise ValueError("Could not extract video frames from output.")
if hasattr(first, "images"):
frames = first.images[0] if first.images else None
first = outputs[0]
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
frames = inner.images[0] if inner.images else None
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images[0]
else:
raise ValueError("Could not extract video frames from output.")
@@ -444,11 +444,11 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
image = Image.new("RGB", (2, 2))
inner = OmniRequestOutput.from_diffusion(request_id="req", images=[image])
outer = OmniRequestOutput.from_stage_output(
inner,
outer = OmniRequestOutput(
request_id="req",
stage_id=0,
final_output_type="image",
request_output=inner,
finished=True,
)
@@ -73,9 +73,11 @@ def test_run_summary_reports_worker_peak_memory():
assert summary["per_run_peak_memory_mb"] == [100.0, 150.0]
def test_request_peak_memory_reads_peak_memory_mb_directly():
result = type("Output", (), {"peak_memory_mb": 123.0})()
assert _request_peak_memory_mb(result) == 123.0
def test_request_peak_memory_prefers_inner_diffusion_output():
inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})()
outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})()
assert _request_peak_memory_mb(outer) == 321.0
def test_metric_guidance_describes_thresholds():
@@ -273,7 +273,7 @@ class TestPipelineStreamingOutputToEntrypoint:
await _wait_for(lambda: len(pipeline.requests) == 1)
assert pipeline.requests[0].request_id.startswith("req-omni-")
assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
assert [output.finished for output in outputs] == [False, True]
assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True]
finally:
await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
@@ -25,7 +25,6 @@ from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")]
@@ -96,8 +95,8 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
for req_output in omni_outputs:
if images := getattr(req_output, "images", None):
return images[0]
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_out = req_output
if hasattr(req_output, "request_output") and req_output.request_output:
stage_out = req_output.request_output
if hasattr(stage_out, "images") and stage_out.images:
return stage_out.images[0]
return None
@@ -24,7 +24,6 @@ from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")]
@@ -123,8 +122,8 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
for req_output in omni_outputs:
if images := getattr(req_output, "images", None):
return images[0]
if isinstance(req_output, OmniRequestOutput) and req_output:
stage_out = req_output
if hasattr(req_output, "request_output") and req_output.request_output:
stage_out = req_output.request_output
if hasattr(stage_out, "images") and stage_out.images:
return stage_out.images[0]
return None
@@ -3,8 +3,8 @@
Daily-Omni settings follow the MiniCPM interleaved AV recipe that reaches
~78% overall accuracy on Daily-Omni (``minicpm-interleave``, ``temperature=0``,
text modalities, and server ``--interleave-mm-strings`` + 1fps / 128-frame
media-io kwargs, also pinned in ``minicpmo_4_5.yaml``).
text modalities with ``use_tts_template``, and server ``--interleave-mm-strings``
+ 1fps / 128-frame media-io kwargs, also pinned in ``minicpmo_4_5.yaml``).
"""
from __future__ import annotations
@@ -40,7 +40,10 @@ _MAX_SEED_TTS_MEAN_WER = 0.05
# Match the validated Daily-Omni client body from daily_omni_bench.sh.
_DAILY_EXTRA_BODY = {
"modalities": ["text"],
"chat_template_kwargs": {"enable_thinking": False},
"chat_template_kwargs": {
"enable_thinking": False,
"use_tts_template": True,
},
}
_SEED_EXTRA_BODY = {
"modalities": ["text", "audio"],
+2 -2
View File
@@ -328,7 +328,7 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag
images = None
cot_text = ""
for out in outputs:
ro = out
ro = getattr(out, "request_output", None)
if ro and getattr(ro, "outputs", None):
cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
if not cot_text:
@@ -467,7 +467,7 @@ def _extract_image(outputs) -> Image.Image:
assert outputs, "Pipeline produced no outputs"
for output in outputs:
images = getattr(output, "images", None)
request_output = output
request_output = getattr(output, "request_output", None)
if not images and request_output is not None:
images = getattr(request_output, "images", None)
if images:
@@ -141,8 +141,8 @@ def _extract_frames(result) -> list[Image.Image]:
if isinstance(result, list):
result = result[0] if result else None
if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result is not None:
result = result
if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None:
result = result.request_output
images = getattr(result, "images", None)
if not images:
raise AssertionError("Pipeline output had no images")
@@ -171,11 +171,11 @@ def _build_text_output(text: str = "This is a test response.") -> OmniRequestOut
metrics=None,
lora_request=None,
)
return OmniRequestOutput.from_stage_output(
request_output,
return OmniRequestOutput(
request_id="test_req_text",
finished=True,
final_output_type="text",
request_output=request_output,
)
@@ -200,11 +200,11 @@ def _build_audio_chat_output(num_samples: int = 24000) -> OmniRequestOutput:
metrics=None,
lora_request=None,
)
return OmniRequestOutput.from_stage_output(
request_output,
return OmniRequestOutput(
request_id="test_req_audio_chat",
finished=True,
final_output_type="audio",
request_output=request_output,
)
@@ -240,11 +240,11 @@ def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
images=[_build_image_output(color="blue")],
finished=True,
)
return OmniRequestOutput.from_stage_output(
request_output,
return OmniRequestOutput(
request_id="test_req_img_chat",
images=[_build_image_output(color="blue")],
finished=True,
final_output_type="image",
request_output=request_output,
)
@@ -515,19 +515,8 @@ def mock_async_omni(
@pytest.fixture
def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni, tmp_path):
def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
"""Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
# Override the STORAGE_MANAGER path to a writable temp directory before
# forking the server subprocess. The default /tmp/storage may be owned
# by a different user (e.g. from a previous CI run) and cause
# PermissionError when the server tries to save generated video files.
from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
old_storage_path = STORAGE_MANAGER.storage_path
new_storage_path = tmp_path / "omni_storage"
new_storage_path.mkdir(exist_ok=True)
STORAGE_MANAGER.storage_path = str(new_storage_path)
parser = TrackingArgumentParser()
subparsers = parser.add_subparsers(dest="command")
cmd = OmniServeCommand()
@@ -566,9 +555,6 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
yield f"http://127.0.0.1:{port}/v1"
# Restore the original storage path after the test.
STORAGE_MANAGER.storage_path = old_storage_path
if server_process.is_alive():
server_process.terminate()
server_process.join(timeout=10)
@@ -294,7 +294,7 @@ class vLLMOmniHttpServerLocal:
"global_steps": self.global_steps,
}
req_output = final_res
req_output = final_res.request_output
if req_output is not None and hasattr(req_output, "finish_reason"):
finish_reason = req_output.finish_reason or "stop"
else:
@@ -53,9 +53,9 @@ def test_audiox_model(omni_runner: OmniRunner) -> None:
assert outputs is not None
first_output = outputs[0]
assert first_output.final_output_type == "audio"
assert isinstance(first_output, OmniRequestOutput) and first_output
assert hasattr(first_output, "request_output") and first_output.request_output
req_out = first_output
req_out = first_output.request_output
assert isinstance(req_out, OmniRequestOutput)
assert req_out.final_output_type == "audio"
assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output
+1 -1
View File
@@ -55,7 +55,7 @@ pytestmark = [
def _extract_text(omni_outputs: list) -> str:
"""Extract generated text from OmniRequestOutput list."""
for req_output in omni_outputs:
ro = req_output
ro = getattr(req_output, "request_output", None)
if ro and getattr(ro, "outputs", None):
return "".join(getattr(o, "text", "") or "" for o in ro.outputs)
return ""
@@ -28,6 +28,7 @@ from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
@@ -175,12 +176,12 @@ def test_cosyvoice3_offline_reference_zero_shot(omni_runner: OmniRunner) -> None
assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
# Code2wav is ``final_output`` (stage 1); ``Omni.generate`` yields its ``request_output`` / completions.
pipeline_out = outputs[0]
pipeline_out = OmniRequestOutput.unwrap_result(outputs[0])
assert pipeline_out.stage_id == 1, f"expected final stage 1, got {pipeline_out.stage_id}"
assert pipeline_out.final_output_type == "audio", (
f"expected audio final_output_type, got {pipeline_out.final_output_type!r}"
)
ro = pipeline_out
ro = pipeline_out.request_output
assert ro is not None, "stage 1 should include request_output"
completions = getattr(ro, "outputs", None) or []
assert completions, "stage 1 request_output should include at least one completion"
@@ -175,7 +175,7 @@ def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None:
# Prefer the latest finished chunk to avoid picking an intermediate stream output.
for stage_output in reversed(matched):
if _is_finished_request_output(stage_output):
if _is_finished_request_output(getattr(stage_output, "request_output", None)):
return stage_output
return matched[-1]
@@ -205,7 +205,7 @@ def _to_token_list(value: Any) -> list[int]:
def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
request_output = stage_output
request_output = getattr(stage_output, "request_output", None)
if request_output is None:
return ""
req_list = request_output if isinstance(request_output, list) else [request_output]
@@ -240,7 +240,7 @@ def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
def _extract_audio(stage_output: Any) -> Any | None:
request_output = stage_output
request_output = getattr(stage_output, "request_output", None)
if request_output is None:
return None
req_list = request_output if isinstance(request_output, list) else [request_output]
@@ -256,7 +256,7 @@ def _extract_audio(stage_output: Any) -> Any | None:
def _extract_image(stage_output: Any) -> Any | None:
request_output = stage_output
request_output = getattr(stage_output, "request_output", None)
if request_output is None:
return None
req_list = request_output if isinstance(request_output, list) else [request_output]
@@ -54,7 +54,7 @@ def _sampling_params() -> OmniDiffusionSamplingParams:
def _first_request_images(outputs) -> list:
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output
req_out = first_output.request_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
return req_out.images
@@ -161,7 +161,7 @@ def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunn
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output
req_out = first_output.request_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
images = req_out.images
@@ -224,7 +224,7 @@ def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunne
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output
req_out = first_output.request_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
images = req_out.images
@@ -288,7 +288,7 @@ def _extract_generated_image(outputs: list[object]) -> Image.Image:
if images := getattr(first_output, "images", None):
return images[0]
request_output = first_output
request_output = getattr(first_output, "request_output", None)
if request_output is not None and (images := getattr(request_output, "images", None)):
return images[0]
@@ -74,7 +74,7 @@ def test_i2t_generates_text(omni: Omni) -> None:
outputs = omni.generate(prompts=[prompt_dict])
assert outputs, "No outputs returned from Omni.generate()"
request_output = outputs[0]
request_output = outputs[0].request_output
assert request_output.outputs, "No completion outputs"
completion = request_output.outputs[0]
@@ -172,7 +172,7 @@ def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner):
found_image = False
for out in outputs:
ro_list = out
ro_list = getattr(out, "request_output", out)
if not isinstance(ro_list, list):
ro_list = [ro_list]
for ro in ro_list:
@@ -137,11 +137,14 @@ def _extract_final_audio_outputs(outputs):
if getattr(item, "final_output_type", None) == "audio":
final_outputs.append(item)
continue
multimodal_output = getattr(item, "multimodal_output", None)
request_output = getattr(item, "request_output", None)
if request_output is None:
continue
multimodal_output = getattr(request_output, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
final_outputs.append(item)
continue
completions = getattr(item, "outputs", None) or []
completions = getattr(request_output, "outputs", None) or []
if any(isinstance(getattr(completion, "multimodal_output", None), Mapping) for completion in completions):
final_outputs.append(item)
return final_outputs
@@ -152,11 +155,16 @@ def _extract_multimodal_output(output) -> Mapping[str, Any]:
if isinstance(multimodal_output, Mapping):
return multimodal_output
completions = getattr(output, "outputs", None) or []
for completion in completions:
multimodal_output = getattr(completion, "multimodal_output", None)
request_output = getattr(output, "request_output", None)
if request_output is not None:
multimodal_output = getattr(request_output, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
return multimodal_output
completions = getattr(request_output, "outputs", None) or []
for completion in completions:
multimodal_output = getattr(completion, "multimodal_output", None)
if isinstance(multimodal_output, Mapping):
return multimodal_output
raise AssertionError("No multimodal audio output found in Ming generate results")
@@ -108,7 +108,7 @@ def _build_request(
def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
"""Run a single request and return (waveform, sample_rate)."""
for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
req_output = stage_outputs
req_output = stage_outputs.request_output
if req_output is not None:
mm = req_output.outputs[0].multimodal_output
assert mm is not None, "Expected multimodal_output to be non-None"
@@ -167,7 +167,7 @@ def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
results = []
# Single-stage model (num_stages=1): one sampling param for all requests.
for stage_outputs in omni_runner.omni.generate(requests, [DEFAULT_SAMPLING]):
req_output = stage_outputs
req_output = stage_outputs.request_output
if req_output is not None:
mm = req_output.outputs[0].multimodal_output
assert mm is not None
@@ -58,7 +58,7 @@ def test_omnivoice_text_to_audio(omni_runner: OmniRunner) -> None:
# Check final output has audio
final_output = outputs[-1]
ro = final_output
ro = final_output.request_output
assert ro is not None, "No request_output"
mm = getattr(ro, "multimodal_output", None)
@@ -42,7 +42,7 @@ def _sampling_params() -> OmniDiffusionSamplingParams:
def _first_request_images(outputs) -> list:
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output
req_out = first_output.request_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
return req_out.images
@@ -62,9 +62,9 @@ def generate_stable_audio_short_clip(
# Audio-output diffusion pipelines (those with ``support_audio_output = True``) now have
# ``final_output_type="audio"`` set on the outer stage metadata as well as the inner request.
assert first_output.final_output_type == "audio"
assert isinstance(first_output, OmniRequestOutput) and first_output
assert hasattr(first_output, "request_output") and first_output.request_output
req_out = first_output
req_out = first_output.request_output
assert isinstance(req_out, OmniRequestOutput)
assert req_out.final_output_type == "audio"
assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output
@@ -63,8 +63,8 @@ def _assert_non_empty_text_output(outputs: list[OmniRequestOutput]) -> None:
text_output = next((o for o in outputs if o.final_output_type == "text"), None)
assert text_output is not None
assert text_output is not None
text_content = text_output.outputs[0].text
assert text_output.request_output is not None
text_content = text_output.request_output.outputs[0].text
assert text_content is not None
assert len(text_content.strip()) > 0
@@ -81,8 +81,8 @@ def test_audio_to_text_and_audio(omni_runner: OmniRunner) -> None:
audio_output = next((o for o in outputs if o.final_output_type == "audio"), None)
if audio_output is not None:
assert audio_output is not None
audio_tensor = audio_output.outputs[0].multimodal_output.get("audio")
assert audio_output.request_output is not None
audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio")
if audio_tensor is not None:
assert audio_tensor.numel() > 0
@@ -53,11 +53,7 @@ def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_
assert result["admission"]["overflow_error_code"] == "resource_exhausted"
# CUDA-only for now: this contract asserts the model speaks while the user is
# still talking, which only holds when the duplex pipeline sustains real-time
# throughput. The current NPU stack runs several times slower than real time,
# so it never reaches a mid-stream decision point.
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -> None:
input_wav = validated_soft_interrupt_wav()
@@ -74,8 +70,13 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -
timeout_s=180.0,
require_audio=True,
no_realtime_pacing=False,
validation_mode="response-required",
min_responses=2,
# Single-GPU co-location cannot reliably finish the fixture's
# two-response soft-interrupt handoff before the WAV ends.
# model-policy + deterministic sampling checks streaming audio
# lifecycle instead of the diagnostic two-response contract.
validation_mode="model-policy",
temperature=0.0,
min_responses=1,
min_audio_deltas_per_response=2,
input_sha256=SOFT_INTERRUPT_SHA256,
expect_followup_response_substring=None,
@@ -87,4 +88,5 @@ def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -
assert result["error_count"] == 0
assert result["response_lifecycle_ok"] is True
assert result["response_audio_contract_ok"] is True
assert result["followup_response_transcript_ok"] is True
assert result["response_before_final_commit"] is True
assert result["enough_responses"] is True
@@ -65,7 +65,7 @@ def test_realtime_duplex_multi_session_script_is_directly_executable():
cwd=MULTI_DEMO_PATH.parents[3],
capture_output=True,
text=True,
timeout=60,
timeout=20,
check=False,
)
@@ -78,7 +78,7 @@ def test_realtime_duplex_demo_pair_script_is_directly_executable():
cwd=PAIR_DEMO_PATH.parents[3],
capture_output=True,
text=True,
timeout=60,
timeout=20,
check=False,
)
@@ -91,7 +91,7 @@ def test_realtime_duplex_soft_interrupt_script_is_directly_executable():
cwd=SOFT_INTERRUPT_DEMO_PATH.parents[3],
capture_output=True,
text=True,
timeout=60,
timeout=20,
check=False,
)
+1 -1
View File
@@ -50,7 +50,7 @@ async def generate(
output_modalities=["text"],
):
stage_id = omni_output.stage_id
out = omni_output
out = omni_output.request_output
if stage_id == 0:
num_tokens = sum(len(output.token_ids) for output in out.outputs)
count = num_tokens
+143 -241
View File
@@ -4,269 +4,171 @@
import pytest
from PIL import Image
from vllm.outputs import CompletionOutput, RequestOutput
from pytest_mock import MockerFixture
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
# ── helpers ────────────────────────────────────────────────────────────────
class TestOmniRequestOutput:
"""Tests for OmniRequestOutput class."""
def _make_text_request_output(
request_id: str = "req-1",
text: str = "hello world",
token_ids: list[int] | None = None,
finished: bool = True,
) -> RequestOutput:
"""Build a minimal vLLM RequestOutput for testing."""
return RequestOutput(
request_id=request_id,
prompt="test prompt",
prompt_token_ids=[1, 2, 3],
prompt_logprobs=None,
outputs=[
CompletionOutput(
index=0,
text=text,
token_ids=token_ids or [10, 11, 12],
cumulative_logprob=-1.0,
logprobs=None,
finish_reason="stop" if finished else None,
stop_reason=None,
)
],
finished=finished,
metrics=None,
lora_request=None,
)
# ── from_stage_output ──────────────────────────────────────────────────────
class TestFromStageOutput:
"""Tests for OmniRequestOutput.from_stage_output()."""
def test_copies_request_output_content(self):
"""Content from a vLLM RequestOutput is flattened onto the new object."""
source = _make_text_request_output(text="copied text")
out = OmniRequestOutput.from_stage_output(source, request_id="req-1", stage_id=0, final_output_type="text")
assert out.prompt == "test prompt"
assert out.prompt_token_ids == [1, 2, 3]
assert len(out.outputs) == 1
assert out.outputs[0].text == "copied text"
assert out.finished is True
assert out.stage_id == 0
assert out.final_output_type == "text"
def test_copies_omni_content(self):
"""Diffusion fields are copied when source is another OmniRequestOutput."""
image = Image.new("RGB", (64, 64), color="blue")
inner = OmniRequestOutput.from_diffusion(request_id="req-diff", images=[image], prompt="a cat")
outer = OmniRequestOutput.from_stage_output(inner, request_id="req-diff", stage_id=0, final_output_type="image")
assert len(outer.images) == 1
assert outer.images[0] is image
assert outer.prompt == "a cat"
assert outer.is_diffusion_output
def test_with_duck_typed_source(self):
"""A duck-typed source (SimpleNamespace-style) copies correctly."""
from types import SimpleNamespace
source = SimpleNamespace(
outputs=[],
prompt="duck prompt",
prompt_token_ids=[7, 8, 9],
finished=False,
def test_from_diffusion(self):
"""Test creating output from diffusion model."""
images = [Image.new("RGB", (64, 64), color="red")]
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=images,
prompt="a cat",
metrics={"steps": 50},
)
out = OmniRequestOutput.from_stage_output(source, request_id="duck-1", stage_id=2)
assert output.request_id == "test-123"
assert output.images == images
assert output.prompt == "a cat"
assert output.metrics == {"steps": 50}
assert output.is_diffusion_output
assert output.num_images == 1
assert out.prompt == "duck prompt"
assert out.prompt_token_ids == [7, 8, 9]
assert out.finished is False
assert out.stage_id == 2
def test_from_pipeline(self, mocker: MockerFixture):
"""Test creating output from pipeline stage."""
mock_request_output = mocker.Mock()
mock_request_output.request_id = "pipeline-123"
mock_request_output.prompt_token_ids = [1, 2, 3]
mock_request_output.outputs = [mocker.Mock()]
mock_request_output.encoder_prompt_token_ids = None
mock_request_output.prompt_logprobs = None
mock_request_output.num_cached_tokens = 10
mock_request_output.kv_transfer_params = None
mock_request_output.multimodal_output = {"image": mocker.Mock()}
def test_request_id_always_copied_from_source(self):
"""request_id is always copied from source, even when passed in kwargs."""
source = _make_text_request_output(request_id="source-id")
out = OmniRequestOutput.from_stage_output(source, request_id="kwargs-id", stage_id=1)
# Source wins — request_id is content, not a control field.
assert out.request_id == "source-id"
def test_passes_control_fields(self):
"""Control fields (metrics, stage_durations, peak_memory_mb) are set."""
source = _make_text_request_output()
out = OmniRequestOutput.from_stage_output(
source,
request_id="r1",
stage_id=3,
metrics={"ttft": 0.5},
stage_durations={"prefill": 0.12},
peak_memory_mb=1024.0,
)
assert out.metrics == {"ttft": 0.5}
assert out.stage_durations == {"prefill": 0.12}
assert out.peak_memory_mb == 1024.0
def test_finished_field_copied(self):
"""The *finished* flag is copied from the source when present."""
source = _make_text_request_output(finished=False)
out = OmniRequestOutput.from_stage_output(source, request_id="r1")
assert out.finished is False
source2 = _make_text_request_output(finished=True)
out2 = OmniRequestOutput.from_stage_output(source2, request_id="r2")
assert out2.finished is True
def test_outputs_field_copied(self):
"""The *outputs* list is copied from the source."""
source = _make_text_request_output(text="hello", token_ids=[1, 2, 3])
out = OmniRequestOutput.from_stage_output(source, request_id="r1")
assert len(out.outputs) == 1
assert out.outputs[0].text == "hello"
assert out.outputs[0].token_ids == [1, 2, 3]
# ── dataclass safety ────────────────────────────────────────────────────────
class TestDataclassSafety:
def test_from_stage_output_is_classmethod(self):
"""from_stage_output must be a classmethod."""
assert callable(OmniRequestOutput.from_stage_output)
# ── msgpack round-trip ─────────────────────────────────────────────────────
class TestMsgpackRoundTrip:
"""Tests that OmniRequestOutput survives msgpack encode → decode without
infinite recursion or data loss."""
def test_round_trip_pipeline_output(self):
"""A pipeline output round-trips without data loss."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
OmniMsgpackEncoder,
)
source = _make_text_request_output(text="round trip")
out = OmniRequestOutput.from_stage_output(
source,
request_id="rt-1",
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
encoder = OmniMsgpackEncoder()
decoder = OmniMsgpackDecoder()
assert output.request_id == "pipeline-123"
assert output.stage_id == 0
assert output.final_output_type == "text"
assert output.is_pipeline_output
encoded = encoder.encode(out)
decoded = decoder.decode(encoded)
def test_prompt_token_ids_property(self, mocker: MockerFixture):
"""Test prompt_token_ids property for streaming compatibility."""
mock_request_output = mocker.Mock()
mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
assert isinstance(decoded, OmniRequestOutput)
assert decoded.request_id == "req-1"
assert decoded.stage_id == 0
assert decoded.final_output_type == "text"
assert decoded.finished is True
assert decoded.outputs[0].text == "round trip"
assert decoded.prompt == "test prompt"
assert decoded.prompt_token_ids == [1, 2, 3]
def test_round_trip_diffusion_output(self):
"""A diffusion output round-trips without data loss."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
OmniMsgpackEncoder,
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
image = Image.new("RGB", (64, 64), color="red")
out = OmniRequestOutput.from_diffusion(
request_id="diff-1",
images=[image],
prompt="a red image",
assert output.prompt_token_ids == [1, 2, 3, 4, 5]
def test_prompt_token_ids_none_when_no_request_output(self):
"""Test prompt_token_ids returns None when no request_output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[],
prompt="a cat",
)
assert output.prompt_token_ids is None
def test_outputs_property(self, mocker: MockerFixture):
"""Test outputs property for chat completion compatibility."""
mock_output = mocker.Mock()
mock_request_output = mocker.Mock()
mock_request_output.outputs = [mock_output]
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.outputs == [mock_output]
def test_outputs_empty_when_no_request_output(self):
"""Test outputs returns empty list when no request_output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[],
prompt="a cat",
)
assert output.outputs == []
def test_encoder_prompt_token_ids_property(self, mocker: MockerFixture):
"""Test encoder_prompt_token_ids property."""
mock_request_output = mocker.Mock()
mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.encoder_prompt_token_ids == [10, 20, 30]
def test_num_cached_tokens_property(self, mocker: MockerFixture):
"""Test num_cached_tokens property."""
mock_request_output = mocker.Mock()
mock_request_output.num_cached_tokens = 42
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
assert output.num_cached_tokens == 42
def test_multimodal_output_property(self, mocker: MockerFixture):
"""Test multimodal_output property."""
mock_request_output = mocker.Mock()
mock_audio = mocker.Mock()
expected_output = {"audio": mock_audio}
mock_request_output.outputs = []
mock_request_output.multimodal_output = expected_output
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="audio",
request_output=mock_request_output,
)
assert output.multimodal_output is expected_output
def test_to_dict_diffusion(self):
"""Test to_dict for diffusion output."""
output = OmniRequestOutput.from_diffusion(
request_id="test-123",
images=[Image.new("RGB", (64, 64), color="red")],
prompt="a cat",
metrics={"steps": 50},
multimodal_output={"audio": None},
)
result = output.to_dict()
encoder = OmniMsgpackEncoder()
decoder = OmniMsgpackDecoder()
assert result["request_id"] == "test-123"
assert result["finished"] is True
assert result["final_output_type"] == "image"
assert result["num_images"] == 1
assert result["prompt"] == "a cat"
encoded = encoder.encode(out)
decoded = decoder.decode(encoded)
def test_to_dict_pipeline(self, mocker: MockerFixture):
"""Test to_dict for pipeline output."""
mock_request_output = mocker.Mock()
mock_request_output.request_id = "pipeline-123"
assert isinstance(decoded, OmniRequestOutput)
assert decoded.request_id == "diff-1"
assert decoded.prompt == "a red image"
assert decoded.num_images == 1
assert decoded.images[0].size == (64, 64)
assert decoded.metrics == {"steps": 50}
def test_round_trip_does_not_recurse(self):
"""Encoding a pipeline output must not trigger infinite recursion."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackEncoder,
output = OmniRequestOutput.from_pipeline(
stage_id=0,
final_output_type="text",
request_output=mock_request_output,
)
result = output.to_dict()
source = _make_text_request_output(text="no recursion")
out = OmniRequestOutput.from_stage_output(source, request_id="rec-1", stage_id=0)
encoder = OmniMsgpackEncoder()
encoded = encoder.encode(out)
assert isinstance(encoded, bytes)
assert len(encoded) > 0
def test_legacy_wire_format_deserialized(self):
"""When the wire format contains a nested ``request_output`` key
(pre-inheritance legacy), it is merged via _copy_content_from."""
from vllm_omni.distributed.omni_connectors.utils.serialization import (
OmniMsgpackDecoder,
)
# Build a legacy-style dict: OmniRequestOutput fields + a nested
# request_output dict containing RequestOutput-like content.
legacy_dict = {
"request_id": "legacy-1",
"finished": True,
"stage_id": 1,
"final_output_type": "text",
"images": [],
"_multimodal_output": {},
"_custom_output": {},
"stage_durations": {},
"peak_memory_mb": 0.0,
# --- nested request_output key (legacy) ---
"request_output": {
"request_id": "legacy-1",
"prompt": "legacy prompt",
"prompt_token_ids": [1, 2],
"prompt_logprobs": None,
"outputs": [],
"finished": True,
"metrics": {},
"lora_request": None,
"encoder_prompt": None,
"encoder_prompt_token_ids": None,
"num_cached_tokens": None,
"kv_transfer_params": None,
},
}
# Manually run the decoder's post_process to test legacy path.
# The decoder's _decode_omni_request_output pops "request_output"
# and merges it via _copy_content_from.
decoder = OmniMsgpackDecoder()
result = decoder._post_process(legacy_dict)
assert isinstance(result, OmniRequestOutput)
assert result.request_id == "legacy-1"
assert result.finished is True
# Content merged from legacy nested key
assert result.prompt == "legacy prompt"
assert result.prompt_token_ids == [1, 2]
assert result["request_id"] == "pipeline-123"
assert result["finished"] is True
assert result["final_output_type"] == "text"
assert result["stage_id"] == 0
@@ -2156,12 +2156,12 @@ def test_direct_listen_decision_survives_inner_completion_metadata():
)
assert decision is not None
output = attach_duplex_output_decision(
OmniRequestOutput.from_stage_output(
inner_output,
OmniRequestOutput(
request_id="duplex-direct-listen",
finished=True,
stage_id=0,
final_output_type=decision.final_output_type,
request_output=inner_output,
metrics={
"stage_metrics": {
"0": {
@@ -3049,12 +3049,12 @@ async def test_minicpmo_auto_response_listen_without_response_does_not_defer_com
)
assert decision is not None
listen_output = attach_duplex_output_decision(
OmniRequestOutput.from_stage_output(
inner_output,
OmniRequestOutput(
request_id=request_id,
finished=True,
stage_id=0,
final_output_type=decision.final_output_type,
request_output=inner_output,
),
decision,
)
@@ -119,6 +119,7 @@ class MockGenerationResult:
def __init__(self, images):
self.images = images
self.request_output = SimpleNamespace(images=images)
self.stage_durations = {}
self.peak_memory_mb = 0.0
@@ -136,7 +137,10 @@ class MockStageResult:
outputs = [SimpleNamespace(text=text, index=0)]
else:
outputs = []
self.outputs = outputs
self.request_output = SimpleNamespace(
outputs=outputs,
images=self.images,
)
self.stage_durations = {}
self.peak_memory_mb = 0.0
@@ -2076,20 +2080,22 @@ def test_extract_images_from_result():
assert all(isinstance(img, Image.Image) for img in images)
assert all(img.size == (64, 64) for img in images)
# Test result with "images" attribute set in __init__
# Test dict path: result.request_output["images"]
class DictRequestOutput:
def __init__(self):
self.images = [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]
self.request_output = {"images": [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]}
result = DictRequestOutput()
images = _extract_images_from_result(result)
assert len(images) == 1
assert isinstance(images[0], Image.Image)
# Test result with "images" attribute from an inner object
# Test attribute path: result.request_output.images
class AttrRequestOutput:
def __init__(self):
self.images = [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]
self.request_output = type(
"obj", (), {"images": [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]}
)()
result = AttrRequestOutput()
images = _extract_images_from_result(result)
@@ -180,10 +180,10 @@ def create_mock_audio_output_for_test(
audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
mock_request_output = MockRequestOutput(request_id=request_id, audio_tensor=audio_tensor)
return OmniRequestOutput.from_stage_output(
mock_request_output,
return OmniRequestOutput(
stage_id=0,
final_output_type="audio",
request_output=mock_request_output,
)
@@ -2292,10 +2292,10 @@ class TestStreamingResponse:
self.prompt_logprobs = None
self.kv_transfer_params = None
return OmniRequestOutput.from_stage_output(
MockRequestOutput(audio_tensor=chunk),
return OmniRequestOutput(
stage_id=0,
final_output_type="audio",
request_output=MockRequestOutput(audio_tensor=chunk),
finished=finished,
)
@@ -3441,10 +3441,10 @@ class TestWAVStreaming:
self.prompt_logprobs = None
self.kv_transfer_params = None
return OmniRequestOutput.from_stage_output(
MockRequestOutput(audio_tensor=chunk),
return OmniRequestOutput(
stage_id=0,
final_output_type="audio",
request_output=MockRequestOutput(audio_tensor=chunk),
finished=finished,
)
@@ -3,7 +3,6 @@ import base64
from types import SimpleNamespace
import pytest
import torch
from fastapi import FastAPI, WebSocket
from pytest_mock import MockerFixture
from starlette.testclient import TestClient
@@ -12,24 +11,11 @@ from starlette.websockets import WebSocketDisconnect
from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
from vllm_omni.model_executor.stage_input_processors.forced_aligner import ALIGNER_WORDS_KEY
from vllm_omni.utils.forced_aligner import WordTimestamp
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _fake_aligner_res(pairs, words):
"""Build a stand-in for the forced-aligner stage's pooling output.
Mirrors what rides the generator in production: ``res.outputs.data`` is a
``[n_words, 2]`` int32 tensor of ``[start_ms, end_ms]`` and the word strings
travel in ``additional_information``. Decoded by ``extract_word_timestamps``.
"""
return SimpleNamespace(
outputs=SimpleNamespace(data=torch.tensor(pairs, dtype=torch.int32), additional_information=None),
additional_information={ALIGNER_WORDS_KEY: list(words)},
)
def _build_test_app(
speech_service=None,
*,
@@ -42,9 +28,9 @@ def _build_test_app(
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav"))
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
speech_service.forced_aligner_enabled = False
speech_service.forced_aligner_config = None
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
for chunk in (b"\x01\x02", b"\x03\x04\x05"):
yield (chunk, 24000) if include_sample_rate else chunk
@@ -249,7 +235,7 @@ class TestStreamingSpeechWebSocket:
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
speech_service.engine_client = mocker.MagicMock()
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = False
speech_service.forced_aligner_config = None
async def mock_prepare_speech_generation(request):
captured_requests.append(request)
@@ -325,13 +311,13 @@ class TestStreamingSpeechWebSocket:
assert error["type"] == "error"
assert "without --forced-aligner" in error["message"]
def test_word_timestamps_emit_pipeline_json_frame(self, mocker: MockerFixture):
def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture):
captured_requests = []
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
speech_service.engine_client = mocker.MagicMock()
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = True
speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
async def mock_prepare_speech_generation(request):
captured_requests.append(request)
@@ -342,15 +328,19 @@ class TestStreamingSpeechWebSocket:
first_chunk = b"\x01" * 1000
second_chunk = b"\x02" * 1000
# The forced-aligner stage rides the same generator: its pooling output
# is surfaced via the ``collect`` channel once the audio has streamed.
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
for chunk in (first_chunk, second_chunk):
yield (chunk, 1000) if include_sample_rate else chunk
if collect is not None:
collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"])
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
# Sentence-level: aligner runs once over the whole sentence audio.
mock_align = mocker.AsyncMock(
return_value=[
WordTimestamp("Hello", 0, 200),
WordTimestamp("world", 200, 900),
]
)
mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
app, _ = _build_test_app(speech_service)
with TestClient(app) as client:
@@ -416,6 +406,7 @@ class TestStreamingSpeechWebSocket:
}
assert captured_requests[0].word_timestamps is True
assert mock_align.await_count == 1
def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture):
# The streaming layer forwards the aligner's (already monotonic,
@@ -424,16 +415,21 @@ class TestStreamingSpeechWebSocket:
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
speech_service.engine_client = mocker.MagicMock()
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = True
speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
chunk = b"\x01" * 1000
yield (chunk, 1000) if include_sample_rate else chunk
if collect is not None:
collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"])
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
mock_align = mocker.AsyncMock(
return_value=[
WordTimestamp("Hello", 0, 1000),
WordTimestamp("world", 1000, 1200),
]
)
mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
app, _ = _build_test_app(speech_service)
with TestClient(app) as client:
@@ -576,7 +572,7 @@ class TestStreamingSpeechWebSocket:
speech_service._generate_pcm_chunks = mocker.AsyncMock()
speech_service.engine_client = mocker.MagicMock()
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = False
speech_service.forced_aligner_config = None
app, _ = _build_test_app(speech_service)
with TestClient(app) as client:
@@ -606,7 +602,7 @@ class TestStreamingSpeechWebSocket:
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {}))
speech_service.engine_client = mocker.MagicMock()
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = False
speech_service.forced_aligner_config = None
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
yield b"\x01\x02"
@@ -700,7 +696,7 @@ class TestStreamingSpeechWebSocket:
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {}))
speech_service.engine_client = mocker.MagicMock()
speech_service.engine_client.abort = mocker.AsyncMock()
speech_service.forced_aligner_enabled = False
speech_service.forced_aligner_config = None
async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
yield b"\x01\x02"
@@ -47,7 +47,7 @@ def _text_result(text: str) -> OmniRequestOutput:
output.text = text
request_output = RequestOutput()
request_output.outputs = [output]
return OmniRequestOutput.from_stage_output(request_output, final_output_type="text")
return OmniRequestOutput(final_output_type="text", request_output=request_output)
def _audio_result(audio_data: Any) -> OmniRequestOutput:
@@ -61,7 +61,7 @@ def _audio_result(audio_data: Any) -> OmniRequestOutput:
output.multimodal_output = {"audio": audio_data}
request_output = RequestOutput()
request_output.outputs = [output]
return OmniRequestOutput.from_stage_output(request_output, final_output_type="audio")
return OmniRequestOutput(final_output_type="audio", request_output=request_output)
class MockWebSocket:
@@ -20,7 +20,6 @@ def test_default_stage_config_includes_cache_backend():
"cache_backend": "cache_dit",
"cache_config": '{"Fn_compute_blocks": 2}',
"vae_use_slicing": True,
"extras": {"ltx2_use_channels_last_3d": True},
"ulysses_degree": 2,
}
)[0]
@@ -30,7 +29,6 @@ def test_default_stage_config_includes_cache_backend():
assert engine_args["cache_backend"] == "cache_dit"
assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
assert engine_args["vae_use_slicing"] is True
assert engine_args["extras"]["ltx2_use_channels_last_3d"] is True
assert engine_args["parallel_config"].ulysses_degree == 2
assert engine_args["model_stage"] == "diffusion"
Loaded 100 of 182 files, more files were not shown because too many files have changed in this diff. Show more