Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
118f646c87 | ||
|
|
7d5a6bc5cb | ||
|
|
8361765b33 | ||
|
|
b3c25a608d |
No files matched your search
@@ -102,7 +102,6 @@ CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_depe
|
||||
SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback"
|
||||
CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN"
|
||||
CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi'
|
||||
CUDA_RUNTIME_CHECK = "python3 .buildkite/cuda/scripts/check_vllm_runtime.py"
|
||||
|
||||
# Bootstrap Buildkite ``if`` expressions.
|
||||
# ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
|
||||
@@ -621,11 +620,11 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
|
||||
if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()):
|
||||
commands = merged.get("commands")
|
||||
if isinstance(commands, list):
|
||||
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, *commands]
|
||||
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands]
|
||||
elif commands is None:
|
||||
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK]
|
||||
merged["commands"] = [CUDA_HF_TOKEN_EXPORT]
|
||||
else:
|
||||
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, commands]
|
||||
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, commands]
|
||||
# Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
|
||||
if "retry" in step:
|
||||
merged["retry"] = step["retry"]
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Verify the pinned vLLM and compiled CUDA imports on a GPU CI worker."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
dockerfile = Path(__file__).resolve().parents[3] / "docker" / "Dockerfile.ci"
|
||||
match = re.search(
|
||||
r"^ARG VLLM_VERSION=([0-9]+\.[0-9]+\.[0-9]+)$",
|
||||
dockerfile.read_text(),
|
||||
re.MULTILINE,
|
||||
)
|
||||
if match is None:
|
||||
raise RuntimeError(f"Missing vLLM release version in {dockerfile}")
|
||||
expected_version = match.group(1)
|
||||
|
||||
import torch
|
||||
import vllm
|
||||
|
||||
assert vllm.__version__ == expected_version, vllm.__version__
|
||||
assert torch.cuda.is_available(), "vLLM runtime check requires a GPU worker"
|
||||
for module in ("flashinfer", "vllm._C_stable_libtorch"):
|
||||
importlib.import_module(module)
|
||||
print(
|
||||
f"Verified vLLM CUDA runtime: vllm={vllm.__version__} "
|
||||
f"expected={expected_version} torch={torch.__version__} "
|
||||
f"cuda={torch.version.cuda} device={torch.cuda.get_device_name()}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -124,30 +124,6 @@ steps:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Omni · Qwen3-Omni Async-Chunk SHM Leak"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_omni_shm_leak.py
|
||||
- tests/helpers/runtime.py
|
||||
- tests/helpers/stage_config.py
|
||||
- vllm_omni/core/sched/omni_ar_scheduler.py
|
||||
- vllm_omni/distributed/omni_connectors/connectors/shm_connector.py
|
||||
- vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py
|
||||
- vllm_omni/engine/orchestrator.py
|
||||
- vllm_omni/engine/stage_engine_core_proc.py
|
||||
- vllm_omni/engine/stage_pool.py
|
||||
- vllm_omni/entrypoints/openai/serving_chat.py
|
||||
- vllm_omni/model_executor/models/qwen3_omni/
|
||||
- vllm_omni/model_executor/models/common/qwen3_code_predictor.py
|
||||
- vllm_omni/model_executor/models/common/snake_activation.py
|
||||
- vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
timeout_in_minutes: 40
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Test"
|
||||
source_file_dependencies: omni_minicpmo_4_5_function
|
||||
timeout_in_minutes: 50
|
||||
|
||||
@@ -48,8 +48,6 @@ steps:
|
||||
- label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
|
||||
source_file_dependencies: omni_qwen3_omni_doc
|
||||
timeout_in_minutes: 90
|
||||
artifact_paths:
|
||||
- audio_*.wav
|
||||
commands:
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and H100 and B200 and cards_2" --run-level "full_model"
|
||||
|
||||
@@ -158,11 +156,6 @@ steps:
|
||||
- label: ":full_moon: TTS · Function Test with L4"
|
||||
source_file_dependencies: tts_qwen3_tts_function
|
||||
timeout_in_minutes: 120
|
||||
env:
|
||||
VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR: .
|
||||
artifact_paths:
|
||||
- failed_speech_*.wav
|
||||
- failed_speech_*.json
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
|
||||
+7
-10
@@ -1,7 +1,10 @@
|
||||
# Use the official release image with its matching CUDA/PyTorch stack.
|
||||
ARG VLLM_VERSION=0.31.0
|
||||
# Target vLLM version: v0.30.0 (tag @ ced6857afa0e).
|
||||
# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer
|
||||
# runtime. Do not reinstall a commit wheel or repair dependencies on top of
|
||||
# an older base image at a release boundary. For a future unreleased target,
|
||||
# restore an explicit wheel pin and validate its complete runtime stack.
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v${VLLM_VERSION}
|
||||
ARG VLLM_BASE_TAG=v0.30.0
|
||||
|
||||
# We use registry-based layer caching for the dependencies,
|
||||
# where the tag is determined based on the hash of the contents
|
||||
@@ -9,7 +12,6 @@ ARG VLLM_BASE_TAG=v${VLLM_VERSION}
|
||||
# everything except the actual project dir, since that should not
|
||||
# change much across branches.
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} AS deps
|
||||
ARG VLLM_VERSION
|
||||
# Keep APP_DIR outside /workspace so k8s Buildkite workspace mounts cannot
|
||||
# shadow the editable install target (#5364). Docker WORKDIR creates /opt if
|
||||
# needed.
|
||||
@@ -34,17 +36,12 @@ RUN apt-get update && \
|
||||
# for CI
|
||||
COPY pyproject.toml setup.py ./
|
||||
COPY requirements/ requirements/
|
||||
# Scope the override to Omni; dependency source builds need their own versions.
|
||||
ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
|
||||
ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
|
||||
# Editable install: deps + project metadata are cached in this layer, while
|
||||
# the package itself is linked to APP_DIR (populated later by COPY --link).
|
||||
# Do NOT use a non-editable install here — that copies an empty stub into
|
||||
# site-packages and breaks ``import vllm_omni.entrypoints``.
|
||||
RUN uv pip install --system -e ".[dev]"
|
||||
# Fail the image build if dependency resolution replaced the selected vLLM.
|
||||
# The CPU builder has no NVIDIA driver; CUDA imports are checked at the start
|
||||
# of each GPU test job by .buildkite/cuda/scripts/check_vllm_runtime.py.
|
||||
RUN python3 -c "import vllm; assert vllm.__version__ == '${VLLM_VERSION}', vllm.__version__"
|
||||
|
||||
# The final stage just copies the working dir on top of the
|
||||
# dependencies image. This uses --link so that we can skip
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# NOTE: this is largely taken from vLLM
|
||||
|
||||
# Exclude .git for CI builds — version is set via VLLM_OMNI_VERSION_OVERRIDE
|
||||
# Exclude .git for CI builds — version is set via SETUPTOOLS_SCM_PRETEND_VERSION
|
||||
.git
|
||||
.venv
|
||||
/build
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.31.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.30.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Keep the ROCm runtime aligned with the vLLM release targeted by
|
||||
# docker/Dockerfile.ci. A stale base image can build successfully while failing
|
||||
# at runtime when vLLM-Omni imports APIs added by the newer vLLM release.
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0
|
||||
|
||||
# Keep the dependency-heavy portion independent from the source tree. The
|
||||
# registry cache is keyed on these inputs, while the final source copy can be
|
||||
@@ -12,11 +12,11 @@ FROM ${BASE_IMAGE} AS deps
|
||||
# BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0.
|
||||
# VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when
|
||||
# USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's
|
||||
# release version. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
|
||||
# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An
|
||||
# intentional custom source pin must update that regression in the same change,
|
||||
# and should return to the stable release as soon as one is available.
|
||||
ARG USE_NIGHTLY_BUILD=0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.31.0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.30.0
|
||||
ARG ARG_PYTORCH_ROCM_ARCH
|
||||
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
|
||||
|
||||
@@ -75,8 +75,7 @@ WORKDIR ${COMMON_WORKDIR}/vllm-omni
|
||||
RUN mkdir -p vllm_omni && touch vllm_omni/__init__.py
|
||||
COPY pyproject.toml setup.py ./
|
||||
COPY requirements/ requirements/
|
||||
# Scope the override to Omni; dependency source builds need their own versions.
|
||||
ENV VLLM_OMNI_VERSION_OVERRIDE=0.0.0
|
||||
ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
|
||||
|
||||
RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
|
||||
--no-cache-dir -e ".[dev]" --no-build-isolation && \
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
This section lists the most common options for running vLLM-Omni.
|
||||
|
||||
For options within a vLLM Engine, please refer to the [vLLM 0.31 configuration guide](https://docs.vllm.ai/en/v0.31.0/configuration/index.html).
|
||||
For options within a vLLM Engine, please refer to the [vLLM 0.30 configuration guide](https://docs.vllm.ai/en/v0.30.0/configuration/index.html).
|
||||
|
||||
Each model defines fixed topology in a registered `PipelineConfig` and runtime
|
||||
overrides in a deploy YAML.
|
||||
|
||||
@@ -186,7 +186,7 @@ their keys only.
|
||||
## Inherited vLLM variables
|
||||
|
||||
vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
|
||||
the [vLLM 0.31 environment-variable reference](https://docs.vllm.ai/en/v0.31.0/configuration/env_vars.html)
|
||||
the [vLLM 0.30 environment-variable reference](https://docs.vllm.ai/en/v0.30.0/configuration/env_vars.html)
|
||||
for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
|
||||
XPU, ModelScope, and FlashInfer workspace settings.
|
||||
|
||||
@@ -214,7 +214,7 @@ collection, examples, bug reports, or logs.
|
||||
|
||||
## Model-specific variables
|
||||
|
||||
The inventory includes 93 variables read by a single model or pipeline family. They are
|
||||
The inventory includes 66 variables read by a single model or pipeline family. They are
|
||||
not listed as public usage options here because doing so would turn implementation
|
||||
escape hatches into an accidental compatibility contract.
|
||||
|
||||
@@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the
|
||||
|
||||
| Disposition | Count | Required outcome |
|
||||
| --- | ---: | --- |
|
||||
| Promote | 64 | Move a stable setting into typed stage or model configuration. |
|
||||
| Promote | 38 | Move a stable setting into typed stage or model configuration. |
|
||||
| Request scope | 5 | Move request-varying behavior into a declared request-option schema. |
|
||||
| External | 0 | Retain only when a supported third-party library owns the contract. |
|
||||
| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. |
|
||||
| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. |
|
||||
| Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. |
|
||||
|
||||
The disposition is a migration target, not a statement that the existing
|
||||
|
||||
@@ -514,50 +514,3 @@ for that comparison; this switch only controls automatic MPS management.
|
||||
|
||||
MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on
|
||||
the intended GPU, and warm the complete pipeline before measuring performance.
|
||||
|
||||
## MOSS-TTS Local 1.5 with Model Runner V2
|
||||
|
||||
`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using
|
||||
the shared runtime introduced for Qwen3-TTS:
|
||||
|
||||
```bash
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
|
||||
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
|
||||
```
|
||||
|
||||
This profile inherits the batching, codec graph buckets and 1-frame/15-frame
|
||||
chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local
|
||||
Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp`
|
||||
capabilities while retaining its V1 `talker_mtp` implementation, sampling
|
||||
defaults and explicit request-seed handling. Codec chunks use the native data
|
||||
plane; the internal Talker retains its final-only orchestrator output policy.
|
||||
|
||||
On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration.
|
||||
The codec retains its own CUDA Graph capture/replay, including batch buckets
|
||||
through 64, but disables Inductor compilation (`compilation_config.mode: 0`)
|
||||
of the stateful decoder to avoid lengthy compilation at startup. This is
|
||||
independent of `enforce_eager`; the profile keeps `enforce_eager: false`.
|
||||
It explicitly selects `cudagraph_mode: FULL`, which works without Inductor;
|
||||
the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear
|
||||
the codec's capture buckets when compilation is disabled.
|
||||
|
||||
For sustained C128 serving on a large-memory CUDA GPU, select the
|
||||
separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128
|
||||
stream slots per stage, a 32 GiB Talker KV budget, and the codec's
|
||||
`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs.
|
||||
The bounded KV budget leaves room for codec state and graphs; it does not
|
||||
guarantee that 128 maximum-length prompts fit simultaneously. See the
|
||||
[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention)
|
||||
for activation, backend comparisons, memory requirements and benchmark commands.
|
||||
|
||||
Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1.
|
||||
NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable
|
||||
MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio;
|
||||
set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and
|
||||
`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM.
|
||||
|
||||
Event-driven orchestration remains independently selectable with
|
||||
`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical
|
||||
when comparing these modes. Model-runner selection does not change the
|
||||
orchestration default or enable experimental reference encoding, chunk ramps,
|
||||
generation-output draining or MPS.
|
||||
@@ -1,7 +1,7 @@
|
||||
# Installation
|
||||
|
||||
!!! important
|
||||
vLLM-Omni is released against the matching upstream vLLM major/minor version. This source checkout targets vLLM 0.31.0. The published vLLM-Omni 0.30.0 wheel instructions use vLLM 0.30.0; choose the installation path that matches your Omni revision.
|
||||
vLLM-Omni is released against the matching upstream vLLM major/minor version. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x; the source-install instructions pin vLLM 0.30.0.
|
||||
|
||||
vLLM-Omni supports the following hardware platforms:
|
||||
|
||||
|
||||
@@ -1,13 +1,11 @@
|
||||
# --8<-- [start:requirements]
|
||||
|
||||
- GPU: compute capability 7.5 or higher (e.g., T4, RTX20xx, A100, L4, H100, etc.)
|
||||
- GPU: compute capability 7.0 or higher (e.g., V100, T4, RTX20xx, A100, L4, H100, etc.)
|
||||
|
||||
# --8<-- [end:requirements]
|
||||
# --8<-- [start:set-up-using-python]
|
||||
|
||||
vLLM-Omni depends on the matching major/minor release of vLLM. This source
|
||||
checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
|
||||
the published 0.30.0 release with its matching vLLM dependency.
|
||||
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.30.x release line uses vLLM 0.30.x.
|
||||
|
||||
!!! note
|
||||
PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
|
||||
@@ -43,17 +41,13 @@ uv pip install 'vllm-omni[demo]==0.30.0'
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, install the final 0.31.0
|
||||
release wheel from its commit-pinned index:
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.31.0 --torch-backend=auto \
|
||||
--extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
|
||||
uv pip install vllm==0.30.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The default 0.31.0 wheel contains CUDA 12.9 binaries. If you need a different
|
||||
CUDA variant or want to reuse an existing PyTorch installation, build vLLM
|
||||
from source instead.
|
||||
The 0.30.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -74,16 +68,16 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.31.0
|
||||
git checkout v0.30.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
# Default CUDA 12.9 wheel for v0.31.0 on x86_64
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6/vllm-0.31.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
#For CUDA 13.0 (the default for v0.30.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.30.0/vllm-0.30.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
VLLM_USE_PRECOMPILED=1 uv pip install --editable . --torch-backend=auto
|
||||
uv pip install --editable .
|
||||
```
|
||||
Install vllm with command below (If you already have PyTorch).
|
||||
```bash
|
||||
@@ -134,13 +128,12 @@ docker run --runtime nvidia --gpus 2 \
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.cuda -t vllm-omni-cuda .
|
||||
```
|
||||
|
||||
The source build defaults to `vllm/vllm-openai:v0.31.0`. To select that
|
||||
upstream release image explicitly:
|
||||
If you want to specify the base vLLM version:
|
||||
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
-f docker/Dockerfile.cuda \
|
||||
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.31.0 \
|
||||
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.30.0 \
|
||||
-t vllm-omni-cuda .
|
||||
```
|
||||
|
||||
|
||||
@@ -10,9 +10,8 @@
|
||||
|
||||
For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images.
|
||||
|
||||
vLLM-Omni depends on the matching major/minor release of vLLM. This source
|
||||
checkout targets vLLM 0.31.0. The pre-built Omni wheel instructions below install
|
||||
the published 0.30.0 release with its matching vLLM dependency.
|
||||
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni
|
||||
0.30.x release line uses vLLM 0.30.x.
|
||||
|
||||
The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The
|
||||
`vllm-omni` package does not install vLLM as a dependency, so non-Docker source
|
||||
@@ -50,14 +49,13 @@ uv pip install onnxruntime-rocm
|
||||
|
||||
#### Installation of vLLM
|
||||
|
||||
If you do not need to modify source code of vLLM, install the final 0.31.0
|
||||
release wheel for ROCm 7.2.3:
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.30.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
|
||||
uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
|
||||
```
|
||||
|
||||
The pre-built 0.31.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The pre-built 0.30.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
|
||||
@@ -77,7 +75,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.31.0
|
||||
git checkout v0.30.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
@@ -89,7 +87,7 @@ python3 setup.py develop
|
||||
#### Build docker image
|
||||
|
||||
The source build defaults to the published upstream base image
|
||||
`vllm/vllm-openai-rocm:v0.31.0`, aligned with this source checkout.
|
||||
`vllm/vllm-openai-rocm:v0.30.0`, aligned with the vLLM release used by CI.
|
||||
Older or custom bases must provide the vLLM APIs checked by the Dockerfile's
|
||||
image-build canary. This upstream base is distinct from the prebuilt
|
||||
`vllm/vllm-omni-rocm` images discussed below; their published-tag availability
|
||||
@@ -104,7 +102,7 @@ To select the upstream ROCm base explicitly, pass `BASE_IMAGE`:
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
-f docker/Dockerfile.rocm \
|
||||
--build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.31.0 \
|
||||
--build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.30.0 \
|
||||
-t vllm-omni-rocm .
|
||||
```
|
||||
|
||||
|
||||
@@ -20,16 +20,12 @@ vLLM-Omni currently recommends using the Docker image setup steps below.
|
||||
#### Build docker image
|
||||
|
||||
```bash
|
||||
export VLLM_VERSION=v0.31.0
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g \
|
||||
--build-arg VLLM_VERSION="$VLLM_VERSION" .
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
|
||||
```
|
||||
|
||||
This layers vLLM-Omni on top of the published `vllm/vllm-openai-xpu:<VLLM_VERSION>`
|
||||
base image, which Docker pulls automatically. The Dockerfile still defaults to
|
||||
v0.30.0; the command above explicitly selects v0.31.0 to match this source
|
||||
checkout. To target a different vLLM release, pass
|
||||
`--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
|
||||
base image (v0.30.0 by default), which Docker pulls automatically. To target a different vLLM release,
|
||||
pass `--build-arg VLLM_VERSION=<tag>`. If that tag has not been published yet, build
|
||||
the base from upstream's own Dockerfile first and point the build at it:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -71,14 +71,10 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
|
||||
|
||||
# --8<-- [start:installation-main]
|
||||
|
||||
You can also install vLLM-Omni from source if you want to use the latest features
|
||||
or bug fixes. This source checkout targets vLLM 0.31.0 and needs a matching
|
||||
vLLM-Ascend runtime. The v0.30.0 images above are for the published Omni v0.30.0
|
||||
release. Check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886)
|
||||
for aligned versions and NPU validation before selecting a development image.
|
||||
You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.30.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU.
|
||||
|
||||
```bash
|
||||
# Inside a container with matching vLLM and vLLM-Ascend, install Omni from source
|
||||
# Inside the vLLM-Ascend v0.30.0 container, install vLLM-Omni from main
|
||||
cd /vllm-workspace
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
|
||||
@@ -19,20 +19,17 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.31.0 --torch-backend=auto \
|
||||
--extra-index-url https://wheels.vllm.ai/db9527a46873454610df6dbedf79a36d6bf1a7f6
|
||||
uv pip install vllm==0.30.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.31.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.31.0/rocm723
|
||||
uv pip install vllm==0.30.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.30.0/rocm723
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
These commands install vLLM 0.31.0 for this source checkout. The CUDA index is
|
||||
pinned to the final release commit. For published vLLM-Omni wheels and additional
|
||||
installation methods, see the [installation guide](installation/README.md).
|
||||
For additional installation methods — please see the [installation guide](installation/README.md).
|
||||
|
||||
!!! note
|
||||
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
||||
|
||||
@@ -130,7 +130,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
### Notes
|
||||
|
||||
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
|
||||
- H200 automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when the CUDA MPS control tool is available; other devices and environments use `cosyvoice3.yaml`. See the [serving defaults and tradeoffs](../online_serving/text_to_speech.md#cosyvoice3). Pass `--deploy-config cosyvoice3.yaml` for the previous profile or `--deploy-config <path>` to override. The example uses the resolved profile's sampling parameters.
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
|
||||
- `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
|
||||
|
||||
---
|
||||
|
||||
@@ -155,22 +155,16 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
|
||||
./cosyvoice3/run_server.sh
|
||||
```
|
||||
|
||||
Streaming is on by default. On CUDA Hopper GPUs with at least 140 GiB of device memory (H200), serving automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when `nvidia-cuda-mps-control` is available on `PATH`: AR32 / codec32, standard sampling, packed Flow, cached ISTFT, HiFT decode graphs, and local CUDA MPS. No `--hf-overrides` is needed. Other devices and environments without the MPS control tool use `cosyvoice3.yaml` with RAS sampling and AR8 / codec8.
|
||||
Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
|
||||
|
||||
The optimized profile is not a universal hardware configuration. Its packed kernels are Hopper-only; the talker's Model Runner V2 is unsupported on NPU/XPU, and CUDA MPS requires a local NVIDIA GPU and its control tool. Larger batches, 25-frame chunks, and graph allocations also require per-device validation on smaller or slower GPUs. Automatic selection retains the generic profile outside the validated device class; an explicit deploy configuration takes precedence.
|
||||
|
||||
The H200 default targets streaming throughput at C32/C64. It changes sampling from RAS to standard (temperature 0.7, top-p 0.8, top-k 20, repetition penalty 1.21); output lengths and quality can differ. HiFT captures a recurring exact shape on its third use, which adds latency to that request, and its graph pool needs memory headroom beyond the engine reservations. This profile is not a claim of optimal first-audio latency at low concurrency. Pass `--deploy-config cosyvoice3.yaml` to retain the previous configuration, or `--deploy-config <path>` to select another profile.
|
||||
|
||||
Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path; use `--deploy-config cosyvoice3.yaml` alongside it when selecting that path on H200.
|
||||
|
||||
Cross-request Stage-1 flow batching is enabled in the H200 profile. For the generic configuration, enable it in the server environment when concurrent requests should share a flow-estimator call:
|
||||
Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call:
|
||||
|
||||
```bash
|
||||
export COSYVOICE3_BATCH_FLOW=1
|
||||
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
|
||||
```
|
||||
|
||||
Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Use `--deploy-config cosyvoice3.yaml` and leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
|
||||
Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
|
||||
|
||||
### CLI client
|
||||
|
||||
@@ -200,7 +194,7 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
|
||||
### Notes
|
||||
|
||||
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
|
||||
- Deploy config auto-loads the device-specific CosyVoice3 default described above. Pass `--deploy-config <path>` to override.
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
|
||||
- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
|
||||
|
||||
---
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import argparse
|
||||
import copy
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
@@ -11,11 +10,13 @@ from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from vllm import SamplingParams
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
|
||||
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
|
||||
from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
|
||||
|
||||
# Upstream zero-shot reference clip
|
||||
ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
|
||||
@@ -58,7 +59,8 @@ def run_e2e():
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path. If unset, selects the bundled CosyVoice3 profile for the device.",
|
||||
help="Override the deploy config path. If unset, auto-loads "
|
||||
"vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
|
||||
)
|
||||
parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
|
||||
parser.add_argument(
|
||||
@@ -103,6 +105,8 @@ def run_e2e():
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
|
||||
sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
|
||||
|
||||
print("Model initialized. Preparing inputs...")
|
||||
ref_audio_path = args.ref_audio or _default_ref_audio()
|
||||
if not os.path.exists(ref_audio_path):
|
||||
@@ -134,7 +138,7 @@ def run_e2e():
|
||||
|
||||
print(f"Generating for prompt: {args.text}")
|
||||
|
||||
config = omni.engine.stage_vllm_configs[0].model_config.hf_config
|
||||
config = CosyVoice3Config()
|
||||
tokenizer = get_qwen_tokenizer(
|
||||
token_path=args.tokenizer,
|
||||
skip_special_tokens=config.skip_special_tokens,
|
||||
@@ -145,10 +149,29 @@ def run_e2e():
|
||||
min_len = int(base_len * config.min_token_text_ratio)
|
||||
max_len = int(base_len * config.max_token_text_ratio)
|
||||
|
||||
# Keep the deploy profile's sampling policy, penalties and required stops.
|
||||
sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
|
||||
sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
|
||||
sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
|
||||
# Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
|
||||
gpt_sampling = SamplingParams(
|
||||
temperature=1.0,
|
||||
top_p=sampling_cfg["top_p"],
|
||||
top_k=sampling_cfg["top_k"],
|
||||
repetition_penalty=2.0,
|
||||
min_tokens=min_len,
|
||||
max_tokens=max_len,
|
||||
stop_token_ids=[sampling_cfg["eos_token_id"]],
|
||||
# allowed_token_ids=list(range(6561+3)),
|
||||
detokenize=False,
|
||||
)
|
||||
# Not used
|
||||
s2mel_sampling = SamplingParams(
|
||||
temperature=1.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
repetition_penalty=2.0,
|
||||
max_tokens=256,
|
||||
detokenize=False,
|
||||
)
|
||||
|
||||
sampling_params_list = [gpt_sampling, s2mel_sampling]
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
if profiler_enabled:
|
||||
|
||||
@@ -1,15 +1,11 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# Launch vLLM-Omni server for CosyVoice3 TTS
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh
|
||||
# CUDA_VISIBLE_DEVICES=0 ./run_server.sh
|
||||
#
|
||||
# Streaming (async-chunk) is on by default. H200 selects the optimized profile;
|
||||
# other devices use vllm_omni/deploy/cosyvoice3.yaml.
|
||||
# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
|
||||
# Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
|
||||
|
||||
set -e
|
||||
|
||||
@@ -35,10 +35,6 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses
|
||||
audio-reference generation. Reference transcripts must match the reference
|
||||
audio; they are not style instructions.
|
||||
|
||||
For the optional CUDA MRV2 runner, see the
|
||||
[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2).
|
||||
The default Local deployment continues to use V1.
|
||||
|
||||
## References
|
||||
|
||||
- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
|
||||
@@ -156,296 +152,3 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
- No `ref_audio` required or accepted for MOSS-SoundEffect.
|
||||
- Input field maps to the `ambient_sound` parameter in the upstream processor.
|
||||
- Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
|
||||
|
||||
## Local 1.5 MRV2 and slot attention
|
||||
|
||||
`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with
|
||||
an explicit deploy profile. The default Local profile continues to use V1.
|
||||
Both profiles below preserve 1-frame initial and 15-frame steady codec chunks
|
||||
and the model's sampling defaults.
|
||||
|
||||
```bash
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
|
||||
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml
|
||||
```
|
||||
|
||||
This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned
|
||||
CUDA graphs without compiling the codec with Inductor.
|
||||
|
||||
For sustained high concurrency on a large-memory CUDA GPU, use:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
|
||||
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \
|
||||
--stage-init-timeout 1200 --init-timeout 1500
|
||||
```
|
||||
|
||||
The high-concurrency profile places both stages on one GPU, sets each stage's
|
||||
capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses
|
||||
mixed FULL CUDA graphs with both its token budget and capture limit set to
|
||||
512, so new prompt chunks remain within graph coverage. It also enables
|
||||
`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and
|
||||
continuation control stay in GPU request slots, gathered through MRV2's
|
||||
batch mapping. Prefill conditioning and explicit per-request sampling seeds
|
||||
retain their original behavior. Published audio codes own their storage, so
|
||||
slot reuse cannot overwrite an in-flight output. Set the override to `false`
|
||||
to compare with the generic Omni model state. The profile selects
|
||||
`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels
|
||||
disabled, and codec CUDA graph buckets through 128. The configuration was
|
||||
validated on one H200. The earlier capacity-256 variant used about 110 GiB of
|
||||
sampled peak GPU memory, including loading and warmup. Reduce the stage capacities,
|
||||
graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec
|
||||
compilation can take several minutes; subsequent starts can reuse the AOT
|
||||
cache. Graph capture is still performed at startup.
|
||||
|
||||
The high-concurrency profile also coalesces ready codec streams with
|
||||
`connectors.shm.extra.generation_min_batch_size: 32` and
|
||||
`generation_max_wait_ms: 12`. The target controls when to stop waiting; it
|
||||
neither caps execution at 32 nor guarantees every codec group has 32 rows.
|
||||
First and final chunks share the bounded window at high concurrency. When
|
||||
fewer than 32 requests remain, dispatch is immediate. Pending outputs are
|
||||
retired before waiting, with the original deadline preserved, so waiting does
|
||||
not prevent the previous batch from releasing its in-flight state. Cancellation
|
||||
and input notifications keep their ordinary scheduler bookkeeping. This option
|
||||
requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to
|
||||
`0` to disable coalescing; other profiles keep the immediate-dispatch default.
|
||||
|
||||
The high-concurrency profile enables two additional Local model-state
|
||||
optimizations under the CUDA stage-0 `hf_overrides`, alongside
|
||||
`mrv2_gpu_slot_state: true`:
|
||||
|
||||
- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines
|
||||
CPU reference-code slices and their destination positions into one pinned
|
||||
upload. Host staging is reused only after the upload event completes;
|
||||
GPU-resident references keep the original path.
|
||||
- `mrv2_direct_tokens: true` returns the already determined text token after
|
||||
the Local audio/stop decision. It avoids constructing and sampling the
|
||||
full text vocabulary. Audio-code and binary-stop sampling are unchanged.
|
||||
Requests requiring distribution metadata or token constraints, and
|
||||
unsupported execution modes, retain the normal sampler.
|
||||
|
||||
Other profiles leave both flags disabled. Set both flags explicitly when
|
||||
comparing their combined effect. Deployment
|
||||
inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also
|
||||
retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as
|
||||
throughput: faster prompt preparation can change competition between the
|
||||
Talker and codec sharing the GPU.
|
||||
|
||||
For the single-H200 capacity-128 throughput baseline, keep both switches on.
|
||||
In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch
|
||||
rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct
|
||||
tokens alone had no reliable gain. The combined path had 685 ms mean first
|
||||
audio versus 679 ms across six control rounds, with lower mean completion time.
|
||||
A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus
|
||||
3.78% for the control. This sample does not establish speaker similarity or
|
||||
exclude small quality changes. At concurrency 256, the combined path had
|
||||
raised mean first-audio latency by about 0.33 s, so other profiles leave the
|
||||
switches off until their latency and quality are checked.
|
||||
|
||||
The optional connector setting `generation_coalescing_policy: idle_wait`
|
||||
waits for an inbox notification when no stream is runnable, all receivers
|
||||
are parked and registered, and no output needs retirement. The first ready
|
||||
arrival starts a fresh coalescing window, preserving the batch budget.
|
||||
Control messages, including cancellations, can wait up to two windows
|
||||
instead of one (24 ms at the 12 ms setting). The default policy remains
|
||||
`fixed`; neither policy changes the global orchestration default.
|
||||
|
||||
The codec backends differ in state access:
|
||||
|
||||
- `sdpa` uses PyTorch attention after gathering and updating the ring cache.
|
||||
- `triton` replaces the attention calculation, retaining the ring-cache
|
||||
gather/copy and explicit mask construction.
|
||||
- `triton_slot` writes surviving K/V directly into request slots, attends
|
||||
directly to the ring, and advances the active slot offsets in three
|
||||
ordered kernels. Graph-padding rows do not advance persistent state.
|
||||
|
||||
The slot path preserves the existing chunk-complete ring semantics, including
|
||||
retaining the final cache-capacity tokens when a chunk exceeds the ring.
|
||||
It does not introduce another cache owner or change request-slot lifetime.
|
||||
Zero-length padding rows skip the attention computation and emit zeros.
|
||||
|
||||
The high-concurrency profile includes intermediate codec batch buckets 6, 12
|
||||
and 24. To compare with power-of-two buckets, change the stage-1
|
||||
`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`.
|
||||
Keep the maximum bucket equal to the state capacity to retain terminal-tail
|
||||
coalescing. Smaller buckets reduce padding work but require more graphs;
|
||||
measure complete serving runs before selecting them for a deployment.
|
||||
|
||||
For a codec output-path comparison, set
|
||||
`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous
|
||||
codec output path; `1` selects GPU output snapshots on a private codec
|
||||
stream. Both configurations use MRV2 and keep streaming audio responses.
|
||||
The option changes input ownership, metadata staging and snapshot handling
|
||||
as well as transfer scheduling, so its timing difference is not just D2H
|
||||
copy time.
|
||||
|
||||
Those chunk-complete semantics truncate the causal window: every ring holds
|
||||
exactly `context` entries and a whole chunk is written before attention runs,
|
||||
so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the
|
||||
tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/
|
||||
400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so
|
||||
its first 80 tokens attend to nothing and no chunk sees the previous one there.
|
||||
Padding a terminal tail to 15 frames therefore decodes it differently from an
|
||||
exact-length decode; the difference is deterministic (identical in fp32 and
|
||||
fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each
|
||||
ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming
|
||||
then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails
|
||||
equal exact tails. It increases state/graph memory and capture time; historical C256 experiments
|
||||
observed about a 2% throughput cost. It is off by default and its perceptual
|
||||
quality benefit has not been established. Terminal tails are padded
|
||||
into the regular 15-frame graph bucket in either mode.
|
||||
The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths
|
||||
retain their existing attention implementation. Event-driven orchestration
|
||||
remains independently selectable and its default is unchanged.
|
||||
|
||||
### Optional progressive chunks
|
||||
|
||||
Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks
|
||||
before steady decoding, for example:
|
||||
|
||||
```yaml
|
||||
connectors:
|
||||
shm:
|
||||
extra:
|
||||
codec_chunk_ramp: [1, 4, 15]
|
||||
```
|
||||
|
||||
The ramp's first entry overrides `initial_codec_chunk_frames`. After the last
|
||||
entry, chunks use `codec_chunk_frames`. Each request advances independently;
|
||||
the final partial chunk is flushed and an empty terminal carries no codec
|
||||
tokens. The codec captures every ramp length, sizes the maximum execution
|
||||
step for the largest entry, and uses the ramp's first entry for its dedicated
|
||||
first-chunk graph. Additional shapes increase startup and graph memory.
|
||||
|
||||
Ramps are opt-in. They can reduce early playback underrun while increasing
|
||||
codec calls, first-audio latency or total generation time. The measured C128
|
||||
throughput presets retain 1→15; their published figures do not establish an
|
||||
additional benefit from enabling ramps. The processor and graph-shape tests
|
||||
run without GPU or model weights:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \
|
||||
tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \
|
||||
-m 'core_model and cpu' --run-level core_model
|
||||
```
|
||||
|
||||
### Low-latency and reference-encoding options
|
||||
|
||||
`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and
|
||||
uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP.
|
||||
Its codec first-chunk fast path uses dedicated graphs and a stream handoff
|
||||
before the regular decoder reuses the request slot. A gate limits contention
|
||||
with regular codec work. Slot waits are bounded at 30 seconds. A failed or
|
||||
stalled decode raises an error rather than replaying a partially advanced
|
||||
slot; a failed output enqueue retains owned PCM for regular delivery.
|
||||
Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch
|
||||
policy; it does not change the client concurrency or stage capacities.
|
||||
|
||||
Reference encoding runs in the API layer, independently of MRV2. Enable
|
||||
reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default
|
||||
eager encoder is retained when the option is absent. The graph path uses the
|
||||
loaded tokenizer's encoder/quantizer, length buckets and optional compilation
|
||||
(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed
|
||||
attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention.
|
||||
Compilation and attention changes need not produce bit-identical codes.
|
||||
|
||||
For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared
|
||||
reference-code storage and, by default, a single encoder host with four
|
||||
workers. Use a dedicated directory per service, checkpoint and encoding
|
||||
configuration. Workers have their own graph resources; all graph captures
|
||||
complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains
|
||||
separate encoders while sharing codes. More workers/graphs consume memory;
|
||||
they do not imply more GPUs. Example after selecting an available GPU and
|
||||
configuring any operator-managed MPS daemon:
|
||||
|
||||
```bash
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
|
||||
VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \
|
||||
VLLM_OMNI_MOSS_REF_GRAPHS=1 \
|
||||
VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \
|
||||
VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \
|
||||
VLLM_OMNI_MOSS_REF_INFLIGHT=4 \
|
||||
VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \
|
||||
VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \
|
||||
--api-server-count 4 \
|
||||
--deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \
|
||||
--stage-init-timeout 1200 --init-timeout 1500
|
||||
```
|
||||
|
||||
This example is a configurable serving profile, not the complete launch
|
||||
configuration of a historical benchmark. In particular, the low-latency YAML
|
||||
uses synchronous stage-1 scheduling; the retained source46 experiment used
|
||||
asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder
|
||||
environment settings. Do not infer a throughput result from the YAML alone.
|
||||
|
||||
### Cold versus hot reference measurements
|
||||
|
||||
A cold reference misses the encoded-reference cache: the request pays for
|
||||
reference parsing/preparation and GPU encoding before speech generation.
|
||||
A hot reference reuses reference codes but still synthesizes new speech;
|
||||
generated audio is not cached. Pre-registering a reference moves encoding
|
||||
into registration and can make subsequent synthesis hot, but registration
|
||||
latency is still part of the first-use cost.
|
||||
|
||||
Model startup/compilation is separate from reference coldness. EN1088 contains
|
||||
repeated references, so its first pass is not uniformly cold. For an all-unique
|
||||
cold pass, use its 666 first-occurrence references, keep order, warm execution
|
||||
with disjoint references, then repeat the same 666 requests. A second pass
|
||||
through independent API-local caches is not guaranteed hot; verify cache
|
||||
coverage before comparing it with a shared-cache result. Report requests/s
|
||||
alongside audio-s/s and retain failed, empty, long and near-silent outputs.
|
||||
|
||||
### Reproduce the serving benchmark
|
||||
|
||||
Use the complete Seed-TTS English test set, including reference audio and text.
|
||||
Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088
|
||||
entries. Run the native benchmark after the server is ready:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000
|
||||
export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2
|
||||
export SEED_TTS_WER_EVAL=0
|
||||
|
||||
for concurrency in 128; do
|
||||
for phase in warm r1 r2; do
|
||||
vllm bench serve --omni \
|
||||
--model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
|
||||
--backend openai-audio-speech --endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \
|
||||
--seed-tts-locale en --disable-shuffle --num-prompts 1088 \
|
||||
--num-warmups 0 --ready-check-timeout-sec 0 \
|
||||
--output-len 256 --max-concurrency "$concurrency" \
|
||||
--request-rate inf --seed 42 \
|
||||
--extra-body '{"task_type":"Base","max_new_tokens":256}' \
|
||||
--save-result --save-detailed --result-dir results/moss-local \
|
||||
--result-filename "c${concurrency}-${phase}.json"
|
||||
done
|
||||
done
|
||||
```
|
||||
|
||||
Discard the complete `warm` pass and combine measured runs as total generated
|
||||
audio seconds divided by total benchmark duration. Require 1088 successes,
|
||||
zero failures and 1088 nonempty-audio metric samples in each run. The output
|
||||
cap matches the benchmark protocol; success alone does not establish speech
|
||||
quality or that every sentence ended before the cap. Dataset seed 42 does
|
||||
not fix an independent sampling seed for every request.
|
||||
|
||||
For an attention-only comparison, copy the high-concurrency YAML beside the
|
||||
original to preserve relative `base_config` resolution. Change its
|
||||
codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`,
|
||||
and compare `codec_attention_backend: triton` against `triton_slot`. Keep all
|
||||
other settings, warmup and client concurrency identical. Restart the server
|
||||
between configurations and use separate result directories. Comparing the
|
||||
compiled slot profile to uncompiled Triton includes both changes.
|
||||
|
||||
Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages,
|
||||
and no model weights:
|
||||
|
||||
```bash
|
||||
python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \
|
||||
tests/model_executor/models/moss_tts/test_streaming_attention.py \
|
||||
-m 'core_model and cuda' --run-level=core_model
|
||||
```
|
||||
@@ -2,7 +2,7 @@
|
||||
# kernels==0.16.1 requires the LayerRepository version=/revision= integration
|
||||
# introduced in transformers 5.13. Keep the tested upper bound: 5.15 has
|
||||
# a reported model-construction regression; later versions need validation.
|
||||
# This range also satisfies vLLM 0.31.0's transformers>=5.10.4,<5.18.0 requirement.
|
||||
# This range also satisfies vLLM 0.30.0's transformers>=5.10.4 requirement.
|
||||
# MOSS-TTS-Nano's list/set weight-loading compatibility is handled by
|
||||
# transformers_keys_to_ignore_compat in model_executor/models/utils.py.
|
||||
transformers >= 5.13.0, < 5.15
|
||||
|
||||
@@ -10,8 +10,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2"
|
||||
AMD_BUILD_SCRIPT = REPO_ROOT / ".buildkite/amd/scripts/build-ci-image.sh"
|
||||
CUDA_RELEASE_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.cuda"
|
||||
CUDA_CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
|
||||
CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci"
|
||||
ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm"
|
||||
ROCM_DOCKERIGNORE = REPO_ROOT / "docker/Dockerfile.rocm.dockerignore"
|
||||
|
||||
@@ -31,18 +30,14 @@ def _line_index(lines: list[str], prefix: str) -> int:
|
||||
return matches[0]
|
||||
|
||||
|
||||
def test_rocm_base_tracks_cuda_vllm_release() -> None:
|
||||
cuda_base = _docker_arg(CUDA_RELEASE_DOCKERFILE, "BASE_IMAGE")
|
||||
cuda_image, _, cuda_release = cuda_base.rpartition(":")
|
||||
def test_rocm_base_tracks_ci_vllm_release() -> None:
|
||||
ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
|
||||
rocm_base = _docker_arg(ROCM_DOCKERFILE, "BASE_IMAGE")
|
||||
|
||||
image_ref, separator, image_tag = rocm_base.rpartition(":")
|
||||
assert separator, f"expected a tagged ROCm base image, got {rocm_base}"
|
||||
assert cuda_image == "vllm/vllm-openai"
|
||||
assert image_ref == "vllm/vllm-openai-rocm"
|
||||
assert image_tag == cuda_release == f"v{_docker_arg(CUDA_CI_DOCKERFILE, 'VLLM_VERSION')}"
|
||||
assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_IMAGE") == cuda_image
|
||||
assert _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_BASE_TAG") == "v${VLLM_VERSION}"
|
||||
assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref
|
||||
assert image_tag == ci_release
|
||||
|
||||
|
||||
def test_rocm_defaults_to_prebuilt_base_image() -> None:
|
||||
@@ -108,11 +103,11 @@ def test_rocm_build_context_excludes_git_history() -> None:
|
||||
assert ".git" in patterns
|
||||
|
||||
|
||||
def test_rocm_source_ref_tracks_cuda_vllm_release() -> None:
|
||||
cuda_release = _docker_arg(CUDA_CI_DOCKERFILE, "VLLM_VERSION")
|
||||
def test_rocm_source_ref_tracks_ci_vllm_release() -> None:
|
||||
ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG")
|
||||
rocm_source_ref = _docker_arg(ROCM_DOCKERFILE, "VLLM_VERSION_OR_COMMIT_HASH")
|
||||
|
||||
assert rocm_source_ref == f"v{cuda_release}"
|
||||
assert rocm_source_ref == ci_release
|
||||
|
||||
|
||||
def test_rocm_dockerfile_contains_vllm_api_canary() -> None:
|
||||
|
||||
@@ -172,11 +172,7 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest
|
||||
env_names = {item["name"] for item in container["env"]}
|
||||
assert "VLLM_CI_HF_TOKEN" in env_names
|
||||
assert "HF_TOKEN" not in env_names
|
||||
assert step["commands"] == [
|
||||
CUDA_HF_TOKEN_EXPORT,
|
||||
"python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
|
||||
"pytest -sv tests/example",
|
||||
]
|
||||
assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"]
|
||||
|
||||
|
||||
def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
|
||||
@@ -243,10 +239,6 @@ def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None:
|
||||
assert "VLLM_CI_HF_TOKEN" in env_names
|
||||
assert "HF_TOKEN" not in env_names
|
||||
assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT
|
||||
assert step["commands"][1:] == [
|
||||
"python3 .buildkite/cuda/scripts/check_vllm_runtime.py",
|
||||
"pytest -sv tests/example",
|
||||
]
|
||||
|
||||
|
||||
def _gpu_limit(step: dict) -> int:
|
||||
|
||||
@@ -571,56 +571,6 @@ class TestPipelineDiscovery:
|
||||
assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
|
||||
assert "qwen3_tts" in OMNI_PIPELINES
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cuda,major,memory_gib,mps,optimized",
|
||||
[
|
||||
(False, None, 0, True, False),
|
||||
(True, None, 141, True, False),
|
||||
(True, 8, 141, True, False),
|
||||
(True, 9, 80, True, False),
|
||||
(True, 10, 180, True, False),
|
||||
(True, 9, 141, False, False),
|
||||
(True, 9, 141, True, True),
|
||||
],
|
||||
)
|
||||
def test_cosyvoice3_device_default_and_explicit_overrides(
|
||||
self, monkeypatch, cuda, major, memory_gib, mps, optimized
|
||||
):
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
|
||||
monkeypatch.setattr(current_omni_platform, "device_name", "cuda" if cuda else "cpu")
|
||||
capability = DeviceCapability(major, 0) if major is not None else None
|
||||
monkeypatch.setattr(current_platform, "get_device_capability", lambda: capability)
|
||||
monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: memory_gib * 1024**3)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.model_executor.models.cosyvoice3.pipeline.shutil.which",
|
||||
lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps else None,
|
||||
)
|
||||
pipeline = resolve_pipeline_config("cosyvoice3")
|
||||
assert pipeline is not None
|
||||
config = VllmOmniConfig.from_pipeline_config(pipeline)
|
||||
talker, codec = config.stage_configs
|
||||
assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == (32 if optimized else 8)
|
||||
assert talker.runtime_config.cuda_mps == codec.runtime_config.cuda_mps == optimized
|
||||
assert talker.model_config.hf_overrides == ({"cosyvoice3_sampling_mode": "standard"} if optimized else None)
|
||||
assert pipeline.stages[0].sampling_constraints["stop_token_ids"] == list(range(6561, 6761))
|
||||
if optimized:
|
||||
assert codec.runtime_config.env["COSYVOICE3_CACHED_ISTFT"] == "1"
|
||||
assert codec.runtime_config.env["COSYVOICE3_HIFT_GRAPH"] == "1"
|
||||
override = VllmOmniConfig.from_pipeline_config(
|
||||
pipeline, cli_overrides={"hf_overrides": {"cosyvoice3_sampling_mode": "ras", "seed": 7}}
|
||||
)
|
||||
assert override.stage_by_id(0).model_config.hf_overrides == {"cosyvoice3_sampling_mode": "ras", "seed": 7}
|
||||
explicit = VllmOmniConfig.from_pipeline_config(
|
||||
pipeline, deploy_config_path=get_deploy_config_path("cosyvoice3.yaml")
|
||||
)
|
||||
assert explicit.stage_by_id(0).scheduler_config.max_num_seqs == 8
|
||||
assert explicit.stage_by_id(0).model_config.hf_overrides is None
|
||||
|
||||
def test_registry_resolver_qwen3_omni_all_stages(self):
|
||||
"""Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
|
||||
pipeline = resolve_pipeline_config(
|
||||
@@ -2978,17 +2928,9 @@ class TestPlatformOverrides:
|
||||
|
||||
assert deploy.stages[0].engine_extras["block_size"] == 128
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename,pipeline_key",
|
||||
[
|
||||
("qwen3_tts_mrv2.yaml", "qwen3_tts"),
|
||||
("moss_tts_local_mrv2.yaml", "moss_tts_local"),
|
||||
("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"),
|
||||
("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
|
||||
def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key):
|
||||
@pytest.mark.parametrize("platform", ["cuda", "npu"])
|
||||
def test_recommended_native_runner_platform_defaults(self, platform):
|
||||
filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
|
||||
deploy = _apply_platform_overrides(deploy, platform=platform)
|
||||
expect_v2 = platform == "cuda"
|
||||
@@ -3004,60 +2946,6 @@ class TestPlatformOverrides:
|
||||
# v2 only engages the native plane on stages declaring support.
|
||||
assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
|
||||
|
||||
def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self):
|
||||
pipeline = resolve_pipeline_config("moss_tts_local")
|
||||
base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name))
|
||||
candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml"))
|
||||
base = load_deploy_config(base_path)
|
||||
candidate = load_deploy_config(candidate_path)
|
||||
assert base.model_runner == "v1"
|
||||
assert candidate.stages == base.stages
|
||||
assert candidate.connectors == base.connectors
|
||||
assert candidate.async_chunk == base.async_chunk
|
||||
cuda = _apply_platform_overrides(candidate, platform="cuda")
|
||||
stages = merge_pipeline_deploy(pipeline, cuda)
|
||||
assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512
|
||||
codec_args = stages[1].yaml_engine_args
|
||||
assert codec_args["max_num_seqs"] == 64
|
||||
assert codec_args["enforce_eager"] is False
|
||||
assert codec_args["compilation_config"] == {
|
||||
"mode": 0,
|
||||
"cudagraph_mode": "FULL",
|
||||
"cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64],
|
||||
"cudagraph_num_of_warmups": 1,
|
||||
}
|
||||
for platform in ("npu", "xpu", "rocm", "musa"):
|
||||
# Platform resolution mutates the deployment; each server loads
|
||||
# a fresh profile before applying its own platform overrides.
|
||||
fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform)
|
||||
original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
|
||||
assert fallback.stages == original.stages
|
||||
for variant in ("moss_tts_delay", "moss_tts_realtime"):
|
||||
assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages)
|
||||
|
||||
def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self):
|
||||
pipeline = resolve_pipeline_config("moss_tts_local")
|
||||
path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml"))
|
||||
deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda")
|
||||
talker, codec = merge_pipeline_deploy(pipeline, deploy)
|
||||
assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128
|
||||
assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True
|
||||
assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True
|
||||
assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True
|
||||
assert talker.yaml_engine_args["max_num_batched_tokens"] == 512
|
||||
assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3
|
||||
assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot"
|
||||
compilation = codec.yaml_engine_args["compilation_config"]
|
||||
assert compilation["mode"] == 3 and compilation["backend"] == "inductor"
|
||||
assert compilation["cudagraph_mode"] == "FULL"
|
||||
assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128]
|
||||
assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False}
|
||||
base_path = Path(get_deploy_config_path("moss_tts_local.yaml"))
|
||||
for platform in ("npu", "xpu", "rocm", "musa"):
|
||||
actual = _apply_platform_overrides(load_deploy_config(path), platform=platform)
|
||||
expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform)
|
||||
assert actual.stages == expected.stages
|
||||
|
||||
@pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
|
||||
def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
|
||||
from unittest.mock import Mock
|
||||
|
||||
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 32,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 95,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 77,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
|
||||
EnvironmentVariableCategory.INTERNAL: 7,
|
||||
}
|
||||
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
|
||||
)
|
||||
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 66,
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 48,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 17,
|
||||
|
||||
@@ -1,258 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Bounded readiness coalescing must not allocate or consume request state."""
|
||||
|
||||
import queue
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm.v1.core.sched.interface import PauseState
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
from vllm_omni.outputs import OmniConnectorOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch):
|
||||
events = [notification("0"), notification("1", "2", "3")]
|
||||
s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])])
|
||||
s._generation_coalescing_policy = "idle_wait"
|
||||
for request in s.requests.values():
|
||||
request.status = RequestStatus.WAITING_FOR_CHUNK
|
||||
assert s._drain_omni_connector_outputs() == events
|
||||
assert clock[0] == pytest.approx(10.006)
|
||||
assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"])
|
||||
def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending):
|
||||
s, clock = scheduler(monkeypatch)
|
||||
s._generation_coalescing_policy = "idle_wait"
|
||||
for request in s.requests.values():
|
||||
request.status = RequestStatus.WAITING_FOR_CHUNK
|
||||
if pending == "in_flight":
|
||||
s.requests["0"].num_in_flight_tokens = 1
|
||||
elif pending == "admission":
|
||||
s.requests["0"].status = RequestStatus.WAITING
|
||||
elif pending == "registration":
|
||||
s.input_coordinator.pending_chunk_registrations = [object()]
|
||||
elif pending == "terminal":
|
||||
s._pending_data_plane_terminal_req_ids = {"0"}
|
||||
else:
|
||||
s.requests.clear()
|
||||
assert s._drain_omni_connector_outputs() == []
|
||||
assert clock[0] == 10
|
||||
assert not s._omni_connector_output_inbox.waits
|
||||
|
||||
|
||||
def test_idle_wait_without_arrivals_is_bounded(monkeypatch):
|
||||
s, clock = scheduler(monkeypatch)
|
||||
s._generation_coalescing_policy = "idle_wait"
|
||||
for request in s.requests.values():
|
||||
request.status = RequestStatus.WAITING_FOR_CHUNK
|
||||
assert s._drain_omni_connector_outputs() == []
|
||||
assert clock[0] == pytest.approx(10.012)
|
||||
|
||||
|
||||
def notification(*ids, terminal=()):
|
||||
return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal))
|
||||
|
||||
|
||||
class Inbox:
|
||||
def __init__(self, clock, initial=(), delayed=()):
|
||||
self.clock = clock
|
||||
self.initial = list(initial)
|
||||
self.delayed = list(delayed)
|
||||
self.waits = []
|
||||
|
||||
def get_nowait(self):
|
||||
if self.initial:
|
||||
return self.initial.pop(0)
|
||||
raise queue.Empty
|
||||
|
||||
def get(self, timeout):
|
||||
self.waits.append(timeout)
|
||||
if not self.delayed or self.delayed[0][0] > timeout:
|
||||
self.clock[0] += timeout
|
||||
raise queue.Empty
|
||||
delay, value = self.delayed.pop(0)
|
||||
self.clock[0] += delay
|
||||
return value
|
||||
|
||||
|
||||
def scheduler(monkeypatch, initial=(), delayed=()):
|
||||
clock = [10.0]
|
||||
monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0])
|
||||
s = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
|
||||
s._generation_max_wait_s = 0.012
|
||||
s._generation_min_batch_size = 4
|
||||
s._pause_state = PauseState.UNPAUSED
|
||||
s.requests = {
|
||||
str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8)
|
||||
}
|
||||
s.running = list(s.requests.values())
|
||||
s.input_coordinator = SimpleNamespace(
|
||||
requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set()
|
||||
)
|
||||
s._latest_omni_connector_output = None
|
||||
s._omni_connector_output_inbox = Inbox(clock, initial, delayed)
|
||||
return s, clock
|
||||
|
||||
|
||||
def test_arrivals_form_target_batch_in_order(monkeypatch):
|
||||
events = [notification("0"), notification("1", "2"), notification("3")]
|
||||
s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])])
|
||||
assert s._drain_omni_connector_outputs() == events
|
||||
assert clock[0] == pytest.approx(10.005)
|
||||
assert not s.input_coordinator.requests_with_ready_chunks
|
||||
assert all(r.num_in_flight_tokens == 0 for r in s.requests.values())
|
||||
|
||||
|
||||
def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch):
|
||||
events = [notification("0", "gone", "1", "2"), notification("0")]
|
||||
s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])])
|
||||
s.requests["1"].is_finished = lambda: True
|
||||
s.requests["2"].is_finished = lambda: True
|
||||
assert s._drain_omni_connector_outputs() == events
|
||||
assert clock[0] == pytest.approx(10.012)
|
||||
assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"])
|
||||
def test_immediate_dispatch(monkeypatch, reason):
|
||||
event = notification("0")
|
||||
s, clock = scheduler(monkeypatch, [event])
|
||||
if reason == "enough":
|
||||
event.chunk_ready_req_ids.update(["1", "2", "3"])
|
||||
elif reason == "empty":
|
||||
event.chunk_ready_req_ids.clear()
|
||||
elif reason == "low":
|
||||
s.requests = {"0": s.requests["0"]}
|
||||
elif reason == "paused":
|
||||
s._pause_state = PauseState.PAUSED_ALL
|
||||
elif reason == "off":
|
||||
s._generation_max_wait_s = 0
|
||||
assert s._drain_omni_connector_outputs() == [event]
|
||||
assert not s._omni_connector_output_inbox.waits
|
||||
assert clock[0] == 10
|
||||
|
||||
|
||||
def test_terminal_arrival_keeps_original_deadline(monkeypatch):
|
||||
events = [notification("0"), notification("1", terminal=["1"])]
|
||||
s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])])
|
||||
assert s._drain_omni_connector_outputs() == events
|
||||
assert clock[0] == pytest.approx(10.012)
|
||||
|
||||
|
||||
def test_pending_ready_and_legacy_notification_preserved(monkeypatch):
|
||||
s, _ = scheduler(monkeypatch)
|
||||
s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"])
|
||||
last = notification("3")
|
||||
s._latest_omni_connector_output = last
|
||||
assert s._drain_omni_connector_outputs() == [last]
|
||||
assert s._latest_omni_connector_output is None
|
||||
assert not s._omni_connector_output_inbox.waits
|
||||
|
||||
|
||||
def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch):
|
||||
s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))])
|
||||
assert len(s._drain_omni_connector_outputs()) == 2
|
||||
assert clock[0] == pytest.approx(10.012)
|
||||
s._omni_connector_output_inbox.initial.append(notification("3"))
|
||||
assert len(s._drain_omni_connector_outputs()) == 2
|
||||
assert clock[0] == pytest.approx(10.024)
|
||||
|
||||
|
||||
def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker):
|
||||
event = notification("0", "cancelled")
|
||||
event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}}
|
||||
s, _ = scheduler(monkeypatch, [event])
|
||||
s._generation_max_wait_s = 0
|
||||
s.waiting = []
|
||||
s.kv_holding_waiting = []
|
||||
s.input_coordinator._async_chunk = True
|
||||
s.input_coordinator.update_request_metadata = mocker.Mock()
|
||||
s.input_coordinator.process_pending_chunks = mocker.Mock()
|
||||
s._consume_pending_connector_output("generation")
|
||||
s.input_coordinator.update_request_metadata.assert_called_once_with(
|
||||
s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation"
|
||||
)
|
||||
s.input_coordinator.process_pending_chunks.assert_called_once()
|
||||
waiting, running, ready, finished = s.input_coordinator.process_pending_chunks.call_args.args
|
||||
assert list(waiting) == []
|
||||
assert running is s.running
|
||||
assert ready == {"0"} and finished == set()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))])
|
||||
def test_invalid_config_rejected(monkeypatch, target, wait):
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler
|
||||
|
||||
def initialize(s, *args, **kwargs):
|
||||
s.max_num_running_reqs = 256
|
||||
s.vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
retains_state_across_chunks=True,
|
||||
stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}},
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(VLLMScheduler, "__init__", initialize)
|
||||
monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None)
|
||||
with pytest.raises(ValueError):
|
||||
OmniGenerationScheduler()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("terminal", [False, True])
|
||||
def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal):
|
||||
event = notification("0", terminal=["0"] if terminal else [])
|
||||
s, clock = scheduler(monkeypatch, [event])
|
||||
s.running = s.running[1:]
|
||||
assert s._drain_omni_connector_outputs() == [event]
|
||||
assert clock[0] == pytest.approx(10.012)
|
||||
|
||||
|
||||
def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch):
|
||||
events = [notification("0"), notification("1")]
|
||||
s, clock = scheduler(monkeypatch, events[:1])
|
||||
s.requests["7"].num_in_flight_tokens = 1
|
||||
assert s._drain_omni_connector_outputs() == events[:1]
|
||||
assert s._generation_defer_batch
|
||||
assert not s._omni_connector_output_inbox.waits
|
||||
s.input_coordinator.requests_with_ready_chunks.add("0")
|
||||
# Simulate EngineCore retiring the previous batch after five milliseconds.
|
||||
clock[0] += 0.005
|
||||
s.requests["7"].num_in_flight_tokens = 0
|
||||
s._omni_connector_output_inbox.initial.append(events[1])
|
||||
assert s._drain_omni_connector_outputs() == events[1:]
|
||||
assert not s._generation_defer_batch
|
||||
assert s._omni_connector_output_inbox.waits == pytest.approx([0.007])
|
||||
assert clock[0] == pytest.approx(10.012)
|
||||
assert s._generation_batch_deadline is None
|
||||
|
||||
|
||||
def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch):
|
||||
s, clock = scheduler(monkeypatch, [notification("0")])
|
||||
s.requests["7"].num_in_flight_tokens = 1
|
||||
s._generation_batch_deadline = clock[0] - 0.001
|
||||
assert len(s._drain_omni_connector_outputs()) == 1
|
||||
assert not s._generation_defer_batch
|
||||
assert not s._omni_connector_output_inbox.waits
|
||||
|
||||
|
||||
def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker):
|
||||
from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler
|
||||
|
||||
request = _chunk_request("pending")
|
||||
s = _make_generation_scheduler(request)
|
||||
s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True)
|
||||
allocate = mocker.spy(s.kv_cache_manager, "allocate_slots")
|
||||
output = s.schedule()
|
||||
assert output.total_num_scheduled_tokens == 0
|
||||
assert not output.num_scheduled_tokens
|
||||
allocate.assert_not_called()
|
||||
assert list(s.waiting) == [request]
|
||||
assert request.num_computed_tokens == request.num_in_flight_tokens == 0
|
||||
assert s.chunk_transfer_adapter.restore_called
|
||||
@@ -1,107 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Exercise generation constructor policy after actual deploy/platform resolution."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config.pipeline_registry import resolve_pipeline_config
|
||||
from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def construct_scheduler(monkeypatch, mocker):
|
||||
# Only upstream cache allocation and connector construction are omitted.
|
||||
# Keep the actual Omni I/O initialization, native decision and constructor.
|
||||
mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter")
|
||||
|
||||
def upstream_init(self, config):
|
||||
self.vllm_config = config
|
||||
self.max_num_running_reqs = config.scheduler_config.max_num_seqs
|
||||
|
||||
monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init)
|
||||
return OmniGenerationScheduler
|
||||
|
||||
|
||||
def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128):
|
||||
return SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
use_v2_model_runner=native,
|
||||
supports_native_mrv2_data_plane=True,
|
||||
retains_state_across_chunks=stateful,
|
||||
async_chunk=True,
|
||||
stage_id=1,
|
||||
stage_connector_config={"extra": extras or {}},
|
||||
),
|
||||
parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp),
|
||||
scheduler_config=SimpleNamespace(max_num_seqs=capacity),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"])
|
||||
@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"])
|
||||
def test_moss_profile_generation_constructor_after_platform_resolution(
|
||||
construct_scheduler, monkeypatch, mocker, profile, platform
|
||||
):
|
||||
# merge_pipeline_deploy owns platform resolution. Mock the detected device
|
||||
# instead of applying an override first and then resolving the worker's
|
||||
# actual platform a second time.
|
||||
monkeypatch.setattr(current_omni_platform, "device_name", platform)
|
||||
deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml"))
|
||||
pipeline = resolve_pipeline_config("moss_tts_local")
|
||||
codec = merge_pipeline_deploy(pipeline, deploy)[1]
|
||||
extras = deploy.connectors["shm"]["extra"]
|
||||
config = _config(
|
||||
native=codec.yaml_engine_args["use_v2_model_runner"],
|
||||
stateful=pipeline.stages[1].retains_state_across_chunks,
|
||||
capacity=codec.yaml_engine_args["max_num_seqs"],
|
||||
extras=extras,
|
||||
)
|
||||
warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning")
|
||||
scheduler = construct_scheduler(config)
|
||||
# The raw extras persist through fallback; the live constructor must cope.
|
||||
assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12
|
||||
if platform == "cuda":
|
||||
assert scheduler._native_data_plane and scheduler.input_coordinator is not None
|
||||
assert scheduler._generation_min_batch_size == 32
|
||||
assert scheduler._generation_max_wait_s == 0.012
|
||||
assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0)
|
||||
warning.assert_not_called()
|
||||
if profile == "low_latency":
|
||||
assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1
|
||||
else:
|
||||
assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None
|
||||
assert scheduler._generation_min_batch_size == 1
|
||||
assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0
|
||||
warning.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)])
|
||||
def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp):
|
||||
config = _config(
|
||||
stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12}
|
||||
)
|
||||
with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"):
|
||||
construct_scheduler(config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extras,message",
|
||||
[
|
||||
({"generation_min_batch_size": 0}, "generation_min_batch_size"),
|
||||
({"generation_min_batch_size": 129}, "generation_min_batch_size"),
|
||||
({"generation_max_wait_ms": -1}, "generation_max_wait_ms"),
|
||||
({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"),
|
||||
({"generation_max_regular_batch": -1}, "generation_max_regular_batch"),
|
||||
({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"),
|
||||
],
|
||||
)
|
||||
def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message):
|
||||
with pytest.raises(ValueError, match=message):
|
||||
construct_scheduler(_config(native=False, extras=extras))
|
||||
@@ -11,9 +11,7 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
from vllm.distributed.aux_output_connector.connector import AuxOutputSchedulerConnector
|
||||
from vllm.v1.core.sched.interface import PauseState
|
||||
from vllm.v1.core.sched.output import CachedRequestData
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
|
||||
@@ -55,14 +53,12 @@ class FakeAdapter:
|
||||
def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
|
||||
scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
|
||||
scheduler.max_num_scheduled_tokens = 8
|
||||
scheduler.max_num_active_reqs = 1
|
||||
scheduler.max_num_running_reqs = 1
|
||||
scheduler._pause_state = PauseState.UNPAUSED
|
||||
scheduler.running = []
|
||||
scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.waiting.add_request(waiting_request)
|
||||
scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.requests = {waiting_request.request_id: waiting_request}
|
||||
scheduler.policy = SchedulingPolicy.FCFS
|
||||
scheduler.chunk_transfer_adapter = FakeAdapter()
|
||||
@@ -90,7 +86,6 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
|
||||
get_manager_metadata=lambda: None,
|
||||
)
|
||||
scheduler.connector = None
|
||||
scheduler.aux_output_connector = None
|
||||
scheduler.ec_connector = None
|
||||
scheduler.prev_step_scheduled_req_ids = set()
|
||||
scheduler._pending_finish_reqs = []
|
||||
@@ -110,45 +105,6 @@ def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
|
||||
return scheduler
|
||||
|
||||
|
||||
def test_generation_scheduler_preserves_auxiliary_output_metadata() -> None:
|
||||
request = Request("aux-contract", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
|
||||
request.block_hashes.append(b"block-hash")
|
||||
scheduler = _make_generation_scheduler(request, use_v2_model_runner=True)
|
||||
scheduler.aux_output_connector = AuxOutputSchedulerConnector()
|
||||
|
||||
output = scheduler._build_generation_scheduler_output(
|
||||
new_reqs_data=[],
|
||||
cached_reqs_data=CachedRequestData.make_empty(),
|
||||
num_scheduled_tokens={request.request_id: 1},
|
||||
scheduled_spec_decode_tokens={},
|
||||
scheduled_encoder_inputs={},
|
||||
num_common_prefix_blocks=[0],
|
||||
)
|
||||
|
||||
metadata = output.aux_output_connector_metadata
|
||||
assert metadata is not None
|
||||
assert metadata.requests == {request.request_id: 0}
|
||||
assert list(metadata.block_hashes[request.request_id]) == request.block_hashes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("runner_slots", "active_limit"), [(1, 2), (2, 1)])
|
||||
def test_generation_admits_kv_holder_before_fresh_requests(runner_slots, active_limit) -> None:
|
||||
fresh = Request("fresh", [1, 2], SamplingParams(max_tokens=2), pooling_params=None)
|
||||
holding = Request("holding", [3, 4], SamplingParams(max_tokens=2), pooling_params=None)
|
||||
holding.num_computed_tokens = 1
|
||||
scheduler = _make_generation_scheduler(fresh)
|
||||
scheduler.max_num_running_reqs = runner_slots
|
||||
scheduler.max_num_active_reqs = active_limit
|
||||
scheduler.requests[holding.request_id] = holding
|
||||
scheduler.kv_holding_waiting.add_request(holding)
|
||||
|
||||
output = scheduler.schedule()
|
||||
|
||||
assert set(output.num_scheduled_tokens) == {holding.request_id}
|
||||
assert list(scheduler.waiting) == [fresh]
|
||||
assert list(scheduler.kv_holding_waiting) == []
|
||||
|
||||
|
||||
def _chunk_request(request_id, **kwargs):
|
||||
defaults = dict(
|
||||
request_id=request_id,
|
||||
@@ -167,13 +123,7 @@ def _chunk_request(request_id, **kwargs):
|
||||
record_event=lambda *args, **kwargs: None,
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return _HashableChunkRequest(**defaults)
|
||||
|
||||
|
||||
class _HashableChunkRequest(SimpleNamespace):
|
||||
# Deferred queues require identity hashing; SimpleNamespace types it as None.
|
||||
__hash__ = object.__hash__ # type: ignore[assignment]
|
||||
__eq__ = object.__eq__
|
||||
return SimpleNamespace(**defaults)
|
||||
|
||||
|
||||
def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None:
|
||||
@@ -188,8 +138,7 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
|
||||
)
|
||||
scheduler.running = []
|
||||
scheduler.waiting = create_request_queue(scheduler.policy)
|
||||
scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = create_request_queue(scheduler.policy)
|
||||
|
||||
completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None)
|
||||
completed.status = RequestStatus.RUNNING
|
||||
@@ -343,7 +292,6 @@ def _express_scheduler():
|
||||
scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True)
|
||||
scheduler.waiting.add_request(first)
|
||||
scheduler.requests[first.request_id] = first
|
||||
scheduler.max_num_active_reqs = 4
|
||||
scheduler.max_num_running_reqs = 4
|
||||
scheduler._native_data_plane = True
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
|
||||
@@ -9,7 +9,7 @@ handing control to upstream, because upstream ``Scheduler.schedule()`` raises
|
||||
admits in a finished state -- which kills the stage's engine core, not just
|
||||
the request.
|
||||
|
||||
Regression for the sweep missing ``kv_holding_waiting``, the third queue
|
||||
Regression for the sweep missing ``skipped_waiting``, the third queue
|
||||
upstream admits from. An aborted duplex session parked there was re-selected
|
||||
by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
|
||||
stage.
|
||||
@@ -49,30 +49,28 @@ class _StubRequest:
|
||||
return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
|
||||
|
||||
|
||||
def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), kv_holding_waiting=(), running=()):
|
||||
def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
|
||||
"""Bare scheduler with just the surface the sweep reads."""
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.waiting = create_request_queue(policy)
|
||||
for req in waiting:
|
||||
scheduler.waiting.add_request(req)
|
||||
scheduler.kv_holding_waiting = create_request_queue(policy)
|
||||
scheduler.deferred_waiting = set()
|
||||
for req in kv_holding_waiting:
|
||||
scheduler.kv_holding_waiting.add_request(req)
|
||||
scheduler.skipped_waiting = create_request_queue(policy)
|
||||
for req in skipped_waiting:
|
||||
scheduler.skipped_waiting.add_request(req)
|
||||
scheduler.running = list(running)
|
||||
return scheduler
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy", _POLICIES)
|
||||
def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
|
||||
def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
policy: SchedulingPolicy,
|
||||
) -> None:
|
||||
"""The public schedule path must remove aborted skipped requests first."""
|
||||
aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
|
||||
scheduler = _make_scheduler(policy, kv_holding_waiting=[aborted])
|
||||
scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.max_num_active_reqs = 8
|
||||
scheduler.max_num_running_reqs = 8
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.num_waiting_for_streaming_input = 0
|
||||
@@ -87,7 +85,7 @@ def test_schedule_sweeps_kv_holding_waiting_before_upstream_selection(
|
||||
sentinel = object()
|
||||
|
||||
def upstream_schedule(_self, _throttle_prefills: bool = False):
|
||||
assert list(scheduler.kv_holding_waiting) == []
|
||||
assert list(scheduler.skipped_waiting) == []
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
|
||||
|
||||
@@ -1,16 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
|
||||
from vllm_omni.engine.serialization import deserialize_additional_information
|
||||
|
||||
@@ -31,7 +25,7 @@ def _make_engine_request() -> EngineCoreRequest:
|
||||
)
|
||||
|
||||
|
||||
def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
|
||||
def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
|
||||
tagged = apply_omni_final_stage_metadata(
|
||||
_make_engine_request(),
|
||||
final_stage_id=0,
|
||||
@@ -43,187 +37,21 @@ def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]:
|
||||
request_id="req",
|
||||
additional_information=tagged.additional_information,
|
||||
)
|
||||
omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request)
|
||||
omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request)
|
||||
result = scheduler._request_omits_kv_transfer_to_next_stage(request)
|
||||
metadata = deserialize_additional_information(tagged.additional_information)
|
||||
return omits_kv, omits_chunk, metadata
|
||||
return result, metadata
|
||||
|
||||
|
||||
def test_stage_zero_request_omits_kv_transfer():
|
||||
omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False)
|
||||
omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
|
||||
|
||||
assert omits_kv
|
||||
assert omits_chunk
|
||||
assert omits_transfer
|
||||
assert "omni_force_kv_transfer" not in metadata
|
||||
|
||||
|
||||
def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
|
||||
omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True)
|
||||
omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
|
||||
|
||||
assert not omits_kv
|
||||
assert omits_chunk
|
||||
assert not omits_transfer
|
||||
assert metadata["omni_final_stage_id"] == 0
|
||||
assert metadata["omni_force_kv_transfer"] is True
|
||||
|
||||
|
||||
class _ChunkRequest(SimpleNamespace):
|
||||
def __hash__(self):
|
||||
return hash(self.request_id)
|
||||
|
||||
def __eq__(self, other):
|
||||
return isinstance(other, _ChunkRequest) and other.request_id == self.request_id
|
||||
|
||||
|
||||
def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False):
|
||||
tagged = apply_omni_final_stage_metadata(
|
||||
_make_engine_request(),
|
||||
final_stage_id=final_stage_id,
|
||||
force_kv_transfer=force_kv_transfer,
|
||||
)
|
||||
request = _ChunkRequest(
|
||||
request_id="req-shm-leak",
|
||||
additional_information=tagged.additional_information,
|
||||
num_in_flight_tokens=1,
|
||||
num_stale_output_tokens=0,
|
||||
sampling_params=SimpleNamespace(num_logprobs=None),
|
||||
has_encoder_inputs=False,
|
||||
pooling_params=None,
|
||||
status=RequestStatus.RUNNING,
|
||||
resumable=False,
|
||||
num_output_placeholders=0,
|
||||
spec_token_ids=[],
|
||||
_output_token_ids=[],
|
||||
num_computed_tokens=1,
|
||||
stop_reason=None,
|
||||
num_nans_in_logits=None,
|
||||
client_index=0,
|
||||
trace_headers=None,
|
||||
)
|
||||
request.is_finished = lambda: False
|
||||
request.get_finished_reason = lambda: None
|
||||
request.take_prefill_stats = lambda: None
|
||||
request.take_events = lambda: None
|
||||
return request
|
||||
|
||||
|
||||
def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None):
|
||||
adapter = mocker.MagicMock()
|
||||
adapter._confirmed_num_computed_tokens.return_value = 1
|
||||
|
||||
sched = mocker.MagicMock()
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.perf_metrics = None
|
||||
sched.defer_block_free = False
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
sched._update_request_with_output.return_value = ([42], True)
|
||||
sched._process_kv_transfer_trigger.return_value = False
|
||||
sched._handle_stopped_request.return_value = True
|
||||
|
||||
def _free_request(req, delay_free_blocks=False):
|
||||
sched._omits_kv_transfer_cache.pop(req.request_id, None)
|
||||
return None, None
|
||||
|
||||
sched._free_request.side_effect = _free_request
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
sched.running = [request]
|
||||
sched.waiting_for_transfer_free = set()
|
||||
sched.transfer_triggered_requests = set()
|
||||
sched.active_kv_transfers = set()
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
bind_omits_transfer_helpers(sched)
|
||||
|
||||
scheduler_output = mocker.MagicMock()
|
||||
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
|
||||
scheduler_output.total_num_scheduled_tokens = 1
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = mocker.MagicMock()
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {request.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
model_runner_output.inter_stage_outputs = inter_stage_outputs
|
||||
|
||||
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
return adapter
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("final_stage_id", "force_kv_transfer", "expect_save"),
|
||||
[
|
||||
pytest.param(0, False, False, id="stage-zero-final"),
|
||||
pytest.param(1, False, True, id="downstream-final"),
|
||||
],
|
||||
)
|
||||
def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save):
|
||||
request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer)
|
||||
adapter = _run_finished_save_step(mocker, request)
|
||||
|
||||
assert adapter.save_async.called is expect_save
|
||||
|
||||
|
||||
def test_stage_zero_final_skips_inter_stage_output(mocker):
|
||||
warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning")
|
||||
request = _make_chunk_request(0)
|
||||
adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}])
|
||||
|
||||
adapter.save_async.assert_not_called()
|
||||
assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list)
|
||||
|
||||
|
||||
def test_streaming_session_update_invalidates_omits_kv_transfer_cache():
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler._omits_kv_transfer_cache = {}
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.kv_holding_waiting = []
|
||||
scheduler.log_stats = False
|
||||
tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
|
||||
|
||||
class Session(SimpleNamespace):
|
||||
__hash__ = object.__hash__
|
||||
|
||||
session = Session(
|
||||
request_id="req",
|
||||
additional_information=tagged.additional_information,
|
||||
status=RequestStatus.WAITING,
|
||||
)
|
||||
assert scheduler._request_omits_chunk_transfer_to_next_stage(session)
|
||||
assert "req" in scheduler._omits_kv_transfer_cache
|
||||
update = SimpleNamespace(
|
||||
additional_information={"omni_final_stage_id": 1},
|
||||
model_intermediate_buffer=None,
|
||||
arrival_time=0.0,
|
||||
sampling_params=None,
|
||||
)
|
||||
|
||||
OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update)
|
||||
|
||||
assert scheduler._omits_kv_transfer_cache == {}
|
||||
assert session.additional_information == {"omni_final_stage_id": 1}
|
||||
assert not scheduler._request_omits_chunk_transfer_to_next_stage(session)
|
||||
assert not scheduler._request_omits_kv_transfer_to_next_stage(session)
|
||||
|
||||
|
||||
def test_omits_cache_misses_when_payload_object_is_replaced():
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler._omits_kv_transfer_cache = {}
|
||||
tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0)
|
||||
request = SimpleNamespace(
|
||||
request_id="req",
|
||||
additional_information=tagged.additional_information,
|
||||
)
|
||||
assert scheduler._request_omits_chunk_transfer_to_next_stage(request)
|
||||
request.additional_information = {"omni_final_stage_id": 1}
|
||||
|
||||
assert not scheduler._request_omits_chunk_transfer_to_next_stage(request)
|
||||
@@ -11,14 +11,12 @@ from types import MethodType, SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.v1.engine import FinishReason
|
||||
from vllm.v1.outputs import LogprobsLists
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.outputs import OmniModelRunnerOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -134,13 +132,11 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
|
||||
scheduler = SimpleNamespace(
|
||||
perf_metrics=None,
|
||||
connector=None,
|
||||
aux_output_connector=None,
|
||||
chunk_transfer_adapter=None,
|
||||
requests={request.request_id: request for request in requests},
|
||||
running=list(requests),
|
||||
waiting=_RequestQueue(),
|
||||
kv_holding_waiting=_RequestQueue(),
|
||||
deferred_waiting=set(),
|
||||
skipped_waiting=_RequestQueue(),
|
||||
structured_output_manager=SimpleNamespace(accept_tokens=lambda _request, _token_ids: True),
|
||||
transfer_triggered_requests=set(),
|
||||
active_kv_transfers=set(),
|
||||
@@ -296,36 +292,6 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
|
||||
np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scored", [False, True])
|
||||
def test_fixed_token_scores_reach_engine_output(scored) -> None:
|
||||
request = _Request("req")
|
||||
request.sampling_params.num_logprobs = None
|
||||
scheduler = _make_scheduler_stub([request])
|
||||
_bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
|
||||
scheduler_output = SimpleNamespace(
|
||||
num_scheduled_tokens={"req": 3},
|
||||
scheduled_spec_decode_tokens={},
|
||||
num_invalid_spec_tokens=0,
|
||||
)
|
||||
scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
|
||||
model_runner_output = OmniModelRunnerOutput(
|
||||
req_ids=["req"],
|
||||
req_id_to_index={"req": 0},
|
||||
sampled_token_ids=[[7]],
|
||||
prompt_logprobs_dict={},
|
||||
prompt_token_id_logprobs_dict={"req": scores} if scored else {},
|
||||
)
|
||||
|
||||
outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
|
||||
(output,) = outputs[0].outputs
|
||||
|
||||
assert output.new_token_ids == [7]
|
||||
if scored:
|
||||
torch.testing.assert_close(output.prompt_token_id_logprobs, scores)
|
||||
else:
|
||||
assert output.prompt_token_id_logprobs is None
|
||||
|
||||
|
||||
def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
sampled_token_ids=[[]],
|
||||
|
||||
@@ -40,7 +40,6 @@ from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
|
||||
|
||||
# isort: on
|
||||
|
||||
@@ -73,8 +72,7 @@ def _replace_streaming_session(session: Request) -> None:
|
||||
sched.num_waiting_for_streaming_input = 0
|
||||
sched.log_stats = False
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.kv_holding_waiting = set()
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = set()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
update = StreamingUpdate(
|
||||
mm_features=None,
|
||||
@@ -105,7 +103,6 @@ def _make_drain_sched(session: Request) -> MagicMock:
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
bind_omits_transfer_helpers(sched)
|
||||
return sched
|
||||
|
||||
|
||||
|
||||
@@ -26,7 +26,6 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS
|
||||
from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
|
||||
OmniChunkTransferAdapter,
|
||||
)
|
||||
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
|
||||
|
||||
# isort: on
|
||||
|
||||
@@ -42,8 +41,7 @@ def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARS
|
||||
sched.num_waiting_for_streaming_input = 0
|
||||
sched.log_stats = False
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.kv_holding_waiting = set()
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = set()
|
||||
sched._free_request_blocks = MagicMock()
|
||||
sched.encoder_cache_manager = MagicMock()
|
||||
sched._inflight_prefills = set()
|
||||
@@ -177,7 +175,6 @@ def _run_resumable_segment_stop(
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
bind_omits_transfer_helpers(sched)
|
||||
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
|
||||
@@ -315,7 +312,6 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
bind_omits_transfer_helpers(sched)
|
||||
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
|
||||
@@ -404,7 +400,6 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None:
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
bind_omits_transfer_helpers(sched)
|
||||
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
|
||||
@@ -1170,8 +1165,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
|
||||
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
|
||||
sched.running = []
|
||||
sched.waiting = queue()
|
||||
sched.kv_holding_waiting = queue(session)
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = queue(session)
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=True,
|
||||
@@ -1183,7 +1177,7 @@ def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
|
||||
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert session in sched.waiting.requests
|
||||
assert session not in sched.kv_holding_waiting.requests
|
||||
assert session not in sched.skipped_waiting.requests
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
|
||||
|
||||
@@ -1210,8 +1204,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
|
||||
receives_chunks=receives_chunks,
|
||||
segment_finished_requests={session.request_id},
|
||||
)
|
||||
sched.kv_holding_waiting = MagicMock()
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = MagicMock()
|
||||
sched._enqueue_waiting_request = MagicMock()
|
||||
|
||||
sched._resume_downstream_chunk_receiver(session)
|
||||
@@ -1219,7 +1212,7 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
|
||||
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
assert sched.num_waiting_for_streaming_input == 1
|
||||
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
|
||||
sched.kv_holding_waiting.remove_requests.assert_not_called()
|
||||
sched.skipped_waiting.remove_requests.assert_not_called()
|
||||
sched._enqueue_waiting_request.assert_not_called()
|
||||
|
||||
|
||||
@@ -1275,8 +1268,7 @@ def _park_session(sched: OmniARScheduler, session: Request) -> None:
|
||||
session.num_computed_tokens = 6
|
||||
session.num_output_placeholders = 0
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.kv_holding_waiting.add_request(session)
|
||||
sched.deferred_waiting.add(session)
|
||||
sched.skipped_waiting.add_request(session)
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
|
||||
|
||||
@@ -1299,7 +1291,7 @@ def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_sessi
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
assert session.request_id not in sched.requests
|
||||
assert len(sched.waiting) == 0
|
||||
assert len(sched.kv_holding_waiting) == 0
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
assert sched.finished_req_ids == {session.request_id}
|
||||
@@ -1407,8 +1399,7 @@ def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
|
||||
sched.max_model_len = max_model_len
|
||||
sched.policy = SchedulingPolicy.FCFS
|
||||
sched.waiting = create_request_queue(sched.policy)
|
||||
sched.kv_holding_waiting = create_request_queue(sched.policy)
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = create_request_queue(sched.policy)
|
||||
sched.running = []
|
||||
sched.requests = {}
|
||||
sched._free_request = MagicMock() # needs a KV manager; not what this covers
|
||||
@@ -1437,7 +1428,7 @@ def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -
|
||||
|
||||
assert finished is True
|
||||
assert len(sched.waiting) == 0
|
||||
assert len(sched.kv_holding_waiting) == 0
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
assert session.resumable is False
|
||||
assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
|
||||
@@ -1456,8 +1447,7 @@ def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
|
||||
finished = sched._handle_stopped_request(session)
|
||||
|
||||
assert finished is False
|
||||
assert list(sched.kv_holding_waiting) == [session]
|
||||
assert session not in sched.deferred_waiting
|
||||
assert len(sched.waiting) == 1
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert not getattr(sched, "_streaming_context_overflow", {})
|
||||
sched._free_request.assert_not_called()
|
||||
@@ -1479,7 +1469,7 @@ def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:
|
||||
|
||||
assert finished is True
|
||||
assert len(sched.waiting) == 0
|
||||
assert len(sched.kv_holding_waiting) == 0
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
assert session.resumable is False
|
||||
sched._free_request.assert_not_called()
|
||||
@@ -1502,7 +1492,7 @@ def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanc
|
||||
finished = sched._handle_stopped_request(session)
|
||||
|
||||
assert finished is True
|
||||
assert len(sched.kv_holding_waiting) == 0
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
sched._free_request.assert_not_called()
|
||||
@@ -1535,7 +1525,6 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
|
||||
sched.running = []
|
||||
sched._native_data_plane = native
|
||||
sched.use_v2_model_runner = native
|
||||
sched.max_num_active_reqs = 8
|
||||
sched.max_num_running_reqs = 8
|
||||
sched.input_coordinator = (
|
||||
SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
|
||||
@@ -1562,7 +1551,7 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
|
||||
observed_limits: list[int] = []
|
||||
|
||||
def fake_schedule(self, _throttle_prefills=False):
|
||||
observed_limits.append(self.max_num_active_reqs)
|
||||
observed_limits.append(self.max_num_running_reqs)
|
||||
return SimpleNamespace(scheduled_new_reqs=[])
|
||||
|
||||
monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule)
|
||||
@@ -1570,4 +1559,4 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
|
||||
sched.schedule()
|
||||
|
||||
assert observed_limits == [7 if native else 8]
|
||||
assert sched.max_num_active_reqs == 8
|
||||
assert sched.max_num_running_reqs == 8
|
||||
@@ -79,19 +79,16 @@ def _make_scheduler(
|
||||
else:
|
||||
hf_config.voxcpm2_runtime_config = runtime_config
|
||||
sched.vllm_config = SimpleNamespace(model_config=model_config)
|
||||
# Upstream ``Scheduler.__init__`` always creates ``kv_holding_waiting``; the
|
||||
# Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
|
||||
# abort sweep at the top of ``schedule()`` reads it.
|
||||
sched.kv_holding_waiting = _MockQueue()
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = _MockQueue()
|
||||
return sched
|
||||
|
||||
|
||||
@pytest.mark.parametrize("queue_name", ["waiting", "kv_holding_waiting"])
|
||||
def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready(queue_name) -> None:
|
||||
def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None:
|
||||
scheduler = _make_scheduler()
|
||||
scheduler.running = [_MockRequest("decode")]
|
||||
scheduler.waiting = _MockQueue()
|
||||
setattr(scheduler, queue_name, _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]))
|
||||
scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)])
|
||||
|
||||
assert scheduler._should_defer_waiting_for_unified_decode_graph()
|
||||
|
||||
@@ -147,17 +144,12 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa
|
||||
|
||||
def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
|
||||
scheduler = _make_scheduler()
|
||||
scheduler.max_num_active_reqs = 8
|
||||
scheduler.max_num_running_reqs = 8
|
||||
scheduler.running = [_MockRequest("decode")]
|
||||
original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
|
||||
deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)
|
||||
original_waiting = _MockQueue([original_waiting_req])
|
||||
scheduler.waiting = original_waiting
|
||||
kv_holder = _MockRequest("kv-holder", status=RequestStatus.WAITING)
|
||||
original_kv_holding = _MockQueue([kv_holder])
|
||||
scheduler.kv_holding_waiting = original_kv_holding
|
||||
preempted_by_upstream = _MockRequest("preempted-by-upstream", status=RequestStatus.PREEMPTED)
|
||||
scheduler.policy = "fcfs"
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
@@ -168,10 +160,7 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
|
||||
def fake_upstream_schedule(self, throttle_prefills: bool = False):
|
||||
assert self.waiting is not original_waiting
|
||||
assert not self.waiting
|
||||
assert self.kv_holding_waiting is not original_kv_holding
|
||||
assert not self.kv_holding_waiting
|
||||
self.waiting.add_request(deferred_by_upstream)
|
||||
self.kv_holding_waiting.add_request(preempted_by_upstream)
|
||||
raise RuntimeError("stop before output wrapping")
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule)
|
||||
@@ -180,6 +169,3 @@ def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> No
|
||||
scheduler.schedule()
|
||||
|
||||
assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req]
|
||||
assert scheduler.waiting is original_waiting
|
||||
assert scheduler.kv_holding_waiting is original_kv_holding
|
||||
assert scheduler.kv_holding_waiting._items == [preempted_by_upstream, kv_holder]
|
||||
@@ -23,12 +23,10 @@ def _scheduler(monkeypatch, waiting, running=()):
|
||||
scheduler.waiting = create_request_queue(scheduler.policy)
|
||||
for request in waiting:
|
||||
scheduler.waiting.add_request(request)
|
||||
scheduler.kv_holding_waiting = create_request_queue(scheduler.policy)
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = create_request_queue(scheduler.policy)
|
||||
scheduler.running = list(running)
|
||||
scheduler.requests = {request.request_id: request for request in [*waiting, *running]}
|
||||
scheduler.max_model_len = 128
|
||||
scheduler.max_num_active_reqs = 2
|
||||
scheduler.max_num_running_reqs = 2
|
||||
scheduler.max_num_scheduled_tokens = 256
|
||||
scheduler.block_size = 16
|
||||
|
||||
@@ -52,11 +52,10 @@ class _SchedulerStub(OmniGenerationScheduler):
|
||||
self.num_waiting_for_streaming_input = 0
|
||||
self.log_stats = log_stats
|
||||
self.chunk_transfer_adapter = _ChunkTransferAdapterStub()
|
||||
self.kv_holding_waiting = _SkippedWaitingStub()
|
||||
self.deferred_waiting: set[Request] = set()
|
||||
self.skipped_waiting = _SkippedWaitingStub()
|
||||
|
||||
def _enqueue_waiting_request(self, session: Request) -> None:
|
||||
raise AssertionError("unexpected enqueue for kv_holding_waiting miss")
|
||||
raise AssertionError("unexpected enqueue for skipped_waiting miss")
|
||||
|
||||
|
||||
class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
|
||||
@@ -135,8 +134,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
||||
sched._handle_stopped_request.side_effect = rearm_resumable_request
|
||||
sched.running = [session]
|
||||
sched.waiting = MagicMock()
|
||||
sched.kv_holding_waiting = MagicMock()
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = MagicMock()
|
||||
sched.structured_output_manager.accept_tokens.return_value = True
|
||||
# Async scheduling can observe the segment boundary in the next schedule()
|
||||
# before this model output is applied and defer the same request for finish.
|
||||
@@ -208,8 +206,7 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete:
|
||||
OmniSchedulerMixin._remove_stopped_requests_from_queues(sched, running, preempted)
|
||||
)
|
||||
sched.waiting = MagicMock()
|
||||
sched.kv_holding_waiting = MagicMock()
|
||||
sched.deferred_waiting = set()
|
||||
sched.skipped_waiting = MagicMock()
|
||||
sched.structured_output_manager.accept_tokens.return_value = False
|
||||
sched._pending_finish_reqs = []
|
||||
sched.recompute_kv_load_failures = False
|
||||
@@ -353,8 +350,6 @@ class _RealignSchedulerStub(OmniSchedulerMixin):
|
||||
self.requests = requests
|
||||
self.running = running
|
||||
self.waiting = waiting
|
||||
self.kv_holding_waiting: list[Request] = []
|
||||
self.deferred_waiting: set[Request] = set()
|
||||
|
||||
|
||||
class TestRealignRequestStatusToQueues:
|
||||
|
||||
@@ -123,21 +123,13 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
|
||||
|
||||
sched = MagicMock()
|
||||
sched.max_num_scheduled_tokens = 8
|
||||
sched.max_num_active_reqs = 8
|
||||
sched.max_num_running_reqs = 8
|
||||
sched._native_data_plane = False
|
||||
sched._retains_state_across_chunks = False
|
||||
# Unset MagicMock attributes are truthy. Keep optional scheduling modes
|
||||
# disabled so this fixture actually schedules its one-shot request.
|
||||
sched._generation_defer_batch = False
|
||||
sched._generation_first_chunk_only = False
|
||||
sched._generation_max_regular_batch = 0
|
||||
sched._pause_state = PauseState.UNPAUSED
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.running = [request]
|
||||
sched.waiting = [] # falsy: the waiting loop never runs
|
||||
sched.kv_holding_waiting = []
|
||||
sched.deferred_waiting = set()
|
||||
sched.policy = SchedulingPolicy.FCFS
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Integration tests for the post-finish defensive purge of
|
||||
``self.running`` placed in
|
||||
``OmniARScheduler.finish_requests`` /
|
||||
@@ -55,8 +52,6 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting):
|
||||
scheduler.requests = requests
|
||||
scheduler.running = running
|
||||
scheduler.waiting = waiting
|
||||
scheduler.kv_holding_waiting = []
|
||||
scheduler.deferred_waiting = set()
|
||||
return scheduler
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
@@ -52,8 +49,6 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
|
||||
scheduler.requests = {}
|
||||
scheduler.running = []
|
||||
scheduler.waiting = []
|
||||
scheduler.kv_holding_waiting = []
|
||||
scheduler.deferred_waiting = set()
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
assert request_ids == ["req-a", "req-b"]
|
||||
@@ -81,7 +76,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.input_coordinator = coordinator
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)}
|
||||
scheduler._omits_kv_transfer_cache = {"req-free": True}
|
||||
scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
|
||||
scheduler.finished_req_ids = set()
|
||||
scheduler.finished_req_ids_dict = None
|
||||
|
||||
@@ -7,10 +7,7 @@ from typing import Any
|
||||
|
||||
import pytest
|
||||
from vllm.config import SchedulerConfig, VllmConfig
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.scheduler import Scheduler
|
||||
from vllm.v1.engine import FinishReason
|
||||
from vllm.v1.request import Request
|
||||
|
||||
from vllm_omni.config.model import OmniModelConfig
|
||||
from vllm_omni.core.sched import omni_scheduler_mixin
|
||||
@@ -21,7 +18,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _Scheduler(OmniSchedulerMixin):
|
||||
_is_blocked_waiting_status = staticmethod(Scheduler._is_blocked_waiting_status)
|
||||
pass
|
||||
|
||||
|
||||
def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
|
||||
@@ -123,11 +120,7 @@ def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_i
|
||||
def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
calls: list[tuple[Any, ...]] = []
|
||||
scheduler = _Scheduler()
|
||||
waiting = Request("waiting", [1], SamplingParams(max_tokens=1), pooling_params=None)
|
||||
holding = Request("holding", [1, 2], SamplingParams(max_tokens=1), pooling_params=None)
|
||||
holding.num_computed_tokens = 1
|
||||
scheduler.waiting = [waiting]
|
||||
scheduler.kv_holding_waiting = [holding]
|
||||
scheduler.waiting = ["waiting"]
|
||||
scheduler.running = ["running"]
|
||||
scheduler.requests = {"request": object()}
|
||||
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
|
||||
@@ -144,16 +137,16 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
scheduler.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=True,
|
||||
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("process", list(waiting), running, scheduler_requests)
|
||||
("process", waiting, running, scheduler_requests)
|
||||
),
|
||||
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("restore-chunks", list(waiting), running, scheduler_requests)
|
||||
("restore-chunks", waiting, running, scheduler_requests)
|
||||
),
|
||||
collect_timed_out_request_ids=_collect_timed_out,
|
||||
collect_failed_send_request_ids=_collect_failed_sends,
|
||||
)
|
||||
scheduler.input_coordinator = SimpleNamespace(
|
||||
restore_queues=lambda waiting, running: calls.append(("restore-full", list(waiting), running))
|
||||
restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running))
|
||||
)
|
||||
|
||||
scheduler._process_pending_omni_inputs("ar")
|
||||
@@ -162,13 +155,13 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
assert calls == [
|
||||
("consume", "ar"),
|
||||
("timeouts",),
|
||||
("process", [holding, waiting], scheduler.running, scheduler.requests),
|
||||
("process", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
# The chunk deadline runs after chunks are applied, so a chunk that
|
||||
# arrived this cycle resets the clock before it is measured (R1.1).
|
||||
("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
|
||||
("failed-sends",),
|
||||
("restore-chunks", [holding, waiting], scheduler.running, scheduler.requests),
|
||||
("restore-full", [holding, waiting], scheduler.running),
|
||||
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
("restore-full", scheduler.waiting, scheduler.running),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -5,13 +5,8 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm import SamplingParams
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator
|
||||
from vllm_omni.outputs import OmniConnectorOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -20,8 +15,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
|
||||
scheduler = OmniSchedulerMixin()
|
||||
scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], []
|
||||
scheduler.kv_holding_waiting = []
|
||||
scheduler._is_blocked_waiting_status = OmniARScheduler._is_blocked_waiting_status
|
||||
coordinator = SimpleNamespace(
|
||||
_async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock()
|
||||
)
|
||||
@@ -39,92 +32,7 @@ def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
|
||||
coordinator.update_request_metadata.assert_called_once_with(
|
||||
scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar"
|
||||
)
|
||||
coordinator.process_pending_chunks.assert_called_once_with(mocker.ANY, [], {"r"}, {"r"})
|
||||
assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
|
||||
coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"})
|
||||
scheduler._consume_pending_connector_output(model_mode="ar")
|
||||
assert coordinator.update_request_metadata.call_count == 1
|
||||
coordinator.process_pending_chunks.assert_called_with(mocker.ANY, [], set(), set())
|
||||
assert list(coordinator.process_pending_chunks.call_args.args[0]) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy", [SchedulingPolicy.FCFS, SchedulingPolicy.PRIORITY])
|
||||
@pytest.mark.parametrize("async_chunk", [False, True])
|
||||
def test_native_input_gate_parks_and_restores_kv_holders(policy, async_chunk):
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.waiting = create_request_queue(policy)
|
||||
scheduler.kv_holding_waiting = create_request_queue(policy)
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.running = []
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = OmniSchedulingCoordinator(
|
||||
stage_id=1, scheduler_max_num_seqs=2, async_chunk=async_chunk
|
||||
)
|
||||
scheduler._init_omni_connector_output_inbox()
|
||||
holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
|
||||
holder.num_computed_tokens = 2
|
||||
fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
|
||||
scheduler.requests = {request.request_id: request for request in (holder, fresh)}
|
||||
scheduler._enqueue_waiting_request(fresh)
|
||||
scheduler._enqueue_waiting_request(holder)
|
||||
|
||||
scheduler._consume_pending_connector_output(model_mode="ar")
|
||||
|
||||
assert not scheduler.waiting and not scheduler.kv_holding_waiting
|
||||
expected_status = RequestStatus.WAITING_FOR_CHUNK if async_chunk else RequestStatus.WAITING_FOR_INPUT
|
||||
assert holder.status == fresh.status == expected_status
|
||||
coordinator = scheduler.input_coordinator
|
||||
registrations = coordinator.pending_chunk_registrations if async_chunk else coordinator.pending_input_registrations
|
||||
assert [entry.request_id for entry in registrations] == ["holder", "fresh"]
|
||||
|
||||
scheduler._restore_omni_wait_queues()
|
||||
assert list(scheduler.kv_holding_waiting) == [holder]
|
||||
assert list(scheduler.waiting) == [fresh]
|
||||
ready_output = (
|
||||
OmniConnectorOutput(chunk_ready_req_ids={"holder"})
|
||||
if async_chunk
|
||||
else OmniConnectorOutput(stage_recv_req_ids={"holder"})
|
||||
)
|
||||
scheduler.enqueue_omni_connector_output(ready_output)
|
||||
scheduler._consume_pending_connector_output(model_mode="ar")
|
||||
|
||||
assert holder.status == RequestStatus.WAITING
|
||||
assert fresh.status == expected_status
|
||||
assert list(scheduler.kv_holding_waiting) == [holder]
|
||||
assert not scheduler.waiting
|
||||
scheduler._restore_omni_wait_queues()
|
||||
assert list(scheduler.kv_holding_waiting) == [holder]
|
||||
assert list(scheduler.waiting) == [fresh]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_chunk", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"status",
|
||||
[
|
||||
RequestStatus.WAITING_FOR_STREAMING_REQ,
|
||||
RequestStatus.WAITING_FOR_REMOTE_KVS,
|
||||
RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR,
|
||||
],
|
||||
)
|
||||
def test_native_input_gate_preserves_upstream_blocked_waits(async_chunk, status):
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.running = []
|
||||
request = Request("blocked", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
|
||||
request.status = status
|
||||
request.num_computed_tokens = 2
|
||||
scheduler.requests = {request.request_id: request}
|
||||
scheduler._enqueue_waiting_request(request)
|
||||
scheduler.input_coordinator = OmniSchedulingCoordinator(
|
||||
stage_id=1, scheduler_max_num_seqs=1, async_chunk=async_chunk
|
||||
)
|
||||
scheduler._init_omni_connector_output_inbox()
|
||||
|
||||
scheduler._consume_pending_connector_output(model_mode="ar")
|
||||
|
||||
assert request.status == status
|
||||
assert list(scheduler.kv_holding_waiting) == [request]
|
||||
assert scheduler.deferred_waiting == {request}
|
||||
assert not scheduler.input_coordinator.pending_chunk_registrations
|
||||
assert not scheduler.input_coordinator.pending_input_registrations
|
||||
coordinator.process_pending_chunks.assert_called_with([], [], set(), set())
|
||||
@@ -59,8 +59,7 @@ def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipp
|
||||
scheduler.requests = requests
|
||||
scheduler.running = running
|
||||
scheduler.waiting = waiting
|
||||
scheduler.kv_holding_waiting = skipped if skipped is not None else []
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = skipped if skipped is not None else []
|
||||
scheduler.num_waiting_for_streaming_input = counter
|
||||
return scheduler
|
||||
|
||||
@@ -98,8 +97,8 @@ def test_finish_requests_clears_counter_left_by_a_stomped_status(
|
||||
assert scheduler.num_waiting_for_streaming_input == 0
|
||||
|
||||
|
||||
def test_resync_counts_parked_requests_in_kv_holding_waiting() -> None:
|
||||
"""``get_num_unfinished_requests`` sums ``waiting`` and ``kv_holding_waiting``,
|
||||
def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
|
||||
"""``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
|
||||
so the resync has to look at both or it under-counts and reintroduces the
|
||||
spin the counter exists to prevent."""
|
||||
parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
|
||||
@@ -129,11 +128,9 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk
|
||||
scheduler.requests = {parked.request_id: parked}
|
||||
scheduler.running = [parked]
|
||||
scheduler.waiting = []
|
||||
scheduler.kv_holding_waiting = []
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = []
|
||||
scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite
|
||||
scheduler.policy = "fcfs"
|
||||
scheduler.max_num_active_reqs = 8
|
||||
scheduler.max_num_running_reqs = 8
|
||||
|
||||
seen: dict[str, int] = {}
|
||||
@@ -164,13 +161,12 @@ def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -
|
||||
scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
|
||||
scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.waiting.add_request(aborted)
|
||||
scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.kv_holding_waiting.add_request(aborted)
|
||||
scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.skipped_waiting.add_request(aborted)
|
||||
scheduler.running = [aborted]
|
||||
|
||||
scheduler._drop_aborted_queued_requests()
|
||||
|
||||
assert list(scheduler.waiting) == []
|
||||
assert list(scheduler.kv_holding_waiting) == []
|
||||
assert list(scheduler.skipped_waiting) == []
|
||||
assert scheduler.running == []
|
||||
@@ -216,40 +216,44 @@ def test_video_chat_stream_invalid_requests(
|
||||
# WS /v1/realtime
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
|
||||
_REALTIME_WS_MODEL_MISMATCH = object()
|
||||
_REALTIME_WS_INVALID_AUDIO_APPEND = object()
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize(
|
||||
"send_frames_spec, ws_error_code, err_message",
|
||||
[
|
||||
pytest.param("{not-json", "invalid_event", "Invalid or unrecognized client event", id="invalid_json"),
|
||||
pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
|
||||
pytest.param(
|
||||
json.dumps({"type": "session.update"}),
|
||||
"invalid_event",
|
||||
"Invalid or unrecognized client event",
|
||||
id="session_update_missing_session",
|
||||
("model", "Missing required field"),
|
||||
id="session_update_missing_model",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "session.update", "session": []}),
|
||||
"invalid_event",
|
||||
"Invalid or unrecognized client event",
|
||||
id="session_update_invalid_session",
|
||||
_REALTIME_WS_MODEL_MISMATCH,
|
||||
"model_not_found",
|
||||
("model", "does not exist"),
|
||||
id="session_update_model_mismatch",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "input_audio_buffer.commit"}),
|
||||
"invalid_request_error",
|
||||
"Input audio buffer is empty",
|
||||
id="commit_without_audio",
|
||||
json.dumps({"type": "input_audio_buffer.commit", "final": False}),
|
||||
"model_not_validated",
|
||||
("session.update", "validate the model"),
|
||||
id="commit_without_model_validation",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "unknown.event"}),
|
||||
"invalid_event",
|
||||
"Invalid or unrecognized client event",
|
||||
"unknown_event",
|
||||
"Unknown event type",
|
||||
id="unknown_event_type",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
|
||||
"invalid_request_error",
|
||||
"Invalid base64 audio data",
|
||||
_REALTIME_WS_INVALID_AUDIO_APPEND,
|
||||
"invalid_audio",
|
||||
"Invalid audio data",
|
||||
id="invalid_audio_append",
|
||||
),
|
||||
],
|
||||
@@ -262,12 +266,27 @@ def test_realtime_invalid_requests(
|
||||
ws_error_code: str,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
send_frames: Any
|
||||
if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
|
||||
send_frames = json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"model": "this-model-is-not-served-on-this-instance",
|
||||
}
|
||||
)
|
||||
elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
|
||||
send_frames = [
|
||||
json.dumps({"type": "session.update", "model": omni_server.model}),
|
||||
json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
|
||||
]
|
||||
else:
|
||||
send_frames = send_frames_spec
|
||||
online_client.send_realtime_ws_request(
|
||||
{
|
||||
"send_frames": send_frames_spec,
|
||||
"send_frames": send_frames,
|
||||
"timeout": 120,
|
||||
"ws_max_size": None,
|
||||
"ws_skip_types": ["session.created", "conversation.created"],
|
||||
"ws_skip_types": ["session.created"],
|
||||
"ws_json_type": "error",
|
||||
"ws_error_code": ws_error_code,
|
||||
"err_message": err_message,
|
||||
|
||||
+1
-3
@@ -329,9 +329,7 @@ def _uneven_cfg_sharded_worker(
|
||||
init_method=init_method,
|
||||
rank=rank,
|
||||
world_size=4,
|
||||
# Spawned workers import torch/vLLM before joining; allow startup
|
||||
# skew without changing the collective deadlock checks below.
|
||||
timeout=datetime.timedelta(seconds=60),
|
||||
timeout=datetime.timedelta(seconds=10),
|
||||
)
|
||||
try:
|
||||
fs_groups = [
|
||||
|
||||
@@ -103,7 +103,7 @@ def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
|
||||
# Independent expected native extras: a crossing item has a negative
|
||||
# relative offset; a block can include the end of one item and the next.
|
||||
first, second = (feature.identifier for feature in row.mm_features)
|
||||
extras = [(("mm", first, 2),), (("mm", first, -2), ("mm", second, 3)), (("mm", second, -1),)]
|
||||
extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
|
||||
parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
|
||||
expected = []
|
||||
for start, keys in zip(range(0, 12, 4), extras, strict=True):
|
||||
|
||||
@@ -31,10 +31,6 @@ def producer(monkeypatch):
|
||||
# vLLM 0.30 send_kv_to_decode also reads the PP size (ce08bb5b34,
|
||||
# #56033), the MLA/Mamba flags and the head-resharding validator.
|
||||
worker.pp_size = 1
|
||||
worker.opaque_packed_storages = set()
|
||||
worker._prepared_transfer_regions = {}
|
||||
worker.region_shared_groups = []
|
||||
worker.region_row_offsets = []
|
||||
worker.use_mla = False
|
||||
worker.kv_cache_config = SimpleNamespace(has_mamba_layers=False)
|
||||
worker._validate_head_resharding_layout = Mock(return_value=None)
|
||||
|
||||
@@ -1253,15 +1253,14 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
speculative_config=None,
|
||||
)
|
||||
|
||||
def select_backend(*, backend, attn_selector_config, num_heads, _run_kv_cache_dtype):
|
||||
# Exercise the real upstream selector, replacing only platform resolution.
|
||||
def select_backend(*, backend, attn_selector_config, num_heads):
|
||||
# Exercise the real 0.29 selector, replacing only platform resolution.
|
||||
# Ulysses has already gathered tokens and sharded heads before the
|
||||
# native kernel; the MoE PCP mapping must not request PCP attention.
|
||||
assert not attn_selector_config.use_pcp
|
||||
assert not attn_selector_config.use_dcp
|
||||
assert attn_selector_config.use_non_causal
|
||||
assert num_heads == 4
|
||||
assert _run_kv_cache_dtype == config.cache_config.cache_dtype
|
||||
selected_backends.append((backend, config.attention_config.backend_per_kind))
|
||||
if selection_fails:
|
||||
raise ValueError("test backend unavailable")
|
||||
|
||||
@@ -14,7 +14,6 @@ from typing import Literal
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
@@ -57,8 +56,6 @@ def _init_worker(
|
||||
"MASTER_PORT": str(master_port),
|
||||
}
|
||||
)
|
||||
if device_kind == "cpu":
|
||||
configure_cpu_collective_worker()
|
||||
backend = "gloo" if device_kind == "cpu" else None
|
||||
init_distributed_environment(backend=backend)
|
||||
|
||||
@@ -230,10 +227,6 @@ def _run_ring_p2p(
|
||||
atol=atol,
|
||||
msg=f"[Rank {local_rank}] Ring P2P data mismatch",
|
||||
)
|
||||
if device_kind == "cpu":
|
||||
# P2P completion only synchronizes neighbors. Keep fast CPU ranks
|
||||
# alive until every rank has finished constructing and using groups.
|
||||
torch.distributed.barrier()
|
||||
finally:
|
||||
destroy_distributed_env()
|
||||
|
||||
@@ -250,7 +243,7 @@ def _spawn_4d_identity(
|
||||
device_kind: DeviceKind,
|
||||
master_port: int,
|
||||
) -> None:
|
||||
start_collective_workers(
|
||||
torch.multiprocessing.spawn(
|
||||
_run_4d_identity,
|
||||
args=(
|
||||
world_size,
|
||||
@@ -264,7 +257,6 @@ def _spawn_4d_identity(
|
||||
master_port,
|
||||
),
|
||||
nprocs=world_size,
|
||||
device_kind=device_kind,
|
||||
)
|
||||
|
||||
|
||||
@@ -280,7 +272,7 @@ def _spawn_5d_identity(
|
||||
device_kind: DeviceKind,
|
||||
master_port: int,
|
||||
) -> None:
|
||||
start_collective_workers(
|
||||
torch.multiprocessing.spawn(
|
||||
_run_5d_identity,
|
||||
args=(
|
||||
world_size,
|
||||
@@ -294,7 +286,6 @@ def _spawn_5d_identity(
|
||||
master_port,
|
||||
),
|
||||
nprocs=world_size,
|
||||
device_kind=device_kind,
|
||||
)
|
||||
|
||||
|
||||
@@ -308,11 +299,10 @@ def _spawn_ring_p2p(
|
||||
device_kind: DeviceKind,
|
||||
master_port: int,
|
||||
) -> None:
|
||||
start_collective_workers(
|
||||
torch.multiprocessing.spawn(
|
||||
_run_ring_p2p,
|
||||
args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
|
||||
nprocs=world_size,
|
||||
device_kind=device_kind,
|
||||
)
|
||||
|
||||
|
||||
@@ -509,9 +499,8 @@ def _run_fused_qkv_matches_three_4d(
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("world_size", [2, 4])
|
||||
def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
|
||||
start_collective_workers(
|
||||
torch.multiprocessing.spawn(
|
||||
_run_fused_qkv_matches_three_4d,
|
||||
args=(world_size, "cpu", 29613),
|
||||
nprocs=world_size,
|
||||
device_kind="cpu",
|
||||
)
|
||||
@@ -19,7 +19,6 @@ from vllm.utils.network_utils import get_file_store_init_method
|
||||
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
|
||||
|
||||
import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
|
||||
from tests.helpers.distributed import collective_context, start_collective_workers
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
@@ -40,25 +39,6 @@ DeviceKind = Literal["cpu", "cuda"]
|
||||
_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _restore_collective_platforms():
|
||||
"""CPU baselines run in the parent process; retain its original platforms."""
|
||||
import vllm.platforms
|
||||
|
||||
from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
|
||||
|
||||
original_platform = vllm.platforms.current_platform
|
||||
original_omni_platforms = {
|
||||
module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state)
|
||||
}
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
vllm.platforms.current_platform = original_platform
|
||||
for module, platform in original_omni_platforms.items():
|
||||
module.current_omni_platform = platform
|
||||
|
||||
|
||||
def _cleanup_distributed() -> None:
|
||||
"""Destroy omni distributed state left dangling after a test.
|
||||
|
||||
@@ -429,10 +409,6 @@ def init_dist(
|
||||
device = _worker_device(local_rank, device_kind)
|
||||
if device_kind == "cuda":
|
||||
current_omni_platform.set_device(device)
|
||||
else:
|
||||
from tests.helpers.distributed import configure_cpu_collective_worker
|
||||
|
||||
configure_cpu_collective_worker()
|
||||
backend = "gloo" if device_kind == "cpu" else None
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
@@ -524,14 +500,13 @@ def isend_irecv_worker(
|
||||
|
||||
|
||||
def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
|
||||
mp_context = collective_context(device_kind)
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
q = manager.Queue()
|
||||
start_collective_workers(
|
||||
torch.multiprocessing.spawn(
|
||||
isend_irecv_worker,
|
||||
args=(pp_size, init_method, device_kind, q),
|
||||
nprocs=pp_size,
|
||||
device_kind=device_kind,
|
||||
)
|
||||
results = {label: tensor for label, tensor in [q.get(), q.get()]}
|
||||
torch.testing.assert_close(
|
||||
@@ -681,16 +656,15 @@ def _run_predict_noise(
|
||||
|
||||
baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
|
||||
|
||||
mp_context = collective_context(device_kind)
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
pp_q = manager.Queue()
|
||||
|
||||
world_size = pp_size * cfg_size
|
||||
start_collective_workers(
|
||||
torch.multiprocessing.spawn(
|
||||
predict_noise_worker,
|
||||
args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
|
||||
nprocs=world_size,
|
||||
device_kind=device_kind,
|
||||
)
|
||||
|
||||
pp_out = pp_q.get()
|
||||
@@ -926,16 +900,15 @@ def _run_scheduler_step(
|
||||
|
||||
baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
|
||||
|
||||
mp_context = collective_context(device_kind)
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
q = manager.Queue()
|
||||
|
||||
world_size = pp_size * cfg_size
|
||||
start_collective_workers(
|
||||
torch.multiprocessing.spawn(
|
||||
scheduler_step_worker,
|
||||
args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
|
||||
nprocs=world_size,
|
||||
device_kind=device_kind,
|
||||
)
|
||||
|
||||
result = q.get()
|
||||
|
||||
@@ -61,10 +61,7 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa
|
||||
|
||||
# The conditioning is prepared once per request and reused by every step.
|
||||
assert prepare_spy.call_count == 1
|
||||
# _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
|
||||
assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5)
|
||||
# The time grid reaches prepare, so the adaLN modulations are computed once.
|
||||
assert prepare_spy.call_args.args[7].numel() == 2
|
||||
assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
|
||||
assert step_spy.call_count == 2
|
||||
capture_spy.assert_not_called()
|
||||
assert not wrapper._cache
|
||||
@@ -87,8 +84,7 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
|
||||
assert padded[4].shape == (1, 100, 4)
|
||||
assert padded[5].shape == (1, 100)
|
||||
assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
|
||||
assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0)
|
||||
assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32)
|
||||
assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
|
||||
assert wrapper.max_graphs == 32
|
||||
|
||||
|
||||
@@ -157,21 +153,18 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
|
||||
inputs["ref"],
|
||||
inputs["ref_mask"],
|
||||
)
|
||||
# sample_latents passes the grid, so the key carries its step count (two steps here).
|
||||
key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2)
|
||||
key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
|
||||
assert key in wrapper._cache
|
||||
|
||||
# The static context was refreshed for this request's conditioning.
|
||||
entry = wrapper._cache[key]
|
||||
grid = torch.tensor([0.0, 0.4], device="cuda")
|
||||
fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid)
|
||||
fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
|
||||
for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
|
||||
if want is None:
|
||||
assert static is None
|
||||
else:
|
||||
torch.testing.assert_close(static, want)
|
||||
torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
|
||||
assert int(entry.static_step) == 1
|
||||
|
||||
assert len(wrapper._cache) == 1
|
||||
|
||||
|
||||
@@ -1,139 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""The AuK DiT fast paths keep the reference numerics.
|
||||
|
||||
Two paths are checked: the grouped conv position embedding as one batched
|
||||
GEMM, and the adaLN modulations precomputed for a whole time grid, eager and
|
||||
under the per-step CUDA graph.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.auk import auk_transformer as T
|
||||
from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid
|
||||
from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_dit(device: str = "cpu") -> AuKTransformer:
|
||||
torch.manual_seed(5)
|
||||
return (
|
||||
AuKTransformer(
|
||||
dim=32,
|
||||
heads=2,
|
||||
dim_head=16,
|
||||
ff_mult=2,
|
||||
latent_dim=4,
|
||||
text_hidden_dim=8,
|
||||
num_layers=2,
|
||||
num_single_layers=2,
|
||||
)
|
||||
.eval()
|
||||
.to(device)
|
||||
)
|
||||
|
||||
|
||||
def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
|
||||
"""The channels-first conv stack as the reference implementation writes it."""
|
||||
keep = None if mask is None else mask.unsqueeze(1)
|
||||
x = x.transpose(1, 2)
|
||||
if keep is not None:
|
||||
x = x.masked_fill(~keep, 0.0)
|
||||
for layer in module.conv1d:
|
||||
x = layer(x)
|
||||
if keep is not None and isinstance(layer, torch.nn.Conv1d):
|
||||
x = x.masked_fill(~keep, 0.0)
|
||||
return x.transpose(1, 2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None:
|
||||
torch.manual_seed(0)
|
||||
conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15)
|
||||
x = torch.randn(2, 37, 64)
|
||||
expected = conv(x.transpose(1, 2)).transpose(1, 2)
|
||||
torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
@pytest.mark.parametrize("masked", [False, True])
|
||||
def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None:
|
||||
torch.manual_seed(1)
|
||||
module = T.ConvPosEmbedding(64).eval()
|
||||
x = torch.randn(2, 37, 64)
|
||||
mask = None
|
||||
if masked:
|
||||
mask = torch.ones(2, 37, dtype=torch.bool)
|
||||
mask[1, 30:] = False
|
||||
torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
@pytest.mark.parametrize("cfg_infer", [False, True])
|
||||
def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None:
|
||||
dit = _make_dit()
|
||||
text = torch.randn(1, 5, 8)
|
||||
ref = torch.randn(1, 3, 4)
|
||||
x = torch.randn(1, 6, 4)
|
||||
grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu")
|
||||
plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer)
|
||||
tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1])
|
||||
assert tabled.modulation is not None and tabled.modulation.shape[0] == 4
|
||||
for i in range(4):
|
||||
expected = dit.step(x, grid[i], plain)
|
||||
torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5)
|
||||
index = torch.tensor(i)
|
||||
torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
def _padded_inputs(device: str) -> dict[str, torch.Tensor]:
|
||||
generator = torch.Generator(device="cpu").manual_seed(3)
|
||||
text = torch.randn(1, 7, 8, generator=generator)
|
||||
c_mask = torch.ones(1, 7, dtype=torch.bool)
|
||||
c_mask[:, 5:] = False
|
||||
text[:, 5:] = 0.0
|
||||
ref = torch.randn(1, 6, 4, generator=generator)
|
||||
ref_mask = torch.ones(1, 6, dtype=torch.bool)
|
||||
ref_mask[:, 4:] = False
|
||||
mask = torch.ones(1, 9, dtype=torch.bool)
|
||||
mask[:, 7:] = False
|
||||
inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask}
|
||||
return {name: value.to(device) for name, value in inputs.items()}
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
|
||||
@torch.inference_mode()
|
||||
@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
|
||||
def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None:
|
||||
dit = _make_dit("cuda")
|
||||
inputs = _padded_inputs("cuda")
|
||||
grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda")
|
||||
x = torch.randn(1, 9, 4, device="cuda")
|
||||
wrapper = AuKCUDAGraphWrapper(dit)
|
||||
common = dict(
|
||||
text=inputs["text"],
|
||||
c_mask=inputs["c_mask"],
|
||||
ref=inputs["ref"],
|
||||
ref_mask=inputs["ref_mask"],
|
||||
cfg_strength=cfg_strength,
|
||||
timesteps=grid[:-1],
|
||||
)
|
||||
for i in range(4):
|
||||
graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common)
|
||||
ctx = dit.prepare(
|
||||
inputs["text"],
|
||||
target_len=9,
|
||||
c_mask=inputs["c_mask"],
|
||||
ref=inputs["ref"],
|
||||
ref_mask=inputs["ref_mask"],
|
||||
cfg_infer=cfg_strength >= 1e-5,
|
||||
timesteps=grid[:-1],
|
||||
)
|
||||
eager = wrapper._step(x, grid[i], ctx, cfg_strength, i)
|
||||
torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4)
|
||||
# One graph serves every step of the grid.
|
||||
assert len(wrapper._cache) == 1
|
||||
assert next(iter(wrapper._cache))[-1] == 4
|
||||
@@ -127,9 +127,7 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non
|
||||
v = dit(x, time=timesteps[i], **inputs)
|
||||
x = x + (timesteps[i + 1] - timesteps[i]) * v
|
||||
|
||||
# sample_latents precomputes the adaLN modulations of the whole grid in one
|
||||
# GEMM per layer, whose rows round differently from one-row GEMMs in fp32.
|
||||
torch.testing.assert_close(got, x, rtol=0, atol=1e-6)
|
||||
torch.testing.assert_close(got, x, rtol=0, atol=0)
|
||||
assert dit.text_cond is None and dit.text_uncond is None
|
||||
|
||||
|
||||
|
||||
@@ -12,8 +12,7 @@ import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.models.auk import auk_vae
|
||||
from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample
|
||||
from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample
|
||||
from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
@@ -247,72 +246,6 @@ def test_tiled_decode_matches_the_whole_decode() -> None:
|
||||
AuKVAEDecodeGraph(vae, tile_frames=60)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@torch.inference_mode()
|
||||
def test_fused_activation_is_skipped_off_cuda() -> None:
|
||||
activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True))
|
||||
x = torch.randn(1, 4, 16)
|
||||
assert not activation._fused_eligible(x)
|
||||
assert activation(x).shape == (1, 4, 16)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
|
||||
@torch.inference_mode()
|
||||
def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda")
|
||||
x = torch.randn(1, 4, 16, device="cuda")
|
||||
assert activation._fused_eligible(x)
|
||||
monkeypatch.setattr(auk_vae, "HAS_TRITON", False)
|
||||
assert not activation._fused_eligible(x)
|
||||
assert activation(x).shape == (1, 4, 16)
|
||||
|
||||
|
||||
def _activation_input(layout: str, length: int) -> torch.Tensor:
|
||||
"""A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer."""
|
||||
if layout == "contiguous":
|
||||
x = torch.randn(2, 6, length, device="cuda")
|
||||
elif layout == "transposed":
|
||||
x = torch.randn(2, length, 6, device="cuda").transpose(1, 2)
|
||||
elif layout == "sliced":
|
||||
x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2]
|
||||
else:
|
||||
raise ValueError(layout)
|
||||
assert x.shape == (2, 6, length)
|
||||
if length > 1: # contiguity of a single-sample view depends on the layout
|
||||
assert x.is_contiguous() == (layout == "contiguous")
|
||||
return x * 3
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA")
|
||||
@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"])
|
||||
@pytest.mark.parametrize("causal", [False, True])
|
||||
# The kernels tile the output in blocks of 1024 samples: the 2x oversampled
|
||||
# signal crosses a block boundary at 512 input samples and the decimated one at
|
||||
# 1024, so both sides of each boundary are covered, plus short and odd lengths.
|
||||
@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025])
|
||||
@torch.inference_mode()
|
||||
def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None:
|
||||
torch.manual_seed(0)
|
||||
activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda")
|
||||
with torch.no_grad():
|
||||
activation.act.alpha.normal_(0.0, 0.3)
|
||||
activation.act.beta.normal_(0.0, 0.3)
|
||||
x = _activation_input(layout, length)
|
||||
original = x.clone()
|
||||
|
||||
eager = activation.downsample(activation.act(activation.upsample(x)))
|
||||
assert activation._fused_eligible(x)
|
||||
fused = activation(x)
|
||||
|
||||
# Same fp32 arithmetic, only the order in which the taps are summed differs.
|
||||
assert fused.shape == eager.shape
|
||||
torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5)
|
||||
# Neither path writes into its (possibly strided) input.
|
||||
torch.testing.assert_close(x, original, atol=0, rtol=0)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
|
||||
@torch.inference_mode()
|
||||
|
||||
@@ -291,18 +291,6 @@ class TestRequestParsing:
|
||||
with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"):
|
||||
build_pipeline(model_config={"max_dit_graphs": max_dit_graphs})
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)]
|
||||
)
|
||||
def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected):
|
||||
pipeline, _ = build_pipeline(model_config=model_config)
|
||||
assert pipeline._ref_cache_size == expected
|
||||
|
||||
@pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None])
|
||||
def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size):
|
||||
with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"):
|
||||
build_pipeline(model_config={"auk_ref_cache_size": size})
|
||||
|
||||
def test_pipeline_declares_audio_output(self, build_pipeline):
|
||||
pipeline, _ = build_pipeline()
|
||||
|
||||
@@ -504,40 +492,6 @@ class TestRequestParsing:
|
||||
assert call["sample"] is True
|
||||
assert isinstance(call["generator"], torch.Generator)
|
||||
|
||||
def test_reference_latents_are_cached_by_content(self, build_pipeline):
|
||||
pipeline, calls = build_pipeline()
|
||||
clip = _silence(1.0)
|
||||
|
||||
for _ in range(3):
|
||||
pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
|
||||
# A different clip is a new entry, the first one still hits.
|
||||
other = (np.full_like(clip[0], 0.25), clip[1])
|
||||
pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1))
|
||||
pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
|
||||
|
||||
assert len(pipeline.vae.encode_calls) == 2
|
||||
assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"]
|
||||
|
||||
def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline):
|
||||
pipeline, _ = build_pipeline()
|
||||
for _ in range(2):
|
||||
pipeline.forward(
|
||||
_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)
|
||||
)
|
||||
assert len(pipeline.vae.encode_calls) == 2
|
||||
|
||||
pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0})
|
||||
for _ in range(2):
|
||||
pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1))
|
||||
assert len(pipeline.vae.encode_calls) == 2
|
||||
|
||||
def test_reference_cache_is_bounded(self, build_pipeline):
|
||||
pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2})
|
||||
for level in (0.1, 0.2, 0.3):
|
||||
clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE)
|
||||
pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1))
|
||||
assert len(pipeline._ref_cache) == 2
|
||||
|
||||
def test_one_request_per_forward(self, build_pipeline):
|
||||
pipeline, _ = build_pipeline()
|
||||
batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1)
|
||||
|
||||
@@ -239,11 +239,11 @@ def test_constructor_rejects_unsupported_mllm_quantization(mock_dependencies, mo
|
||||
)
|
||||
def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_config, quantize_mllm):
|
||||
from transformers import AutoModel
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8Static128BlockSym
|
||||
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
|
||||
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
||||
|
||||
od_config = OmniDiffusionConfig(
|
||||
model="dummy-boogu",
|
||||
@@ -287,14 +287,15 @@ def test_constructor_routes_mllm_quantization(mock_dependencies, quantization_co
|
||||
mock_dependencies["mllm_config_loader"].assert_not_called()
|
||||
mock_dependencies["mllm_builder"].assert_not_called()
|
||||
assert [source.subfolder for source in pipeline.weights_sources] == ["transformer"]
|
||||
assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], DiffusionFp8Config)
|
||||
assert isinstance(mock_dependencies["transformer_cls"].call_args.kwargs["quant_config"], Fp8Config)
|
||||
|
||||
|
||||
def test_constructor_preserves_mllm_ignored_layers(mock_dependencies):
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
|
||||
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
|
||||
quant_config = DiffusionFp8Config(
|
||||
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
|
||||
|
||||
quant_config = Fp8Config(
|
||||
ignored_layers=["mllm.language_model.layers.0.self_attn.q_proj", "transformer.blocks.0.attn.to_q"]
|
||||
)
|
||||
od_config = OmniDiffusionConfig(
|
||||
|
||||
@@ -16,13 +16,17 @@ class _RecordingBlock(nn.Module):
|
||||
|
||||
|
||||
def test_flux2_propagates_online_quant_config_to_all_transformer_blocks(monkeypatch):
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
|
||||
from vllm_omni.diffusion.models.flux2 import flux2_transformer
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
monkeypatch.setattr(flux2_transformer, "Flux2TransformerBlock", _RecordingBlock)
|
||||
monkeypatch.setattr(flux2_transformer, "Flux2SingleTransformerBlock", _RecordingBlock)
|
||||
|
||||
quant_config = build_quant_config("fp8")
|
||||
quant_config = Fp8Config(
|
||||
is_checkpoint_fp8_serialized=False,
|
||||
activation_scheme="dynamic",
|
||||
)
|
||||
model = flux2_transformer.Flux2Transformer2DModel(
|
||||
in_channels=4,
|
||||
out_channels=4,
|
||||
|
||||
@@ -71,12 +71,12 @@ def _small_od_config(*, class_name=None):
|
||||
|
||||
def test_fp8_scope_and_prefix_propagation(monkeypatch):
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
is_layer_skipped,
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
|
||||
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
||||
|
||||
monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
|
||||
@@ -92,7 +92,7 @@ def test_fp8_scope_and_prefix_propagation(monkeypatch):
|
||||
"blocks.0.adaln_proj.linear",
|
||||
"final_layer.adaln_proj.linear",
|
||||
}
|
||||
fp8_config = DiffusionFp8Config(ignored_layers=sorted(ignored_layers))
|
||||
fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
|
||||
model = h3.MiniMaxH3DiTModel(
|
||||
_small_od_config(),
|
||||
quant_config=fp8_config,
|
||||
|
||||
@@ -15,74 +15,6 @@ def test_build_quant_config_fp8():
|
||||
assert config is not None
|
||||
assert config.get_name() == "fp8"
|
||||
assert config.activation_scheme == "dynamic"
|
||||
assert not config.is_checkpoint_fp8_serialized
|
||||
|
||||
|
||||
def test_build_quant_config_fp8_uses_native_online_quantization():
|
||||
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
config = build_quant_config("fp8", ignored_layers=["proj_out"])
|
||||
assert isinstance(config, OnlineQuantizationConfig)
|
||||
assert config.args.linear.weight == kFp8StaticTensorSym
|
||||
assert config.args.moe.weight == kFp8StaticTensorSym
|
||||
assert config.ignored_layers == ["proj_out"]
|
||||
|
||||
|
||||
def test_build_quant_config_upstream_online_fp8_shorthand():
|
||||
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import kFp8StaticTensorSym
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
config = build_quant_config("fp8_per_tensor", ignore=["proj_out"])
|
||||
assert isinstance(config, OnlineQuantizationConfig)
|
||||
assert config.args.linear.weight == kFp8StaticTensorSym
|
||||
assert config.ignored_layers == ["proj_out"]
|
||||
|
||||
|
||||
def test_online_fp8_rejects_static_activation():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="activation_scheme='dynamic'"):
|
||||
build_quant_config("fp8", activation_scheme="static")
|
||||
|
||||
|
||||
def test_build_online_quant_config_preserves_explicit_native_args():
|
||||
from vllm.config.quantization import QuantizationConfigArgs, QuantSpec
|
||||
from vllm.model_executor.layers.quantization.online.base import OnlineQuantizationConfig
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
args = QuantizationConfigArgs(linear=QuantSpec(weight="fp8_per_tensor_static"), ignore=["proj_out"])
|
||||
config = build_quant_config("online", args=args)
|
||||
assert isinstance(config, OnlineQuantizationConfig)
|
||||
assert config.args is args
|
||||
assert config.ignored_layers == ["proj_out"]
|
||||
|
||||
|
||||
def test_build_online_quant_config_requires_args():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
with pytest.raises(ValueError, match="requires quantization config arguments"):
|
||||
build_quant_config("online")
|
||||
|
||||
|
||||
def test_serialized_checkpoint_replaces_online_fp8_config():
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
|
||||
|
||||
config = resolve_quant_config_from_disk(
|
||||
build_quant_config("fp8"),
|
||||
{"quant_method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
|
||||
)
|
||||
assert isinstance(config, Fp8Config)
|
||||
assert config.is_checkpoint_fp8_serialized
|
||||
assert config.activation_scheme == "static"
|
||||
|
||||
|
||||
def test_build_quant_config_none():
|
||||
@@ -107,7 +39,7 @@ def test_build_quant_config_invalid():
|
||||
def test_build_quant_config_dict():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
config = build_quant_config({"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"})
|
||||
config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
|
||||
assert config is not None
|
||||
assert config.get_name() == "fp8"
|
||||
assert config.activation_scheme == "static"
|
||||
@@ -116,7 +48,7 @@ def test_build_quant_config_dict():
|
||||
def test_build_quant_config_dict_not_mutated():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
original = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
|
||||
original = {"method": "fp8", "activation_scheme": "static"}
|
||||
copy = original.copy()
|
||||
build_quant_config(original)
|
||||
assert original == copy
|
||||
@@ -167,7 +99,7 @@ def test_build_quant_config_per_component_inner_dict_not_mutated():
|
||||
"""Inner component dicts should not be mutated by build_quant_config."""
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
inner = {"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"}
|
||||
inner = {"method": "fp8", "activation_scheme": "static"}
|
||||
original = inner.copy()
|
||||
build_quant_config({"transformer": inner, "vae": None})
|
||||
assert inner == original
|
||||
@@ -183,18 +115,20 @@ def test_flat_dict_not_misdetected_as_per_component():
|
||||
|
||||
|
||||
def test_build_quant_config_passthrough():
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
|
||||
fp8 = DiffusionFp8Config()
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
|
||||
assert build_quant_config(fp8) is fp8
|
||||
|
||||
|
||||
def test_component_config_routing():
|
||||
from vllm_omni.quantization import ComponentQuantizationConfig
|
||||
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
|
||||
fp8 = DiffusionFp8Config()
|
||||
from vllm_omni.quantization import ComponentQuantizationConfig
|
||||
|
||||
fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
|
||||
config = ComponentQuantizationConfig(
|
||||
component_configs={"transformer": fp8, "vae": None},
|
||||
)
|
||||
@@ -206,10 +140,11 @@ def test_component_config_routing():
|
||||
|
||||
|
||||
def test_component_config_with_default():
|
||||
from vllm_omni.quantization import ComponentQuantizationConfig
|
||||
from vllm_omni.quantization.fp8_config import DiffusionFp8Config
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
|
||||
fp8 = DiffusionFp8Config()
|
||||
from vllm_omni.quantization import ComponentQuantizationConfig
|
||||
|
||||
fp8 = Fp8Config(is_checkpoint_fp8_serialized=False, activation_scheme="dynamic")
|
||||
config = ComponentQuantizationConfig(
|
||||
component_configs={"vae": None},
|
||||
default_config=fp8,
|
||||
@@ -232,7 +167,7 @@ def test_integration_dict():
|
||||
|
||||
config = OmniDiffusionConfig(
|
||||
model="test",
|
||||
quantization_config={"method": "fp8", "is_checkpoint_fp8_serialized": True, "activation_scheme": "static"},
|
||||
quantization_config={"method": "fp8", "activation_scheme": "static"},
|
||||
)
|
||||
assert config.quantization_config is not None
|
||||
assert config.quantization_config.get_name() == "fp8"
|
||||
|
||||
@@ -1,72 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Exercise online FP8 loading and inference through the native CUDA path."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
from vllm.config import set_current_vllm_config
|
||||
from vllm.distributed import (
|
||||
destroy_distributed_environment,
|
||||
destroy_model_parallel,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
|
||||
from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.torch_utils import set_default_torch_dtype
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("method", ["fp8", "fp8_per_tensor"])
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_online_fp8_loads_bf16_weights_and_runs_native_kernel(method):
|
||||
ignored_key = "ignored_layers" if method == "fp8" else "ignore"
|
||||
config = build_quant_config(method, **{ignored_key: ["skip"]})
|
||||
od_config = OmniDiffusionConfig(model="test", dtype=torch.bfloat16, quantization_config=config)
|
||||
vllm_config = create_diffusion_vllm_config(torch.device("cuda:0"), od_config)
|
||||
|
||||
try:
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
local_rank=0,
|
||||
distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
|
||||
backend="nccl",
|
||||
)
|
||||
with set_current_vllm_config(vllm_config), set_default_torch_dtype(torch.bfloat16), torch.device("cuda:0"):
|
||||
initialize_model_parallel(tensor_model_parallel_size=1, pipeline_model_parallel_size=1)
|
||||
layer = ReplicatedLinear(
|
||||
128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="linear", disable_tp=True
|
||||
)
|
||||
assert isinstance(layer.quant_method, Fp8PerTensorOnlineLinearMethod)
|
||||
assert layer.weight.is_meta
|
||||
|
||||
generator = torch.Generator(device="cuda").manual_seed(42)
|
||||
weights = torch.randn(128, 128, generator=generator, dtype=torch.bfloat16) / 128**0.5
|
||||
inputs = torch.randn(32, 128, generator=generator, dtype=torch.bfloat16)
|
||||
expected = F.linear(inputs, weights).float()
|
||||
|
||||
# This is the loader installed by the native online quantizer.
|
||||
# It materializes and quantizes the BF16 checkpoint in one pass.
|
||||
layer.weight.weight_loader(layer.weight, weights)
|
||||
assert layer.weight.dtype == torch.float8_e4m3fn
|
||||
actual = layer(inputs)[0].float()
|
||||
relative_mse = (actual - expected).square().mean() / expected.square().mean()
|
||||
assert relative_mse.item() < 0.002
|
||||
|
||||
ignored = ReplicatedLinear(
|
||||
128, 128, bias=False, params_dtype=torch.bfloat16, quant_config=config, prefix="skip", disable_tp=True
|
||||
)
|
||||
assert isinstance(ignored.quant_method, UnquantizedLinearMethod)
|
||||
assert ignored.weight.dtype == torch.bfloat16
|
||||
finally:
|
||||
destroy_model_parallel()
|
||||
destroy_distributed_environment()
|
||||
@@ -11,12 +11,10 @@ import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
|
||||
from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
|
||||
from tests.helpers.omni_scheduler import bind_omits_transfer_helpers
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
|
||||
@@ -97,30 +95,6 @@ def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatc
|
||||
assert not adapter.code_prompt_token_ids
|
||||
|
||||
|
||||
def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender):
|
||||
adapter, connector = shm_sender
|
||||
ext_id = f"release_{uuid.uuid4().hex}"
|
||||
request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id)
|
||||
sibling_key = f"{ext_id}_1_0_0"
|
||||
assert connector.put("0", "1", sibling_key, "sibling")[0]
|
||||
for _ in range(3):
|
||||
adapter.save_async(None, request)
|
||||
adapter._send_single_request(adapter._pending_save_reqs.popleft())
|
||||
# The consumer read only the first chunk before its own request ended.
|
||||
assert connector.get("0", "1", f"{ext_id}_0_0") is not None
|
||||
|
||||
adapter.release_shm_resources(ext_id)
|
||||
# The release is queued behind sends; nothing is unlinked on the caller.
|
||||
assert f"{ext_id}_0_1" in connector._pending_keys
|
||||
while adapter._pending_save_reqs:
|
||||
adapter._send_single_request(adapter._pending_save_reqs.popleft())
|
||||
|
||||
assert adapter._reclaimed_shm_total == 2
|
||||
assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys)
|
||||
assert connector.get("0", "1", f"{ext_id}_0_1") is None
|
||||
assert connector.get("0", "1", sibling_key)[0] == "sibling"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fail_terminal", [False, True])
|
||||
def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal):
|
||||
adapter, connector = shm_sender
|
||||
@@ -636,45 +610,6 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
|
||||
return _build
|
||||
|
||||
|
||||
@pytest.mark.parametrize("active_stream_window", [0, 2])
|
||||
def test_scheduler_chunk_gate_preserves_kv_holding_queue(build_adapter, active_stream_window):
|
||||
adapter, _ = build_adapter(active_stream_window=active_stream_window)
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.running = []
|
||||
scheduler.chunk_transfer_adapter = adapter
|
||||
scheduler.input_coordinator = None
|
||||
holder = Request("holder", [1, 2, 3], SamplingParams(max_tokens=1), pooling_params=None)
|
||||
holder.num_computed_tokens = 2
|
||||
fresh = Request("fresh", [1], SamplingParams(max_tokens=1), pooling_params=None)
|
||||
scheduler.requests = {request.request_id: request for request in (holder, fresh)}
|
||||
scheduler._enqueue_waiting_request(fresh)
|
||||
scheduler._enqueue_waiting_request(holder)
|
||||
|
||||
scheduler._process_pending_omni_inputs(model_mode="ar")
|
||||
|
||||
assert not scheduler.waiting and not scheduler.kv_holding_waiting
|
||||
assert list(adapter.waiting_for_chunk_waiting_requests) == [holder, fresh]
|
||||
assert holder.status == fresh.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
scheduler._restore_omni_wait_queues()
|
||||
assert list(scheduler.kv_holding_waiting) == [holder]
|
||||
assert list(scheduler.waiting) == [fresh]
|
||||
|
||||
adapter._finished_load_reqs.add(holder.request_id)
|
||||
scheduler._process_pending_omni_inputs(model_mode="ar")
|
||||
|
||||
assert holder.status == RequestStatus.WAITING
|
||||
assert fresh.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert list(scheduler.kv_holding_waiting) == [holder]
|
||||
assert not scheduler.waiting
|
||||
assert adapter.requests_with_ready_chunks == {holder.request_id}
|
||||
scheduler._restore_omni_wait_queues()
|
||||
assert list(scheduler.kv_holding_waiting) == [holder]
|
||||
assert list(scheduler.waiting) == [fresh]
|
||||
|
||||
|
||||
def _dequeue_load_entry(adapter, request):
|
||||
"""Register and take the entry as the background recv loop would."""
|
||||
adapter.load_async(request)
|
||||
@@ -2485,8 +2420,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
|
||||
scheduler.requests = {request.request_id: request}
|
||||
scheduler.running = [request] if placement == "running" else []
|
||||
scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
|
||||
scheduler.kv_holding_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
|
||||
scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
|
||||
|
||||
freed = []
|
||||
@@ -2510,7 +2444,7 @@ def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
|
||||
assert not adapter.waiting_for_chunk_running_requests
|
||||
assert scheduler.running == []
|
||||
assert list(scheduler.waiting) == []
|
||||
assert list(scheduler.kv_holding_waiting) == []
|
||||
assert list(scheduler.skipped_waiting) == []
|
||||
assert scheduler.num_waiting_for_streaming_input == 0
|
||||
|
||||
fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
|
||||
@@ -2533,8 +2467,7 @@ def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapt
|
||||
scheduler.requests = {request.request_id: request}
|
||||
scheduler.running = []
|
||||
scheduler.waiting = DummyWaitingQueue()
|
||||
scheduler.kv_holding_waiting = DummyWaitingQueue()
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = DummyWaitingQueue()
|
||||
scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
|
||||
|
||||
assert (
|
||||
@@ -2608,7 +2541,6 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
|
||||
|
||||
scheduler = mocker.MagicMock()
|
||||
scheduler._first_chunk_express = False
|
||||
scheduler.aux_output_connector = None
|
||||
scheduler._express_min_slack_s = 0.0
|
||||
scheduler.chunk_transfer_adapter = adapter_mock
|
||||
scheduler.connector = None
|
||||
@@ -2688,10 +2620,10 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
|
||||
adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
|
||||
adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
|
||||
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
|
||||
scheduler = mocker.MagicMock()
|
||||
scheduler.aux_output_connector = None
|
||||
scheduler.chunk_transfer_adapter = adapter_mock
|
||||
bind_omits_transfer_helpers(scheduler)
|
||||
scheduler.connector = None
|
||||
scheduler.perf_metrics = None
|
||||
scheduler.log_stats = False
|
||||
@@ -2788,8 +2720,6 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
|
||||
sched.requests = {}
|
||||
sched.running = []
|
||||
sched.waiting = []
|
||||
sched.kv_holding_waiting = []
|
||||
sched.deferred_waiting = set()
|
||||
|
||||
with patch.object(VLLMScheduler, "finish_requests", _super_finish):
|
||||
OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
|
||||
@@ -2805,8 +2735,7 @@ def _parked_sender_scheduler(adapter, session):
|
||||
scheduler.requests = {session.request_id: session}
|
||||
scheduler.running = []
|
||||
scheduler.waiting = DummyWaitingQueue()
|
||||
scheduler.kv_holding_waiting = DummyWaitingQueue([session])
|
||||
scheduler.deferred_waiting = set()
|
||||
scheduler.skipped_waiting = DummyWaitingQueue([session])
|
||||
scheduler.num_waiting_for_streaming_input = 1
|
||||
scheduler.finished_req_ids_dict = {}
|
||||
|
||||
|
||||
@@ -213,7 +213,7 @@ class TestCleanup:
|
||||
connector.put("s0", "s1", "cleanup_req_42", data)
|
||||
assert "cleanup_req_42" in connector._pending_keys
|
||||
|
||||
assert connector.cleanup("cleanup_req_42") is True
|
||||
connector.cleanup("cleanup_req_42")
|
||||
assert "cleanup_req_42" not in connector._pending_keys
|
||||
|
||||
result = connector.get("s0", "s1", "cleanup_req_42", metadata=None)
|
||||
@@ -224,7 +224,7 @@ class TestCleanup:
|
||||
connector.put("s0", "s1", "consumed_req_99", data)
|
||||
connector.get("s0", "s1", "consumed_req_99", metadata=None)
|
||||
|
||||
assert connector.cleanup("consumed_req_99") is False
|
||||
connector.cleanup("consumed_req_99")
|
||||
assert "consumed_req_99" not in connector._pending_keys
|
||||
|
||||
def test_close_cleans_all_pending(self, connector):
|
||||
@@ -296,24 +296,6 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch):
|
||||
assert key not in connector._pending_keys
|
||||
|
||||
|
||||
@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_0", "_suffix"])
|
||||
def test_cleanup_prefix_only_matches_own_chunk_keys(connector, suffix):
|
||||
request_id = f"prefix_{uuid.uuid4().hex}"
|
||||
own_keys = [f"{request_id}_0_{i}" for i in range(3)]
|
||||
sibling_key = f"{request_id}{suffix}_0_0"
|
||||
for key in own_keys:
|
||||
assert connector.put("0", "1", key, key)[0]
|
||||
assert connector.put("0", "1", sibling_key, "sibling")[0]
|
||||
# A consumed chunk is no longer counted as reclaimed.
|
||||
assert connector.get("0", "1", own_keys[0])[0] == own_keys[0]
|
||||
|
||||
assert connector.cleanup_prefix(f"{request_id}_0_") == 2
|
||||
for key in own_keys:
|
||||
assert key not in connector._pending_keys
|
||||
assert connector.get("0", "1", key) is None
|
||||
assert connector.get("0", "1", sibling_key)[0] == "sibling"
|
||||
|
||||
|
||||
# ── Arrival wakeups ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -1,26 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sana_model_parallel(tmp_path):
|
||||
"""Construct native parallel layers with real single-rank GPU groups."""
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
destroy_distributed_env,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
local_rank=0,
|
||||
distributed_init_method=f"file://{tmp_path / 'sana-distributed'}",
|
||||
)
|
||||
try:
|
||||
initialize_model_parallel()
|
||||
yield
|
||||
finally:
|
||||
destroy_distributed_env()
|
||||
@@ -202,7 +202,6 @@ def test_sana_video_transformer_matches_frozen_golden(
|
||||
variant: str,
|
||||
_hardware,
|
||||
accuracy_artifact_root: Path,
|
||||
sana_model_parallel,
|
||||
) -> None:
|
||||
del _hardware
|
||||
from vllm.config import LoadConfig
|
||||
|
||||
@@ -189,7 +189,6 @@ def _cfg_prediction(
|
||||
def test_sana_video_quick_pipeline_stages_match_diffusers(
|
||||
variant: str,
|
||||
_hardware,
|
||||
sana_model_parallel,
|
||||
) -> None:
|
||||
"""Compare the first denoising step without replacing native attention.
|
||||
|
||||
|
||||
@@ -1,175 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""AuK encoder stage: FULL CUDA graph replay matches eager, request by request.
|
||||
|
||||
``auk.yaml`` captures the encoder prefill as FULL CUDA graphs padded to a few
|
||||
token buckets. A replay reuses static input buffers and the graph pool, so a
|
||||
request could in principle read conditioning left behind by an earlier, longer
|
||||
prompt. This test runs the same request sequence through an eager encoder and
|
||||
through the graphed one and compares the text condition stage 0 hands to
|
||||
stage 1 (the fused thinker hidden states) for every request. The sequence
|
||||
alternates prompts of different lengths inside one capture bucket, crosses
|
||||
into a larger bucket, and then returns to the first prompt.
|
||||
|
||||
Like ``test_auk.py`` it needs an assembled checkpoint directory:
|
||||
|
||||
VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk_encoder_cudagraph.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import bisect
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.model_executor.stage_input_processors import auk as auk_input_processor
|
||||
from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params
|
||||
|
||||
MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR"
|
||||
|
||||
# Text-only instruct prompts: the token count, and so the capture bucket,
|
||||
# follows the text length. SHORT_A and SHORT_B differ in length but share a
|
||||
# bucket; LONG lands in a larger one. The test checks these properties on the
|
||||
# recorded condition lengths, so a tokenizer change fails loudly instead of
|
||||
# silently shrinking the coverage.
|
||||
_TEMPLATE = 'Generate speech based on the following description: "{desc}". The content to speak is: "{text}".'
|
||||
SHORT_A = _TEMPLATE.format(
|
||||
desc="A calm young woman speaking warmly.",
|
||||
text="Welcome back, how was your day? I kept some soup warm for you, and the kettle is on if you want tea.",
|
||||
)
|
||||
SHORT_B = _TEMPLATE.format(
|
||||
desc="An older man with a deep, slightly raspy voice, speaking slowly.",
|
||||
text=(
|
||||
"The train to the coast leaves at seven, so please be on time. Bring a warm coat, "
|
||||
"some cash for the ferry, and the blue umbrella from the hall, because the forecast says rain."
|
||||
),
|
||||
)
|
||||
LONG = _TEMPLATE.format(
|
||||
desc=(
|
||||
"A cheerful radio host in her thirties, bright and energetic, with clear articulation, "
|
||||
"a slight smile in her voice, and a brisk but relaxed pace, recorded close to the microphone"
|
||||
),
|
||||
text=(
|
||||
"Good morning and welcome to the weekend edition. Today we are talking about small gardens, "
|
||||
"the plants that thrive on a sunny balcony, and the neighbours who turned an empty lot into "
|
||||
"a place where the whole street now meets on Saturday afternoons to trade seedlings and stories. "
|
||||
"Stay with us, because after the news we will open the phone lines and hear from you."
|
||||
),
|
||||
)
|
||||
SEQUENCE = [SHORT_A, SHORT_B, SHORT_A, LONG, SHORT_B, SHORT_A]
|
||||
|
||||
# The DiT output is not inspected; keep stage 1 as cheap as it goes (it also
|
||||
# runs eager in both configs, see _deploy).
|
||||
GEN_SECONDS = 1.0
|
||||
NFE = 1
|
||||
|
||||
# bf16 hidden states. Graph replay pads the prefill to the bucket size but runs
|
||||
# the same kernels on the real tokens, so the two paths agree to bf16 rounding.
|
||||
# Measured on H20: cosine >= 0.99998 and max-abs error <= 2.9e-3 of the
|
||||
# condition's max magnitude for every request; the bounds leave about 3x
|
||||
# headroom for other GPUs.
|
||||
MIN_COSINE = 0.9999
|
||||
MAX_REL_ABS = 1e-2
|
||||
|
||||
_model_dir = os.environ.get(MODEL_DIR_ENV)
|
||||
pytestmark = [
|
||||
pytest.mark.slow,
|
||||
pytest.mark.tts,
|
||||
pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"),
|
||||
]
|
||||
|
||||
|
||||
def _deploy(tmp_path: Path, *, graph_encoder: bool) -> str:
|
||||
"""``auk.yaml`` with stage 1 eager, and the encoder stage eager or graphed as shipped.
|
||||
|
||||
Only the encoder differs between the two runs. The conditions compared here
|
||||
are recorded before stage 1 runs, so its DiT compile, graph capture and
|
||||
codec bucket warm-up would only add startup time; eager stage 1 skips them
|
||||
and still takes the real stage handoff.
|
||||
"""
|
||||
with open(get_deploy_config_path("auk.yaml")) as f:
|
||||
deploy = yaml.safe_load(f)
|
||||
for stage in deploy["stages"]:
|
||||
if stage["stage_id"] == 1:
|
||||
stage["enforce_eager"] = True
|
||||
elif not graph_encoder:
|
||||
stage["enforce_eager"] = True
|
||||
stage.pop("compilation_config", None)
|
||||
path = tmp_path / f"auk_{'graph' if graph_encoder else 'eager'}_encoder.yaml"
|
||||
path.write_text(yaml.safe_dump(deploy))
|
||||
return str(path)
|
||||
|
||||
|
||||
def _graph_capture_sizes() -> list[int]:
|
||||
with open(get_deploy_config_path("auk.yaml")) as f:
|
||||
deploy = yaml.safe_load(f)
|
||||
stage = next(s for s in deploy["stages"] if s["stage_id"] == 0)
|
||||
assert stage.get("enforce_eager") is False, "auk.yaml no longer graphs the encoder; update this test"
|
||||
return sorted(stage["compilation_config"]["cudagraph_capture_sizes"])
|
||||
|
||||
|
||||
def _run_sequence(deploy_config: str, monkeypatch: pytest.MonkeyPatch) -> list[torch.Tensor]:
|
||||
"""Run SEQUENCE one request at a time and return each request's text condition."""
|
||||
conditions: list[torch.Tensor] = []
|
||||
original = auk_input_processor.encoder2dit
|
||||
|
||||
def recording_encoder2dit(*args: Any, **kwargs: Any) -> dict[str, Any] | None:
|
||||
result = original(*args, **kwargs)
|
||||
if result is not None:
|
||||
conditions.append(result["prompt_embeds"].detach().to("cpu", torch.float32).clone())
|
||||
return result
|
||||
|
||||
# The orchestrator runs in this process and resolves the stage input
|
||||
# function by module attribute at startup, so the patch must precede Omni.
|
||||
with monkeypatch.context() as patch:
|
||||
patch.setattr(auk_input_processor, "encoder2dit", recording_encoder2dit)
|
||||
with OmniRunner(str(Path(_model_dir).resolve()), deploy_config=deploy_config) as runner:
|
||||
for text in SEQUENCE:
|
||||
prompt = auk_prompt(text, None, gen_seconds=GEN_SECONDS)
|
||||
outputs = runner.omni.generate(prompt, auk_sampling_params(nfe=NFE, cfg=2.0, seed=0))
|
||||
assert outputs, "no outputs returned"
|
||||
assert len(conditions) == len(SEQUENCE), f"recorded {len(conditions)} conditions for {len(SEQUENCE)} requests"
|
||||
return conditions
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_encoder_graph_replay_matches_eager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
capture_sizes = _graph_capture_sizes()
|
||||
eager = _run_sequence(_deploy(tmp_path, graph_encoder=False), monkeypatch)
|
||||
graphed = _run_sequence(_deploy(tmp_path, graph_encoder=True), monkeypatch)
|
||||
|
||||
lengths = {text: cond.shape[0] for text, cond in zip(SEQUENCE, eager)}
|
||||
for i, (ref, got) in enumerate(zip(eager, graphed)):
|
||||
assert got.shape == ref.shape, f"request {i}: shape {tuple(got.shape)} != eager {tuple(ref.shape)}"
|
||||
assert torch.isfinite(got).all(), f"request {i}: non-finite condition"
|
||||
cosine = torch.nn.functional.cosine_similarity(got.flatten(), ref.flatten(), dim=0).item()
|
||||
rel_abs = ((got - ref).abs().max() / ref.abs().max()).item()
|
||||
assert cosine >= MIN_COSINE and rel_abs <= MAX_REL_ABS, (
|
||||
f"request {i} ({lengths[SEQUENCE[i]]} tokens): cosine {cosine:.6f}, max-abs error {rel_abs:.2e} "
|
||||
"of the condition's magnitude"
|
||||
)
|
||||
|
||||
# The sequence must exercise what it claims to: two different lengths in
|
||||
# one bucket, a second bucket, and every prompt inside the graphed range.
|
||||
assert max(lengths.values()) <= capture_sizes[-1], f"{lengths} exceeds the largest capture size"
|
||||
bucket = {text: capture_sizes[bisect.bisect_left(capture_sizes, n)] for text, n in lengths.items()}
|
||||
assert lengths[SHORT_A] != lengths[SHORT_B] and bucket[SHORT_A] == bucket[SHORT_B], (lengths, bucket)
|
||||
assert bucket[LONG] > bucket[SHORT_A], (lengths, bucket)
|
||||
|
||||
# Repeats of a prompt reproduce its first condition exactly, eager or graphed:
|
||||
# nothing from the requests in between leaks into the replay.
|
||||
for run in (eager, graphed):
|
||||
first: dict[str, torch.Tensor] = {}
|
||||
for text, cond in zip(SEQUENCE, run):
|
||||
if text in first:
|
||||
torch.testing.assert_close(cond, first[text], atol=0, rtol=0)
|
||||
else:
|
||||
first[text] = cond
|
||||
@@ -23,8 +23,7 @@ BASE_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"})
|
||||
def _get_diffusion_feature_cases(model: str):
|
||||
return [
|
||||
pytest.param(
|
||||
# The unsharded transformer alone exceeds the L4's 22 GiB.
|
||||
OmniServerParams(model=model, server_args=["--enable-layerwise-offload"]),
|
||||
OmniServerParams(model=model),
|
||||
id="base",
|
||||
marks=BASE_FEATURE_MARKS,
|
||||
),
|
||||
|
||||
@@ -28,7 +28,7 @@ import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.advanced_model,
|
||||
@@ -84,12 +84,7 @@ tts_server_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=modify_stage_config(
|
||||
get_deploy_config_path("moss_tts_realtime.yaml"),
|
||||
# Leave room on the single L4 for the codec weights, reference
|
||||
# encoder, and speech-serving activations in addition to KV.
|
||||
updates={"stages": {0: {"gpu_memory_utilization": 0.35}}},
|
||||
),
|
||||
stage_config_path=get_deploy_config_path("moss_tts_realtime.yaml"),
|
||||
server_args=["--disable-log-stats"],
|
||||
),
|
||||
id="moss_tts_realtime",
|
||||
|
||||
@@ -21,14 +21,7 @@ models = ["Qwen/Qwen2.5-Omni-7B"]
|
||||
|
||||
# Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
|
||||
# platforms: section in vllm_omni/deploy/ci/qwen2_5_omni.yaml.
|
||||
stage_configs = [
|
||||
modify_stage_config(
|
||||
get_deploy_config_path("ci/qwen2_5_omni.yaml"),
|
||||
# This suite requests two cards. Share the Talker's card with
|
||||
# Code2Wav and leave room for its whole-sequence DiT workspace.
|
||||
updates={"stages": {1: {"gpu_memory_utilization": 0.3}, 2: {"devices": "1"}}},
|
||||
)
|
||||
]
|
||||
stage_configs = [modify_stage_config(get_deploy_config_path("ci/qwen2_5_omni.yaml"))]
|
||||
|
||||
# Create parameter combinations for model and stage config
|
||||
test_params = [
|
||||
|
||||
@@ -1,155 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Live /dev/shm leak check for Qwen3-Omni ``async_chunk`` (stage-0-final).
|
||||
|
||||
The default L2 Qwen3-Omni job uses ``ci/qwen3_omni_moe.yaml`` with
|
||||
``async_chunk: false`` plus ``--no-async-chunk``, so it never exercises the
|
||||
thinker→talker SharedMemoryConnector path.
|
||||
|
||||
A text-only request is tagged ``omni_final_stage_id=0``, so nothing downstream
|
||||
ever ``get()``s what stage-0 ``put()``s for it. Before #7245 the scheduler still
|
||||
``put()`` a finished marker on finish and nothing unlinked it, so every such
|
||||
request left a ``SharedMemoryConnector`` segment plus lockfile in ``/dev/shm``.
|
||||
Stage-0-final finishes now skip the ``put()`` and the orchestrator reclaims any
|
||||
undrained segment on request cleanup.
|
||||
|
||||
This module starts the production async-chunk deploy and asserts the three
|
||||
text-only request ids leave no connector lockfile or POSIX segment in
|
||||
``/dev/shm``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
|
||||
_MODEL = os.environ.get("VLLM_OMNI_TEST_MODEL", "Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
_SHM_DIR = Path("/dev/shm")
|
||||
_NUM_TEXT_REQUESTS = 3
|
||||
_SETTLE_S = 15.0
|
||||
|
||||
# CI overlay keeps seq/token limits small; flip async_chunk back on so stage-0
|
||||
# still uses SharedMemoryConnector. Production yaml is already async_chunk:true.
|
||||
_ASYNC_CHUNK_DEPLOY = modify_stage_config(
|
||||
get_deploy_config_path("ci/qwen3_omni_moe.yaml"),
|
||||
updates={"async_chunk": True},
|
||||
)
|
||||
|
||||
test_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=_MODEL,
|
||||
stage_config_path=_ASYNC_CHUNK_DEPLOY,
|
||||
use_stage_cli=True,
|
||||
),
|
||||
id="async_chunk",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _system_prompt() -> dict:
|
||||
return {
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"You are Qwen, a virtual human developed by the Qwen Team, "
|
||||
"Alibaba Group, capable of perceiving auditory and visual inputs, "
|
||||
"as well as generating text and speech."
|
||||
),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _shm_used_mib() -> str:
|
||||
try:
|
||||
st = os.statvfs(_SHM_DIR)
|
||||
used = (st.f_blocks - st.f_bavail) * st.f_frsize
|
||||
return f"{used / (1024 * 1024):.1f}MiB"
|
||||
except OSError:
|
||||
return "unknown"
|
||||
|
||||
|
||||
def _is_entry_for_request(name: str, request_id: str) -> bool:
|
||||
"""Match a segment or ``shm_`` lockfile keyed by the ``{request_id}-{suffix}`` orchestrator id."""
|
||||
return name.removeprefix("shm_").startswith((f"{request_id}-", f"{request_id}_"))
|
||||
|
||||
|
||||
def _entries_for_request_ids(request_ids: list[str]) -> list[str]:
|
||||
if not _SHM_DIR.is_dir():
|
||||
pytest.skip("/dev/shm is not available")
|
||||
found: list[str] = []
|
||||
with os.scandir(_SHM_DIR) as entries:
|
||||
for entry in entries:
|
||||
if any(_is_entry_for_request(entry.name, request_id) for request_id in request_ids):
|
||||
found.append(entry.name)
|
||||
return found
|
||||
|
||||
|
||||
def _wait_no_entries_for_request_ids(request_ids: list[str], timeout_s: float = _SETTLE_S) -> list[str]:
|
||||
"""Reclaim is fire-and-forget; poll until these ids are gone or *timeout_s*."""
|
||||
deadline = time.monotonic() + timeout_s
|
||||
time.sleep(1.0)
|
||||
leftover: list[str] = _entries_for_request_ids(request_ids)
|
||||
while leftover:
|
||||
if time.monotonic() >= deadline:
|
||||
return leftover
|
||||
time.sleep(2.0)
|
||||
leftover = _entries_for_request_ids(request_ids)
|
||||
return leftover
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.")
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_only_async_chunk_does_not_leak_shm(omni_server, online_client) -> None:
|
||||
"""Stage-0-final (text-only) requests must not leave connector SHM behind."""
|
||||
print(f"[shm-leak] before requests: /dev/shm_used={_shm_used_mib()}", flush=True)
|
||||
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=_system_prompt(),
|
||||
content_text="What is the capital of China? Answer in 20 words.",
|
||||
)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"modalities": ["text"],
|
||||
"key_words": {"text": ["beijing"]},
|
||||
}
|
||||
responses = online_client.send_omni_request(request_config, request_num=_NUM_TEXT_REQUESTS)
|
||||
produced = [resp.text_content for resp in responses if getattr(resp, "text_content", None)]
|
||||
assert len(produced) == _NUM_TEXT_REQUESTS, (
|
||||
f"Need {_NUM_TEXT_REQUESTS} text completions to exercise stage-0 finish; "
|
||||
f"got {len(produced)} non-empty texts. A preprocess/engine miss would "
|
||||
f"leave /dev/shm unchanged and hide a leak."
|
||||
)
|
||||
request_ids = [resp.request_id for resp in responses]
|
||||
assert all(request_ids) and len(request_ids) == _NUM_TEXT_REQUESTS, (
|
||||
f"Need chat completion ids to assert per-request SHM cleanup; got {request_ids}"
|
||||
)
|
||||
|
||||
leftover = _wait_no_entries_for_request_ids(request_ids)
|
||||
print(
|
||||
f"[shm-leak] after {_NUM_TEXT_REQUESTS} text-only requests: "
|
||||
f"/dev/shm_used={_shm_used_mib()} leftover={leftover} request_ids={request_ids}",
|
||||
flush=True,
|
||||
)
|
||||
assert not leftover, (
|
||||
f"SharedMemoryConnector leaked /dev/shm entries for {leftover} after stage-0-final request(s) {request_ids}"
|
||||
)
|
||||
@@ -20,11 +20,7 @@ import pytest
|
||||
from omegaconf import OmegaConf
|
||||
from pydantic import ValidationError
|
||||
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.exceptions import VLLMValidationError
|
||||
from vllm.v1.engine.input_processor import InputProcessor
|
||||
from vllm.v1.worker.gpu.sample.logit_bias import LogitBiasState
|
||||
|
||||
from tests.helpers.mock import patch_hf_snapshot_download
|
||||
from vllm_omni.config.model import OmniModelConfig
|
||||
@@ -36,36 +32,6 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorMode
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_v2", [False, True])
|
||||
def test_stage_runner_selects_matching_input_processor_validator(monkeypatch, use_v2):
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.model_executor.models.audex.cfg import AudexCFGLogitsProcessor
|
||||
|
||||
# Upstream's process-wide preference must not override a stage's deploy choice.
|
||||
monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", not use_v2)
|
||||
config = VllmConfig()
|
||||
model_config = object.__new__(OmniModelConfig)
|
||||
model_config.use_v2_model_runner = use_v2
|
||||
model_config.runner_type = "generate"
|
||||
model_config.logits_processors = [LogitBiasState if use_v2 else AudexCFGLogitsProcessor]
|
||||
config.model_config = model_config
|
||||
assert config.use_v2_model_runner is use_v2
|
||||
processor = object.__new__(InputProcessor)
|
||||
processor.vllm_config, processor.model_config = config, model_config
|
||||
validator = processor._build_logits_processors_params_validator()
|
||||
validator(SamplingParams(extra_args={"cfg_role": "cond", "cfg_scale": 1.5}))
|
||||
if not use_v2:
|
||||
with pytest.raises(VLLMValidationError, match="cfg_role"):
|
||||
validator(SamplingParams(extra_args={"cfg_role": "invalid"}))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_v2", [False, True])
|
||||
def test_plain_vllm_config_retains_upstream_runner_preference(monkeypatch, use_v2):
|
||||
monkeypatch.setattr("vllm.envs.VLLM_USE_V2_MODEL_RUNNER", use_v2)
|
||||
assert VllmConfig().use_v2_model_runner is use_v2
|
||||
|
||||
|
||||
def test_sync_config_is_omni():
|
||||
"""Ensure create_model_config gives the right type."""
|
||||
cfg = OmniEngineArgs().create_model_config()
|
||||
|
||||
@@ -8,7 +8,6 @@ from __future__ import annotations
|
||||
from dataclasses import fields
|
||||
|
||||
import pytest
|
||||
from vllm.config import AuxOutputConfig
|
||||
|
||||
from vllm_omni.engine.arg_utils import (
|
||||
SHARED_FIELDS,
|
||||
@@ -50,26 +49,3 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
|
||||
"""
|
||||
blacklist = internal_blacklist_keys()
|
||||
assert blacklist == orchestrator_field_names() - SHARED_FIELDS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"aux_args",
|
||||
[
|
||||
{"enable_return_routed_experts": True},
|
||||
{"aux_output_config": {"enable_return_routed_experts": True}},
|
||||
{"aux_output_config": AuxOutputConfig(enable_return_routed_experts=True)},
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("use_v2_model_runner", [False, True])
|
||||
@pytest.mark.parametrize("is_npu", [False, True])
|
||||
def test_aux_outputs_require_a_supported_omni_stage_runner(aux_args, use_v2_model_runner, is_npu, monkeypatch):
|
||||
# Upstream runner selection reads this environment variable; Omni workers
|
||||
# select their runner from the separate per-stage flag instead.
|
||||
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
|
||||
monkeypatch.setattr("vllm_omni.engine.arg_utils.current_omni_platform.is_npu", lambda: is_npu)
|
||||
if use_v2_model_runner or is_npu:
|
||||
args = OmniEngineArgs(model="unused", use_v2_model_runner=use_v2_model_runner, **aux_args)
|
||||
assert args.aux_output_config.enabled
|
||||
else:
|
||||
with pytest.raises(ValueError, match="use_v2_model_runner=True for this Omni stage"):
|
||||
OmniEngineArgs(model="unused", **aux_args)
|
||||
@@ -168,37 +168,6 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
|
||||
output_processor.add_request.assert_not_called()
|
||||
|
||||
|
||||
def test_build_add_request_message_preserves_kv_hints_after_omni_payload_upgrade(mocker):
|
||||
from vllm.v1.engine.input_processor import InputProcessor
|
||||
from vllm.v1.kv_hints import KvHintsEnvelope
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
engine.default_sampling_params_list = [params]
|
||||
engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
|
||||
engine.supported_tasks = ("generate",)
|
||||
processor = mocker.create_autospec(InputProcessor, instance=True)
|
||||
hints = KvHintsEnvelope("1", "hint-message", [])
|
||||
|
||||
def process(**kwargs):
|
||||
request = _make_engine_core_request()
|
||||
request.kv_hints = kwargs["kv_hints"]
|
||||
return request
|
||||
|
||||
processor.process_inputs.side_effect = process
|
||||
engine.input_processor = processor
|
||||
msg = engine._build_add_request_message(
|
||||
request_id="req-1",
|
||||
prompt={"prompt_token_ids": [1], "additional_information": {"speaker": "voice"}},
|
||||
sampling_params_list=[params],
|
||||
kv_hints=hints,
|
||||
)
|
||||
assert processor.process_inputs.call_args.kwargs["kv_hints"] is hints
|
||||
assert isinstance(msg.prompt, OmniEngineCoreRequest)
|
||||
assert msg.prompt.kv_hints is hints
|
||||
assert msg.prompt.additional_information is not None
|
||||
|
||||
|
||||
def test_build_add_request_message_injects_global_id_before_prompt_transform(mocker: MockerFixture):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
params = SamplingParams(max_tokens=8)
|
||||
|
||||
@@ -268,21 +268,11 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
|
||||
"""Regression test for forward references in make_async_mp_client."""
|
||||
import vllm_omni.engine.stage_engine_core_client as client_mod
|
||||
|
||||
# Snapshot the module namespace so a reload cannot strand other test
|
||||
# modules that captured the pre-reload class objects (isinstance checks
|
||||
# against the stale class would then fail). ``importlib.reload`` mutates
|
||||
# the module dict in place; restoring it leaves the runtime exactly as it
|
||||
# was before this test ran.
|
||||
saved = dict(client_mod.__dict__)
|
||||
try:
|
||||
importlib.reload(client_mod)
|
||||
importlib.reload(client_mod)
|
||||
|
||||
assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
|
||||
"StageEngineCoreClient | DPLBStageEngineCoreClient"
|
||||
)
|
||||
finally:
|
||||
client_mod.__dict__.clear()
|
||||
client_mod.__dict__.update(saved)
|
||||
assert client_mod.StageEngineCoreClientBase.make_async_mp_client.__annotations__["return"] == (
|
||||
"StageEngineCoreClient | DPLBStageEngineCoreClient"
|
||||
)
|
||||
|
||||
|
||||
def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch):
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Deterministic tests for the CFG-companion output lifecycle.
|
||||
|
||||
Regression suite for the companion-output race behind the nightly Bagel
|
||||
@@ -94,9 +94,6 @@ class _FakePool:
|
||||
def release_bindings(self, request_ids):
|
||||
pass
|
||||
|
||||
async def release_request_resources(self, request_ids):
|
||||
pass
|
||||
|
||||
|
||||
def _make_orchestrator(num_stages: int = 2) -> Orchestrator:
|
||||
orch = Orchestrator.__new__(Orchestrator)
|
||||
|
||||
@@ -13,16 +13,13 @@ substitutable with the base class — including base-style calls that pass
|
||||
|
||||
import inspect
|
||||
|
||||
import msgspec
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.kv_hints import KvHintAction, KvHintsEnvelope
|
||||
from vllm.v1.request import Request
|
||||
|
||||
from vllm_omni.engine import OmniEngineCoreRequest, PromptEmbedsPayload
|
||||
from vllm_omni.engine import PromptEmbedsPayload
|
||||
from vllm_omni.request import OmniRequest, OmniStreamingUpdate
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -114,26 +111,3 @@ def test_model_intermediate_buffer_round_trips_to_streaming_update():
|
||||
assert update is not None
|
||||
assert request.model_intermediate_buffer is buffer
|
||||
assert update.model_intermediate_buffer is buffer
|
||||
|
||||
|
||||
def test_kv_hints_survive_omni_upgrade_wire_transport_and_scheduler_conversion():
|
||||
hints = KvHintsEnvelope("1", "hint-message", [KvHintAction("a", "prefetch", "1", {"blocks": [7]})])
|
||||
core = EngineCoreRequest(
|
||||
request_id="hint-request",
|
||||
prompt_token_ids=[1, 2],
|
||||
mm_features=None,
|
||||
sampling_params=SamplingParams(max_tokens=2),
|
||||
pooling_params=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
session_id="session",
|
||||
kv_hints=hints,
|
||||
)
|
||||
upgraded = OmniEngineCoreRequest.from_request(core)
|
||||
decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(upgraded), type=OmniEngineCoreRequest)
|
||||
scheduled = OmniRequest.from_engine_core_request(decoded, block_hasher=None)
|
||||
assert decoded.session_id == "session"
|
||||
assert scheduled.kv_hints == hints
|
||||
assert scheduled.session_id == "session"
|
||||
@@ -2,7 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for OmniRequestOutput class."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
@@ -158,7 +157,6 @@ class TestMsgpackRoundTrip:
|
||||
)
|
||||
|
||||
source = _make_text_request_output(text="round trip")
|
||||
source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
|
||||
out = OmniRequestOutput.from_stage_output(
|
||||
source,
|
||||
request_id="rt-1",
|
||||
@@ -180,7 +178,6 @@ class TestMsgpackRoundTrip:
|
||||
assert decoded.outputs[0].text == "round trip"
|
||||
assert decoded.prompt == "test prompt"
|
||||
assert decoded.prompt_token_ids == [1, 2, 3]
|
||||
np.testing.assert_array_equal(decoded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
|
||||
|
||||
def test_round_trip_diffusion_output(self):
|
||||
"""A diffusion output round-trips without data loss."""
|
||||
@@ -305,16 +302,6 @@ class TestRequestOutputFieldParity:
|
||||
assert omni.num_cached_tokens is None
|
||||
assert omni.num_cache_creation_tokens is None
|
||||
|
||||
def test_fixed_token_scores_survive_stage_output_wrapping(self) -> None:
|
||||
source = _make_text_request_output()
|
||||
source.prompt_token_id_logprobs = np.array([[-0.5, -1.5], [-0.2, -2.2]], dtype=np.float32)
|
||||
|
||||
omni = OmniRequestOutput.from_stage_output(source, stage_id=0)
|
||||
forwarded = OmniRequestOutput.from_stage_output(omni, stage_id=1)
|
||||
|
||||
np.testing.assert_array_equal(forwarded.prompt_token_id_logprobs, source.prompt_token_id_logprobs)
|
||||
assert OmniRequestOutput(request_id="unscored").prompt_token_id_logprobs is None
|
||||
|
||||
def test_transfer_param_fields_are_copied_from_the_stage_output(self) -> None:
|
||||
"""Declaring the fields is not enough — they must also be copied.
|
||||
|
||||
|
||||
@@ -1282,7 +1282,6 @@ async def test_cache_reset_rpc_failure_is_nonfatal_and_visits_all_replicas(metho
|
||||
[
|
||||
("reset_mm_cache", {}, True),
|
||||
("sleep", {}, True),
|
||||
("release_kv_cache_memory", {}, True),
|
||||
("pause_scheduler", {"clear_cache": True}, True),
|
||||
("pause_scheduler", {"clear_cache": False}, False),
|
||||
("reset_prefix_cache", {}, False),
|
||||
|
||||
@@ -7,9 +7,8 @@ from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.v1.engine import EngineCoreRequest, FinishReason
|
||||
from vllm.v1.engine.logprobs import LogprobsProcessor
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
from vllm.v1.engine import FinishReason
|
||||
|
||||
from vllm_omni.engine import OmniEngineCoreOutput
|
||||
from vllm_omni.engine.output_modality import OutputModality
|
||||
@@ -17,7 +16,6 @@ from vllm_omni.engine.output_processor import (
|
||||
MultimodalOutputProcessor,
|
||||
OmniRequestState,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -111,37 +109,6 @@ def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload():
|
||||
assert not state.mm_accumulated.is_empty
|
||||
|
||||
|
||||
def test_fixed_token_scores_survive_output_processing_and_stage_wrapping():
|
||||
processor, state = _make_processor(OutputModality.TEXT, _Detokenizer())
|
||||
request = EngineCoreRequest(
|
||||
request_id="r",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
mm_features=None,
|
||||
sampling_params=SamplingParams(prompt_logprob_token_ids=[4, 5]),
|
||||
pooling_params=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
data_parallel_rank=None,
|
||||
)
|
||||
state.logprobs_processor = LogprobsProcessor.from_new_request(None, request)
|
||||
scores = torch.tensor([[-0.5, -1.5], [-0.2, -2.2]])
|
||||
processed = processor.process_outputs(
|
||||
[
|
||||
OmniEngineCoreOutput(
|
||||
request_id="r",
|
||||
new_token_ids=[42],
|
||||
prompt_token_id_logprobs=scores,
|
||||
finish_reason=FinishReason.STOP,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
(raw_output,) = processed.request_outputs
|
||||
final_output = OmniRequestOutput.from_stage_output(raw_output, stage_id=0)
|
||||
torch.testing.assert_close(torch.from_numpy(final_output.prompt_token_id_logprobs), scores)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("streaming", [False, True])
|
||||
def test_generation_stage_mm_delivered_only_with_terminal(streaming):
|
||||
# Non-terminal chunks accumulate; the terminal output carries the audio.
|
||||
|
||||
@@ -97,6 +97,7 @@ _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
|
||||
"enable_multithread_weight_load",
|
||||
"env",
|
||||
"has_sampling_extra_args",
|
||||
"log_level",
|
||||
"log_stats",
|
||||
"model_arch",
|
||||
"model_subdir",
|
||||
|
||||
@@ -5,19 +5,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from vllm import SamplingParams
|
||||
from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
|
||||
from vllm_omni.engine.stage_engine_core_client import (
|
||||
DPLBStageEngineCoreClient,
|
||||
StageEngineCoreClient,
|
||||
StageEngineCoreClientBase,
|
||||
)
|
||||
from vllm_omni.engine.stage_init_utils import StageMetadata
|
||||
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -38,46 +30,3 @@ def test_check_health_raises_when_resources_engine_dead():
|
||||
client = _make_client(engine_dead=True)
|
||||
with pytest.raises(EngineDeadError, match="engine core is dead"):
|
||||
client.check_health()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dp_size, external_lb", [(1, False), (2, False), (2, True)])
|
||||
@pytest.mark.parametrize("renderer", [None, object()])
|
||||
def test_stage_client_forwards_renderer_to_the_actual_upstream_constructor(dp_size, external_lb, renderer):
|
||||
config = SimpleNamespace(
|
||||
model_config=None,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=dp_size, data_parallel_external_lb=external_lb),
|
||||
)
|
||||
metadata = StageMetadata(
|
||||
stage_id=0,
|
||||
stage_type="llm",
|
||||
engine_output_type="text",
|
||||
is_comprehension=True,
|
||||
requires_multimodal_data=False,
|
||||
engine_input_source=[],
|
||||
final_output=True,
|
||||
final_output_type="text",
|
||||
default_sampling_params=SamplingParams(),
|
||||
custom_process_input_func=None,
|
||||
model_stage=None,
|
||||
runtime_cfg=None,
|
||||
)
|
||||
internal_lb = dp_size > 1 and not external_lb
|
||||
parent = DPLBAsyncMPClient if internal_lb else AsyncMPClient
|
||||
# Keep Omni construction/factory routing real; intercept only upstream's
|
||||
# subprocess/socket boundary, using its current signature via autospec.
|
||||
with patch.object(parent, "__init__", autospec=True, return_value=None) as init:
|
||||
addresses = {"input_address": "tcp://192.0.2.1:1"}
|
||||
client = StageEngineCoreClientBase.make_async_mp_client(
|
||||
config, object, metadata=metadata, client_addresses=addresses, renderer=renderer
|
||||
)
|
||||
assert isinstance(client, DPLBStageEngineCoreClient if internal_lb else StageEngineCoreClient)
|
||||
init.assert_called_once_with(
|
||||
client,
|
||||
config,
|
||||
object,
|
||||
log_stats=False,
|
||||
client_addresses=addresses,
|
||||
client_count=1,
|
||||
client_index=0,
|
||||
renderer=renderer,
|
||||
)
|
||||
@@ -2,7 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@@ -26,15 +25,11 @@ def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
|
||||
expected,
|
||||
):
|
||||
executor = UniProcExecutor.__new__(UniProcExecutor)
|
||||
executor.vllm_config = SimpleNamespace(
|
||||
parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)
|
||||
)
|
||||
executor.driver_worker = mocker.Mock()
|
||||
runner = executor.driver_worker.worker.model_runner
|
||||
runner.model.first_frame_decoder = object() if first_decoder else None
|
||||
runner.model.stream_decoder = object() if stream_decoder else None
|
||||
runner.model.stream_first_audio = stream_first_audio
|
||||
runner._omni_data_plane = None
|
||||
assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
|
||||
assert runner.model_state.set_first_audio_sink.call_count == int(expected)
|
||||
|
||||
@@ -59,64 +54,3 @@ def test_preprocess_add_request_preserves_omni_fields():
|
||||
assert current_wave == 3
|
||||
assert result.external_req_id == "external"
|
||||
assert result.additional_information == {"conditioning": "payload"}
|
||||
|
||||
|
||||
def test_codec_uses_generic_first_audio_binding(monkeypatch):
|
||||
import queue
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
|
||||
from vllm_omni.engine import stage_engine_core_proc as module
|
||||
|
||||
calls: dict[str, Any] = {}
|
||||
|
||||
def hook(*args):
|
||||
return False
|
||||
|
||||
def bind(sink):
|
||||
calls["sink"] = sink
|
||||
return hook
|
||||
|
||||
plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value))
|
||||
model = SimpleNamespace(bind_first_chunk_fast_path=bind)
|
||||
runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane)
|
||||
executor = SimpleNamespace(
|
||||
vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
|
||||
driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
|
||||
)
|
||||
monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
|
||||
outputs: queue.Queue = queue.Queue()
|
||||
scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)})
|
||||
assert module._bind_first_audio_sink(executor, outputs, scheduler)
|
||||
assert calls["hook"] is hook
|
||||
calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
|
||||
client, batch = outputs.get_nowait()
|
||||
assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output
|
||||
|
||||
|
||||
def test_first_audio_binding_preserves_talker_marker(monkeypatch):
|
||||
import queue
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
|
||||
from vllm_omni.engine import stage_engine_core_proc as module
|
||||
|
||||
calls: dict[str, Any] = {}
|
||||
runner = SimpleNamespace(
|
||||
model=SimpleNamespace(first_frame_decoder=object()),
|
||||
model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)),
|
||||
)
|
||||
executor = SimpleNamespace(
|
||||
vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)),
|
||||
driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)),
|
||||
)
|
||||
monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace)
|
||||
outputs: queue.Queue = queue.Queue()
|
||||
scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)})
|
||||
assert module._bind_first_audio_sink(executor, outputs, scheduler)
|
||||
calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000))
|
||||
assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY]
|
||||
executor.vllm_config.parallel_config.tensor_parallel_size = 2
|
||||
assert not module._bind_first_audio_sink(executor, outputs, scheduler)
|
||||
@@ -83,27 +83,6 @@ def test_collective_rpc_control_method_reraises_worker_error():
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_kv_memory_release_uses_engine_core_and_propagates_pause_errors():
|
||||
async def run() -> None:
|
||||
release = AsyncMock(return_value=None)
|
||||
pool, client = _make_pool(release_kv_cache_memory_async=release)
|
||||
await pool.collective_rpc(0, "release_kv_cache_memory", timeout=2.0)
|
||||
release.assert_awaited_once_with()
|
||||
client.collective_rpc_async.assert_not_awaited()
|
||||
|
||||
release.side_effect = RuntimeError("requires a completed pause first")
|
||||
with pytest.raises(RuntimeError, match="completed pause"):
|
||||
await pool.collective_rpc(0, "release_kv_cache_memory")
|
||||
|
||||
del client.release_kv_cache_memory_async
|
||||
result = await pool.collective_rpc(0, "release_kv_cache_memory")
|
||||
assert result["supported"] is False
|
||||
client.collective_rpc_async.assert_not_awaited()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_collective_rpc_reset_caches_use_control_helpers():
|
||||
async def run() -> None:
|
||||
@@ -175,46 +154,6 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails():
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_release_request_resources_skips_without_async_chunk():
|
||||
async def run() -> None:
|
||||
call = AsyncMock()
|
||||
pool = StagePool(
|
||||
0,
|
||||
[SimpleNamespace(call_utility_async=call)], # type: ignore[list-item]
|
||||
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)),
|
||||
)
|
||||
|
||||
await pool.release_request_resources(["req-1"])
|
||||
call.assert_not_awaited()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch):
|
||||
async def run() -> None:
|
||||
async def hang(*_args):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
live = AsyncMock()
|
||||
pool = StagePool(
|
||||
0,
|
||||
[SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)], # type: ignore[list-item]
|
||||
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)),
|
||||
)
|
||||
monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2)
|
||||
|
||||
release = asyncio.create_task(pool.release_request_resources(["req-1"]))
|
||||
await asyncio.sleep(0.05)
|
||||
live.assert_awaited_once_with("omni_release_request_resources", ["req-1"])
|
||||
assert not release.done()
|
||||
|
||||
await asyncio.wait_for(release, timeout=1.0)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"])
|
||||
@pytest.mark.parametrize("failure", ["error", "timeout", "missing"])
|
||||
|
||||
@@ -1,85 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""The Local-v1.5 continuation prompt fast path equals the processor layout."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _Processor:
|
||||
"""Segment API of the model's processor with a deterministic tokenizer."""
|
||||
|
||||
def __init__(self):
|
||||
self.model_config = SimpleNamespace(
|
||||
n_vq=3,
|
||||
im_start_token_id=1,
|
||||
im_end_token_id=2,
|
||||
audio_start_token_id=3,
|
||||
audio_pad_token_id=1024,
|
||||
audio_assistant_slot_token_id=5,
|
||||
)
|
||||
self.encoded = []
|
||||
|
||||
def _encode_text(self, text):
|
||||
self.encoded.append(text)
|
||||
return [100 + ord(c) % 50 for c in text]
|
||||
|
||||
def _user_prompt_prefix_ids(self):
|
||||
return [1] + self._encode_text("user\n") + self._encode_text("<ref>")
|
||||
|
||||
def _user_prompt_after_reference_ids(self, language, fields):
|
||||
return self._encode_text(f"lang={fields.get('language', language)}")
|
||||
|
||||
def _assistant_prompt_prefix_ids(self):
|
||||
return self._encode_text("</user>") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant")
|
||||
|
||||
def reference(self, text, language, codes):
|
||||
# Mirrors the processor's _build_continuation_codes + _pad (one item).
|
||||
fields = {} if language is None else {"language": language}
|
||||
ids = (
|
||||
self._user_prompt_prefix_ids()
|
||||
+ self._encode_text("None")
|
||||
+ self._user_prompt_after_reference_ids(language, fields)
|
||||
+ self._encode_text(text)
|
||||
+ self._assistant_prompt_prefix_ids()
|
||||
+ [3]
|
||||
)
|
||||
text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long)
|
||||
text_rows[:, 0] = torch.tensor(ids)
|
||||
audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long)
|
||||
audio_rows[:, 0] = 5
|
||||
audio_rows[:, 1:] = codes
|
||||
unified = torch.cat([text_rows, audio_rows])
|
||||
return unified[:, 0].tolist(), unified[:, 1:].contiguous()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("language", [None, "English"])
|
||||
def test_matches_processor_and_caches_constant_segments(language):
|
||||
proc = _Processor()
|
||||
codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32)
|
||||
kwargs = {"text": "ref words target words"}
|
||||
if language is not None:
|
||||
kwargs["language"] = language
|
||||
ids, audio = _local_continuation_prompt(proc, kwargs, [codes])
|
||||
expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long())
|
||||
assert ids == expected_ids
|
||||
assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio)
|
||||
|
||||
proc.encoded.clear()
|
||||
_local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes])
|
||||
assert proc.encoded == ["other"]
|
||||
|
||||
|
||||
def test_unsupported_layouts_fall_back():
|
||||
proc = _Processor()
|
||||
codes = torch.zeros((2, 3), dtype=torch.long)
|
||||
assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None
|
||||
assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None
|
||||
assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None
|
||||
assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None
|
||||
@@ -691,45 +691,6 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker)
|
||||
audio_tokenizer.to.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"variant,graphs_env,prepares",
|
||||
[("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)],
|
||||
)
|
||||
def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch):
|
||||
if graphs_env is None:
|
||||
monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env)
|
||||
adapter = _moss_adapter_with_stage_configs([], mocker)
|
||||
adapter._moss_variant = variant
|
||||
encoder = mocker.Mock()
|
||||
mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
|
||||
asyncio.run(adapter.warmup())
|
||||
assert encoder.prepare.called is prepares
|
||||
|
||||
|
||||
def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch):
|
||||
monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
|
||||
adapter = _moss_adapter_with_stage_configs([], mocker)
|
||||
encoder = mocker.Mock()
|
||||
encoder.prepare.side_effect = RuntimeError("capture failed")
|
||||
mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
|
||||
asyncio.run(adapter.warmup())
|
||||
encoder.prepare.assert_called_once_with()
|
||||
|
||||
|
||||
def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch):
|
||||
monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1")
|
||||
from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError
|
||||
|
||||
adapter = _moss_adapter_with_stage_configs([], mocker)
|
||||
encoder = mocker.Mock()
|
||||
encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling")
|
||||
mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder)
|
||||
with pytest.raises(SharedReferenceEncoderStartupError):
|
||||
asyncio.run(adapter.warmup())
|
||||
|
||||
|
||||
def test_qwen3_tts_metadata():
|
||||
assert Qwen3TTSAdapter.backend == "ar"
|
||||
assert issubclass(Qwen3TTSAdapter, ARTTSAdapter)
|
||||
|
||||
@@ -10,7 +10,6 @@ import pytest
|
||||
from vllm.engine.protocol import StreamingInput
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
||||
from vllm.v1.kv_hints import KvHintsEnvelope
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
@@ -139,39 +138,6 @@ def test_generate_forwards_lora_request_to_engine():
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("streaming", [None, "chunk", "empty"])
|
||||
def test_generate_forwards_kv_hints_on_initial_stage_submission(streaming):
|
||||
async def run():
|
||||
submissions = []
|
||||
|
||||
async def add(**kwargs):
|
||||
submissions.append(kwargs)
|
||||
|
||||
omni = get_async_omni_instance(fake_add_request=add)
|
||||
omni.engine.add_streaming_update_async = _noop
|
||||
omni.engine.stage_configs = [SimpleNamespace(is_comprehension=True, stage_type="llm")]
|
||||
omni.engine.stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(is_encoder_decoder=False))]
|
||||
hints = KvHintsEnvelope("1", "message", [])
|
||||
params = SamplingParams(max_tokens=2, output_kind=RequestOutputKind.DELTA)
|
||||
|
||||
async def chunks():
|
||||
if streaming == "chunk":
|
||||
yield StreamingInput(prompt={"prompt_token_ids": [1]})
|
||||
|
||||
async for _ in omni.generate(
|
||||
prompt=chunks() if streaming else {"prompt_token_ids": [1]},
|
||||
request_id="hints",
|
||||
sampling_params_list=[params],
|
||||
kv_hints=hints,
|
||||
):
|
||||
pass
|
||||
assert len(submissions) == 1
|
||||
assert submissions[0]["kv_hints"] is hints
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize(
|
||||
"req_ids,cancel_prefix,expected_cancel_count",
|
||||
|
||||
@@ -194,56 +194,6 @@ def test_reset_encoder_cache_forwards_to_ar_stages():
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_release_kv_memory_requires_pause_and_wakes_only_ar_kv_allocations():
|
||||
async def run() -> None:
|
||||
omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
|
||||
with pytest.raises(RuntimeError, match="completed pause"):
|
||||
await omni.release_kv_cache_memory()
|
||||
omni.collective_rpc.assert_not_awaited()
|
||||
|
||||
omni._paused = True
|
||||
await omni.release_kv_cache_memory()
|
||||
omni._clear_frontend_mm_cache.assert_awaited_once_with()
|
||||
omni.collective_rpc.assert_awaited_once_with(
|
||||
method="release_kv_cache_memory",
|
||||
args=(),
|
||||
kwargs=None,
|
||||
stage_ids=[0, 2],
|
||||
timeout=CACHE_RESET_TIMEOUT_S,
|
||||
)
|
||||
assert omni._stage_sleeping_tags == {0: {"kv_cache"}, 2: {"kv_cache"}}
|
||||
assert not omni._level2_sleeping
|
||||
|
||||
omni.collective_rpc.reset_mock()
|
||||
await omni.wake_up(stage_ids=[0, 2], tags=["kv_cache"])
|
||||
omni.collective_rpc.assert_awaited_once_with(
|
||||
method="wake_up", args=(), kwargs={"tags": ["kv_cache"]}, stage_ids=[0, 2]
|
||||
)
|
||||
assert not omni._sleeping_tags
|
||||
assert omni._paused
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_release_kv_memory_skips_diffusion_and_reports_engine_failures():
|
||||
async def run() -> None:
|
||||
diffusion = _make_omni(stage_types=["diffusion"])
|
||||
await diffusion.release_kv_cache_memory()
|
||||
diffusion.collective_rpc.assert_not_awaited()
|
||||
assert not diffusion._paused
|
||||
|
||||
omni = _make_omni(stage_types=["llm"])
|
||||
omni._paused = True
|
||||
omni.collective_rpc.return_value = [{"supported": False, "error": "requires all executor memory resident"}]
|
||||
with pytest.raises(RuntimeError, match="all executor memory resident"):
|
||||
await omni.release_kv_cache_memory()
|
||||
assert not omni._sleeping_tags
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_reset_mm_cache_clears_frontend_then_forwards_to_ar_stages():
|
||||
async def run() -> None:
|
||||
|
||||
@@ -288,9 +288,6 @@ def _setup_log_mocks(monkeypatch):
|
||||
def build_and_log_summary(self):
|
||||
return "Fake summary"
|
||||
|
||||
def log_timing_summary(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.entrypoints.omni.OrchestratorAggregator",
|
||||
_FakeOrchestratorAggregator,
|
||||
|
||||
@@ -4,11 +4,9 @@
|
||||
"""Assertion and response validation helpers for tests."""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import wave
|
||||
@@ -1001,31 +999,6 @@ def _assert_transcript_matches(
|
||||
)
|
||||
|
||||
|
||||
def _retain_failed_speech_audio(
|
||||
audio_bytes: bytes | None, transcript: str, expected_text: Any, error: AssertionError
|
||||
) -> None:
|
||||
"""Retain the exact ASR input when a CI speech-content assertion fails."""
|
||||
directory = os.environ.get("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR")
|
||||
if not directory or not audio_bytes:
|
||||
return
|
||||
stem = f"failed_speech_{hashlib.sha256(audio_bytes).hexdigest()[:16]}"
|
||||
root = Path(directory)
|
||||
try:
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
(root / f"{stem}.wav").write_bytes(audio_bytes)
|
||||
(root / f"{stem}.json").write_text(
|
||||
json.dumps(
|
||||
{"transcript": transcript, "expected_text": str(expected_text), "assertion": str(error)},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
print(f"Retained failed speech ASR input: {root / f'{stem}.wav'}")
|
||||
except OSError as exc:
|
||||
# Diagnostic I/O must never replace or suppress the content failure.
|
||||
print(f"Could not retain failed speech ASR input: {exc}")
|
||||
|
||||
|
||||
def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str | None = None) -> None:
|
||||
"""Validate speech API results from :class:`~tests.helpers.client.OmniResponse`.
|
||||
|
||||
@@ -1090,21 +1063,14 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
|
||||
if expected_text:
|
||||
print(f"audio content is: {transcript}")
|
||||
print(f"input text is: {expected_text}")
|
||||
transcription_audio = _speech_audio_for_transcription(
|
||||
getattr(response, "audio_bytes", None), request_config
|
||||
_assert_transcript_matches(
|
||||
transcript,
|
||||
_speech_audio_for_transcription(getattr(response, "audio_bytes", None), request_config),
|
||||
expected_text,
|
||||
threshold=0.9,
|
||||
escalation_model=request_config.get("transcript_escalation_model"),
|
||||
language=request_config.get("transcript_language"),
|
||||
)
|
||||
try:
|
||||
_assert_transcript_matches(
|
||||
transcript,
|
||||
transcription_audio,
|
||||
expected_text,
|
||||
threshold=0.9,
|
||||
escalation_model=request_config.get("transcript_escalation_model"),
|
||||
language=request_config.get("transcript_language"),
|
||||
)
|
||||
except AssertionError as exc:
|
||||
_retain_failed_speech_audio(transcription_audio, transcript, expected_text, exc)
|
||||
raise
|
||||
_assert_preset_voice_gender_from_audio(
|
||||
response.audio_bytes,
|
||||
request_config.get("voice"),
|
||||
|
||||
@@ -97,7 +97,6 @@ def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[
|
||||
class OmniResponse:
|
||||
"""Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``)."""
|
||||
|
||||
request_id: str | None = None
|
||||
text_content: str | None = None
|
||||
audio_data: list[str] | None = None
|
||||
audio_content: str | None = None
|
||||
@@ -422,8 +421,6 @@ class OnlineOmniClient:
|
||||
text_content = ""
|
||||
audio_data = []
|
||||
for chunk in chat_completion:
|
||||
if result.request_id is None:
|
||||
result.request_id = getattr(chunk, "id", None)
|
||||
for choice in chunk.choices:
|
||||
content = getattr(getattr(choice, "delta", None), "content", None)
|
||||
modality = getattr(chunk, "modality", None)
|
||||
@@ -456,7 +453,6 @@ class OnlineOmniClient:
|
||||
"""Wall clock from *before* ``chat.completions.create`` through response parse + local decode."""
|
||||
result = OmniResponse()
|
||||
try:
|
||||
result.request_id = getattr(chat_completion, "id", None)
|
||||
audio_data = None
|
||||
text_content = None
|
||||
for choice in chat_completion.choices:
|
||||
|
||||
@@ -1,19 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Import preload for the CPU-only collective forkserver.
|
||||
|
||||
This module is imported in the forkserver process, before it starts workers.
|
||||
Its CUDA visibility does not alter the pytest process or CUDA parity workers.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
||||
|
||||
import torch # noqa: E402
|
||||
|
||||
# Retain the expensive worker-module imports across CPU process groups.
|
||||
from tests.diffusion.distributed import test_comm, test_pipeline_parallel # noqa: E402, F401
|
||||
|
||||
if torch.cuda.is_initialized():
|
||||
raise RuntimeError("CPU collective forkserver must not initialize CUDA")
|
||||
@@ -1,62 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Process and platform setup for isolated collective test workers."""
|
||||
|
||||
from collections.abc import Callable
|
||||
from multiprocessing.context import BaseContext
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def collective_context(device_kind: str) -> BaseContext:
|
||||
"""Reuse imports in a clean CPU forkserver; CUDA workers use spawn."""
|
||||
if device_kind == "cpu" and "forkserver" in torch.multiprocessing.get_all_start_methods():
|
||||
torch.multiprocessing.set_forkserver_preload(["tests.helpers.cpu_collective_server"])
|
||||
return torch.multiprocessing.get_context("forkserver")
|
||||
return torch.multiprocessing.get_context("spawn")
|
||||
|
||||
|
||||
def start_collective_workers(
|
||||
worker: Callable[..., None], *, args: tuple[object, ...], nprocs: int, device_kind: str
|
||||
) -> None:
|
||||
context = collective_context(device_kind)
|
||||
if context.get_start_method() == "spawn":
|
||||
torch.multiprocessing.spawn(worker, args=args, nprocs=nprocs)
|
||||
else:
|
||||
torch.multiprocessing.start_processes(worker, args=args, nprocs=nprocs, start_method=context.get_start_method())
|
||||
|
||||
|
||||
def configure_cpu_collective_worker() -> None:
|
||||
"""Select CPU platforms inside a spawned worker that only uses Gloo.
|
||||
|
||||
Upstream coordinators resolve local ranks through the selected
|
||||
platform's visible devices. CPU ranks can outnumber visible GPUs, so they
|
||||
must use the CPU platform even when the test host has CUDA installed.
|
||||
The worker still creates real process groups and runs real collectives.
|
||||
"""
|
||||
import vllm.platforms
|
||||
from vllm.platforms.cpu import CpuPlatform
|
||||
|
||||
from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
|
||||
from vllm_omni.platforms.interface import UnspecifiedOmniPlatform
|
||||
|
||||
class CpuCollectivePlatform(UnspecifiedOmniPlatform):
|
||||
dist_backend = "gloo"
|
||||
|
||||
@classmethod
|
||||
def get_device_count(cls) -> int:
|
||||
# CPU tensors share one logical device regardless of process rank.
|
||||
return 1
|
||||
|
||||
@classmethod
|
||||
def set_device(cls, device: torch.device) -> None:
|
||||
assert device.type == "cpu"
|
||||
|
||||
@classmethod
|
||||
def synchronize(cls) -> None:
|
||||
pass
|
||||
|
||||
vllm.platforms.current_platform = CpuPlatform()
|
||||
platform = CpuCollectivePlatform()
|
||||
for module in (comm, group_coordinator, parallel_state):
|
||||
module.current_omni_platform = platform
|
||||
@@ -1,28 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Scheduler test helpers shared by AR / chunk-transfer unit tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import MethodType
|
||||
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
|
||||
|
||||
def bind_omits_transfer_helpers(sched) -> None:
|
||||
"""Bind the real omit helpers so a MagicMock scheduler does not skip save_async.
|
||||
|
||||
MagicMock is truthy, so an unbound ``_request_omits_chunk_transfer_to_next_stage``
|
||||
silently skips the chunk put.
|
||||
"""
|
||||
sched._omits_kv_transfer_cache = {}
|
||||
sched._omni_final_stage_flags = MethodType(OmniARScheduler._omni_final_stage_flags, sched)
|
||||
sched._request_omits_kv_transfer_to_next_stage = MethodType(
|
||||
OmniARScheduler._request_omits_kv_transfer_to_next_stage,
|
||||
sched,
|
||||
)
|
||||
sched._request_omits_chunk_transfer_to_next_stage = MethodType(
|
||||
OmniARScheduler._request_omits_chunk_transfer_to_next_stage,
|
||||
sched,
|
||||
)
|
||||
+12
-58
@@ -659,14 +659,6 @@ class OmniRunner:
|
||||
self.seed = seed
|
||||
self._prompt_len_estimate_cache: dict[str, Any] = {}
|
||||
self.omni: Any = None
|
||||
self._runner_process = psutil.Process()
|
||||
self._initial_child_ids: set[tuple[int, float]] = set()
|
||||
self._owned_engine_processes: dict[int, float] = {}
|
||||
for child in self._runner_process.children(recursive=True):
|
||||
try:
|
||||
self._initial_child_ids.add((child.pid, child.create_time()))
|
||||
except psutil.Error:
|
||||
pass
|
||||
try:
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
@@ -678,14 +670,12 @@ class OmniRunner:
|
||||
deploy_config=deploy_config,
|
||||
**kwargs,
|
||||
)
|
||||
self._remember_engine_processes(startup=True)
|
||||
startup_s = time.perf_counter() - startup_t0
|
||||
if log_stats:
|
||||
print(f"OmniRunner startup took {startup_s:.3f}s (model={model_name})", flush=True)
|
||||
except BaseException:
|
||||
# ``with OmniRunner(...)`` never reaches ``__enter__``/``__exit__``
|
||||
# when construction fails after worker processes have started.
|
||||
self._remember_engine_processes(startup=True)
|
||||
self.__exit__(None, None, None)
|
||||
raise
|
||||
|
||||
@@ -909,51 +899,18 @@ class OmniRunner:
|
||||
def stop_profile(self, stages: list[int] | None = None) -> list[Any]:
|
||||
return self.omni.stop_profile(stages=stages)
|
||||
|
||||
def _remember_engine_processes(self, *, startup: bool = False) -> None:
|
||||
"""Track startup workers, then only descendants of those workers.
|
||||
|
||||
Existing children may belong to another runner or a shared transcriber.
|
||||
Record identities before close() can reparent surviving workers.
|
||||
"""
|
||||
try:
|
||||
if startup:
|
||||
candidates = self._runner_process.children(recursive=True)
|
||||
else:
|
||||
candidates = []
|
||||
for pid, created in list(self._owned_engine_processes.items()):
|
||||
try:
|
||||
proc = psutil.Process(pid)
|
||||
if proc.create_time() == created:
|
||||
candidates.extend(proc.children(recursive=True))
|
||||
except psutil.Error:
|
||||
pass
|
||||
for proc in candidates:
|
||||
try:
|
||||
identity = (proc.pid, proc.create_time())
|
||||
if identity in self._initial_child_ids:
|
||||
continue
|
||||
if startup and any(
|
||||
(parent.pid, parent.create_time()) in self._initial_child_ids for parent in proc.parents()
|
||||
):
|
||||
continue
|
||||
cmdline = " ".join(proc.cmdline()).lower()
|
||||
if "enginecore" in cmdline or "enginecore" in proc.name().lower():
|
||||
self._owned_engine_processes[identity[0]] = identity[1]
|
||||
except psutil.Error:
|
||||
pass
|
||||
except psutil.Error as error:
|
||||
print(f"Warning: could not inspect runner workers: {error}")
|
||||
|
||||
def _cleanup_process(self):
|
||||
try:
|
||||
keywords = ["enginecore"]
|
||||
matched = []
|
||||
for pid, created in self._owned_engine_processes.items():
|
||||
for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
|
||||
try:
|
||||
proc = psutil.Process(pid)
|
||||
if proc.create_time() == created:
|
||||
print(f"Found owned vllm process: PID={proc.pid}")
|
||||
cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
|
||||
name = proc.name().lower()
|
||||
if any(k in cmdline for k in keywords) or any(k in name for k in keywords):
|
||||
print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
|
||||
matched.append(proc)
|
||||
except psutil.Error:
|
||||
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
||||
pass
|
||||
for proc in matched:
|
||||
try:
|
||||
@@ -981,14 +938,11 @@ class OmniRunner:
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self._remember_engine_processes()
|
||||
try:
|
||||
omni = getattr(self, "omni", None)
|
||||
if omni is not None and hasattr(omni, "close"):
|
||||
omni.close()
|
||||
finally:
|
||||
self._cleanup_process()
|
||||
cleanup_test_environment()
|
||||
omni = getattr(self, "omni", None)
|
||||
if omni is not None and hasattr(omni, "close"):
|
||||
omni.close()
|
||||
self._cleanup_process()
|
||||
cleanup_test_environment()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -2,9 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import io
|
||||
import json
|
||||
import wave
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
@@ -251,85 +249,6 @@ def test_pcm_transcript_escalation_uses_same_valid_wav_and_rejects_two_wrong_tra
|
||||
_assert_wav_wraps_pcm(wav_bytes, pcm_bytes, 24_000)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("response_format", ["wav", "pcm"])
|
||||
@pytest.mark.parametrize("escalation_model", [None, "large-v3"])
|
||||
def test_failed_speech_capture_retains_exact_asr_input_and_original_failure(
|
||||
monkeypatch, tmp_path, response_format, escalation_model
|
||||
):
|
||||
pcm_bytes = _pcm_sine()
|
||||
wav = io.BytesIO()
|
||||
with wave.open(wav, "wb") as writer:
|
||||
writer.setnchannels(1)
|
||||
writer.setsampwidth(2)
|
||||
writer.setframerate(24_000)
|
||||
writer.writeframes(pcm_bytes)
|
||||
audio_bytes = wav.getvalue() if response_format == "wav" else pcm_bytes
|
||||
transcribed = []
|
||||
|
||||
def transcribe(audio_bytes, **kwargs):
|
||||
transcribed.append(audio_bytes)
|
||||
return "this is unrelated speech"
|
||||
|
||||
directory = tmp_path / "failed_audio"
|
||||
monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
|
||||
monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", transcribe)
|
||||
config = {
|
||||
"input": "please read this requested sentence accurately",
|
||||
"response_format": response_format,
|
||||
"transcript_pcm_sample_rate": 24_000,
|
||||
"transcript_escalation_model": escalation_model,
|
||||
}
|
||||
response = OmniResponse(success=True, audio_bytes=audio_bytes, audio_format=f"audio/{response_format}")
|
||||
with pytest.raises(AssertionError, match="Transcript doesn't match input") as error:
|
||||
assert_audio_speech_response(response, config, "full_model")
|
||||
|
||||
clips = list(directory.glob("*.wav"))
|
||||
assert len(clips) == 1
|
||||
assert all(clip == clips[0].read_bytes() for clip in transcribed)
|
||||
_assert_wav_wraps_pcm(clips[0].read_bytes(), pcm_bytes, 24_000)
|
||||
metadata = json.loads(clips[0].with_suffix(".json").read_text())
|
||||
assert metadata["expected_text"] == config["input"]
|
||||
assert metadata["transcript"] == "this is unrelated speech"
|
||||
assert metadata["assertion"] == str(error.value)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("capture_enabled", [False, True])
|
||||
@pytest.mark.parametrize("content_failure", [False, True])
|
||||
def test_speech_capture_is_opt_in_and_does_not_write_successes(monkeypatch, tmp_path, capture_enabled, content_failure):
|
||||
directory = tmp_path / "failed_audio"
|
||||
if capture_enabled:
|
||||
monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
|
||||
else:
|
||||
monkeypatch.delenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", raising=False)
|
||||
expected = "please read this requested sentence accurately"
|
||||
monkeypatch.setattr(
|
||||
assertions,
|
||||
"convert_audio_bytes_to_text",
|
||||
lambda *_args, **_kwargs: "this is unrelated speech" if content_failure else expected,
|
||||
)
|
||||
with pytest.raises(AssertionError, match="Transcript doesn't match input") if content_failure else nullcontext():
|
||||
assert_audio_speech_response(
|
||||
OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
|
||||
{"input": expected, "response_format": "wav"},
|
||||
"full_model",
|
||||
)
|
||||
assert directory.exists() is (capture_enabled and content_failure)
|
||||
assert len(list(tmp_path.rglob("*.wav"))) == int(capture_enabled and content_failure)
|
||||
|
||||
|
||||
def test_speech_capture_io_failure_does_not_mask_content_failure(monkeypatch, tmp_path):
|
||||
directory = tmp_path / "blocked"
|
||||
directory.write_text("a file cannot be an output directory")
|
||||
monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
|
||||
monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "this is unrelated speech")
|
||||
with pytest.raises(AssertionError, match="Transcript doesn't match input"):
|
||||
assert_audio_speech_response(
|
||||
OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
|
||||
{"input": "please read this requested sentence accurately", "response_format": "wav"},
|
||||
"full_model",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pcm_bytes",
|
||||
[
|
||||
|
||||
@@ -1,15 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Tests for OmniRunner startup rollback and worker ownership."""
|
||||
"""Tests for OmniRunner startup rollback."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Generator
|
||||
|
||||
import psutil
|
||||
import pytest
|
||||
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
@@ -43,114 +38,3 @@ def test_runner_init_rolls_back_on_omni_startup_failure(monkeypatch: pytest.Monk
|
||||
assert cleaned.count("process") == 1
|
||||
# Once at the start of ``__init__``, once from the constructor rollback.
|
||||
assert cleaned.count("env") == 2
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def live_workers() -> Generator[list[subprocess.Popen[bytes]], None, None]:
|
||||
workers: list[subprocess.Popen[bytes]] = []
|
||||
try:
|
||||
yield workers
|
||||
finally:
|
||||
for worker in workers:
|
||||
if worker.poll() is None:
|
||||
worker.kill()
|
||||
worker.wait(timeout=10)
|
||||
if worker.stdout is not None:
|
||||
worker.stdout.close()
|
||||
|
||||
|
||||
def _start_engine_worker(
|
||||
workers: list[subprocess.Popen[bytes]], *, ignore_terminate: bool = False
|
||||
) -> subprocess.Popen[bytes]:
|
||||
script = "import signal, time; "
|
||||
if ignore_terminate:
|
||||
script += "signal.signal(signal.SIGTERM, signal.SIG_IGN); "
|
||||
script += "print('ready', flush=True); time.sleep(90)"
|
||||
worker = subprocess.Popen([sys.executable, "-c", script, "enginecore"], stdout=subprocess.PIPE)
|
||||
workers.append(worker)
|
||||
assert worker.stdout is not None
|
||||
assert worker.stdout.readline() == b"ready\n"
|
||||
return worker
|
||||
|
||||
|
||||
@pytest.mark.parametrize("close_raises", [False, True])
|
||||
def test_runner_cleanup_preserves_unrelated_engine_workers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
live_workers: list[subprocess.Popen[bytes]],
|
||||
close_raises: bool,
|
||||
) -> None:
|
||||
existing = _start_engine_worker(live_workers)
|
||||
owned: list[subprocess.Popen[bytes]] = []
|
||||
late_unrelated: list[subprocess.Popen[bytes]] = []
|
||||
monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
|
||||
|
||||
class _OwnedOmni:
|
||||
def __init__(self, *args: object, **kwargs: object) -> None:
|
||||
owned.append(_start_engine_worker(live_workers))
|
||||
|
||||
def close(self) -> None:
|
||||
late_unrelated.append(_start_engine_worker(live_workers))
|
||||
if close_raises:
|
||||
raise RuntimeError("close failed")
|
||||
|
||||
monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _OwnedOmni)
|
||||
if close_raises:
|
||||
with pytest.raises(RuntimeError, match="close failed"), OmniRunner("fake-model"):
|
||||
pass
|
||||
else:
|
||||
with OmniRunner("fake-model"):
|
||||
pass
|
||||
|
||||
assert owned[0].poll() is not None
|
||||
assert existing.poll() is None
|
||||
assert late_unrelated[0].poll() is None
|
||||
|
||||
|
||||
def test_runner_startup_failure_cleans_only_new_workers(
|
||||
monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
|
||||
) -> None:
|
||||
unrelated = _start_engine_worker(live_workers)
|
||||
owned: list[subprocess.Popen[bytes]] = []
|
||||
monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
|
||||
|
||||
class _FailingOmni:
|
||||
def __init__(self, *args: object, **kwargs: object) -> None:
|
||||
owned.append(_start_engine_worker(live_workers))
|
||||
raise RuntimeError("startup failed")
|
||||
|
||||
monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _FailingOmni)
|
||||
with pytest.raises(RuntimeError, match="startup failed"):
|
||||
OmniRunner("fake-model")
|
||||
|
||||
assert owned[0].poll() is not None
|
||||
assert unrelated.poll() is None
|
||||
|
||||
|
||||
def test_runner_cleanup_rejects_stale_process_identity(live_workers: list[subprocess.Popen[bytes]]) -> None:
|
||||
unrelated = _start_engine_worker(live_workers)
|
||||
runner = object.__new__(OmniRunner)
|
||||
runner._owned_engine_processes = {unrelated.pid: psutil.Process(unrelated.pid).create_time() + 1}
|
||||
|
||||
runner._cleanup_process()
|
||||
|
||||
assert unrelated.poll() is None
|
||||
|
||||
|
||||
def test_runner_cleanup_kills_owned_worker_ignoring_terminate(
|
||||
monkeypatch: pytest.MonkeyPatch, live_workers: list[subprocess.Popen[bytes]]
|
||||
) -> None:
|
||||
owned: list[subprocess.Popen[bytes]] = []
|
||||
monkeypatch.setattr("tests.helpers.runtime.cleanup_test_environment", lambda: None)
|
||||
|
||||
class _StubbornOmni:
|
||||
def __init__(self, *args: object, **kwargs: object) -> None:
|
||||
owned.append(_start_engine_worker(live_workers, ignore_terminate=True))
|
||||
|
||||
def close(self) -> None:
|
||||
pass
|
||||
|
||||
monkeypatch.setattr("vllm_omni.entrypoints.omni.Omni", _StubbornOmni)
|
||||
with OmniRunner("fake-model"):
|
||||
assert owned[0].poll() is None
|
||||
|
||||
assert owned[0].poll() is not None
|
||||
@@ -1,17 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from importlib import import_module
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from transformers import BatchFeature
|
||||
from vllm.config.multimodal import MultiModalDummyOptions
|
||||
from vllm.multimodal.parse import AudioProcessorItems, ImageProcessorItems, MultiModalDataItems
|
||||
from vllm.multimodal.processing import BaseMultiModalProcessor
|
||||
|
||||
from vllm_omni.inputs.mm_processor import OmniDummyInputsBuilder, OmniMultiModalProcessor
|
||||
from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -28,14 +24,10 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts,
|
||||
|
||||
processor = object.__new__(Processor)
|
||||
processor.dummy_inputs = SimpleNamespace(get_dummy_text=Mock(return_value="dummy"))
|
||||
processor._get_hf_mm_data = Mock(return_value=({"audios": ["wave"]}, {"embedding": "preserve"}))
|
||||
processor._call_hf_processor = Mock(return_value=BatchFeature({"features": "processed"}))
|
||||
|
||||
class AudioItems(AudioProcessorItems):
|
||||
def get_passthrough_data(self):
|
||||
return {"embedding": "preserve"}
|
||||
|
||||
items = MultiModalDataItems({"audio": AudioItems(["wave"]), "image": ImageProcessorItems([])} if counts else {})
|
||||
items.select = Mock(wraps=items.select)
|
||||
items = Mock()
|
||||
items.get_all_counts.return_value = counts
|
||||
kwargs = {"sampling_rate": 16000}
|
||||
if original_text is not None:
|
||||
kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text
|
||||
@@ -44,83 +36,9 @@ def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts,
|
||||
items.select.assert_called_once_with({k for k, v in counts.items() if v > 0})
|
||||
processor._call_hf_processor.assert_called_once_with(
|
||||
"dummy" if original_text is None else original_text,
|
||||
{"audios": ["wave"]} if counts else {},
|
||||
{"audios": ["wave"]},
|
||||
{"sampling_rate": 16000},
|
||||
{},
|
||||
)
|
||||
assert dict(result) == ({"features": "processed", "embedding": "preserve"} if counts else {"features": "processed"})
|
||||
assert dict(result) == {"features": "processed", "embedding": "preserve"}
|
||||
assert kwargs == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"module_name,builder_name,processor_name,expected_kwargs",
|
||||
[
|
||||
(
|
||||
"cosyvoice3",
|
||||
"CosyVoice3DummyInputsBuilder",
|
||||
"CosyVoice3MultiModalProcessor",
|
||||
{"prompt_text": "Testing my voices. Why should I not?"},
|
||||
),
|
||||
(
|
||||
"glm_tts",
|
||||
"GLMTTSDummyInputsBuilder",
|
||||
"GLMTTSMultiModalProcessor",
|
||||
{"prompt_text": "This is the reference voice."},
|
||||
),
|
||||
(
|
||||
"omnivoice",
|
||||
"OmniVoiceDummyInputsBuilder",
|
||||
"OmniVoiceMultiModalProcessor",
|
||||
{"ref_text": "Testing voice cloning."},
|
||||
),
|
||||
("mimo_audio", "MiMoAudioLLMDummyInputsBuilder", "MiMoAudioLLMMultiModalProcessor", {}),
|
||||
],
|
||||
)
|
||||
def test_dummy_budget_preserves_model_conditioning(module_name, builder_name, processor_name, expected_kwargs):
|
||||
"""The v0.31 budget path must reach each model's custom builder hook."""
|
||||
module = import_module(f"vllm_omni.model_executor.models.{module_name}.{module_name}")
|
||||
builder = object.__new__(getattr(module, builder_name))
|
||||
assert isinstance(builder, OmniDummyInputsBuilder)
|
||||
items = MultiModalDataItems({})
|
||||
builder.info = SimpleNamespace(ctx=SimpleNamespace(tokenizer=None), parse_mm_data=Mock(return_value=items))
|
||||
builder.get_dummy_mm_data = Mock(return_value={})
|
||||
processor = object.__new__(getattr(module, processor_name))
|
||||
processor.dummy_inputs = builder
|
||||
|
||||
inputs = processor.get_dummy_inputs(128, {"audio": 1}, MultiModalDummyOptions())
|
||||
|
||||
assert inputs.mm_data_items is items
|
||||
assert inputs.hf_processor_mm_kwargs == expected_kwargs
|
||||
assert inputs.prompt == ("<|empty|>" if module_name == "mimo_audio" else [])
|
||||
|
||||
|
||||
def test_dummy_budget_keeps_upstream_path_for_standard_builders(monkeypatch):
|
||||
class Processor(OmniMultiModalProcessor):
|
||||
def _get_mm_fields_config(self, *args):
|
||||
return {}
|
||||
|
||||
def _get_prompt_updates(self, *args):
|
||||
return []
|
||||
|
||||
processor = object.__new__(Processor)
|
||||
processor.dummy_inputs = SimpleNamespace()
|
||||
expected = object()
|
||||
upstream = Mock(return_value=expected)
|
||||
monkeypatch.setattr(BaseMultiModalProcessor, "get_dummy_inputs", upstream)
|
||||
options = MultiModalDummyOptions()
|
||||
|
||||
assert processor.get_dummy_inputs(128, {"audio": 1}, options) is expected
|
||||
upstream.assert_called_once_with(128, {"audio": 1}, options)
|
||||
|
||||
|
||||
def test_voxtral_dummy_budget_uses_audio_token_prompt():
|
||||
from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import VoxtralTTSMultiModalProcessor
|
||||
|
||||
processor = object.__new__(VoxtralTTSMultiModalProcessor)
|
||||
expected = SimpleNamespace(prompt=[1, 25, 24, 35])
|
||||
builder = SimpleNamespace(get_dummy_processor_inputs=Mock(return_value=expected))
|
||||
processor.dummy_inputs = builder
|
||||
options = MultiModalDummyOptions()
|
||||
|
||||
assert processor.get_dummy_inputs(128, {"audio": 1}, options) is expected
|
||||
builder.get_dummy_processor_inputs.assert_called_once_with(128, {"audio": 1}, options)
|
||||
@@ -8,9 +8,7 @@ from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.exceptions import VLLMValidationError
|
||||
from vllm.inputs import mm_input, tokens_input
|
||||
from vllm.multimodal.inputs import PlaceholderRange
|
||||
from vllm.inputs import tokens_input
|
||||
from vllm.pooling_params import PoolingParams
|
||||
from vllm.renderers import BaseRenderer
|
||||
from vllm.renderers.params import TokenizeParams
|
||||
@@ -117,35 +115,6 @@ def test_omni_renderer_is_a_real_subclass_created_once():
|
||||
assert not hasattr(instance, "_renderer")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("async_mode", [False, True])
|
||||
@pytest.mark.parametrize("no_media_kwargs", [False, True])
|
||||
@pytest.mark.parametrize("limit", [2, 4])
|
||||
def test_renderer_applies_expanded_prompt_bound_without_splitting_media(renderer, async_mode, no_media_kwargs, limit):
|
||||
renderer._process_multimodal.return_value = mm_input(
|
||||
list(range(6)), {}, {"image": ["hash"]}, {"image": [PlaceholderRange(offset=1, length=2)]}
|
||||
)
|
||||
prompt = {"prompt_token_ids": [1, 2, 3], "additional_information": {"speaker": 1}}
|
||||
if no_media_kwargs:
|
||||
prompt["mm_processor_kwargs"] = {"target_h": 512}
|
||||
else:
|
||||
prompt["multi_modal_data"] = {"image": "image"}
|
||||
params = TokenizeParams(max_total_tokens=128, truncate_prompt_tokens=limit, truncation_side="right")
|
||||
|
||||
def render():
|
||||
if async_mode:
|
||||
return asyncio.run(renderer.render_cmpl_async([prompt], params))
|
||||
return renderer.render_cmpl([prompt], params)
|
||||
|
||||
if limit == 2:
|
||||
with pytest.raises(VLLMValidationError, match="would split or drop"):
|
||||
render()
|
||||
else:
|
||||
(result,) = render()
|
||||
assert result["prompt_token_ids"] == [0, 1, 2, 3]
|
||||
assert result["mm_placeholders"]["image"] == [PlaceholderRange(offset=1, length=2)]
|
||||
assert result["additional_information"] == {"speaker": 1}
|
||||
|
||||
|
||||
def test_concrete_renderer_overrides_are_not_suppressed():
|
||||
"""The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides."""
|
||||
|
||||
|
||||
@@ -111,7 +111,7 @@ class TestFailureCounterWiring:
|
||||
obj.request_states["req-1"] = SimpleNamespace(
|
||||
metrics=SimpleNamespace(
|
||||
e2e_done=set(),
|
||||
log_timing_summary=lambda: None,
|
||||
build_and_log_summary=lambda: None,
|
||||
),
|
||||
failure_recorded=False,
|
||||
)
|
||||
@@ -127,7 +127,7 @@ class TestFailureCounterWiring:
|
||||
obj.request_states["req-1"] = SimpleNamespace(
|
||||
metrics=SimpleNamespace(
|
||||
e2e_done=set(),
|
||||
log_timing_summary=lambda: None,
|
||||
build_and_log_summary=lambda: None,
|
||||
),
|
||||
failure_recorded=False,
|
||||
)
|
||||
@@ -426,8 +426,6 @@ def _make_scheduler_shell() -> OmniARScheduler:
|
||||
obj.requests = {}
|
||||
obj.running = []
|
||||
obj.waiting = []
|
||||
obj.kv_holding_waiting = []
|
||||
obj.deferred_waiting = set()
|
||||
obj._kv_wait_start_ts = {}
|
||||
obj._omni_kv_config = None
|
||||
return obj
|
||||
|
||||
Loaded 100 of 234 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user